From 684d4bd3175deb029389325a4f6b221283454159 Mon Sep 17 00:00:00 2001 From: "Daisuke Majima (MLBoy)" Date: Thu, 3 Sep 2026 00:52:50 +0900 Subject: [PATCH 001/190] Skip impure ops in constant_prop_pass (#22418) ## Summary `constant_prop_pass` folds every `call_function` node whose arguments are all constants. Ops that draw from the RNG take only sizes as arguments, so they qualify: a model returning `x + torch.rand(4)` came out of the pass with the draw frozen into `_prop_tensor_constant0`, and returned the same value on every call. @JakeStevens spotted this while reviewing #22391 (repro there). The pass already runs in the Qualcomm and Samsung backends and in `quant_fusion_pass`, so the fix stands on its own. Skip nodes that `torch.fx.Node.is_impure()` reports as impure. That covers ops tagged `nondeterministic_seeded` (`rand`, `randn`, `bernoulli`, `dropout`, ...), mutable schemas and side-effectful functions, and is the same check `eliminate_dead_code` uses to decide what it must keep. ## Test plan New `test_constant_prop_pass_skips_nondeterministic_ops` in `exir/tests/test_passes.py`: after the pass one `aten.rand` node remains, no constant was added, and two calls give different outputs. It fails on main with `0 != 1`. `python -m unittest executorch.exir.tests.test_passes -k constant_prop`: 15 tests pass (torch 2.13.0, macOS arm64). --- exir/passes/constant_prop_pass.py | 4 ++++ exir/tests/test_passes.py | 25 +++++++++++++++++++++++++ 2 files changed, 29 insertions(+) diff --git a/exir/passes/constant_prop_pass.py b/exir/passes/constant_prop_pass.py index 11640d875c0..ea8ee1ad3a9 100644 --- a/exir/passes/constant_prop_pass.py +++ b/exir/passes/constant_prop_pass.py @@ -146,6 +146,10 @@ def get_propagated_const_tensor_dict( node.op != "call_function" or node.target is memory.alloc or node.target in all_skip_targets + # Ops with side effects (RNG draws, mutation) have to run at + # runtime. `aten.rand` has no tensor inputs, so without this check + # it would be folded into a single frozen draw. + or node.is_impure() ): continue diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 0b586bd44cd..54211a490a3 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -2475,6 +2475,31 @@ def forward(self, x): # 1 constant: a (= self.w @ self.cst) self.assertEqual(1, len(pass_result.constants)) + def test_constant_prop_pass_skips_nondeterministic_ops(self) -> None: + """ + Ops that draw from the RNG take no tensor inputs, so they look constant + to the pass. They have to stay in the graph: folding one would freeze a + single random draw into the program. + """ + + class RandomAdd(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + torch.rand(4) + + x = torch.zeros(4) + edge = to_edge(export(RandomAdd(), (x,), strict=True)) + new_ep = constant_prop_pass(edge.exported_program()) + + rand_nodes = [ + node + for node in new_ep.graph.nodes + if node.target == exir_ops.edge.aten.rand.default + ] + self.assertEqual(len(rand_nodes), 1) + self.assertEqual(len(new_ep.constants), 0) + module = new_ep.module() + self.assertFalse(torch.equal(module(x), module(x))) + def test_constant_prop_pass_zero_stride_tensors(self) -> None: """ Test that constant propagation correctly handles tensors with zero strides From c65ad53ad6861de3d9a76f31ddaf882a7115176f Mon Sep 17 00:00:00 2001 From: Ben Mehlow Date: Wed, 2 Sep 2026 12:39:14 -0400 Subject: [PATCH 002/190] Add oncall to all executorch build files missing one (#20496) Differential Revision: D109546129 Pull Request resolved: https://github.com/pytorch/executorch/pull/20496 --- BUCK | 1 + backends/aoti/BUCK | 2 ++ backends/aoti/slim/c10/core/BUCK | 2 ++ backends/aoti/slim/c10/core/test/BUCK | 2 ++ backends/aoti/slim/c10/macros/BUCK | 2 ++ backends/aoti/slim/core/BUCK | 2 ++ backends/aoti/slim/core/test/BUCK | 2 ++ backends/aoti/slim/factory/BUCK | 2 ++ backends/aoti/slim/util/BUCK | 2 ++ backends/aoti/slim/util/test/BUCK | 2 ++ backends/arm/BUCK | 2 ++ backends/arm/_passes/BUCK | 2 ++ backends/arm/debug/BUCK | 2 ++ backends/arm/operator_support/BUCK | 2 ++ backends/arm/operators/BUCK | 2 ++ backends/arm/quantizer/BUCK | 2 ++ backends/arm/tosa/BUCK | 2 ++ backends/arm/tosa/dialect/BUCK | 2 ++ backends/mediatek/BUCK | 2 ++ backends/mediatek/_passes/BUCK | 2 ++ backends/mediatek/quantizer/BUCK | 2 ++ backends/qualcomm/debugger/BUCK | 2 ++ backends/xnnpack/quantizer/BUCK | 2 ++ examples/apple/coreml/llama/BUCK | 2 ++ examples/apple/coreml/scripts/BUCK | 2 ++ exir/backend/test/demos/BUCK | 2 ++ shim_et/BUCK | 2 ++ shim_et/third-party/nlohmann-json/BUCK | 2 ++ shim_et/third-party/re2/BUCK | 2 ++ 29 files changed, 57 insertions(+) create mode 100644 BUCK diff --git a/BUCK b/BUCK new file mode 100644 index 00000000000..fd09ad5ebd3 --- /dev/null +++ b/BUCK @@ -0,0 +1 @@ +oncall("executorch") diff --git a/backends/aoti/BUCK b/backends/aoti/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/BUCK +++ b/backends/aoti/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/aoti/slim/c10/core/BUCK b/backends/aoti/slim/c10/core/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/slim/c10/core/BUCK +++ b/backends/aoti/slim/c10/core/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/aoti/slim/c10/core/test/BUCK b/backends/aoti/slim/c10/core/test/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/c10/core/test/BUCK +++ b/backends/aoti/slim/c10/core/test/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/c10/macros/BUCK b/backends/aoti/slim/c10/macros/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/c10/macros/BUCK +++ b/backends/aoti/slim/c10/macros/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/core/BUCK b/backends/aoti/slim/core/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/core/BUCK +++ b/backends/aoti/slim/core/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/core/test/BUCK b/backends/aoti/slim/core/test/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/core/test/BUCK +++ b/backends/aoti/slim/core/test/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/factory/BUCK b/backends/aoti/slim/factory/BUCK index 981043e51d7..6b124bd857b 100644 --- a/backends/aoti/slim/factory/BUCK +++ b/backends/aoti/slim/factory/BUCK @@ -1,4 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("targets.bzl", "define_common_targets") +oncall("executorch") + fbcode_target(_kind = define_common_targets,) diff --git a/backends/aoti/slim/util/BUCK b/backends/aoti/slim/util/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/slim/util/BUCK +++ b/backends/aoti/slim/util/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/aoti/slim/util/test/BUCK b/backends/aoti/slim/util/test/BUCK index 77871de4469..f91c46c0f20 100644 --- a/backends/aoti/slim/util/test/BUCK +++ b/backends/aoti/slim/util/test/BUCK @@ -1,3 +1,5 @@ load("targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/backends/arm/BUCK b/backends/arm/BUCK index eb8c6bc7590..f6746f29951 100644 --- a/backends/arm/BUCK +++ b/backends/arm/BUCK @@ -9,6 +9,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "constants", diff --git a/backends/arm/_passes/BUCK b/backends/arm/_passes/BUCK index 45aca62d95f..401cbd8b18e 100644 --- a/backends/arm/_passes/BUCK +++ b/backends/arm/_passes/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "core", diff --git a/backends/arm/debug/BUCK b/backends/arm/debug/BUCK index 8374fab364e..e87e419763f 100644 --- a/backends/arm/debug/BUCK +++ b/backends/arm/debug/BUCK @@ -2,6 +2,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "schema", diff --git a/backends/arm/operator_support/BUCK b/backends/arm/operator_support/BUCK index 9152adf2b5b..aee02389128 100644 --- a/backends/arm/operator_support/BUCK +++ b/backends/arm/operator_support/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "operator_support", diff --git a/backends/arm/operators/BUCK b/backends/arm/operators/BUCK index 2555aa78035..5d71cf151e5 100644 --- a/backends/arm/operators/BUCK +++ b/backends/arm/operators/BUCK @@ -2,6 +2,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "node_visitor", diff --git a/backends/arm/quantizer/BUCK b/backends/arm/quantizer/BUCK index 632a38523e2..c9c2fcdb203 100644 --- a/backends/arm/quantizer/BUCK +++ b/backends/arm/quantizer/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + # Exposed through __init__.py fbcode_target( _kind = runtime.python_library, diff --git a/backends/arm/tosa/BUCK b/backends/arm/tosa/BUCK index b7073c97a15..ccf4f461a8b 100644 --- a/backends/arm/tosa/BUCK +++ b/backends/arm/tosa/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "schemas", srcs = ["schemas/__init__.py"], diff --git a/backends/arm/tosa/dialect/BUCK b/backends/arm/tosa/dialect/BUCK index b3a484bb8d9..3c1a070f82a 100644 --- a/backends/arm/tosa/dialect/BUCK +++ b/backends/arm/tosa/dialect/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "core", srcs = [ diff --git a/backends/mediatek/BUCK b/backends/mediatek/BUCK index fa68ea9c2be..c194e3335a9 100644 --- a/backends/mediatek/BUCK +++ b/backends/mediatek/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "preprocess", srcs = [ diff --git a/backends/mediatek/_passes/BUCK b/backends/mediatek/_passes/BUCK index 931c337cf4e..4f8b78983b7 100644 --- a/backends/mediatek/_passes/BUCK +++ b/backends/mediatek/_passes/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "passes", srcs = [ diff --git a/backends/mediatek/quantizer/BUCK b/backends/mediatek/quantizer/BUCK index 847bfc547b2..1136a5c3888 100644 --- a/backends/mediatek/quantizer/BUCK +++ b/backends/mediatek/quantizer/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "quantizer", diff --git a/backends/qualcomm/debugger/BUCK b/backends/qualcomm/debugger/BUCK index 28b5e68a879..ac2e13cebd0 100644 --- a/backends/qualcomm/debugger/BUCK +++ b/backends/qualcomm/debugger/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "utils", diff --git a/backends/xnnpack/quantizer/BUCK b/backends/xnnpack/quantizer/BUCK index 72dd65c00f8..8257d561197 100644 --- a/backends/xnnpack/quantizer/BUCK +++ b/backends/xnnpack/quantizer/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target( _kind = runtime.python_library, name = "xnnpack_quantizer", diff --git a/examples/apple/coreml/llama/BUCK b/examples/apple/coreml/llama/BUCK index 8604a50d9f7..ac9cd6eaf3a 100644 --- a/examples/apple/coreml/llama/BUCK +++ b/examples/apple/coreml/llama/BUCK @@ -4,6 +4,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "no load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +oncall("executorch") + fbcode_target(_kind = runtime.python_library, name = "llama_transformer", srcs = [ diff --git a/examples/apple/coreml/scripts/BUCK b/examples/apple/coreml/scripts/BUCK index 42a97ea893f..2685fe53140 100644 --- a/examples/apple/coreml/scripts/BUCK +++ b/examples/apple/coreml/scripts/BUCK @@ -3,6 +3,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "no # targets.bzl. This file can contain fbcode-only targets. load("@fbcode_macros//build_defs:python_binary.bzl", "python_binary") +oncall("executorch") + fbcode_target(_kind = python_binary, name = "extract_coreml_models", srcs = [ diff --git a/exir/backend/test/demos/BUCK b/exir/backend/test/demos/BUCK index 8404c5982c6..b022a9f8fcc 100644 --- a/exir/backend/test/demos/BUCK +++ b/exir/backend/test/demos/BUCK @@ -1,6 +1,8 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +oncall("executorch") + fbcode_target(_kind = python_unittest, name = "test_delegate_aten_mode", srcs = [ diff --git a/shim_et/BUCK b/shim_et/BUCK index a1a9bdaf65d..2020ed0a73e 100644 --- a/shim_et/BUCK +++ b/shim_et/BUCK @@ -5,6 +5,8 @@ load("@prelude//toolchains:genrule.bzl", "system_genrule_toolchain") load("@prelude//toolchains:python.bzl", "system_python_bootstrap_toolchain", "system_python_toolchain") load("@prelude//toolchains:remote_test_execution.bzl", "remote_test_execution_toolchain") +oncall("executorch") + # Although the non-Android toolchains below are present in shim/BUCK, it appears that we # have to duplicate them here or builds won't work. system_cxx_toolchain( diff --git a/shim_et/third-party/nlohmann-json/BUCK b/shim_et/third-party/nlohmann-json/BUCK index c0b4f27eb52..0a42614a385 100644 --- a/shim_et/third-party/nlohmann-json/BUCK +++ b/shim_et/third-party/nlohmann-json/BUCK @@ -1,3 +1,5 @@ load(":targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() diff --git a/shim_et/third-party/re2/BUCK b/shim_et/third-party/re2/BUCK index c0b4f27eb52..0a42614a385 100644 --- a/shim_et/third-party/re2/BUCK +++ b/shim_et/third-party/re2/BUCK @@ -1,3 +1,5 @@ load(":targets.bzl", "define_common_targets") +oncall("executorch") + define_common_targets() From 457a2a8b9f7d103765d73752c5d2efc6b2e8c8bc Mon Sep 17 00:00:00 2001 From: Nikhil Viswanath Sivakumar <68182521+nil-is-all@users.noreply.github.com> Date: Wed, 2 Sep 2026 13:31:34 -0400 Subject: [PATCH 003/190] Enhance workflow to label external PRs (#22381) Updated the workflow to label external PRs with a community label and refined the exclusion list for authors and bots. We already have a label - *community: contribution*. Just updating the workflow which adds community contributed PRs to the Project Board to have this label. --- .../workflows/add-unanswered-to-project.yml | 86 +++++++++++++------ 1 file changed, 59 insertions(+), 27 deletions(-) diff --git a/.github/workflows/add-unanswered-to-project.yml b/.github/workflows/add-unanswered-to-project.yml index 4a5702b2826..acbdb906b2b 100644 --- a/.github/workflows/add-unanswered-to-project.yml +++ b/.github/workflows/add-unanswered-to-project.yml @@ -11,11 +11,12 @@ on: pull_request: paths: - .github/workflows/add-unanswered-to-project.yml + jobs: add_to_project: runs-on: ubuntu-latest steps: - - name: Add open issues and open, non-draft PRs to org project (excluding certain authors and bots) + - name: Add open issues and open, non-draft PRs to org project and label external PRs (excluding certain authors and bots) uses: actions/github-script@v7 with: github-token: ${{ secrets.ET_EXT_CONTRIB }} @@ -26,36 +27,36 @@ jobs: // List of authors to exclude const excludedAuthors = new Set([ - "nil-is-all", "tanvirislam-meta", "cbilgin", "kimishpatel", "psiddh", "digantdesai", "SS-JIA", "ahmtox", "mcr229", - "shoumikhin", "manuelcandales", "metascroy", "cccclai", "rohansjoshi", "kirklandsign", "abhinaykukkadapu", "JacobSzwejbka", - "Conarnar", "rascani", "kiymetakdemir", "JCNTH", "lucylq", "larryliu0820", "BujSet", "Gasoonjia", "Juntian777", "guangy10", - "jackzhxng", "GregoryComer", "leafs1", "swolchok", "mergennachin", "tarun292", "byjlw", "jathu", "Jack-Khuu", "georgehong", - "zhenyan-zhang-meta", "silverguo", "harishs88ss", "AlannaBurke", "Doggeral", "laithsakka", "Reubend", "dbort", "huydhn", "mcremon-meta", - "trivedivivek", "angelayi", "helunwencser", "hsharma35", "zhxchen17", "iseeyuan", "svekars", "nathanaelsee", "dulinriley", - "jerryzh168", "cmodi-meta", "bigfootjon", "sxu", "ydwu4", "Riandy", "tugsbayasgalan", "bsoyluoglu", "yangw-dev", + "nil-is-all", "tanvirislam-meta", "cbilgin", "kimishpatel", "psiddh", "digantdesai", "SS-JIA", "ahmtox", "mcr229", + "shoumikhin", "manuelcandales", "metascroy", "cccclai", "rohansjoshi", "kirklandsign", "abhinaykukkadapu", "JacobSzwejbka", + "Conarnar", "rascani", "kiymetakdemir", "JCNTH", "lucylq", "larryliu0820", "BujSet", "Gasoonjia", "Juntian777", "guangy10", + "jackzhxng", "GregoryComer", "leafs1", "swolchok", "mergennachin", "tarun292", "byjlw", "jathu", "Jack-Khuu", "georgehong", + "zhenyan-zhang-meta", "silverguo", "harishs88ss", "AlannaBurke", "Doggeral", "laithsakka", "Reubend", "dbort", "huydhn", "mcremon-meta", + "trivedivivek", "angelayi", "helunwencser", "hsharma35", "zhxchen17", "iseeyuan", "svekars", "nathanaelsee", "dulinriley", + "jerryzh168", "cmodi-meta", "bigfootjon", "sxu", "ydwu4", "Riandy", "tugsbayasgalan", "bsoyluoglu", "yangw-dev", "YIWENX14", "namanahuja", "yushangdi", "limintang", "pianpwk", "viveknayakatmeta", "andreanicastro", "JakeStevens", "gmagogsfm", "zonglinpeng", "eigen-k", "derekxu", "salilsdesai", "skrtskrtfb", "pssrawat", "r-barnes", "kalpit-meta-1", "Will-MingLun-Li", "KapJI", "piyengar", "j-bahr", "BoyuanFeng", "fgasperij", "DariusHolmgren", "sammarden-meta", "kushrast", "meta-emilian", "Rittzz", "jeanschmidt", "copyrightly", "mikekgfb", "vmpuri", - "zonglinpengmeta", "maggiemoss", "aorenste", "hoangminhle98", "Solumin", "meyering", "rchen152", "AishwaryaSivaraman", - "migeed-z", "ebgraham", "Esteb37", "nausicaasnow", "Camyll", "ezyang", "huiyujie", "dltn", "cjhopman", "blackm00n", - "agunapal", "SamGondelman", "Ninja91", "ivayloen", "DrJessop", "rodrigos01meta", "akrieger", "cmt0", "yiming0416", - "ethansfng", "ThomasJannaud", "nirvanagth", "marcinkwiatkowski", "3l1", "omerjerk", "nitish2112", "yipjustin", - "ejnguyen", "andrewor14", "phaiting", "mgiordy", "LeeOHzzZ", "adicatana", "Polyomino", "ezrilow", "navsud", - "michaelmaitland", "RahulC7", "seyeong-han", "thdusdl1219", "jaejunku", "felixweilbach", "apullin", "trviv", "junluan01", - "mvartani-meta", "abeakkas", "elpdumont", "corporateshark", "bdemirb", "GeorgeTzoupis", "AdithyaReddy9", "drinkmorewaterr", - "aliafzal", "YifanShenSZ", "RdoubleA", "Olivia-liu", "Abhi-hpp", "Vysarat","azad-meta", "junpi", - "pytorchbot", "pytorchmergebot", "pytorchupdatebot", "facebook-github-bot", "app/dependabot", - "Erik-Lundell", "zingo", "AdrianLundell", "oscarandersson8218", "per", "Sebastian-Larsson", "SaoirseARM", "robell", - "mansnils", "martinlsm", "freddan80", "YufengShi-dudu", "tom-arm", "perheld", "Jerry-Ge", "gggekov", "fumchin", "wwwind", - "benkli01", "Tessil", "maddun01", "Michiel-Olieslagers", "armwaheed", "agrima1304", "emmakujala", "annietllnd", - "MatthiasHertel80", "AlexTawseArm", "jmahbs", "morgolock", "Christoffer-JL", "ArmRyan", "xingguo01", "tgonzalezorlandoarm", + "zonglinpengmeta", "maggiemoss", "aorenste", "hoangminhle98", "Solumin", "meyering", "rchen152", "AishwaryaSivaraman", + "migeed-z", "ebgraham", "Esteb37", "nausicaasnow", "Camyll", "ezyang", "huiyujie", "dltn", "cjhopman", "blackm00n", + "agunapal", "SamGondelman", "Ninja91", "ivayloen", "DrJessop", "rodrigos01meta", "akrieger", "cmt0", "yiming0416", + "ethansfng", "ThomasJannaud", "nirvanagth", "marcinkwiatkowski", "3l1", "omerjerk", "nitish2112", "yipjustin", + "ejnguyen", "andrewor14", "phaiting", "mgiordy", "LeeOHzzZ", "adicatana", "Polyomino", "ezrilow", "navsud", + "michaelmaitland", "RahulC7", "seyeong-han", "thdusdl1219", "jaejunku", "felixweilbach", "apullin", "trviv", "junluan01", + "mvartani-meta", "abeakkas", "elpdumont", "corporateshark", "bdemirb", "GeorgeTzoupis", "AdithyaReddy9", "drinkmorewaterr", + "aliafzal", "YifanShenSZ", "RdoubleA", "Olivia-liu", "Abhi-hpp", "Vysarat","azad-meta", "junpi", + "pytorchbot", "pytorchmergebot", "pytorchupdatebot", "facebook-github-bot", "app/dependabot", + "Erik-Lundell", "zingo", "AdrianLundell", "oscarandersson8218", "per", "Sebastian-Larsson", "SaoirseARM", "robell", + "mansnils", "martinlsm", "freddan80", "YufengShi-dudu", "tom-arm", "perheld", "Jerry-Ge", "gggekov", "fumchin", "wwwind", + "benkli01", "Tessil", "maddun01", "Michiel-Olieslagers", "armwaheed", "agrima1304", "emmakujala", "annietllnd", + "MatthiasHertel80", "AlexTawseArm", "jmahbs", "morgolock", "Christoffer-JL", "ArmRyan", "xingguo01", "tgonzalezorlandoarm", "chizkiyahu", "sarah-blades", "itsMarco-G", "usamahz", "Rob-Hughes-Arm", "swha815", "FabulousSuperDude", - "haowhsu-quic", "shewu-quic", "winskuo-quic", "chunit-quic", "DannyYuyang-quic", "chuntl", "thchenqti", "jethroqti", - "chenweng-quic", "qti-horodnic", "qti-mmadhava", "quic-boyuc", "zhaoxul-qti", - "cymbalrush", "DenisVieriu97", "billmguo", - "StrycekSimon", "jirioc", "robert-kalmar", "skywall", "MartinPavella", "roman-janik-nxp", "novak-vaclav", "irtrukhina", - "neuropilot-captain", "dijopaul", "cad-rlc", "cad-audio", "ynimmaga", "daniil-lyakhov", + "haowhsu-quic", "shewu-quic", "winskuo-quic", "chunit-quic", "DannyYuyang-quic", "chuntl", "thchenqti", "jethroqti", + "chenweng-quic", "qti-horodnic", "qti-mmadhava", "quic-boyuc", "zhaoxul-qti", + "cymbalrush", "DenisVieriu97", "billmguo", + "StrycekSimon", "jirioc", "robert-kalmar", "skywall", "MartinPavella", "roman-janik-nxp", "novak-vaclav", "irtrukhina", + "neuropilot-captain", "dijopaul", "cad-rlc", "cad-audio", "ynimmaga", "daniil-lyakhov", "emmanuel-ferdman", "cavusmustafa", "anzr299", "suryasidd", "Jiseong-oh", "alexdean08", // explicitly include the dependabot bot login seen in PRs "dependabot[bot]" @@ -69,7 +70,10 @@ jobs: // Labels on PRs to exclude from being added to the project const excludedPrLabels = new Set(["fb-exported", "meta-exported"]); - + + // Label applied to PRs that pass every external-contributor check below + const communityLabel = "community: contribution"; + // Simple cache for user -> boolean (member of excluded org) const orgsCache = new Map(); const companyCache = new Map(); @@ -125,6 +129,31 @@ jobs: return false; } } + + function hasCommunityLabel(item) { + if (!item || !item.labels) return false; + return item.labels.some(l => l && l.name && l.name.toLowerCase() === communityLabel); + } + + async function addCommunityLabel(pr) { + if (hasCommunityLabel(pr)) { + console.log(`PR #${pr.number} already has "${communityLabel}"`); + return; + } + try { + await github.rest.issues.addLabels({ + owner, + repo, + issue_number: pr.number, + labels: [communityLabel] + }); + console.log(`Labeled PR #${pr.number} with "${communityLabel}"`); + } catch (error) { + // Labeling is best-effort: a failure here must not stop the rest of the run. + console.log(`Error labeling PR #${pr.number}: ${error.message}`); + } + } + async function addItem(contentId, type, number) { try { await github.graphql(` @@ -156,6 +185,7 @@ jobs: filter: 'all' } ); + for (const issue of issues) { if (issue.pull_request) { console.log(`Skipping PR #${issue.number} (listed in issues)`); @@ -185,6 +215,7 @@ jobs: state: 'open', } ); + for (const pr of prs) { if (pr.draft) { console.log(`Skipping PR #${pr.number} (draft)`); @@ -207,6 +238,7 @@ jobs: continue; } await addItem(pr.node_id, 'pr', pr.number); + await addCommunityLabel(pr); } } catch (error) { core.setFailed(`Workflow failed: ${error.message}`); From a6b115bff6948db23e2f528372db08ff6390834f Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 2 Sep 2026 11:14:50 -0700 Subject: [PATCH 004/190] Fix MLX C++20 errors (#22424) --- backends/mlx/CMakeLists.txt | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 2a3d1546755..58f36465357 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -452,6 +452,14 @@ target_compile_options(mlxdelegate PRIVATE ${_common_compile_options}) # Core tensor headers carry pre-existing narrowing conversions that trip Xcode's # -Wshorten-64-to-32 -Werror; suppress it here as XNNPACK and abseil already do. target_compile_options(mlxdelegate PRIVATE -Wno-shorten-64-to-32) +# MLX headers use C++20 defaulted comparison operators while we build C++17. +# They normally arrive via -isystem from the imported mlx target, which +# suppresses the warning. But this backend also installs them for downstream +# consumers, and every -I is searched before any -isystem, so once an install +# has run they resolve out of ${CMAKE_INSTALL_PREFIX}/include instead and +# -Werror turns fatal. That is why an incremental rebuild used to need a +# cmake-out wipe: a clean tree has nothing installed yet. +target_compile_options(mlxdelegate PRIVATE -Wno-c++20-extensions) if(EXECUTORCH_MLX_ENABLE_SANITIZERS) target_link_options(mlxdelegate PRIVATE ${_mlx_sanitizer_link_options}) endif() From 4747ab71f67ce9d0c64f5a18af5eb624b796989b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 2 Sep 2026 12:26:34 -0700 Subject: [PATCH 005/190] Run attention on MLX wherever the fused kernel can compute it (#22419) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### The problem Attention on rank-3 tensors exports without complaint and then fails when you run the model: ``` MLX execute failed: [scaled_dot_product_attention] input with shape (2,16,64) expected to be rank 4 ``` PyTorch accepts rank 2, 3, 4 and 5 here. The fused kernel takes rank 4 only, and has other requirements besides. The handler matched the operator by name and checked none of them, so it claimed calls the kernel cannot run. ### Adapting the shapes that can be adapted `ExpandDimsNode` → `SdpaNode` → `SqueezeNode`, as suggested in review, so rank 3 keeps the fused kernel. Rank 2 needs it too and takes two added dimensions. The dimensions go at the front. For a rank-3 input the first dimension is already the head one, so inserting in the middle moves it into the batch slot, which misaligns masks and grouped heads. Admitting a lower rank also reaches code that was only ever given rank 4. The grouped key and value unwrap looks for a repeat on dimension 1, which is the head dimension at rank 4 and the key sequence below it, so at rank 3 it would absorb a repeat carrying real keys. Without a mask the result still matches, because a duplicated key and its value give the same weighted sum, which is what makes this easy to miss. With a causal mask it is off by 5.1. Both unwraps now run only at rank 4. ### Three shapes left to decompose Decomposed calls still run on this backend, which the tests assert rather than assume. | Call | Today | Here | | --- | --- | --- | | Rank 5 and above | fails at execute | decomposes, 3.6e-07 | | Unequal batch at rank 4 | fails at execute | decomposes, 2.4e-07 | | Causal at rank 2 or 3, query shorter than key | not reachable | decomposes, 4.2e-07 | The causal one needs a word. Torch anchors a causal mask at the top left and MLX at the bottom right, so the two disagree whenever the lengths differ, silently. Rank 4 already reaches the kernel today and keeps its current behaviour: fixing that needs either a change to what the speech example computes or a fix to an off-by-one in the decomposed path, so it is filed separately as #22426. What this change does is avoid opening two more ranks onto the same trap. ### Declining a call is not free Preservation from decomposition is requested per operator, so one claimed call keeps the operator whole for a declined one in the same graph, and that call is then neither lowered nor decomposed: ``` RuntimeError: Missing out variants: {'aten::scaled_dot_product_attention'} ``` Give the whole operator back when any of its calls is unsupported. The framework does offer a per-node filter that would keep the other calls fused. It is deliberately not used, and the docstring says why: on an attention block that reshapes its output, with a declined call in the same graph, that path fails to lower at all with `Cannot view a tensor with shape (1,16,4,16) and strides (1024,16,256,1)`. It trades the fusion cost for a hard failure on a very ordinary shape. The cost of the coarser choice is real and written down: one declined call unfuses the operator's other calls in that graph. This half is not specific to attention. Two calls to `torch.roll` in one graph, one supported and one not, already fail at execute today for the same reason. ### Test plan Fifteen partitioner tests. **Eleven fail before this change**, all fifteen pass after. They assert on serialized nodes, so they check which path a call took rather than only that it answered, and the rejection cases also assert the work stayed on this backend, so they cannot pass when nothing is delegated at all. A rank-3 case was added to the operator suite too, which runs the compiled runtime. Ran every combination on an Apple Silicon Mac against eager: plain, causal, explicit scale, float mask, boolean mask, grouped-query attention, batch size 1, float16 and bfloat16, ranks 2 through 5. - rank 3 matches to 3.6e-07 and stays fused in all seven variants; rank 2 to 3.6e-07 - rank 5 decomposes at 3.6e-07 where it previously failed; unequal batch from failing to 2.4e-07 - the grouped key case measured both ways: absorbed and correct at rank 4, and off by 5.1 at rank 3 with a causal mask if absorbed, which is what the guard prevents - a zero key head count used to divide by zero and abort the whole export; it now declines that node - a mixed supported and unsupported graph exports and runs; two rolls in one graph go from failing at execute to matching eager exactly Exported the speech example and compared it against the same export without this change: **identical serialized node counts and an identical greedy token sequence over twelve decode steps**, so that model is unaffected. Ran the neighbouring backend test files, 71 tests, all passing. Co-authored-by: PyTorch Bot --- backends/mlx/partitioner.py | 17 ++ backends/mlx/patterns.py | 131 ++++++++++++++- backends/mlx/test/test_ops.py | 24 +++ backends/mlx/test/test_partitioner.py | 231 +++++++++++++++++++++++++- 4 files changed, 395 insertions(+), 8 deletions(-) diff --git a/backends/mlx/partitioner.py b/backends/mlx/partitioner.py index 7814e883588..a82f54cda00 100644 --- a/backends/mlx/partitioner.py +++ b/backends/mlx/partitioner.py @@ -133,6 +133,18 @@ def ops_to_not_decompose( handler that rejects the 6-arg edge form, for instance). Preserving an op the handler then rejects is worse than not preserving it, because the op neither decomposes into something delegatable nor lowers itself. + + A target is only preserved when every node carrying it is supported. One + unsupported node is enough to give the whole operator back to decomposition, + because keeping it would leave that node neither lowered nor decomposed and + export would stop. + + The second return value is a per-node filter, which would keep the supported + calls fused and decompose only the rest. It is deliberately not used: it puts + the program on exir's EDGE_DO_NOT_DECOMP path, which fails on an ordinary + attention block that reshapes its output, a shape this backend has to lower. + The cost of the coarser choice is that one declined call also unfuses the + operator's other calls in that graph. """ from executorch.backends.mlx.builder.program_builder import MLXProgramBuilder @@ -157,6 +169,7 @@ def ops_to_not_decompose( # Collect ops for nodes that are actually supported do_not_decompose: list[torch._ops.OpOverload] = [] + declined: set[torch._ops.OpOverload] = set() for node in ep.graph.nodes: if node.op == "call_function" and isinstance( @@ -166,6 +179,10 @@ def ops_to_not_decompose( if info is not None and info.supported: if node.target not in do_not_decompose: do_not_decompose.append(node.target) + else: + declined.add(node.target) + + do_not_decompose = [op for op in do_not_decompose if op not in declined] self._not_decompose_cache = (weakref.ref(ep), do_not_decompose) diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 1760dc63b9a..5a8525eacce 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -44,6 +44,7 @@ AddIntNode, AddNode, AsTypeNode, + ExpandDimsNode, IndexCopyNode, IntOrVid, ModIntNode, @@ -52,10 +53,12 @@ SdpaNode, SliceNode, SliceUpdateNode, + SqueezeNode, SubtractIntNode, SymSizeNode, ) from torch.export.exported_program import ExportedProgram +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.node import Node @@ -527,6 +530,61 @@ def _try_unwrap_repeat_kv(cls, node: Node) -> Optional[Tuple[Node, List[Node]]]: body = [e for e in entries if e is not None] return base, body + @classmethod + def _kernel_can_compute(cls, sdpa_node: Node) -> bool: + """Whether the fused kernel can compute this call faithfully. + + Its preconditions are narrower than what PyTorch accepts, and a claimed call + that the kernel cannot compute fails loudly at execute rather than falling + back, so declining it here is what sends it to decomposition instead. + """ + q, k, v, attn_mask, _, is_causal, _, _ = cls._parse_sdpa_args_and_kwargs( + sdpa_node + ) + operand_vals = [ + operand.meta.get("val") if isinstance(operand, Node) else None + for operand in (q, k, v) + ] + if any(val is None for val in operand_vals): + return False + + # Ranks 2 and 3 are lifted to 4 on emission. Beyond 4 the leading dimensions + # would have to fold together, and a fold pairs the wrong operands as soon as + # one of them broadcasts a batch the others do not. + ranks = {val.dim() for val in operand_vals} + if len(ranks) != 1 or not 2 <= next(iter(ranks)) <= 4: + return False + + # Torch anchors a causal mask at the top left and MLX at the bottom right, so + # the two agree only when the query and key lengths are equal. Rank 4 already + # reaches the kernel today and is left alone; the lower ranks are newly lifted + # here, so do not open a path that returns wrong values with no error. + if ( + next(iter(ranks)) < 4 + and is_causal + and not statically_known_true( + operand_vals[0].shape[-2] == operand_vals[1].shape[-2] + ) + ): + return False + + # The kernel requires the batch sizes to match and rejects the call otherwise, + # so a broadcast batch has to be decomposed rather than fused. + if next(iter(ranks)) == 4 and any( + not statically_known_true(val.shape[0] == operand_vals[0].shape[0]) + for val in operand_vals[1:] + ): + return False + + if attn_mask is not None: + mask_val = ( + attn_mask.meta.get("val") if isinstance(attn_mask, Node) else None + ) + if mask_val is None or mask_val.dim() > 4: + return False + + return True + @classmethod def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler"]: sdpa_node = head @@ -535,15 +593,23 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler" ): return None + if not cls._kernel_can_compute(sdpa_node): + return None + q, k, v, _, _, _, _, _ = cls._parse_sdpa_args_and_kwargs(sdpa_node) - # Detect grouped kv attention pattern with repeat_interleave before SDPA + # Detect grouped kv attention pattern with repeat_interleave before SDPA. + # Both unwraps below key on dim 1, which is the head dimension only at rank 4. + # At a lower rank dim 1 is the key sequence, and absorbing a repeat there + # drops keys, which a causal mask then turns into a wrong answer. + is_rank4 = q.meta["val"].dim() == 4 if isinstance(q, Node) else False is_grouped_kv = False k_base = k v_base = v body: List[Node] = [] if ( - match_target(k, torch.ops.aten.repeat_interleave.self_int) + is_rank4 + and match_target(k, torch.ops.aten.repeat_interleave.self_int) and has_single_user(k) and (len(k.args) == 3) and (len(k.kwargs) == 0) @@ -563,7 +629,7 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler" # Detect HuggingFace repeat_kv pattern: # unsqueeze(dim=2) → expand → clone → view - if not is_grouped_kv: + if is_rank4 and not is_grouped_kv: k_unwrap = cls._try_unwrap_repeat_kv(k) v_unwrap = cls._try_unwrap_repeat_kv(v) if k_unwrap is not None and v_unwrap is not None: @@ -572,6 +638,27 @@ def maybe_create(cls, ep: ExportedProgram, head: Node) -> Optional["SDPAHandler" is_grouped_kv = True body = k_body + v_body + # Checked after the unwrapping above, because grouped-query attention reaches + # the kernel with its original head counts. MLX pairs heads only when the key + # and value agree and the query is a whole multiple of them. + kernel_vals = [ + node.meta.get("val") if isinstance(node, Node) else None + for node in (q, k_base, v_base) + ] + if any(val is None for val in kernel_vals): + return None + q_heads, k_heads, v_heads = ( + 1 if val.dim() == 2 else val.shape[-3] for val in kernel_vals + ) + # A zero head count would make the multiple test below divide by zero, and a + # raise here aborts the whole export rather than declining this one node. + if not statically_known_true(k_heads > 0): + return None + if not statically_known_true(k_heads == v_heads): + return None + if not statically_known_true(q_heads % k_heads == 0): + return None + head = sdpa_node if not is_grouped_kv: body = [] @@ -593,19 +680,49 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: assert dropout_p == 0.0, "SDPA with dropout is not supported" q, k, v, attn_mask = P.slot_map([q, k, v, attn_mask]) + # Add the dimensions the kernel is missing at the front, never in the middle. + # For a rank-3 input the first dimension is already the head one, so inserting + # there would move it into the batch slot and misalign masks and grouped heads. + input_nodes = (self.q_node, self.k_node, self.v_node) + inputs = [q, k, v] + for i, input_node in enumerate(input_nodes): + for _ in range(4 - input_node.meta["val"].dim()): + _, expanded = P.make_tmp_slot() + P.emit( + ExpandDimsNode( + x=P.slot_to_tid(inputs[i]), + out=P.slot_to_tid(expanded), + axis=0, + ) + ) + inputs[i] = expanded + + output_rank = n.meta["val"].dim() + out = P.make_or_get_slot(n) + sdpa_out = out + if output_rank < 4: + _, sdpa_out = P.make_tmp_slot() P.emit( SdpaNode( - q=P.slot_to_tid(q), - k=P.slot_to_tid(k), - v=P.slot_to_tid(v), - out=P.slot_to_tid(out), + q=P.slot_to_tid(inputs[0]), + k=P.slot_to_tid(inputs[1]), + v=P.slot_to_tid(inputs[2]), + out=P.slot_to_tid(sdpa_out), scale=scale, mask=P.slot_to_tid(attn_mask) if attn_mask else None, causal=is_causal, ) ) + if output_rank < 4: + P.emit( + SqueezeNode( + x=P.slot_to_tid(sdpa_out), + out=P.slot_to_tid(out), + dims=list(range(4 - output_rank)), + ) + ) return out diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 452ae37c40b..9ecdf5a8e58 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -6275,6 +6275,30 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: return (q, k, v) +@register_test +class SDPARank3Test(OpTestCase): + """Attention on rank-3 tensors, which PyTorch accepts and the fused kernel does not. + + The node counts are the point of the test: they assert the fused kernel is still + used, rather than the operator having been decomposed into primitives. + """ + + name = "sdpa_rank3" + rtol = 1e-3 + atol = 1e-3 + expected_node_counts = { + "SdpaNode": 1, + "ExpandDimsNode": 3, + "SqueezeNode": 1, + } + + def create_model(self) -> nn.Module: + return SDPAModel() + + def create_inputs(self) -> Tuple[torch.Tensor, ...]: + return tuple(torch.randn(2, 16, 64) for _ in range(3)) + + class CustomSDPAModel(nn.Module): """ Test model for mlx::custom_sdpa with KVCache. diff --git a/backends/mlx/test/test_partitioner.py b/backends/mlx/test/test_partitioner.py index 4a5833aa656..3b82e306b1f 100644 --- a/backends/mlx/test/test_partitioner.py +++ b/backends/mlx/test/test_partitioner.py @@ -9,12 +9,16 @@ Tests for the MLX partitioner. """ +import tempfile import unittest +from pathlib import Path import torch import torch.nn as nn from executorch.backends.mlx.partitioner import MLXPartitioner -from executorch.exir import EdgeCompileConfig, to_edge +from executorch.backends.mlx.test.test_utils import get_mlx_node_counts +from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower +from executorch.runtime import Runtime from torch.export import export @@ -41,5 +45,230 @@ def forward(self, x): self.assertIn("to_edge_transform_and_lower", str(ctx.exception)) +def _lower(model, inputs): + return to_edge_transform_and_lower( + export(model, inputs, strict=False), + partitioner=[MLXPartitioner()], + ).to_executorch() + + +def _delegate_count(program) -> int: + return sum( + 1 + for node in program.exported_program().graph_module.graph.nodes + if node.op == "call_function" and "executorch_call_delegate" in str(node.target) + ) + + +def _run(model, inputs): + """Lower, execute, and return the node counts, the delegate count and the error. + + The delegate count is returned so a test can tell "decomposed onto this backend" + apart from "not lowered here at all", which a node count alone cannot show. + """ + with torch.no_grad(): + ref = model(*inputs) + program = _lower(model, inputs) + delegates = _delegate_count(program) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + counts = get_mlx_node_counts(path) + method = Runtime.get().load_program(path).load_method("forward") + out = method.execute(list(inputs))[0] + return counts, delegates, (out - ref).abs().max().item() + + +class Sdpa(nn.Module): + def __init__(self, is_causal: bool = False): + super().__init__() + self.is_causal = is_causal + + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=self.is_causal + ) + + +class GroupedSdpa(nn.Module): + """Grouped key/value attention, where the repeat is unwrapped before the kernel.""" + + def __init__(self, dim: int, is_causal: bool = False): + super().__init__() + self.dim = dim + self.is_causal = is_causal + + def forward(self, q, k, v): + k = k.repeat_interleave(2, dim=self.dim) + v = v.repeat_interleave(2, dim=self.dim) + return torch.nn.functional.scaled_dot_product_attention( + q, k, v, is_causal=self.is_causal + ) + + +class TestMLXPartitionerSdpaShapes(unittest.TestCase): + """The fused kernel takes rank 4, so other ranks are adapted or left alone.""" + + def test_rank4_is_unchanged(self): + counts, _, err = _run( + Sdpa(), tuple(torch.randn(2, 4, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 0) + self.assertEqual(counts.get("SqueezeNode", 0), 0) + self.assertLess(err, 1e-4) + + def test_rank3_is_lifted_once(self): + counts, _, err = _run(Sdpa(), tuple(torch.randn(2, 16, 64) for _ in range(3))) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 3) + self.assertEqual(counts.get("SqueezeNode", 0), 1) + self.assertLess(err, 1e-4) + + def test_rank2_is_lifted_twice(self): + counts, _, err = _run(Sdpa(), tuple(torch.randn(16, 64) for _ in range(3))) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 6) + self.assertEqual(counts.get("SqueezeNode", 0), 1) + self.assertLess(err, 1e-4) + + def test_rank5_is_decomposed_on_this_backend(self): + # Folding the leading dimensions pairs the wrong operands once one of them + # broadcasts a batch, so this decomposes rather than fusing. + counts, delegates, err = _run( + Sdpa(), tuple(torch.randn(2, 2, 4, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + def test_unequal_batch_is_decomposed_on_this_backend(self): + counts, delegates, err = _run( + Sdpa(), + ( + torch.randn(2, 4, 16, 64), + torch.randn(1, 4, 16, 64), + torch.randn(1, 4, 16, 64), + ), + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + def test_zero_head_count_declines_instead_of_raising(self): + # The head multiple test would divide by zero here, and raising from the + # matcher aborts the whole export rather than declining this one node. Only + # lowering is checked: a zero-size operand is not executable either way. + program = _lower( + Sdpa(), + ( + torch.randn(1, 4, 8, 16), + torch.randn(1, 0, 8, 16), + torch.randn(1, 0, 8, 16), + ), + ) + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "model.pte" + path.write_bytes(program.buffer) + self.assertEqual(get_mlx_node_counts(path).get("SdpaNode", 0), 0) + + +class TestMLXPartitionerGroupedKeys(unittest.TestCase): + """The grouped key/value unwrap reads dim 1 as the head, which holds at rank 4.""" + + def test_rank4_head_repeat_is_absorbed(self): + counts, _, err = _run( + GroupedSdpa(dim=1), + ( + torch.randn(2, 4, 16, 64), + torch.randn(2, 2, 16, 64), + torch.randn(2, 2, 16, 64), + ), + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("RepeatNode", 0), 0) + self.assertLess(err, 1e-4) + + def test_rank3_sequence_repeat_is_kept(self): + # At rank 3 dim 1 is the key sequence, so absorbing the repeat would drop + # half the keys. Without a mask that still sums correctly, which is what + # makes it easy to miss; with a causal mask it is wrong by whole units. + counts, _, err = _run( + GroupedSdpa(dim=1, is_causal=True), + (torch.randn(2, 16, 64), torch.randn(2, 8, 64), torch.randn(2, 8, 64)), + ) + self.assertEqual(counts.get("RepeatNode", 0), 2) + self.assertLess(err, 1e-4) + + +class TestMLXPartitionerSdpaCausal(unittest.TestCase): + """MLX anchors a causal mask at the bottom right and torch at the top left.""" + + def test_equal_lengths_stay_fused(self): + counts, _, err = _run( + Sdpa(is_causal=True), tuple(torch.randn(1, 4, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertLess(err, 1e-4) + + def test_rank3_equal_lengths_are_lifted(self): + counts, _, err = _run( + Sdpa(is_causal=True), tuple(torch.randn(2, 16, 64) for _ in range(3)) + ) + self.assertEqual(counts.get("SdpaNode", 0), 1) + self.assertEqual(counts.get("ExpandDimsNode", 0), 3) + self.assertLess(err, 1e-4) + + def test_rank3_unequal_lengths_are_not_lifted(self): + # The two conventions disagree here and the disagreement is silent, so a + # shape this backend could not previously reach is not opened up. + counts, delegates, err = _run( + Sdpa(is_causal=True), + (torch.randn(2, 6, 64), torch.randn(2, 16, 64), torch.randn(2, 16, 64)), + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + def test_rank2_unequal_lengths_are_not_lifted(self): + counts, _, err = _run( + Sdpa(is_causal=True), + (torch.randn(6, 64), torch.randn(16, 64), torch.randn(16, 64)), + ) + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertLess(err, 1e-4) + + +class TestMLXPartitionerMixedSupport(unittest.TestCase): + """An operator is preserved from decomposition per operator, not per call.""" + + def test_supported_and_unsupported_calls_in_one_graph(self): + class Mixed(nn.Module): + def forward(self, a, b): + x = torch.nn.functional.scaled_dot_product_attention(a, a, a) + y = torch.nn.functional.scaled_dot_product_attention(b, b, b) + return x.sum() + y.sum() + + # Without giving the whole operator back, the rank-5 call would be neither + # lowered nor decomposed and this would raise a missing out variant. + counts, delegates, err = _run( + Mixed().eval(), + (torch.randn(1, 4, 16, 64), torch.randn(2, 2, 4, 16, 64)), + ) + # The cost of the coarse choice: the supported call is unfused as well. + self.assertEqual(counts.get("SdpaNode", 0), 0) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-3) + + def test_mixed_support_outside_attention(self): + class TwoRolls(nn.Module): + def forward(self, x): + return torch.roll(x, 1, dims=0).sum() + torch.roll(x, 1).sum() + + _, delegates, err = _run(TwoRolls().eval(), (torch.randn(4, 8),)) + self.assertGreater(delegates, 0) + self.assertLess(err, 1e-4) + + if __name__ == "__main__": unittest.main() From d5e4ae962e9e9ed3a6f129d5bb9f40eedd70249b Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 2 Sep 2026 13:54:49 -0700 Subject: [PATCH 006/190] Give causal attention the mask PyTorch asks for (#22443) Fixes #22426. ### The problem Attention with `is_causal=True` returns wrong values, with nothing raised, whenever the query is shorter than the keys: | shape | error against eager | | --- | --- | | query 1, keys 16 | **2.7 to 3.4** | | query 16, keys 16 | 3.6e-07 | The two libraries anchor a causal mask at opposite corners. PyTorch puts it at the top left, so the query at row `i` attends to keys `0..i`. MLX puts it at the bottom right, aligning the last query with the last key. They agree only when the lengths are equal, and nothing checked that. This is the decode step of any model that generates a token at a time, so it is not a corner case. ### The fix Slice the keys and values to the query length when the lengths are not provably equal, then let the kernel apply its own mask to what is now a square problem. That is what PyTorch computes: a row past the last query never reads a later key, so dropping those keys changes nothing. Measured over query lengths 1 to 63 against key lengths 2 to 512, the sliced form and the original agree **exactly** in float32. A slice before attention is also what the cache-aware path in this file already does, so no mask tensor is built and the fused kernel is kept. A causal query **longer** than its keys is left to decompose instead. PyTorch clamps each row to the keys that exist, and neither the kernel's own mask nor a slice reproduces that. ### Why the speech example changes too This backend already knew about the conflict. `mlx::custom_sdpa` exists for cached attention and takes a start position so the kernel gets a square problem, and seven models here use it. The speech example does not. It slices the cache itself and then asks for a causal mask over the result, which in PyTorch means **the new token sees only the first cached key**. It has been getting attention over the whole cache because the kernel read its request the other way round, so it depends on the behaviour this change corrects. Moving it onto the same path as the other models makes it ask for what it wants. Worth stating plainly: that example's own eager math was already 3.3 away from what it intended at a decode step, so this is not a regression being introduced, it is one being removed. ### Test plan Two cases added to the operator suite, a query of 1 and of 6 against 32 keys, which run the compiled runtime and compare against eager. **Both fail before this change on numerics and pass after.** Ran the shapes on an Apple Silicon Mac against eager: - query 1, 2 and 6 against longer keys, float32 and bfloat16: now exact or within 4.2e-07, previously off by around 3 - equal lengths and non-causal calls untouched, fused kernel kept in every case, no mask tensor built - a causal query longer than its keys decomposes and matches eager to 3.6e-07, previously off by 2.7 - equal-length causal attention serializes to a **byte-identical program** before and after, so that path is bit-for-bit unchanged Exported the speech model and compared its logits against HuggingFace's own decoder at a cached decode step: unchanged from before this change, and the top token still matches. Ran the neighbouring backend test files, all passing. Co-authored-by: PyTorch Bot --- .../mlx/examples/whisper/export_whisper.py | 19 ++++++----- backends/mlx/patterns.py | 33 +++++++++++++++++++ backends/mlx/test/test_ops.py | 13 ++++++-- 3 files changed, 55 insertions(+), 10 deletions(-) diff --git a/backends/mlx/examples/whisper/export_whisper.py b/backends/mlx/examples/whisper/export_whisper.py index 97d3a22bc79..84b45042ef5 100644 --- a/backends/mlx/examples/whisper/export_whisper.py +++ b/backends/mlx/examples/whisper/export_whisper.py @@ -122,14 +122,17 @@ def forward( # Update KV cache k_cache, v_cache = self.kv_cache.update(pos_int, k, v) - # Explicit windowing: slice cache to valid positions - end_pos = pos_int + T - k_win = k_cache[:, :, :end_pos, :] - v_win = v_cache[:, :, :end_pos, :] - - # SDPA with causal mask - attn_out = F.scaled_dot_product_attention( - q, k_win, v_win, attn_mask=None, is_causal=True, scale=self.scale + # The cache-aware op slices the cache to start_pos + query length itself, and + # applies the causal mask the kernel wants for a cached decode step. Passing + # a hand-sliced window with is_causal instead means something different: torch + # would let the new token see only the first cached key. + attn_out = torch.ops.mlx.custom_sdpa( + q, + k_cache, + v_cache, + start_pos=pos_int, + is_causal=True, + scale=self.scale, ) # Reshape back diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index 5a8525eacce..f2899baf6d4 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -583,6 +583,14 @@ def _kernel_can_compute(cls, sdpa_node: Node) -> bool: if mask_val is None or mask_val.dim() > 4: return False + # A causal query longer than its keys is left to decompose. Torch clamps each + # row to the keys that exist, and neither the kernel's own mask nor slicing the + # keys to the query length reproduces that. + if is_causal and not statically_known_true( + operand_vals[0].shape[-2] <= operand_vals[1].shape[-2] + ): + return False + return True @classmethod @@ -704,6 +712,31 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: if output_rank < 4: _, sdpa_out = P.make_tmp_slot() + # MLX anchors its causal mask at the bottom right and torch at the top left, so + # the flag alone only means the same thing when the lengths are equal. Slicing + # the keys and values to the query length makes the problem square, where the + # two conventions agree, and it is what torch computes: a query at row i attends + # to keys 0..i, so keys past the last query row are never read. + q_len = self.q_node.meta["val"].shape[-2] + k_len = self.k_node.meta["val"].shape[-2] + if is_causal and not statically_known_true(q_len == k_len): + _, rows = P.slot_manager.make_tmp_value_slot() + P.emit( + SymSizeNode(a=P.slot_to_tid(inputs[0]), dim=2, out=P.slot_to_vid(rows)) + ) + for i in (1, 2): + _, sliced = P.make_tmp_slot() + P.emit( + SliceNode( + x=P.slot_to_tid(inputs[i]), + out=P.slot_to_tid(sliced), + axis=IntOrVid.from_literal(2), + start=IntOrVid.from_literal(0), + stop=P.to_int_or_vid(rows), + ) + ) + inputs[i] = sliced + P.emit( SdpaNode( q=P.slot_to_tid(inputs[0]), diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 9ecdf5a8e58..bfc2e931e67 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -6212,6 +6212,7 @@ def __init__( is_causal: bool = False, use_mask: bool = False, use_bool_mask: bool = False, + kv_seq_len: Optional[int] = None, ): self.batch_size = batch_size self.num_heads = num_heads @@ -6221,12 +6222,15 @@ def __init__( self.is_causal = is_causal self.use_mask = use_mask self.use_bool_mask = use_bool_mask + self.kv_seq_len = kv_seq_len if kv_seq_len is not None else seq_len parts = ["sdpa"] if num_kv_heads is not None: parts.append(f"gqa{num_kv_heads}") if is_causal: parts.append("causal") + if self.kv_seq_len != seq_len: + parts.append(f"q{seq_len}kv{self.kv_seq_len}") if use_mask: parts.append("mask") if use_bool_mask: @@ -6241,6 +6245,11 @@ def get_test_configs(cls) -> List["SDPATest"]: cls(num_kv_heads=4), cls(use_mask=True), cls(use_bool_mask=True), # Test boolean mask conversion + # A decode step against a longer key cache. MLX anchors its causal mask at + # the bottom right and torch at the top left, so they only agree when the + # lengths match. + cls(is_causal=True, seq_len=1, kv_seq_len=32), + cls(is_causal=True, seq_len=6, kv_seq_len=32), ] def create_model(self) -> nn.Module: @@ -6256,8 +6265,8 @@ def create_model(self) -> nn.Module: def create_inputs(self) -> Tuple[torch.Tensor, ...]: q = torch.randn(self.batch_size, self.num_heads, self.seq_len, self.head_dim) kv_heads = self.num_kv_heads if self.num_kv_heads else self.num_heads - k = torch.randn(self.batch_size, kv_heads, self.seq_len, self.head_dim) - v = torch.randn(self.batch_size, kv_heads, self.seq_len, self.head_dim) + k = torch.randn(self.batch_size, kv_heads, self.kv_seq_len, self.head_dim) + v = torch.randn(self.batch_size, kv_heads, self.kv_seq_len, self.head_dim) if self.use_mask: # Additive float mask: 0 = attend, -inf = masked From e39bda831ab44c82e4054970650199d1f43af8ba Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 2 Sep 2026 14:29:08 -0700 Subject: [PATCH 007/190] Add metrics to runner (#22283) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds timing instrumentation to the batched runner, which previously reported nothing. metrics.h defines two tiers: GenerationMetrics, published on GenerationHandle beside finish_reason(), covering the submit→first-step→first-token→end timeline, token counts, and an inter-token latency summary; and EngineMetrics, owned by the runner and read after shutdown(), covering step counts and latency, decode/prefill sequences and tokens, admitted-versus-ready decode sequences, and TTFT rollup. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- extension/llm/batching/CMakeLists.txt | 2 +- extension/llm/batching/metrics.cpp | 140 +++++++ extension/llm/batching/metrics.h | 416 ++++++++++++++++++++ extension/llm/batching/runner.cpp | 369 +++++++++++++++-- extension/llm/batching/runner.h | 10 + extension/llm/batching/test/fake_executor.h | 26 +- extension/llm/batching/test/runner_test.cpp | 333 ++++++++++++++++ extension/llm/batching/types.h | 2 +- 8 files changed, 1264 insertions(+), 34 deletions(-) create mode 100644 extension/llm/batching/metrics.cpp create mode 100644 extension/llm/batching/metrics.h diff --git a/extension/llm/batching/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt index 6d00f601b27..2175964352b 100644 --- a/extension/llm/batching/CMakeLists.txt +++ b/extension/llm/batching/CMakeLists.txt @@ -14,7 +14,7 @@ if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) endif() -add_library(extension_llm_batching runner.cpp) +add_library(extension_llm_batching runner.cpp metrics.cpp) # std::optional, std::function, and std::future in the public headers. target_compile_features(extension_llm_batching PUBLIC cxx_std_17) target_include_directories( diff --git a/extension/llm/batching/metrics.cpp b/extension/llm/batching/metrics.cpp new file mode 100644 index 00000000000..365e2e1bb79 --- /dev/null +++ b/extension/llm/batching/metrics.cpp @@ -0,0 +1,140 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// The report formatters, kept out of metrics.h so the stream headers they need +// stay out of every translation unit that merely reads a counter. + +#include + +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +namespace detail { + +inline std::string fixed(double v, int places) { + std::ostringstream os; + os << std::fixed << std::setprecision(places) << v; + return os.str(); +} + +inline std::string ms(std::int64_t microseconds) { + return fixed(static_cast(microseconds) / 1000.0, 2); +} + +} // namespace detail + +// Returned rather than printed: these headers stay free of ExecuTorch's +// logging, so the caller decides where it goes. +std::string format_report(const GenerationMetrics& m) { + std::ostringstream os; + os << "session " << m.sid << ": " << m.n_prompt_tokens << " prompt + " + << m.n_generated_tokens << " generated tokens\n" + << " ttft " << detail::ms(m.ttft_us()) << " ms = queue " + << detail::ms(m.queue_wait_us()) << " + prefill " + << detail::ms(m.prefill_span_us()) << " (" + << detail::fixed(m.prefill_tokens_per_sec(), 1) << " tok/s, " + << m.n_prefill_steps + << (m.n_prefill_steps == 1 ? " step)\n" : " steps)\n"); + if (m.itl_count > 0) { + // The decode span and its rate are omitted: both follow from the mean + // below, which is decode time over decode tokens by construction. + os << " decode " << detail::ms(m.itl_sum_us / m.itl_count) + << " ms/token, max " << detail::ms(m.itl_max_us) << " -> " + << detail::fixed(m.decode_tokens_per_sec(), 1) << " tok/s over " + << m.n_decode_steps << " steps\n"; + } + os << " total " << detail::ms(m.e2e_us()) << " ms\n"; + return os.str(); +} + +std::string format_report(const EngineMetrics& m) { + std::ostringstream os; + const double wall = m.wall_us(); + const auto share = [wall](std::int64_t part) { + return wall > 0.0 ? detail::fixed(100.0 * part / wall, 0) + : std::string("0"); + }; + // An empty bucket has no mean. Printing 0.00 ms would read as a measurement + // rather than an absence, and the count of zero is the informative part. + const auto bucket = [](double mean_us, std::uint64_t n) { + return n > 0 ? detail::ms(static_cast(mean_us)) + " ms" + : std::string("--"); + }; + + os << "engine\n" + << " run " << detail::fixed(wall / 1e6, 2) << " s wall, " + << m.steps << " steps"; + if (m.steps_failed > 0) { + os << " (" << m.steps_failed << " failed)"; + } + os << ", " << detail::fixed(100.0 * m.idle_fraction(), 1) << "% idle\n"; + if (m.init_us > 0) { + os << " init " << detail::ms(m.init_us) + << " ms executor setup, before the wall above\n"; + } + os << " generations " << m.generations_completed << " of " + << m.generations_started << ": " << m.finished_stop_token << " stop, " + << m.finished_token_limit << " limit, " << m.finished_cancelled + << " cancelled, " << m.finished_failed << " failed\n" + << " sessions " << m.peak_concurrent_generations + << " peak concurrent, " << m.sessions_refused << " refused at capacity\n" + << "\n" + // A partition, so the three shares add up. The gap between the first two + // is what packing prefill alongside decode cost the decodes. + << " step time decode-only " + << bucket(m.mean_decode_only_step_us(), m.decode_only_steps()) << " x" + << m.decode_only_steps() << " (" << share(m.decode_only_latency_sum_us) + << "%)\n" + << " mixed " + << bucket(m.mean_mixed_step_us(), m.mixed_steps()) << " x" + << m.mixed_steps() << " (" << share(m.mixed_latency_sum_us) << "%)\n" + << " prefill-only " + << bucket(m.mean_prefill_only_step_us(), m.prefill_only_steps()) << " x" + << m.prefill_only_steps() << " (" << share(m.prefill_only_latency_sum_us) + << "%)\n" + << "\n" + << " decode " << detail::fixed(m.decode_only_tokens_per_sec(), 1) + << " tok/s on decode-only steps at " + << detail::fixed(m.mean_decode_step_sessions(), 2) << " sessions/step\n" + << " " + << detail::fixed(m.mean_decode_us_per_token() / 1000.0, 2) + << " ms/token per session across all " << m.steps_with_decode + << " decode steps\n" + << " context " << detail::fixed(m.mean_context_per_step(), 0) + << " tokens resident/step, max " << m.context_max << " per session\n" + << "\n" + << " scheduler admitted " + << detail::fixed(m.mean_admitted_decode_sessions(), 2) << " of " + << detail::fixed(m.mean_ready_decode_sessions(), 2) + << " ready decode sessions (" + << detail::fixed(100.0 * m.admitted_ratio(), 1) << "%)\n" + << " ttft mean " << detail::ms(m.mean_ttft_us()) << " ms, min " + << detail::ms(m.min_ttft_us()) << " ms, max " << detail::ms(m.ttft_max_us) + << " ms over " << m.ttft_count + << " generations\n" + // Wall-clock rates: these move with the prompt-to-generation mix, so they + // describe the workload as much as the engine. + << " delivered " << m.total_generated_tokens << " generated -> " + << detail::fixed(m.generated_tokens_per_sec(), 1) << " tok/s wall\n" + << " processed " << m.model_input_tokens() << " tokens -> " + << detail::fixed(m.model_input_tokens_per_sec(), 1) << " tok/s (" + << m.decode_tokens_total << " decode + " << m.prefill_tokens_total + << " prompt)\n"; + return os.str(); +} + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/metrics.h b/extension/llm/batching/metrics.h new file mode 100644 index 00000000000..312c62d5e1c --- /dev/null +++ b/extension/llm/batching/metrics.h @@ -0,0 +1,416 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// What a batched run measured, in two tiers. +// +// GenerationMetrics covers one generation and is published on its handle, so a +// caller reads it beside finish_reason(). EngineMetrics covers the engine and +// is owned by the runner. +// +// There is deliberately no session tier: Session::position() already reports a +// session's context, and a session's generations are recovered by grouping on +// GenerationMetrics::sid. Keeping generations separate is the point -- a second +// generation on a warm session has a different profile from a cold one, and an +// average over the two hides exactly that. +// +// Free of ExecuTorch runtime types, like the rest of these headers: no ET_LOG, +// no exceptions, and nothing here allocates outside the format_report calls. + +#include +#include +#include +#include + +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +// Monotonic: these are durations, and a wall-clock adjustment mid-run would +// otherwise produce negative ones. +using MetricsClock = std::chrono::steady_clock; +using MetricsTime = MetricsClock::time_point; + +ET_EXPERIMENTAL inline std::int64_t us_between( + MetricsTime from, + MetricsTime to) { + return std::chrono::duration_cast(to - from) + .count(); +} + +// A default-constructed time_point means the event never happened, which is +// distinct from it happening at time zero: the clock's epoch is the process's, +// so no real event lands there. +ET_EXPERIMENTAL inline bool stamped(MetricsTime t) { + return t.time_since_epoch().count() != 0; +} + +// One generation's timeline and counts. The finish reason is not duplicated +// here; GenerationHandle::finish_reason() already carries it, and two copies +// would be free to disagree. +struct ET_EXPERIMENTAL GenerationMetrics { + SessionId sid = 0; + + // generate_async() was called. Taken on the caller's thread, before the + // request is queued. + MetricsTime t_submit{}; + // The first batch this generation appeared in. Everything between here and + // t_submit is time the scheduler did not pick it. + MetricsTime t_first_step{}; + MetricsTime t_first_token{}; + MetricsTime t_end{}; + + std::int64_t n_prompt_tokens = 0; + std::int64_t n_generated_tokens = 0; + std::int32_t n_prefill_steps = 0; + std::int32_t n_decode_steps = 0; + + // Caller-visible inter-token latency, excluding the gap to the first token, + // which is TTFT. The first token in a later callback gets the elapsed gap; + // further tokens in that callback get zero-time samples. Summary rather than + // samples: exact mean at four scalars, where a per-token vector would cost + // 8 KB on a long generation. + std::int64_t itl_count = 0; + std::int64_t itl_sum_us = 0; + std::int64_t itl_min_us = std::numeric_limits::max(); + std::int64_t itl_max_us = 0; + + // Time to first token: what a caller waits before anything appears. + std::int64_t ttft_us() const { + return stamped(t_submit) && stamped(t_first_token) + ? us_between(t_submit, t_first_token) + : 0; + } + + // The queueing share of ttft_us(). Large means the scheduler was busy, not + // that prefill was slow. + std::int64_t queue_wait_us() const { + return stamped(t_submit) && stamped(t_first_step) + ? us_between(t_submit, t_first_step) + : 0; + } + + // The compute share of ttft_us(). Includes other sessions' work in the steps + // this one's prefill was spread across, so read it with n_prefill_steps. + std::int64_t prefill_span_us() const { + return stamped(t_first_step) && stamped(t_first_token) + ? us_between(t_first_step, t_first_token) + : 0; + } + + // Generation proper, after the first token. + std::int64_t decode_span_us() const { + return stamped(t_first_token) && stamped(t_end) + ? us_between(t_first_token, t_end) + : 0; + } + + std::int64_t e2e_us() const { + return stamped(t_submit) && stamped(t_end) ? us_between(t_submit, t_end) + : 0; + } + + double itl_mean_us() const { + return itl_count > 0 ? static_cast(itl_sum_us) / itl_count : 0.0; + } + + // Zero rather than the sentinel when no gap was ever sampled, as happens to + // any generation that produced a single token. Mirrors min_ttft_us(). + std::int64_t min_itl_us() const { + return itl_count > 0 ? itl_min_us : 0; + } + + // What this one caller saw, which is not the engine's aggregate rate. + double decode_tokens_per_sec() const { + if (itl_count == 0) { + return 0.0; + } + return itl_sum_us > 0 ? 1e6 * static_cast(itl_count) / itl_sum_us + : std::numeric_limits::infinity(); + } + + // Prompt tokens over the wall time this generation's prefill took. Like the + // decode rate above it is what this caller experienced, so it includes any + // work sharing those steps -- a prompt that waits behind another session's + // chunks reports a lower rate, which is what that caller actually got. + double prefill_tokens_per_sec() const { + const std::int64_t span = prefill_span_us(); + return span > 0 ? 1e6 * static_cast(n_prompt_tokens) / span : 0.0; + } +}; + +// The engine's own view, accumulated on the engine thread and read once it has +// stopped. +struct ET_EXPERIMENTAL EngineMetrics { + std::uint64_t steps = 0; + std::uint64_t steps_failed = 0; + + // Summed over steps, so dividing by `steps` gives the mean. Sessions, not + // tokens: a prefill chunk is one session and many tokens. + std::uint64_t decode_sessions_total = 0; + std::uint64_t prefill_sessions_total = 0; + // Generations that could have decoded in this step: alive and past their + // first token. Prefilling generations are excluded because they are not + // waiting on a decode slot, they are doing their own work. Against + // decode_sessions_total this says how much of the eligible work the scheduler + // ran, without naming any scheduler's limits. + std::uint64_t ready_total = 0; + + std::int64_t step_latency_sum_us = 0; + std::int64_t step_latency_max_us = 0; + + // Every token the engine processed. Task::is_decode classifies each one + // exactly, so these are complete. + // + // There is deliberately no rate over every step that held them. A step + // mixing both kinds runs them in one forward pass over one weight read, so + // no share of its latency belongs to either, and dividing by the time of + // those steps would make packing prefill alongside decode -- which raises + // total throughput -- look like a decode regression. The rate that is safe + // to publish is decode_only_tokens_per_sec(), taken over steps that held no + // prefill at all, where the attribution is not in question. + std::uint64_t decode_tokens_total = 0; + std::uint64_t prefill_tokens_total = 0; + // The part of decode_tokens_total that ran on decode-only steps, so it can + // be divided by their time. The comparable figure to a fixed-batch decode + // benchmark, which measures exactly this shape of step. + std::uint64_t decode_only_tokens = 0; + + // Steps holding at least one task of each kind. A step can hold both, so + // these overlap by mixed_steps(). + std::uint64_t steps_with_decode = 0; + std::uint64_t steps_with_prefill = 0; + + // step_latency_sum_us split by what the step held. A partition, unlike the + // two counts above: every step is exactly one of these three, so they sum to + // step_latency_sum_us. Comparing the first two prices the scheduler's choice + // to pack prefill alongside decode, which no single blended mean can show. + std::int64_t decode_only_latency_sum_us = 0; + std::int64_t mixed_latency_sum_us = 0; + std::int64_t prefill_only_latency_sum_us = 0; + + // Committed session length summed across the batch, once per step: what + // attention had to carry. Generic -- executor.h defines a session's length, + // saying nothing about how the state is stored -- and the largest reason one + // decode step costs more than another. + std::int64_t context_sum = 0; + std::int64_t context_max = 0; // longest single session seen in any step + + // Step latency charged once to every decode session in the step. A latency + // may be counted for several sessions because each of them really did wait + // it -- unlike time as a cost, waiting is not divided up. Includes mixed + // steps, where a decode stuck behind a prefill chunk waited the whole thing. + std::int64_t decode_session_time_sum_us = 0; + + // Opens the executor refused for being at capacity, and the most generations + // seen installed at once. The peak is sampled once per executed step, so a + // generation that began and ended between two steps is not counted. + std::uint64_t sessions_refused = 0; + std::uint64_t peak_concurrent_generations = 0; + + std::uint64_t generations_started = 0; + std::uint64_t generations_completed = 0; + std::uint64_t finished_stop_token = 0; + std::uint64_t finished_token_limit = 0; + std::uint64_t finished_cancelled = 0; + std::uint64_t finished_failed = 0; + + // Over generations that reached a first token, which is not every + // completion: one cancelled or failed during prefill has no TTFT to report. + std::uint64_t ttft_count = 0; + std::int64_t ttft_sum_us = 0; + std::int64_t ttft_min_us = std::numeric_limits::max(); + std::int64_t ttft_max_us = 0; + + std::int64_t total_prompt_tokens = 0; + std::int64_t total_generated_tokens = 0; + + MetricsTime t_first_step{}; + MetricsTime t_last_step{}; + + // Executor::initialize(), timed here because nothing else can: it runs on the + // engine thread after the constructor has already returned, so a caller has + // no two points to measure between. Outside wall_us(), which starts at the + // first step -- folding it in would hide one-time setup inside the run. + std::int64_t init_us = 0; + + double wall_us() const { + return stamped(t_first_step) && stamped(t_last_step) + ? static_cast(us_between(t_first_step, t_last_step)) + : 0.0; + } + + // Mean decode sessions per step, and the mean that were eligible. Raw + // counts: normalising against a scheduler's decode limit would tie these to + // one scheduler, and against a workload smaller than that limit it would + // report idle capacity that no prompt existed to fill. + // + // Both divide by every step, so the pair is comparable. For the width of an + // actual decode forward, which is what a batched matmul sees, use + // mean_decode_step_sessions(). + double mean_admitted_decode_sessions() const { + return steps > 0 ? static_cast(decode_sessions_total) / steps : 0.0; + } + + double mean_ready_decode_sessions() const { + return steps > 0 ? static_cast(ready_total) / steps : 0.0; + } + + // Decode sessions per step that actually held a decode. Each contributes one + // token, so this is also the decode token width of the forward. + double mean_decode_step_sessions() const { + return steps_with_decode > 0 + ? static_cast(decode_sessions_total) / steps_with_decode + : 0.0; + } + + // Every session in the step, both kinds. How full the forward pass was, + // which is the batching question; the decode figures above are the + // scheduling one. + double mean_step_sessions() const { + return steps > 0 + ? static_cast(decode_sessions_total + prefill_sessions_total) / + steps + : 0.0; + } + + // Below 1 the scheduler is holding eligible work back rather than running + // out of it. Slightly under 1 is normal: a generation is briefly eligible + // but unqueued between its output being handled and its continuation being + // submitted. + double admitted_ratio() const { + return ready_total > 0 + ? static_cast(decode_sessions_total) / ready_total + : 0.0; + } + + // What the engine processed. Distinct from the generation tier's totals, + // which count what callers were given: a generation's first token comes out + // of a prefill step, and prompt tokens are processed but never emitted. + std::uint64_t model_input_tokens() const { + return decode_tokens_total + prefill_tokens_total; + } + + std::int64_t total_tokens() const { + return total_prompt_tokens + total_generated_tokens; + } + + double mean_step_tokens() const { + return steps > 0 ? static_cast(model_input_tokens()) / steps : 0.0; + } + + // Wall time the engine spent outside execute(): waiting for work, draining + // commands, running callbacks. + double idle_fraction() const { + const double wall = wall_us(); + return wall > 0.0 ? 1.0 - static_cast(step_latency_sum_us) / wall + : 0.0; + } + + // What callers were given, per second. The comparable figure when swapping + // executors: a speculative one feeds one token per decode step and returns + // several, so processed tokens would stay flat while this rises. + double generated_tokens_per_sec() const { + const double wall = wall_us(); + return wall > 0.0 ? 1e6 * static_cast(total_generated_tokens) / wall + : 0.0; + } + + // What the engine put through the model, per second. Against the rate above + // it shows how much output each processed token bought. + double model_input_tokens_per_sec() const { + const double wall = wall_us(); + return wall > 0.0 ? 1e6 * static_cast(model_input_tokens()) / wall + : 0.0; + } + + // Steps that held both kinds at once. Implied by the overlap rather than + // counted: every step holds decode, prefill, or both. + std::uint64_t mixed_steps() const { + const std::uint64_t overlap = steps_with_decode + steps_with_prefill; + return overlap > steps ? overlap - steps : 0; + } + + double mean_ttft_us() const { + return ttft_count > 0 ? static_cast(ttft_sum_us) / ttft_count : 0.0; + } + + // The sentinel is never shown: with no samples there is no minimum. + // Mean wall time a decode session waited per step it took part in, mixed + // steps included. The engine-side counterpart to per-generation inter-token + // latency, and a cross-check on it. Weighted by sessions, so it is not + // directly comparable to the unweighted bucket means above -- read it as + // what sessions experienced, not as what a step cost. + double mean_decode_us_per_token() const { + return decode_sessions_total > 0 + ? static_cast(decode_session_time_sum_us) / + decode_sessions_total + : 0.0; + } + + // Steps holding only one kind. Derived, because every step is decode-only, + // prefill-only, or mixed. + std::uint64_t decode_only_steps() const { + const std::uint64_t mixed = mixed_steps(); + return steps_with_decode > mixed ? steps_with_decode - mixed : 0; + } + + std::uint64_t prefill_only_steps() const { + const std::uint64_t mixed = mixed_steps(); + return steps_with_prefill > mixed ? steps_with_prefill - mixed : 0; + } + + double mean_decode_only_step_us() const { + const std::uint64_t n = decode_only_steps(); + return n > 0 ? static_cast(decode_only_latency_sum_us) / n : 0.0; + } + + double mean_mixed_step_us() const { + const std::uint64_t n = mixed_steps(); + return n > 0 ? static_cast(mixed_latency_sum_us) / n : 0.0; + } + + double mean_prefill_only_step_us() const { + const std::uint64_t n = prefill_only_steps(); + return n > 0 ? static_cast(prefill_only_latency_sum_us) / n : 0.0; + } + + // Decode throughput measured only where no prefill shared the forward, so + // every microsecond in the denominator was spent on these tokens. The one + // rate here that compares across engines, since a fixed-batch decode + // benchmark runs exactly this shape of step. + double decode_only_tokens_per_sec() const { + return decode_only_latency_sum_us > 0 ? 1e6 * + static_cast(decode_only_tokens) / decode_only_latency_sum_us + : 0.0; + } + + // Committed tokens the batch carried, averaged over steps. + double mean_context_per_step() const { + return steps > 0 ? static_cast(context_sum) / steps : 0.0; + } + + std::int64_t min_ttft_us() const { + return ttft_count > 0 ? ttft_min_us : 0; + } +}; +// Human-readable reports. Defined in metrics.cpp so a runner that never formats +// one does not pull and into its translation units -- +// everything above is trivially inline and allocation-free, these are not. +ET_EXPERIMENTAL std::string format_report(const GenerationMetrics& m); +ET_EXPERIMENTAL std::string format_report(const EngineMetrics& m); + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index 0df1a8f6ae7..2ea0e69ba92 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -66,6 +66,7 @@ struct GenerationHandleState { CompletionPhase phase = CompletionPhase::Pending; std::optional reason; std::string error_message; + GenerationMetrics metrics; std::atomic cancelled{false}; }; @@ -120,11 +121,14 @@ class TerminalCompletion { ~TerminalCompletion() { if (state_) { - finish(TerminalOutcome::failed({})); + finish(TerminalOutcome::failed({}), GenerationMetrics{}); } } - void finish(TerminalOutcome outcome) { + // Metrics ride alongside the outcome rather than inside it: they describe + // the generation, not the reason it ended. Published here, with the reason + // and under the same lock, so a handle never shows one without the other. + void finish(TerminalOutcome outcome, const GenerationMetrics& metrics) { assert(state_); auto state = std::move(state_); { @@ -132,6 +136,7 @@ class TerminalCompletion { assert(state->phase == CompletionPhase::DeliveringCallback); state->reason = outcome.reason; state->error_message = std::move(outcome.error_message); + state->metrics = metrics; state->phase = CompletionPhase::Done; } state->cv.notify_all(); @@ -184,7 +189,8 @@ void finalize_terminal( if (!callback_result.succeeded) { outcome = TerminalOutcome::failed(std::move(callback_result.error_message)); } - completion->finish(std::move(outcome)); + // Rejected before admission, so there is no timeline to report. + completion->finish(std::move(outcome), GenerationMetrics{}); } // --- GenerationHandle ------------------------------------------------------ @@ -236,6 +242,17 @@ std::string GenerationHandle::error_message() const { return state_->error_message; } +// Published with the reason, which lands only after the terminal callback +// returns. Reading this from inside that callback yields an empty snapshot, +// exactly as finish_reason() yields nullopt there. +GenerationMetrics GenerationHandle::metrics() const { + if (!state_) { + return {}; + } + std::lock_guard lock(state_->mutex); + return state_->metrics; +} + // Everything the runner owns. Held by shared_ptr from both Runner and every // Session, so a Session outliving its Runner finds a stopped object rather // than a dangling one. @@ -270,6 +287,11 @@ class RunnerImpl : public std::enable_shared_from_this { GenConfig config, GenerationCallback on_update); + // Engine-thread data, so only stable once that thread is joined. + EngineMetrics metrics() const { + return metrics_; + } + private: enum class Lifecycle { Running, Stopping, Stopped }; @@ -281,6 +303,10 @@ class RunnerImpl : public std::enable_shared_from_this { // Shared with every handle, so cancelling needs no route back to the // runner and works after it is gone. std::shared_ptr state; + GenerationMetrics m; + // Engine-side only. An inter-token gap needs the previous delivery, and + // the published metrics keep only the summary, not the last timestamp. + MetricsTime last_token_at{}; }; // Start-only data. The sampling policy is installed on the executor at @@ -433,14 +459,27 @@ class RunnerImpl : public std::enable_shared_from_this { CallbackResult dispatch_update_( const Generation& generation, GenerationUpdate update); - void deliver_claimed_terminal_( + FinishReason deliver_claimed_terminal_( const Generation& generation, TerminalCompletion completion, TerminalOutcome outcome); // Invoke the terminal callback and publish state for a detached generation. - void complete_generation_(Generation generation, TerminalOutcome outcome); - void complete_request_(GenerationRequest request, TerminalOutcome outcome); + // + // `on_engine_thread` is false only on the post-shutdown path out of + // generate_async(), which runs on the caller's thread. The engine may still + // be draining there, so that path publishes to the handle but must leave the + // engine's own counters alone. + void complete_generation_( + Generation generation, + TerminalOutcome outcome, + bool on_engine_thread); + void complete_request_( + GenerationRequest request, + TerminalOutcome outcome, + bool on_engine_thread); + // Engine thread only: rolls one finished generation into metrics_. + void record_completion_(const GenerationMetrics& m, FinishReason reason); std::optional detach_active_generation_(SessionId session); void complete_active_generation_(SessionId session, TerminalOutcome outcome); void fail_active_generation_after_callback_( @@ -470,6 +509,7 @@ class RunnerImpl : public std::enable_shared_from_this { // Kept after records close to enforce executor IDs are lifetime-unique. std::unordered_set issued_session_ids_; TaskId next_tid_ = 1; + EngineMetrics metrics_; std::thread engine_; }; @@ -557,6 +597,10 @@ void Runner::shutdown() { impl_->shutdown(); } +EngineMetrics Runner::metrics() const { + return impl_->metrics(); +} + std::future> Runner::open_session_async() { return impl_->open_session_async(); } @@ -627,7 +671,10 @@ void RunnerImpl::run_() { // Before the loop and before any command is answered, so a caller is never // handed a session for an executor that did not come up, and so one-time // setup is not charged to whichever generation happened to go first. - if (!executor_.initialize()) { + const MetricsTime init_start = MetricsClock::now(); + const bool ready = executor_.initialize(); + metrics_.init_us = us_between(init_start, MetricsClock::now()); + if (!ready) { // Stop without running work. The drain below still answers whatever was // queued while this was starting, so no caller is left waiting. std::lock_guard lock(control_mutex_); @@ -673,7 +720,9 @@ void RunnerImpl::run_() { for (auto& entry : open_sessions) { if (entry.second) { complete_generation_( - std::move(*entry.second), TerminalOutcome::cancelled()); + std::move(*entry.second), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); } executor_.close_session(entry.first); } @@ -714,7 +763,13 @@ void RunnerImpl::process_pending_commands_() { } void RunnerImpl::process_command_(OpenCommand command) { - auto sid = is_running_() ? executor_.open_session() : std::nullopt; + const bool running = is_running_(); + auto sid = running ? executor_.open_session() : std::nullopt; + if (running && !sid) { + // At capacity. Counted apart from a refusal caused by the runner stopping, + // which says nothing about how loaded the executor was. + ++metrics_.sessions_refused; + } bool newly_issued = false; bool published = false; if (sid) { @@ -750,7 +805,9 @@ void RunnerImpl::process_command_(CloseCommand command) { } if (retired->active_generation) { complete_generation_( - std::move(*retired->active_generation), TerminalOutcome::cancelled()); + std::move(*retired->active_generation), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); } executor_.close_session(command.session); } @@ -789,9 +846,166 @@ bool RunnerImpl::execute_one_batch_() { return false; } + // Composition is read here, before to_batch_input moves the Inputs out and + // drops is_decode with the rest of the scheduling fields. + // + // Decode tasks are one sequence each. A session's prefill can arrive as + // several chunks which are not necessarily adjacent -- DecodeFirstScheduler + // rotates, taking one chunk per session per pass, so two sessions prefilling + // together interleave as A, B, A, B. Track every session seen rather than + // comparing with the previous one, which would count each chunk as a new + // sequence and charge the step to the generation repeatedly. + std::uint64_t decode_sessions = 0; + std::uint64_t prefill_sessions = 0; + std::uint64_t decode_tokens = 0; + std::uint64_t prefill_tokens = 0; + std::vector prefilling; // small: bounded by the batch width + const MetricsTime step_start = MetricsClock::now(); + for (const Task& task : tasks) { + bool first_chunk = false; + if (task.is_decode) { + ++decode_sessions; + decode_tokens += task.input.size; + } else { + prefill_tokens += task.input.size; + first_chunk = + std::find(prefilling.begin(), prefilling.end(), task.input.sid) == + prefilling.end(); + if (first_chunk) { + ++prefill_sessions; + prefilling.push_back(task.input.sid); + } + } + // Charge the step to the generation once, however many chunks it brought. + if (!task.is_decode && !first_chunk) { + continue; + } + auto session = sessions_.find(task.input.sid); + if (session == sessions_.end() || !session->second.active_generation) { + continue; + } + GenerationMetrics& m = session->second.active_generation->m; + if (!stamped(m.t_first_step)) { + m.t_first_step = step_start; + } + if (task.is_decode) { + ++m.n_decode_steps; + } else { + ++m.n_prefill_steps; + } + } + // Generations eligible for a decode slot, admitted or not. Past their first + // token, so a generation still prefilling is not counted as held back when + // it is simply busy elsewhere. Against decode_sessions this is what + // separates a scheduler holding work back from there being no work. + // + // The same pass records how many generations are installed at once, so the + // batch widths above can be read against the concurrency that was available. + std::uint64_t live_generations = 0; + for (const auto& entry : sessions_) { + const auto& generation = entry.second.active_generation; + if (!generation) { + continue; + } + ++live_generations; + if (stamped(generation->m.t_first_token)) { + ++metrics_.ready_total; + } + } + metrics_.peak_concurrent_generations = + std::max(metrics_.peak_concurrent_generations, live_generations); + + // Committed context the batch carries into the forward: what attention has + // to cover. Read from the input positions, which say nothing about how the + // executor stores the state. + // + // Counted once per session, at the end of its furthest chunk. A prompt split + // across several chunks of one step is one context, not one per chunk, and + // scheduler.h lets a session appear more than once for exactly that reason. + std::vector> context_ends; + for (const Task& task : tasks) { + const auto end = static_cast(task.input.position) + + static_cast(task.input.offset) + + static_cast(task.input.size); + metrics_.context_max = std::max(metrics_.context_max, end); + auto entry = std::find_if( + context_ends.begin(), context_ends.end(), [&](const auto& seen) { + return seen.first == task.input.sid; + }); + if (entry == context_ends.end()) { + context_ends.emplace_back(task.input.sid, end); + } else { + entry->second = std::max(entry->second, end); + } + } + for (const auto& entry : context_ends) { + metrics_.context_sum += entry.second; + } + + // The forward and the batch handed to it, timed apart from step_start. The + // scans above are measurement, and folding them into the buckets would bias + // exactly the numbers metrics.h offers for comparison against other engines; + // they are also the parts that grow with concurrency, so the bias would not + // be constant. step_start still marks the step for the generation timeline. + const MetricsTime exec_start = MetricsClock::now(); BatchInput batch = to_batch_input(tasks); BatchOutput out; const bool ok = executor_.execute(batch, out); + const MetricsTime step_end = MetricsClock::now(); + + const std::int64_t latency = us_between(exec_start, step_end); + ++metrics_.steps; + metrics_.decode_sessions_total += decode_sessions; + metrics_.prefill_sessions_total += prefill_sessions; + // Only what the model is known to have taken in. A failed execute leaves + // what it processed unknown -- that is why the batch is condemned and its + // sessions poisoned -- so counting the attempt as throughput would credit + // work that may never have happened. The time is still counted below, + // because it was really spent, and steps_failed records the attempt. + if (ok) { + metrics_.decode_tokens_total += decode_tokens; + metrics_.prefill_tokens_total += prefill_tokens; + } + metrics_.step_latency_sum_us += latency; + metrics_.step_latency_max_us = + std::max(metrics_.step_latency_max_us, latency); + if (!stamped(metrics_.t_first_step)) { + metrics_.t_first_step = step_start; + } + metrics_.t_last_step = step_end; + if (decode_tokens > 0) { + ++metrics_.steps_with_decode; + // Charged once per session: each of them waited this whole step. + metrics_.decode_session_time_sum_us += + latency * static_cast(decode_sessions); + } + if (prefill_tokens > 0) { + ++metrics_.steps_with_prefill; + } + // Exactly one of the three, so the sums partition step_latency_sum_us. The + // three step-count conditions below use the same two predicates, so the + // counts partition `steps` too. + assert( + (decode_tokens > 0 || prefill_tokens > 0) && + "a non-empty batch carries tokens of at least one kind"); + if (decode_tokens > 0 && prefill_tokens > 0) { + metrics_.mixed_latency_sum_us += latency; + } else if (decode_tokens > 0) { + metrics_.decode_only_latency_sum_us += latency; + // Attributable, because no prefill shared this forward. Gated on `ok` for + // the same reason as decode_tokens_total: a failed execute leaves what the + // model consumed unknown. The latency above is still counted, since it was + // really spent and the three buckets have to partition the total. + if (ok) { + metrics_.decode_only_tokens += decode_tokens; + } + } else { + metrics_.prefill_only_latency_sum_us += latency; + } + if (!ok) { + ++metrics_.steps_failed; + } + if (!is_running_()) { return true; // discard an in-flight result after the stop boundary } @@ -937,6 +1151,10 @@ GenerationHandle RunnerImpl::generate_async( request.generation.stop_tokens = std::move(config.stop_tokens); request.generation.on_update = std::move(on_update); request.generation.state = state; + // The caller's thread, before the request is queued: the wait a caller sees + // starts here, not when the engine gets round to it. + request.generation.m.sid = session; + request.generation.m.t_submit = MetricsClock::now(); auto handle = GenerationHandle(state); bool admitted = false; @@ -954,7 +1172,10 @@ GenerationHandle RunnerImpl::generate_async( // After shutdown nothing drains the inbox, so complete synchronously instead // of admitting a start that can never report completion. - complete_request_(std::move(request), TerminalOutcome::cancelled()); + complete_request_( + std::move(request), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/false); return handle; } @@ -1005,24 +1226,40 @@ std::shared_ptr> RunnerImpl::build_initial_delta_( } void RunnerImpl::start_generation_(GenerationRequest request) { + // Counted on arrival at the engine, not on successful install: every path + // below ends in a completion, so deferring this would let completions + // exceed starts. + ++metrics_.generations_started; if (!is_running_() || request.generation.state->cancelled.load()) { - complete_request_(std::move(request), TerminalOutcome::cancelled()); + complete_request_( + std::move(request), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); return; } auto session = sessions_.find(request.session); if (session == sessions_.end()) { complete_request_( - std::move(request), TerminalOutcome::failed("session is not open")); + std::move(request), + TerminalOutcome::failed("session is not open"), + /*on_engine_thread=*/true); return; } auto& record = session->second; if (auto rejection = validate_generation_start_(request, record)) { - complete_request_(std::move(request), std::move(*rejection)); + complete_request_( + std::move(request), std::move(*rejection), /*on_engine_thread=*/true); return; } const auto start_position = record.position(); + // The caller's own delta, captured before build_initial_delta_ moves it and + // before any carried token is prepended. The generation tier counts what + // callers gave; tokens actually fed to the model are EngineMetrics' + // model_input_tokens(). + request.generation.m.n_prompt_tokens = + static_cast(request.delta->size()); auto delta = build_initial_delta_(request, record); executor_.set_sampling(request.session, request.sampling, request.seed); @@ -1036,7 +1273,10 @@ void RunnerImpl::start_generation_(GenerationRequest request) { } if (!installed) { // Sampling began before the stop transition, but no task was submitted. - complete_request_(std::move(request), TerminalOutcome::cancelled()); + complete_request_( + std::move(request), + TerminalOutcome::cancelled(), + /*on_engine_thread=*/true); return; } @@ -1178,6 +1418,34 @@ std::optional RunnerImpl::prepare_output_( auto interpreted = interpret_output_(generation, *output); generation.remaining_tokens = interpreted.remaining_tokens; + // Counted here rather than at delivery: this is the one place that sees the + // emitted run with the generation in hand, and it covers both the + // completing and the continuing branch below, which move the tokens away. + if (!interpreted.emitted_tokens.empty()) { + const auto emitted = + static_cast(interpreted.emitted_tokens.size()); + const MetricsTime now = MetricsClock::now(); + generation.m.n_generated_tokens += emitted; + if (!stamped(generation.m.t_first_token)) { + generation.m.t_first_token = now; + // The first token's wait is TTFT. Further tokens in the same run have + // zero caller-visible latency between them. + const std::int64_t intra_burst = emitted - 1; + if (intra_burst > 0) { + generation.m.itl_count += intra_burst; + generation.m.itl_min_us = 0; + } + } else { + const std::int64_t gap = us_between(generation.last_token_at, now); + generation.m.itl_count += emitted; + generation.m.itl_sum_us += gap; + generation.m.itl_min_us = std::min( + generation.m.itl_min_us, emitted > 1 ? std::int64_t{0} : gap); + generation.m.itl_max_us = std::max(generation.m.itl_max_us, gap); + } + generation.last_token_at = now; + } + // The transcript grows by what the caller keeps, capped by what the executor // committed. The last emitted token lands only when it is fed back. record.advance(interpreted.committed_tokens); @@ -1282,7 +1550,10 @@ CallbackResult RunnerImpl::dispatch_update_( return invoke_callback(generation.on_update, std::move(update)); } -void RunnerImpl::deliver_claimed_terminal_( +// Returns the reason actually published, which is not the one passed in when +// the terminal callback throws. The engine counts that final reason, so the +// tallies agree with what the handle reports. +FinishReason RunnerImpl::deliver_claimed_terminal_( const Generation& generation, TerminalCompletion completion, TerminalOutcome outcome) { @@ -1293,12 +1564,16 @@ void RunnerImpl::deliver_claimed_terminal_( if (!callback_result.succeeded) { outcome = TerminalOutcome::failed(std::move(callback_result.error_message)); } - completion.finish(std::move(outcome)); + const FinishReason reason = outcome.reason; + completion.finish(std::move(outcome), generation.m); + return reason; } void RunnerImpl::complete_generation_( Generation generation, - TerminalOutcome outcome) { + TerminalOutcome outcome, + bool on_engine_thread) { + generation.m.t_end = MetricsClock::now(); std::optional completion; { // The claim is taken under the runner lock. User code still runs only after @@ -1319,14 +1594,54 @@ void RunnerImpl::complete_generation_( } completion.emplace(std::move(*claimed)); } - deliver_claimed_terminal_( + const FinishReason reason = deliver_claimed_terminal_( generation, std::move(*completion), std::move(outcome)); + if (on_engine_thread) { + record_completion_(generation.m, reason); + } +} + +void RunnerImpl::record_completion_( + const GenerationMetrics& m, + FinishReason reason) { + ++metrics_.generations_completed; + switch (reason) { + case FinishReason::StopToken: + ++metrics_.finished_stop_token; + break; + case FinishReason::NewTokenLimit: + ++metrics_.finished_token_limit; + break; + case FinishReason::Cancelled: + ++metrics_.finished_cancelled; + break; + case FinishReason::Failed: + ++metrics_.finished_failed; + break; + } + metrics_.total_prompt_tokens += m.n_prompt_tokens; + metrics_.total_generated_tokens += m.n_generated_tokens; + // Zero for a generation that never reached a first token. Counted + // separately from completions so the mean divides by the samples it has, + // and so the minimum stays untouched when there are none. + // Gated on the event, not on it having taken measurable time: a first token + // in the same microsecond as the submit is still a first token, and stamped() + // is what "happened" means everywhere else here. + if (stamped(m.t_first_token)) { + const std::int64_t ttft = m.ttft_us(); + ++metrics_.ttft_count; + metrics_.ttft_sum_us += ttft; + metrics_.ttft_min_us = std::min(metrics_.ttft_min_us, ttft); + metrics_.ttft_max_us = std::max(metrics_.ttft_max_us, ttft); + } } void RunnerImpl::complete_request_( GenerationRequest request, - TerminalOutcome outcome) { - complete_generation_(std::move(request.generation), std::move(outcome)); + TerminalOutcome outcome, + bool on_engine_thread) { + complete_generation_( + std::move(request.generation), std::move(outcome), on_engine_thread); } std::optional RunnerImpl::detach_active_generation_( @@ -1353,7 +1668,8 @@ void RunnerImpl::complete_active_generation_( if (!active) { return; } - complete_generation_(std::move(*active), std::move(outcome)); + complete_generation_( + std::move(*active), std::move(outcome), /*on_engine_thread=*/true); } void RunnerImpl::fail_active_generation_after_callback_( @@ -1366,11 +1682,18 @@ void RunnerImpl::fail_active_generation_after_callback_( if (!active) { return; } + active->m.t_end = MetricsClock::now(); auto completion = TerminalCompletion::try_claim(active->state); if (!completion) { return; } - completion->finish(TerminalOutcome::failed(std::move(error_message))); + completion->finish( + TerminalOutcome::failed(std::move(error_message)), active->m); + // This path claims the terminal itself instead of going through + // complete_generation_, so it has to count its own completion. Without this + // the engine would report fewer completions than starts, and precisely for + // the generations that failed most interestingly. + record_completion_(active->m, FinishReason::Failed); } } // namespace batching diff --git a/extension/llm/batching/runner.h b/extension/llm/batching/runner.h index b26d943fbe7..a3af2fde5c7 100644 --- a/extension/llm/batching/runner.h +++ b/extension/llm/batching/runner.h @@ -40,6 +40,7 @@ #include #include +#include #include #include #include // ET_EXPERIMENTAL @@ -147,6 +148,10 @@ class ET_EXPERIMENTAL GenerationHandle { // when valid() && done(); empty when no diagnostic is available. std::string error_message() const; + // This generation's timeline and counts, complete once done(). Empty on a + // default-constructed handle. + GenerationMetrics metrics() const; + private: friend class RunnerImpl; friend class Session; @@ -246,6 +251,11 @@ class ET_EXPERIMENTAL Runner { // from its callback. void shutdown(); + // What the engine measured. Read it after shutdown(): the counters are the + // engine thread's, so joining it is what makes them stable and visible. A + // call before then returns a torn snapshot. + EngineMetrics metrics() const; + private: std::shared_ptr impl_; }; diff --git a/extension/llm/batching/test/fake_executor.h b/extension/llm/batching/test/fake_executor.h index 1f51f416fed..a4d90cdb5ac 100644 --- a/extension/llm/batching/test/fake_executor.h +++ b/extension/llm/batching/test/fake_executor.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include #include #include @@ -17,6 +18,7 @@ #include #include #include +#include #include #include @@ -51,6 +53,9 @@ class FakeExecutor : public Executor { }; bool initialize() override { + if (initialize_delay.count() > 0) { + std::this_thread::sleep_for(initialize_delay); + } std::lock_guard lock(mutex_); ++initialize_calls_; return !fail_initialize; @@ -132,17 +137,19 @@ class FakeExecutor : public Executor { int capacity = 8; // Refuse to come up, so the runner should admit no work at all. bool fail_initialize = false; + // Stand in for real setup, so its cost is large enough to assert on. + std::chrono::milliseconds initialize_delay{0}; // Batch index from which execute() starts failing. Negative never fails. int fail_batches_from = -1; // Once a session has produced emit_before_stop tokens, every later one is // stop_token. Counted per session across the whole run, so a stop can be - // placed part way into a multi-token decode. A negative stop_token disables - // this. - Token stop_token = -1; + // placed part way into a multi-token decode. Unset disables this; a sentinel + // cannot, since Token is unsigned here and every value is a valid token. + std::optional stop_token; int emit_before_stop = 0; - // Tokens a decode step produces. 1 is a plain executor; more simulates a - // speculative one answering with the run it accepted plus the model's own - // next token. Prefill always produces one whatever this is. + // Tokens an output-producing step returns. Values above 1 simulate a + // speculative executor answering with an accepted run plus the next token. + std::size_t tokens_per_prefill = 1; std::size_t tokens_per_decode = 1; // Malformed answers. An Output carries only the tokens an input produced, so // the only ways to break the contract are to produce none, or to answer for @@ -240,7 +247,8 @@ class FakeExecutor : public Executor { // Task::is_decode. Good enough for a fake: the runner only ever feeds one // token to continue. std::vector produce(const Input& input) { - const std::size_t n = input.size == 1 ? tokens_per_decode : 1; + const std::size_t n = + input.size == 1 ? tokens_per_decode : tokens_per_prefill; std::vector produced; produced.reserve(n); for (std::size_t i = 0; i < n; ++i) { @@ -251,8 +259,8 @@ class FakeExecutor : public Executor { Token next_token(SessionId session) { const int n = ++produced_[session]; - if (stop_token >= 0 && n > emit_before_stop) { - return stop_token; + if (stop_token && n > emit_before_stop) { + return *stop_token; } auto it = sampling_.find(session); if (it == sampling_.end()) { diff --git a/extension/llm/batching/test/runner_test.cpp b/extension/llm/batching/test/runner_test.cpp index f021f2c429c..7e258802400 100644 --- a/extension/llm/batching/test/runner_test.cpp +++ b/extension/llm/batching/test/runner_test.cpp @@ -29,6 +29,7 @@ #include using executorch::extension::llm::batching::DecodeFirstScheduler; +using executorch::extension::llm::batching::EngineMetrics; using executorch::extension::llm::batching::Executor; using executorch::extension::llm::batching::FinishReason; using executorch::extension::llm::batching::GenConfig; @@ -1017,6 +1018,27 @@ TEST(GenerationTest, ImmediateStopTokenIsDeliveredAndWinsOverBudget) { << "the prompt produces STOP, so no continuation is scheduled"; } +TEST(SpeculativeTest, FirstBurstRecordsIntraBurstTokenLatencies) { + FakeExecutor executor; + executor.tokens_per_prefill = 3; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(2), config(3), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + + const auto metrics = handle.metrics(); + EXPECT_EQ(metrics.n_generated_tokens, 3); + EXPECT_EQ(metrics.itl_count, 2); + EXPECT_EQ(metrics.itl_sum_us, 0); + EXPECT_EQ(metrics.itl_min_us, 0); + EXPECT_EQ(metrics.itl_max_us, 0); + EXPECT_EQ( + metrics.decode_tokens_per_sec(), std::numeric_limits::infinity()); +} + // Regression: the runner counts a step's produced tokens as far as the // executor committed them, so the next turn resumes past them rather than on // top of them. The last token of a run is not committed until it is fed back, @@ -1663,6 +1685,316 @@ TEST(ShutdownTest, ConcurrentAdmissionDoesNotStrandCallers) { EXPECT_EQ(stranded.load(), 0); } +// --- engine-tier metrics --------------------------------------------------- +// +// EngineMetrics is only stable once the engine thread is joined, so every test +// here shuts the runner down before reading it. + +namespace { + +// The four terminal reasons partition the completions, so this is the tally +// that reveals a terminal path which forgot to account for itself. +void expect_balanced(const EngineMetrics& m, std::uint64_t expected_starts) { + EXPECT_EQ(m.generations_started, expected_starts); + EXPECT_EQ(m.generations_completed, m.generations_started) + << "every generation the engine started must also be counted as done"; + EXPECT_EQ( + m.finished_stop_token + m.finished_token_limit + m.finished_cancelled + + m.finished_failed, + m.generations_completed) + << "the per-reason tallies must partition the completions"; +} + +} // namespace + +TEST(EngineMetricsTest, CountsStepsSequencesAndTokens) { + FakeExecutor executor; + // Chunk 8 with a 20-token budget, so a 10-token prompt arrives as 8 + 2 and + // both chunks still fit in one batch. + Fixture fixture(executor, 4, 8); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(10), config(3), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + ASSERT_EQ(updates->finish(), FinishReason::NewTokenLimit); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + // One prefill step carrying both chunks, then one decode step per token + // after the first, which the prefill produced. + EXPECT_EQ(m.steps, 3u); + EXPECT_EQ(m.steps_failed, 0u); + EXPECT_EQ(m.steps_with_prefill, 1u); + EXPECT_EQ(m.steps_with_decode, 2u); + // Charged once per session per step, not once per chunk. + EXPECT_EQ(m.prefill_sessions_total, 1u); + EXPECT_EQ(m.decode_sessions_total, 2u); + EXPECT_EQ(m.prefill_tokens_total, 10u); + EXPECT_EQ(m.decode_tokens_total, 2u); + // The prompt plus the tokens fed back. The last token delivered is still + // pending, so it was never an input. + EXPECT_EQ(m.model_input_tokens(), 12u); + EXPECT_EQ(m.total_prompt_tokens, 10); + EXPECT_EQ(m.total_generated_tokens, 3); + EXPECT_EQ(m.ttft_count, 1u); + // The one sample is the minimum and the whole sum. Stated as a relation + // rather than "> 0" so it holds however fast the machine is. + EXPECT_EQ(m.min_ttft_us(), m.ttft_sum_us); + EXPECT_EQ(m.ttft_max_us, m.ttft_sum_us); + expect_balanced(m, 1); + EXPECT_EQ(m.finished_token_limit, 1u); +} + +TEST(EngineMetricsTest, StartsAndCompletionsBalanceAcrossMixedOutcomes) { + FakeExecutor executor; + executor.stop_token = 999; + executor.emit_before_stop = 1; + Fixture fixture(executor); + + // Runs to a stop token. + Session stopping = open(fixture.runner); + GenConfig stop_config = config(50); + stop_config.stop_tokens = {999}; + auto stopped = std::make_shared(); + generate(stopping, tokens(2), stop_config, stopped); + ASSERT_TRUE(stopped->wait()); + + // Rejected by validation once it reaches the engine. + Session invalid = open(fixture.runner); + auto rejected = std::make_shared(); + generate(invalid, tokens(2), config(0), rejected); + ASSERT_TRUE(rejected->wait()); + + // Cancelled explicitly. + Session cancelling = open(fixture.runner); + auto cancelled = std::make_shared(); + GenerationHandle doomed = + generate(cancelling, tokens(2), config(100000), cancelled); + doomed.cancel(); + ASSERT_TRUE(cancelled->wait()); + + stopping = Session{}; + invalid = Session{}; + cancelling = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + expect_balanced(m, 3); + EXPECT_EQ(m.finished_stop_token, 1u); + EXPECT_EQ(m.finished_failed, 1u) << "the rejected budget is a failure"; + EXPECT_EQ(m.finished_cancelled, 1u); + EXPECT_EQ(m.finished_token_limit, 0u); +} + +TEST(EngineMetricsTest, FailedBatchIsCountedWithoutCreditingItsTokens) { + FakeExecutor executor; + executor.fail_batches_from = 0; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + generate(session, tokens(4), config(2), updates); + ASSERT_TRUE(updates->wait()); + ASSERT_EQ(updates->finish(), FinishReason::Failed); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_EQ(m.steps, 1u); + EXPECT_EQ(m.steps_failed, 1u); + // The sequence was in the batch, but a failed execute leaves what the model + // consumed unknown, so its tokens are not credited as throughput. + EXPECT_EQ(m.prefill_sessions_total, 1u); + EXPECT_EQ(m.prefill_tokens_total, 0u); + EXPECT_EQ(m.model_input_tokens(), 0u); + EXPECT_EQ(m.total_generated_tokens, 0); + // Never reached a first token, so it contributes no TTFT sample and leaves + // the minimum untouched. + EXPECT_EQ(m.ttft_count, 0u); + EXPECT_EQ(m.min_ttft_us(), 0); + expect_balanced(m, 1); + EXPECT_EQ(m.finished_failed, 1u); +} + +TEST(EngineMetricsTest, SingleTokenGenerationReportsZeroMinimumItl) { + FakeExecutor executor; + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + GenerationHandle handle = generate(session, tokens(2), config(1), updates); + ASSERT_TRUE(updates->wait()); + handle.wait(); + + // One token means no gap was ever sampled. The raw field still holds its + // sentinel; the accessor is what callers should read. + const auto m = handle.metrics(); + EXPECT_EQ(m.n_generated_tokens, 1); + EXPECT_EQ(m.itl_count, 0); + EXPECT_EQ(m.min_itl_us(), 0); + EXPECT_EQ(m.itl_mean_us(), 0.0); + EXPECT_EQ(m.itl_min_us, std::numeric_limits::max()); +} + +#if ET_HAS_EXCEPTIONS +// Regression: this path claims the terminal itself rather than going through +// complete_generation_, so it once published nothing and counted nothing, +// leaving completions permanently behind starts. +TEST(EngineMetricsTest, CallbackExceptionIsCountedAsACompletedFailure) { + FakeExecutor executor; + Fixture fixture(executor); + Session session = open(fixture.runner); + + GenerationHandle failed = session.generate_async( + tokens(2), config(4), [](const GenerationUpdate& update) { + if (!update.finish_reason) { + throw std::runtime_error("callback failed"); + } + }); + failed.wait(); + ASSERT_EQ(failed.finish_reason(), FinishReason::Failed); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + expect_balanced(m, 1); + EXPECT_EQ(m.finished_failed, 1u); + // The tokens delivered before the callback threw still count as generated. + EXPECT_GT(m.total_generated_tokens, 0); +} +#endif + +// The three latency buckets are a partition of step_latency_sum_us: every step +// holds decode, prefill, or both, and its whole latency lands in exactly one. +// Without this, a step kind added later could quietly go unaccounted. +TEST(EngineMetricsTest, LatencyBucketsPartitionTotalStepTime) { + FakeExecutor executor; + // Chunk 2 with a 10-token budget, so a 6-token prompt needs several chunks + // and the run produces decode-only, prefill-only, and mixed steps. + Fixture fixture(executor, 2, 2); + Session first = open(fixture.runner); + Session second = open(fixture.runner); + + auto first_updates = std::make_shared(); + generate(first, tokens(6), config(6), first_updates); + auto second_updates = std::make_shared(); + generate(second, tokens(6), config(6), second_updates); + ASSERT_TRUE(first_updates->wait()); + ASSERT_TRUE(second_updates->wait()); + + first = Session{}; + second = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_EQ( + m.decode_only_latency_sum_us + m.mixed_latency_sum_us + + m.prefill_only_latency_sum_us, + m.step_latency_sum_us); + EXPECT_EQ( + m.decode_only_steps() + m.mixed_steps() + m.prefill_only_steps(), + m.steps); + // Only the attributable subset feeds the decode-only rate. + EXPECT_LE(m.decode_only_tokens, m.decode_tokens_total); + expect_balanced(m, 2); +} + +TEST(EngineMetricsTest, ReportsContextConcurrencyAndRefusals) { + FakeExecutor executor; + executor.capacity = 2; + Fixture fixture(executor); + Session first = open(fixture.runner); + Session second = open(fixture.runner); + + // A third open has nowhere to go, which is invisible without the counter. + auto refused = fixture.runner.open_session_async(); + ASSERT_EQ(refused.wait_for(kTimeout), std::future_status::ready); + EXPECT_FALSE(refused.get().has_value()); + + auto first_updates = std::make_shared(); + generate(first, tokens(4), config(3), first_updates); + auto second_updates = std::make_shared(); + generate(second, tokens(4), config(3), second_updates); + ASSERT_TRUE(first_updates->wait()); + ASSERT_TRUE(second_updates->wait()); + + first = Session{}; + second = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_EQ(m.sessions_refused, 1u); + EXPECT_GE(m.peak_concurrent_generations, 1u); + EXPECT_LE(m.peak_concurrent_generations, 2u); + // Each session ends holding its prompt plus the tokens it was fed back, and + // context is summed across the batch, so the mean is at least one prompt. + EXPECT_GT(m.mean_context_per_step(), 0.0); + EXPECT_GE(m.context_max, 4); + // Counts, not the rate derived from them: a fake executor's whole run is a + // few microseconds, so us_between can floor it to 0 and the rate accessor + // then reports exactly 0.0. That says nothing about the metric. + EXPECT_GT(m.decode_only_steps(), 0u); + EXPECT_GT(m.decode_only_tokens, 0u); + EXPECT_GT(m.mean_decode_step_sessions(), 0.0); +} + +// Regression: context is a property of a session, not of a task. A prompt that +// arrives as several chunks of one step is one context; summing per chunk +// counted it once per chunk and inflated the total. +TEST(EngineMetricsTest, MultiChunkPromptCountsContextOncePerSession) { + FakeExecutor executor; + // Chunk 2 with a 6-token budget, so a 6-token prompt fits as 3 chunks in a + // single step. max_new_tokens 1 ends the generation on that step's output, + // leaving exactly one step to account for. + Fixture fixture(executor, 2, 2); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + generate(session, tokens(6), config(1), updates); + ASSERT_TRUE(updates->wait()); + ASSERT_EQ(updates->finish(), FinishReason::NewTokenLimit); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + ASSERT_EQ(m.steps, 1u) << "the whole prompt should fit in one step"; + // The session ends that step holding 6 tokens. Per-chunk summing would give + // 2 + 4 + 6 = 12. + EXPECT_EQ(m.context_sum, 6); + EXPECT_EQ(m.context_max, 6); + EXPECT_DOUBLE_EQ(m.mean_context_per_step(), 6.0); +} + +TEST(EngineMetricsTest, InitializationIsTimedAndKeptOutOfTheWall) { + FakeExecutor executor; + // Large enough to dwarf the fake's steps, so the comparison below is not a + // race between two similar durations. + executor.initialize_delay = std::chrono::milliseconds(50); + Fixture fixture(executor); + Session session = open(fixture.runner); + auto updates = std::make_shared(); + + generate(session, tokens(2), config(2), updates); + ASSERT_TRUE(updates->wait()); + + session = Session{}; + fixture.runner.shutdown(); + const EngineMetrics m = fixture.runner.metrics(); + + EXPECT_GE(m.init_us, 40000) << "setup should be measured, not dropped"; + // wall_us() starts at the first step. Were setup folded into it, the wall + // could not be shorter than setup. + EXPECT_GT(static_cast(m.init_us), m.wall_us()) + << "one-time setup must stay outside the run window"; +} + // --- executor initialization ----------------------------------------------- TEST(InitializeTest, RunsOnceBeforeAnythingElseIsAskedOfTheExecutor) { @@ -1695,6 +2027,7 @@ TEST(InitializeTest, FailureStopsTheRunnerAndRefusesSessions) { fixture.runner.shutdown(); EXPECT_EQ(executor.initialize_calls(), 1); + EXPECT_EQ(fixture.runner.metrics().steps, 0u); EXPECT_TRUE(executor.seen().empty()) << "no batch should have run"; EXPECT_TRUE(executor.opened().empty()) << "a session must not be opened on an executor that failed to start"; diff --git a/extension/llm/batching/types.h b/extension/llm/batching/types.h index 3bfa39b74c4..211b8a4be0f 100644 --- a/extension/llm/batching/types.h +++ b/extension/llm/batching/types.h @@ -26,7 +26,7 @@ namespace extension { namespace llm { namespace batching { -using Token = std::int64_t; +using Token = std::uint64_t; using SessionId = std::int64_t; using Position = std::int32_t; // Wide enough that a monotonically issued id cannot wrap in any realistic From 2abf9a67580710d41795d0f9b127c87dcfd960f2 Mon Sep 17 00:00:00 2001 From: Jacob Stevens Date: Wed, 2 Sep 2026 18:08:38 -0400 Subject: [PATCH 008/190] Fix nondeterministic mypy CI dependencies and lint error exposed (#22482) Pin the mypy job to the repository's Torch 2.13 stack before installing timm, and include the Torch pin in its pip cache key. This makes it deterministic instead of depending on order of resolution. Adapt Arm shape and dtype checks to the SymInt and SymBool annotations exposed by newer Torch releases. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- .github/workflows/lint.yml | 12 +++++++++--- backends/arm/_passes/rewrite_conv_pass.py | 2 +- backends/arm/tosa/dialect/ops/resize.py | 2 ++ backends/arm/tosa/partitioner.py | 20 +++++++++++++++++--- 4 files changed, 29 insertions(+), 7 deletions(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 3393f2dbfc6..ea9fa1dfbe1 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -41,16 +41,22 @@ jobs: with: python-version: '3.11' cache: 'pip' - cache-dependency-path: requirements-lintrunner.txt + cache-dependency-path: | + requirements-lintrunner.txt + torch_pin.py - name: Install dependencies run: | - pip install torch --index-url https://download.pytorch.org/whl/cpu + TORCH_VERSION=$(python -c "from torch_pin import TORCH_VERSION; print(TORCH_VERSION)") + pip install \ + "torch==${TORCH_VERSION}" \ + torchvision \ + torchaudio \ + --index-url https://download.pytorch.org/whl/cpu pip install lintrunner==0.12.7 lintrunner-adapters==0.14.1 pip install -r requirements-lintrunner.txt USE_CPP=0 pip install --no-build-isolation third-party/ao pip install pytest numpy parameterized huggingface_hub transformers timm expecttest types-requests - pip install torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu - name: Generate mypy stubs for C++ bindings run: | diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index 69d4d76c9cf..6fb686e7ac3 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -528,7 +528,7 @@ def _is_direct_int32_rescale(node: torch.fx.Node) -> bool: node.op == "call_function" and node.target == exir_ops.backend.tosa.RESCALE.default and len(node.args) > 1 - and node.args[1] == torch.int32 + and node.args[1] is torch.int32 ) def _get_direct_int32_rescale_users( diff --git a/backends/arm/tosa/dialect/ops/resize.py b/backends/arm/tosa/dialect/ops/resize.py index 18ebe6c6210..35db1ae66f5 100644 --- a/backends/arm/tosa/dialect/ops/resize.py +++ b/backends/arm/tosa/dialect/ops/resize.py @@ -96,6 +96,8 @@ def RESIZE( validation_error = get_tosa_resize_output_hw_validation_error(output_hw) if validation_error is not None: raise TosaValueError(validation_error, op="RESIZE") + OH: int | torch.SymInt + OW: int | torch.SymInt if output_hw is None: scale_y_n, scale_y_d, scale_x_n, scale_x_d = scale offset_y, offset_x = offset diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 57aa1db2777..5e5ea1d8423 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -50,6 +50,7 @@ from executorch.exir.graph_module import get_cond_while_submodules from torch.export.exported_program import ExportedProgram from torch.fx import GraphModule +from torch.fx.experimental.symbolic_shapes import statically_known_true from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition from torch.fx.passes.operator_support import any_chain, OperatorSupportBase @@ -152,9 +153,22 @@ def _is_noop_as_strided_copy(node: torch.fx.Node) -> bool: input_tensor = get_first_fake_tensor(ensure_type(torch.fx.Node, node.args[0])) output_tensor = get_first_fake_tensor(node) return ( - input_tensor.shape == output_tensor.shape - and input_tensor.stride() == output_tensor.stride() - and input_tensor.storage_offset() == output_tensor.storage_offset() + len(input_tensor.shape) == len(output_tensor.shape) + and all( + statically_known_true(input_dim == output_dim) + for input_dim, output_dim in zip( + input_tensor.shape, output_tensor.shape + ) + ) + and all( + statically_known_true(input_stride == output_stride) + for input_stride, output_stride in zip( + input_tensor.stride(), output_tensor.stride() + ) + ) + and statically_known_true( + input_tensor.storage_offset() == output_tensor.storage_offset() + ) ) From 902fcf5e3f5ca7a29e259f7a7d6f25c5b57ee97a Mon Sep 17 00:00:00 2001 From: Jacob Stevens Date: Wed, 2 Sep 2026 18:15:58 -0400 Subject: [PATCH 009/190] Fix SLEEF preprocessor macro name to match ATen vec headers Differential Revision: D118475079 Pull Request resolved: https://github.com/pytorch/executorch/pull/22442 --- kernels/optimized/lib_defs.bzl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/kernels/optimized/lib_defs.bzl b/kernels/optimized/lib_defs.bzl index 2ea329a8baa..29d27d97b28 100644 --- a/kernels/optimized/lib_defs.bzl +++ b/kernels/optimized/lib_defs.bzl @@ -25,16 +25,16 @@ def get_vec_preprocessor_flags(): # various ovr_configs are not available in oss preprocessor_flags = select({ "ovr_config//os:linux-x86_64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "ovr_config//os:iphoneos-arm64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "ovr_config//os:macos-arm64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "ovr_config//os:android-arm64": [ - "-DET_BUILD_ARM_VEC256_WITH_SLEEF", + "-DAT_BUILD_ARM_VEC256_WITH_SLEEF", ] if not runtime.is_oss else [], "DEFAULT": [], }) From 33ed3c55192d1a8ce78aaa66d89b36455be5812f Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 2 Sep 2026 15:22:42 -0700 Subject: [PATCH 010/190] Read the query length at build time where it is known (#22494) Follow-up polish on the causal slicing added in #22443. No behaviour change for any shape that worked before. ### Read the query length at build time where it is known The slice that shortens the keys for a causal call always emitted a node to read the query length at run time. For a decode step that length is a constant known when the program is built, so the node was pure overhead on exactly the path this slicing exists to speed up. There is already a helper for this, used by the cached attention handler a few hundred lines down in the same file: it returns a literal for a concrete dimension and emits a size node only for a symbolic one. | Call | Before | After | | --- | --- | --- | | query 1 against 32 keys, static | size node, 2336 bytes | no size node, **2288 bytes** | | query 6 against 32 keys, static | size node, 2336 bytes | no size node, **2288 bytes** | | dynamic query length | size node | size node, unchanged | ### Two comments promised less than the code delivered The guard that declines causal attention also declines a query length that **cannot be compared** at build time, two unrelated dynamic dimensions for instance, and the comment above it described only the case where the query is genuinely longer. It now says both, because the second is the case a reader is more likely to meet. ### A half-wired test knob The key length knob added to the attention test reached the query, key and value but not either mask branch, so a case combining a mask with an unequal key length built a mask of the wrong width. No shipped case does that today, so nothing was failing, but the next person to use it would have hit it. ### Test plan Causal attention at every shape the operator suite covers, run against eager on an Apple Silicon Mac: query 1, 2 and 6 against longer keys, equal lengths, non-causal, float32 and bfloat16. All match to 4.8e-07 or exactly, 8 of 8 cases. A dynamic query length bounded by the keys keeps its size node, exports, and matches eager to 1.2e-07. Building the test mask both ways in eager torch: with the key length it runs, and with the query length it raises the size mismatch this removes. Exported the speech model: unchanged, 8 fused attention nodes and no size nodes. Ran the neighbouring backend test files, all passing. Co-authored-by: PyTorch Bot --- backends/mlx/patterns.py | 18 ++++++++++-------- backends/mlx/test/test_ops.py | 8 ++++---- 2 files changed, 14 insertions(+), 12 deletions(-) diff --git a/backends/mlx/patterns.py b/backends/mlx/patterns.py index f2899baf6d4..9ff360fdb8f 100644 --- a/backends/mlx/patterns.py +++ b/backends/mlx/patterns.py @@ -21,6 +21,7 @@ from executorch.backends.mlx.builder.op_helpers import ( emit_quantized_biases, emit_quantized_gather, + emit_shape, emit_stop_position, mlx_qparams_supported, parse_dequant_int4_node, @@ -583,9 +584,11 @@ def _kernel_can_compute(cls, sdpa_node: Node) -> bool: if mask_val is None or mask_val.dim() > 4: return False - # A causal query longer than its keys is left to decompose. Torch clamps each - # row to the keys that exist, and neither the kernel's own mask nor slicing the - # keys to the query length reproduces that. + # Causal attention is only taken when the query is known to be no longer than + # the keys. Torch clamps a longer query to the keys that exist, and neither the + # kernel's own mask nor slicing the keys reproduces that. A query length that + # cannot be compared at build time, two unrelated dynamic dimensions for + # instance, is also declined, because the relation has to hold for every call. if is_causal and not statically_known_true( operand_vals[0].shape[-2] <= operand_vals[1].shape[-2] ): @@ -720,10 +723,9 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: q_len = self.q_node.meta["val"].shape[-2] k_len = self.k_node.meta["val"].shape[-2] if is_causal and not statically_known_true(q_len == k_len): - _, rows = P.slot_manager.make_tmp_value_slot() - P.emit( - SymSizeNode(a=P.slot_to_tid(inputs[0]), dim=2, out=P.slot_to_vid(rows)) - ) + # A literal when the query length is known at build time, which is the + # decode case this exists for, and a size node only when it is symbolic. + rows = emit_shape(P, self.q_node, inputs[0])[-2] for i in (1, 2): _, sliced = P.make_tmp_slot() P.emit( @@ -732,7 +734,7 @@ def __call__(self, P: MLXProgramBuilder, n: Node) -> Slot: out=P.slot_to_tid(sliced), axis=IntOrVid.from_literal(2), start=IntOrVid.from_literal(0), - stop=P.to_int_or_vid(rows), + stop=rows, ) ) inputs[i] = sliced diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index bfc2e931e67..1f8459564c3 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -6270,16 +6270,16 @@ def create_inputs(self) -> Tuple[torch.Tensor, ...]: if self.use_mask: # Additive float mask: 0 = attend, -inf = masked - mask = torch.zeros(self.batch_size, 1, self.seq_len, self.seq_len) - mask[:, :, :, : self.seq_len // 4] = float("-inf") + mask = torch.zeros(self.batch_size, 1, self.seq_len, self.kv_seq_len) + mask[:, :, :, : self.kv_seq_len // 4] = float("-inf") return (q, k, v, mask) elif self.use_bool_mask: # Boolean mask: True = attend, False = masked # This tests that the backend correctly converts bool -> additive format mask = torch.ones( - self.batch_size, 1, self.seq_len, self.seq_len, dtype=torch.bool + self.batch_size, 1, self.seq_len, self.kv_seq_len, dtype=torch.bool ) - mask[:, :, :, : self.seq_len // 4] = False # Mask out first quarter + mask[:, :, :, : self.kv_seq_len // 4] = False # Mask out first quarter return (q, k, v, mask) return (q, k, v) From 5bcf795a78c0feeb963e41c3fd2b04d9322ae8c0 Mon Sep 17 00:00:00 2001 From: Jiseong-oh Date: Thu, 3 Sep 2026 13:17:37 +0900 Subject: [PATCH 011/190] Optimize Enn runner (#18735) ### Summary * Open Model with zero-copy using dmabufheap * Improve model loading with zero-copy * Remove unncessary comments or headers. Change file name cc @SS-JIA @digantdesai @kimishpatel --------- Signed-off-by: Jiseong oh Signed-off-by: Jiseong Oh Co-authored-by: Hoon Choi --- backends/samsung/runtime/CMakeLists.txt | 8 +- .../runtime/enn_api_implementation.cpp | 41 +-- .../samsung/runtime/enn_api_implementation.h | 20 +- backends/samsung/runtime/enn_backend.cpp | 1 + backends/samsung/runtime/enn_executor.cpp | 34 ++- backends/samsung/runtime/enn_executor.h | 2 +- .../runtime/enn_shared_memory_manager.cpp | 119 ++++++++ .../runtime/enn_shared_memory_manager.h | 43 +++ .../extension/exynos_file_data_loader.cpp | 265 +++++++++++++++++ .../extension/exynos_file_data_loader.h | 88 ++++++ .../runtime/extension/test/CMakeLists.txt | 20 ++ .../test/exynos_file_data_loader_test.cpp | 266 ++++++++++++++++++ .../executor_runner/enn_executor_runner.cpp | 30 +- 13 files changed, 886 insertions(+), 51 deletions(-) create mode 100644 backends/samsung/runtime/enn_shared_memory_manager.cpp create mode 100644 backends/samsung/runtime/enn_shared_memory_manager.h create mode 100644 backends/samsung/runtime/extension/exynos_file_data_loader.cpp create mode 100644 backends/samsung/runtime/extension/exynos_file_data_loader.h create mode 100644 backends/samsung/runtime/extension/test/CMakeLists.txt create mode 100644 backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp diff --git a/backends/samsung/runtime/CMakeLists.txt b/backends/samsung/runtime/CMakeLists.txt index deb93f31bc8..aec4a71de5d 100644 --- a/backends/samsung/runtime/CMakeLists.txt +++ b/backends/samsung/runtime/CMakeLists.txt @@ -7,7 +7,7 @@ # logging target_sources( enn_logging - PUBLIC ${CMAKE_CURRENT_LIST_DIR}/logging.h + PUBLIC $ PRIVATE ${CMAKE_CURRENT_LIST_DIR}/logging.cpp ) @@ -17,6 +17,12 @@ if(${ANDROID}) enn_backend PRIVATE ${CMAKE_CURRENT_LIST_DIR}/enn_backend.cpp ${CMAKE_CURRENT_LIST_DIR}/enn_executor.cpp + ${CMAKE_CURRENT_LIST_DIR}/enn_shared_memory_manager.cpp ${CMAKE_CURRENT_LIST_DIR}/enn_api_implementation.cpp + ${CMAKE_CURRENT_LIST_DIR}/extension/exynos_file_data_loader.cpp ) + + if(BUILD_TESTING) + add_subdirectory(${CMAKE_CURRENT_LIST_DIR}/extension/test) + endif() endif() diff --git a/backends/samsung/runtime/enn_api_implementation.cpp b/backends/samsung/runtime/enn_api_implementation.cpp index bbef883fd10..0cdf810d5c4 100644 --- a/backends/samsung/runtime/enn_api_implementation.cpp +++ b/backends/samsung/runtime/enn_api_implementation.cpp @@ -33,31 +33,35 @@ void* loadApiFunction(void* handle, const char* name, bool optional) { return fn; } -std::mutex EnnApi::instance_mutex_; - EnnApi* EnnApi::getEnnApiInstance() { - std::lock_guard lgd(instance_mutex_); static EnnApi enn_api; - if (!enn_api.getInitialize()) { - auto status = enn_api.loadApiLib(); - if (status == Error::Ok) { - ENN_LOG_INFO("Loading ENN API library Completed.") - enn_api.initialize_ = true; - } else { - ENN_LOG_ERROR("Failed to load enn api library. %s", dlerror()); - } - } return &enn_api; } +EnnApi::EnnApi() { + auto status = loadApiLib(); + if (status != Error::Ok) { + ET_LOG(Error, "Failed to load enn api library. %s", dlerror()); + return; + } + auto ret = EnnInitialize(); + if (ret != ENN_RET_SUCCESS) { + ET_LOG(Error, "EnnInitialize failed: %d", static_cast(ret)); + unloadApiLib(); + return; + } + ET_LOG(Info, "Loading ENN API library Completed."); + initialize_ = true; +} + EnnApi::~EnnApi() { - std::lock_guard lgd(instance_mutex_); - if (getInitialize()) { + if (initialize_) { + EnnDeinitialize(); unloadApiLib(); } } -bool EnnApi::getInitialize() const { +bool EnnApi::isInitialized() const { return initialize_; } @@ -87,13 +91,18 @@ Error EnnApi::loadApiLib() { ENN_LOAD_API_FUNC(libenn_public_api_, EnnBufferCommit, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnGetBuffersInfo, this); ENN_LOAD_API_FUNC(libenn_public_api_, EnnReleaseBuffers, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnCreateBuffer, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnReleaseBuffer, this); + ENN_LOAD_API_FUNC( + libenn_public_api_, EnnGetFileDescriptorFromEnnBuffer, this); + ENN_LOAD_API_FUNC(libenn_public_api_, EnnOpenModelFromFd, this); return Error::Ok; } Error EnnApi::unloadApiLib() { if (dlclose(libenn_public_api_) != 0) { - ENN_LOG_ERROR("Failed to close ENN API library. %s", dlerror()); + ET_LOG(Error, "Failed to close ENN API library. %s", dlerror()); return Error::Internal; } libenn_public_api_ = nullptr; diff --git a/backends/samsung/runtime/enn_api_implementation.h b/backends/samsung/runtime/enn_api_implementation.h index e6f8df01f7a..3e0bf1c5f2a 100644 --- a/backends/samsung/runtime/enn_api_implementation.h +++ b/backends/samsung/runtime/enn_api_implementation.h @@ -28,6 +28,7 @@ class EnnApi { ~EnnApi(); static EnnApi* getEnnApiInstance(); + bool isInitialized() const; EnnReturn (*EnnInitialize)(void); EnnReturn (*EnnSetPreferencePerfMode)(const uint32_t val); @@ -37,6 +38,7 @@ class EnnApi { const char* va, const uint32_t size, EnnModelId* model_id); + EnnReturn (*EnnOpenModelFromFd)(int _fd, EnnModelId* model_id); EnnReturn (*EnnSetFastIpc)(void); EnnReturn (*EnnUnsetFastIpc)(void); EnnReturn (*EnnExecuteModelFastIpc)( @@ -67,6 +69,13 @@ class EnnApi { NumberOfBuffersInfo* buffers_info); EnnReturn ( *EnnReleaseBuffers)(EnnBufferPtr* buffers, const int32_t numOfBuffers); + EnnReturn (*EnnCreateBuffer)( + const uint32_t req_size, + const uint32_t ion_flag, + EnnBufferPtr* out); + EnnReturn (*EnnReleaseBuffer)(EnnBufferPtr buf); + EnnReturn ( + *EnnGetFileDescriptorFromEnnBuffer)(EnnBufferPtr buffer, int32_t* fd); private: static std::mutex instance_mutex_; @@ -75,8 +84,7 @@ class EnnApi { void* libenn_public_api_ = nullptr; static std::atomic ref_count_; - EnnApi() = default; - bool getInitialize() const; + EnnApi(); Error loadApiLib(); Error unloadApiLib(); }; @@ -120,6 +128,14 @@ typedef EnnReturn (*EnnGetBuffersInfo_fn)( NumberOfBuffersInfo* buffers_info); typedef EnnReturn ( *EnnReleaseBuffers_fn)(EnnBufferPtr* buffers, const int32_t numOfBuffers); +typedef EnnReturn (*EnnCreateBuffer_fn)( + const uint32_t req_size, + const uint32_t ion_flag, + EnnBufferPtr* out); +typedef EnnReturn (*EnnReleaseBuffer_fn)(EnnBufferPtr buf); +typedef EnnReturn ( + *EnnGetFileDescriptorFromEnnBuffer_fn)(EnnBufferPtr buffer, int32_t* fd); +typedef EnnReturn (*EnnOpenModelFromFd_fn)(int _fd, EnnModelId* model_id); } // namespace enn } // namespace executor diff --git a/backends/samsung/runtime/enn_backend.cpp b/backends/samsung/runtime/enn_backend.cpp index 44838342013..81484afc44c 100644 --- a/backends/samsung/runtime/enn_backend.cpp +++ b/backends/samsung/runtime/enn_backend.cpp @@ -9,6 +9,7 @@ #include #include #include + #include #include #include diff --git a/backends/samsung/runtime/enn_executor.cpp b/backends/samsung/runtime/enn_executor.cpp index f6c1b08a8a6..99a2ce11e5e 100644 --- a/backends/samsung/runtime/enn_executor.cpp +++ b/backends/samsung/runtime/enn_executor.cpp @@ -7,12 +7,12 @@ * */ #include +#include #include #include -#include -#include -#include +#include +#include #include namespace torch { @@ -21,21 +21,39 @@ namespace enn { Error EnnExecutor::initialize(const char* binary_buf_addr, size_t buf_size) { EXYNOS_ATRACE_FUNCTION_LINE(); + auto sm_instance = executorch::backends::enn::shared_memory_manager:: + SharedMemoryManager::getInstance(); const EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); - auto ret = enn_api_inst->EnnInitialize(); ET_CHECK_OR_RETURN_ERROR( - ret == ENN_RET_SUCCESS, Internal, "Enn initialize failed."); + enn_api_inst->isInitialized(), Internal, "Enn initialize failed."); + EnnReturn ret = ENN_RET_SUCCESS; - ET_LOG(Info, "Start to open model %p, %ld", binary_buf_addr, buf_size); - ret = enn_api_inst->EnnOpenModelFromMemory( - binary_buf_addr, buf_size, &model_id_); + ET_LOG(Info, "Start to open model %p, %zu", binary_buf_addr, buf_size); + EnnBufferPtr shared_buffer = nullptr; + if (sm_instance->query(&shared_buffer, binary_buf_addr, buf_size)) { + int32_t fd; + if (shared_buffer->va == binary_buf_addr && + !enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(shared_buffer, &fd)) { + ret = enn_api_inst->EnnOpenModelFromFd(fd, &model_id_); + if (ret == ENN_RET_SUCCESS) { + ET_LOG(Info, "Opened model from file descriptor, so fd is closed"); + sm_instance->free(shared_buffer->va); + } + } + } + if (!model_id_) { + ET_LOG(Info, "Open model from memory"); + ret = enn_api_inst->EnnOpenModelFromMemory( + binary_buf_addr, buf_size, &model_id_); + } ET_CHECK_OR_RETURN_ERROR( ret == ENN_RET_SUCCESS, Internal, "Failed to load Enn model from buffer %d", (int)ret); ET_LOG(Info, "Open successfully."); + NumberOfBuffersInfo buffers_info; ret = enn_api_inst->EnnAllocateAllBuffersWithSessionId( model_id_, &alloc_buffer_, &buffers_info, 0, true); diff --git a/backends/samsung/runtime/enn_executor.h b/backends/samsung/runtime/enn_executor.h index 902b420a036..68ee1f56847 100644 --- a/backends/samsung/runtime/enn_executor.h +++ b/backends/samsung/runtime/enn_executor.h @@ -40,7 +40,7 @@ class EnnExecutor { ~EnnExecutor(); private: - EnnModelId model_id_; + EnnModelId model_id_ = 0ULL; EnnBufferPtr* alloc_buffer_ = nullptr; int32_t num_of_inputs_ = 0; int32_t num_of_outputs_ = 0; diff --git a/backends/samsung/runtime/enn_shared_memory_manager.cpp b/backends/samsung/runtime/enn_shared_memory_manager.cpp new file mode 100644 index 00000000000..9babf34de9f --- /dev/null +++ b/backends/samsung/runtime/enn_shared_memory_manager.cpp @@ -0,0 +1,119 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + */ +#include + +#include +#include + +#include +#include +#include + +namespace executorch { +namespace backends { +namespace enn { +namespace shared_memory_manager { + +using torch::executor::enn::EnnApi; + +static std::mutex instance_mutex_; + +SharedMemoryManager* SharedMemoryManager::getInstance() { + // Touch the EnnApi singleton first so that it outlives this instance: the + // destructor below releases buffers through the ENN API. + EnnApi::getEnnApiInstance(); + static SharedMemoryManager instance; + return &instance; +} + +SharedMemoryManager::~SharedMemoryManager() { + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + for (auto& buffer : buffers_) { + if (enn_api_inst->EnnReleaseBuffer(buffer)) { + ET_LOG(Error, "Failed to destroy buffer: %p", buffer->va); + } + } + buffers_.clear(); +} + +void* SharedMemoryManager::alloc(const size_t size) { + if (size > std::numeric_limits::max()) { + ET_LOG( + Error, "Requested size %zu exceeds ENN buffer limit (uint32_t)", size); + return nullptr; + } + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + EnnBufferPtr bufferPtr; + auto ret = enn_api_inst->EnnCreateBuffer(size, 0, &bufferPtr); + if (ret) { + ET_LOG(Error, "Buffer Creation Error"); + return nullptr; + } + buffers_.emplace_back(bufferPtr); + return bufferPtr->va; +} + +bool SharedMemoryManager::query( + EnnBufferPtr* out, + const void* ptr, + const size_t size) { + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + for (const auto& buffer : buffers_) { + if (buffer->va <= ptr && + ptr < static_cast(buffer->va) + buffer->size) { + int32_t fd; + auto ret = enn_api_inst->EnnGetFileDescriptorFromEnnBuffer(buffer, &fd); + if (ret) { + ET_LOG( + Info, + "va: %p, size: %zu is in LUT, but failed to get FileDescriptor", + ptr, + size); + return false; + } + *out = buffer; + return true; + } + } + ET_LOG(Info, "va: %p, size: %zu is not in LUT", ptr, size); + *out = nullptr; + return false; +} + +void SharedMemoryManager::free(void* ptr) { + std::lock_guard lgd(instance_mutex_); + auto enn_api_inst = EnnApi::getEnnApiInstance(); + for (auto it = buffers_.begin(); it != buffers_.end(); ++it) { + if ((*it)->va == ptr) { + if (enn_api_inst->EnnReleaseBuffer(*it)) { + ET_LOG( + Error, + "Failed to destroy buffer: %p, keeping tracked for retry", + ptr); + return; + } + ET_LOG( + Info, + "va(%p), size(%" PRIu32 "), offset(%" PRIu32 ") is erased from LUT", + ptr, + (*it)->size, + (*it)->offset); + buffers_.erase(it); + return; + } + } +} + +} // namespace shared_memory_manager +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/enn_shared_memory_manager.h b/backends/samsung/runtime/enn_shared_memory_manager.h new file mode 100644 index 00000000000..301b18f8cd5 --- /dev/null +++ b/backends/samsung/runtime/enn_shared_memory_manager.h @@ -0,0 +1,43 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + */ +#pragma once + +#include + +#include +#include + +namespace executorch { +namespace backends { +namespace enn { +namespace shared_memory_manager { + +class SharedMemoryManager { + public: + static SharedMemoryManager* getInstance(); + + SharedMemoryManager() = default; + ~SharedMemoryManager(); + SharedMemoryManager(const SharedMemoryManager&) = delete; + SharedMemoryManager& operator=(const SharedMemoryManager&) = delete; + SharedMemoryManager(SharedMemoryManager&&) = delete; + SharedMemoryManager& operator=(SharedMemoryManager&&) = delete; + + void* alloc(const size_t size); + void free(void* ptr); + bool query(EnnBufferPtr* out, const void* ptr, const size_t size); + + private: + std::vector buffers_; +}; + +} // namespace shared_memory_manager +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/extension/exynos_file_data_loader.cpp b/backends/samsung/runtime/extension/exynos_file_data_loader.cpp new file mode 100644 index 00000000000..e7bd252f866 --- /dev/null +++ b/backends/samsung/runtime/extension/exynos_file_data_loader.cpp @@ -0,0 +1,265 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include +#include +#include + +// Some platforms (e.g. Xtensa) do not support pread() that we use to read the +// file at different offsets simultaneously from multiple threads not affecting +// each other. We list them below and use a workaround for them. +#if defined(__xtensa__) || defined(__hexagon__) +#define ET_HAVE_PREAD 0 +#endif // defined(__xtensa__) || defined(__hexagon__) + +#ifndef ET_HAVE_PREAD +#define ET_HAVE_PREAD 1 +#endif // !ET_HAVE_PREAD + +using executorch::backends::enn::shared_memory_manager::SharedMemoryManager; +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; + +namespace executorch { +namespace backends { +namespace enn { + +namespace { + +/** + * Returns true if the value is an integer power of 2. + */ +bool is_power_of_2(size_t value) { + return value > 0 && (value & ~(value - 1)) == value; +} + +/** + * FreeableBuffer::FreeFn-compatible callback. + * + * `data` is the original buffer pointer. `context` and `size` are unused. + */ +void FreeSegment(ET_UNUSED void* context, void* data, ET_UNUSED size_t size) { + SharedMemoryManager::getInstance()->free(data); +} + +} // namespace + +ExynosFileDataLoader::~ExynosFileDataLoader() { + // file_name_ can be nullptr if this instance was moved from, but freeing a + // null pointer is safe. + std::free(const_cast(file_name_)); + // fd_ can be -1 if this instance was moved from, but closing a negative fd is + // safe (though it will return an error). + if (fd_ == -1) { + return; + } + ::close(fd_); +} + +Result ExynosFileDataLoader::from( + const char* file_name, + size_t alignment) { + ET_CHECK_OR_RETURN_ERROR( + is_power_of_2(alignment), + InvalidArgument, + "Alignment %zu is not a power of 2", + alignment); + + const size_t page_size = static_cast(sysconf(_SC_PAGESIZE)); + ET_CHECK_OR_RETURN_ERROR( + alignment <= page_size, + InvalidArgument, + "Alignment %zu exceeds page size %zu; ENN shared memory buffers " + "are only guaranteed to be page-aligned", + alignment, + page_size); + + ET_CHECK_OR_RETURN_ERROR( + file_name != nullptr, InvalidArgument, "File name cannot be empty."); + + // Use open() instead of fopen() to avoid the layer of buffering that + // fopen() does. We will be reading large portions of the file in one shot, + // so buffering does not help. + int fd = ::open(file_name, O_RDONLY); + if (fd < 0) { + ET_LOG( + Error, "Failed to open %s: %s (%d)", file_name, strerror(errno), errno); + return Error::AccessFailed; + } + + // Cache the file size. + struct stat st; + int err = ::fstat(fd, &st); + if (err < 0) { + ET_LOG( + Error, + "Could not get length of %s: %s (%d)", + file_name, + ::strerror(errno), + errno); + ::close(fd); + return Error::AccessFailed; + } + size_t file_size = st.st_size; + + // Copy the filename so we can print better debug messages if reads fail. + const char* file_name_copy = ::strdup(file_name); + + if (file_name_copy == nullptr) { + ET_LOG(Error, "strdup(%s) failed", file_name); + ::close(fd); + return Error::MemoryAllocationFailed; + } + + return ExynosFileDataLoader(fd, file_size, alignment, file_name_copy); +} + +Result ExynosFileDataLoader::load( + size_t offset, + size_t size, + ET_UNUSED const DataLoader::SegmentInfo& segment_info) const { + ET_CHECK_OR_RETURN_ERROR( + // Probably had its value moved to another instance. + fd_ >= 0, + InvalidState, + "Uninitialized"); + ET_CHECK_OR_RETURN_ERROR( + size <= file_size_ && offset <= file_size_ - size, + InvalidArgument, + "File %s: offset %zu + size %zu > file_size_ %zu", + file_name_, + offset, + size, + file_size_); + + // Don't bother allocating/freeing for empty segments. + if (size == 0) { + return FreeableBuffer(nullptr, 0, /*free_fn=*/nullptr); + } + + auto* shared_memory = SharedMemoryManager::getInstance(); + void* buffer = shared_memory->alloc(size); + if (buffer == nullptr) { + ET_LOG( + Error, + "Reading from %s at offset %zu: alloc(%zu) failed", + file_name_, + offset, + size); + return Error::MemoryAllocationFailed; + } + + auto err = load_into(offset, size, segment_info, buffer); + if (err != Error::Ok) { + shared_memory->free(buffer); + return err; + } + + return FreeableBuffer(buffer, size, FreeSegment); +} + +Result ExynosFileDataLoader::size() const { + ET_CHECK_OR_RETURN_ERROR( + // Probably had its value moved to another instance. + fd_ >= 0, + InvalidState, + "Uninitialized"); + return file_size_; +} + +ET_NODISCARD Error ExynosFileDataLoader::load_into( + size_t offset, + size_t size, + ET_UNUSED const SegmentInfo& segment_info, + void* buffer) const { + ET_CHECK_OR_RETURN_ERROR( + // Probably had its value moved to another instance. + fd_ >= 0, + InvalidState, + "Uninitialized"); + ET_CHECK_OR_RETURN_ERROR( + size <= file_size_ && offset <= file_size_ - size, + InvalidArgument, + "File %s: offset %zu + size %zu > file_size_ %zu", + file_name_, + offset, + size, + file_size_); + ET_CHECK_OR_RETURN_ERROR( + buffer != nullptr, InvalidArgument, "Provided buffer cannot be null"); + + // Read the data into the aligned address. + size_t needed = size; + uint8_t* buf = reinterpret_cast(buffer); + + // Make a duplicate fd if pread() is not available and we have to seek(). + // Cannot use the standard dup() or fcntl() calls because the returned + // duplicate will share the underlying file record and affect the original fd + // when seeking on multiple threads simultaneously. + const auto dup_fd = ET_HAVE_PREAD ? fd_ : ::open(file_name_, O_RDONLY); + + while (needed > 0) { + // Reads on macOS will fail with EINVAL if size > INT32_MAX. + const auto chunk_size = std::min( + needed, static_cast(std::numeric_limits::max())); + const auto nread = +#if ET_HAVE_PREAD + ::pread(dup_fd, buf, chunk_size, offset); +#else + (::lseek(dup_fd, offset, SEEK_SET) == (off_t)-1) + ? -1 + : ::read(dup_fd, buf, chunk_size); +#endif + if (nread < 0 && errno == EINTR) { + // Interrupted by a signal; zero bytes read. + continue; + } + if (nread <= 0) { + // nread == 0 means EOF, which we shouldn't see if we were able to read + // the full amount. nread < 0 means an error occurred. + ET_LOG( + Error, + "Reading from %s: failed to read %zu bytes at offset %zu: %s", + file_name_, + size, + offset, + nread == 0 ? "EOF" : strerror(errno)); + if (!ET_HAVE_PREAD) { + ::close(dup_fd); + } + return Error::AccessFailed; + } + needed -= nread; + buf += nread; + offset += nread; + } + if (!ET_HAVE_PREAD) { + ::close(dup_fd); + } + return Error::Ok; +} + +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/extension/exynos_file_data_loader.h b/backends/samsung/runtime/extension/exynos_file_data_loader.h new file mode 100644 index 00000000000..d33453e92bf --- /dev/null +++ b/backends/samsung/runtime/extension/exynos_file_data_loader.h @@ -0,0 +1,88 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +#include +#include +#include + +namespace executorch { +namespace backends { +namespace enn { + +/** + * Loads a file from disk into buffers backed by ENN shared memory (dmabuf), so + * that segments handed to the ENN backend can be opened without an extra copy. + * + * Mirrors executorch::extension::FileDataLoader, which cannot be reused + * directly because it is final and always allocates from the heap. + */ +class ExynosFileDataLoader final : public executorch::runtime::DataLoader { + public: + // `alignment` must not exceed the system page size: ENN shared memory + // buffers are only guaranteed to be page-aligned, and larger requests are + // rejected rather than silently under-aligned. + static executorch::runtime::Result from( + const char* file_name, + size_t alignment = alignof(std::max_align_t)); + + ExynosFileDataLoader(ExynosFileDataLoader&& rhs) noexcept + : file_name_(rhs.file_name_), + file_size_(rhs.file_size_), + alignment_(rhs.alignment_), + fd_(rhs.fd_) { + const_cast(rhs.file_name_) = nullptr; + const_cast(rhs.file_size_) = 0; + const_cast(rhs.alignment_) = {}; + const_cast(rhs.fd_) = -1; + } + + ~ExynosFileDataLoader() override; + + ET_NODISCARD + executorch::runtime::Result load( + size_t offset, + size_t size, + const DataLoader::SegmentInfo& segment_info) const override; + + ET_NODISCARD executorch::runtime::Result size() const override; + + ET_NODISCARD executorch::runtime::Error load_into( + size_t offset, + size_t size, + ET_UNUSED const SegmentInfo& segment_info, + void* buffer) const override; + + private: + ExynosFileDataLoader( + int fd, + size_t file_size, + size_t alignment, + const char* file_name) + : file_name_(file_name), + file_size_(file_size), + alignment_{alignment}, + fd_(fd) {} + + // Not safely copyable. + ExynosFileDataLoader(const ExynosFileDataLoader&) = delete; + ExynosFileDataLoader& operator=(const ExynosFileDataLoader&) = delete; + ExynosFileDataLoader& operator=(ExynosFileDataLoader&&) = delete; + + const char* const file_name_; // Owned by the instance. + const size_t file_size_; + const std::align_val_t alignment_; + const int fd_; // Owned by the instance. +}; + +} // namespace enn +} // namespace backends +} // namespace executorch diff --git a/backends/samsung/runtime/extension/test/CMakeLists.txt b/backends/samsung/runtime/extension/test/CMakeLists.txt new file mode 100644 index 00000000000..139a2f4959e --- /dev/null +++ b/backends/samsung/runtime/extension/test/CMakeLists.txt @@ -0,0 +1,20 @@ +# Copyright (c) 2025 Samsung Electronics Co. LTD +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# logging + +cmake_minimum_required(VERSION 3.19) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../../..) + +include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + +set(_test_srcs exynos_file_data_loader_test.cpp) + +et_cxx_test( + exynos_file_data_loader_test SOURCES ${_test_srcs} EXTRA_LIBS enn_backend + enn_logging +) diff --git a/backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp b/backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp new file mode 100644 index 00000000000..777c5665acc --- /dev/null +++ b/backends/samsung/runtime/extension/test/exynos_file_data_loader_test.cpp @@ -0,0 +1,266 @@ +/* + * Copyright (c) 2025 Samsung Electronics Co. LTD + * All rights reserved + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + * + */ +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +using namespace ::testing; +using executorch::backends::enn::ExynosFileDataLoader; +using executorch::backends::enn::shared_memory_manager::SharedMemoryManager; +using executorch::extension::testing::TempFile; +using executorch::runtime::DataLoader; +using executorch::runtime::Error; +using executorch::runtime::FreeableBuffer; +using executorch::runtime::Result; +using torch::executor::enn::EnnApi; + +class ExynosFileDataLoaderTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + // Constructing the singleton initializes the ENN API. + EnnApi::getEnnApiInstance(); + } + + size_t alignment() const { + return GetParam(); + } +}; + +TEST_P(ExynosFileDataLoaderTest, InBoundsLoadsSucceed) { + uint8_t data[256]; + for (int i = 0; i < sizeof(data); ++i) { + data[i] = i; + } + TempFile tf(data, sizeof(data)); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // size() should succeed and reflect the total size. + Result size = fdl->size(); + ASSERT_EQ(size.error(), Error::Ok); + EXPECT_EQ(*size, sizeof(data)); + + // Load the first bytes of the data. + { + Result fb = fdl->load( + /*offset=*/0, + /*size=*/8, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + EXPECT_EQ(fb->size(), 8); + EXPECT_EQ( + 0, + std::memcmp( + fb->data(), + "\x00\x01\x02\x03" + "\x04\x05\x06\x07", + fb->size())); + + // Freeing should release the buffer and clear out the segment. + fb->Free(); + EXPECT_EQ(fb->size(), 0); + EXPECT_EQ(fb->data(), nullptr); + + // Safe to call multiple times. + fb->Free(); + } + + // Load the last few bytes of the data, a different size than the first time. + { + Result fb = fdl->load( + /*offset=*/sizeof(data) - 3, + /*size=*/3, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + EXPECT_EQ(fb->size(), 3); + EXPECT_EQ(0, std::memcmp(fb->data(), "\xfd\xfe\xff", fb->size())); + } + + // Loading all of the data succeeds. + { + Result fb = fdl->load( + /*offset=*/0, + /*size=*/sizeof(data), + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + EXPECT_EQ(fb->size(), sizeof(data)); + EXPECT_EQ(0, std::memcmp(fb->data(), data, fb->size())); + } + + // Loading zero-sized data succeeds, even at the end of the data. + { + Result fb = fdl->load( + /*offset=*/sizeof(data), + /*size=*/0, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_EQ(fb->size(), 0); + } +} + +TEST_P(ExynosFileDataLoaderTest, OutOfBoundsLoadFails) { + // Create a temp file; contents don't matter. + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // Loading beyond the end of the data should fail. + { + Result fb = fdl->load( + /*offset=*/0, + /*size=*/sizeof(data) + 1, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + EXPECT_NE(fb.error(), Error::Ok); + } + + // Loading zero bytes still fails if it's past the end of the data. + { + Result fb = fdl->load( + /*offset=*/sizeof(data) + 1, + /*size=*/0, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + EXPECT_NE(fb.error(), Error::Ok); + } +} + +TEST_P(ExynosFileDataLoaderTest, FromMissingFileFails) { + // Wrapping a file that doesn't exist should fail. + Result fdl = ExynosFileDataLoader::from( + "/tmp/FILE_DOES_NOT_EXIST_EXECUTORCH_EXYNOS_LOADER_TEST"); + EXPECT_NE(fdl.error(), Error::Ok); +} + +TEST_P(ExynosFileDataLoaderTest, FromEmptyFilePathFails) { + // Nullptr should fail + Result fdl = ExynosFileDataLoader::from(nullptr); + EXPECT_NE(fdl.error(), Error::Ok); +} + +TEST_P(ExynosFileDataLoaderTest, BadAlignmentFails) { + // Create a temp file; contents don't matter. + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + + // Creating a loader with default alignment works fine. + { + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str()); + ASSERT_EQ(fdl.error(), Error::Ok); + } + + // Bad alignments fail. + const std::vector bad_alignments = {0, 3, 5, 17}; + for (size_t bad_alignment : bad_alignments) { + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), bad_alignment); + ASSERT_EQ(fdl.error(), Error::InvalidArgument); + } +} + +// Tests that the move ctor works. +TEST_P(ExynosFileDataLoaderTest, MoveCtor) { + // Create a loader. + std::string contents = "FILE_CONTENTS"; + TempFile tf(contents); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + EXPECT_EQ(fdl->size().get(), contents.size()); + + // Move it into another instance. + ExynosFileDataLoader dl2(std::move(*fdl)); + + // Old loader should now be invalid. + EXPECT_EQ( + fdl->load( + 0, + 0, + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)) + .error(), + Error::InvalidState); + EXPECT_EQ(fdl->size().error(), Error::InvalidState); + + // New loader should point to the file. + EXPECT_EQ(dl2.size().get(), contents.size()); + Result fb = dl2.load( + /*offset=*/0, + contents.size(), + DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EXPECT_ALIGNED(fb->data(), alignment()); + ASSERT_EQ(fb->size(), contents.size()); + EXPECT_EQ(0, std::memcmp(fb->data(), contents.data(), fb->size())); +} + +// TODO: Allocation failure test +TEST_P(ExynosFileDataLoaderTest, EnnQueryFailure) { + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // Allocate valid buffer + Result fb = fdl->load( + 0, 8, DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + // Test query with invalid pointer + EnnBufferPtr out; + void* invalid_ptr = reinterpret_cast(0xDEADBEEF); + bool found = SharedMemoryManager::getInstance()->query(&out, invalid_ptr, 8); + EXPECT_FALSE(found); + EXPECT_EQ(out, nullptr); +} + +TEST_P(ExynosFileDataLoaderTest, EnnFreeFailure) { + uint8_t data[256] = {}; + TempFile tf(data, sizeof(data)); + Result fdl = + ExynosFileDataLoader::from(tf.path().c_str(), alignment()); + ASSERT_EQ(fdl.error(), Error::Ok); + + // Allocate and free with invalid ENN state + { + Result fb = fdl->load( + 0, 8, DataLoader::SegmentInfo(DataLoader::SegmentInfo::Type::Program)); + ASSERT_EQ(fb.error(), Error::Ok); + EnnApi::getEnnApiInstance()->EnnDeinitialize(); + fb->Free(); // Free should fail gracefully + EnnApi::getEnnApiInstance()->EnnInitialize(); + } +} + +// Run all ExynosFileDataLoaderTests multiple times, varying the return value of +// `GetParam()` based on the `testing::Values` list. The tests will interpret +// the value as "alignment". +INSTANTIATE_TEST_SUITE_P( + VariedSegments, + ExynosFileDataLoaderTest, + testing::Values( + 1, + 4, + alignof(std::max_align_t), + 2 * alignof(std::max_align_t), + 128, + 1024)); \ No newline at end of file diff --git a/examples/samsung/executor_runner/enn_executor_runner.cpp b/examples/samsung/executor_runner/enn_executor_runner.cpp index de168be9d7c..bdbff74088c 100644 --- a/examples/samsung/executor_runner/enn_executor_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_runner.cpp @@ -17,8 +17,8 @@ */ #include +#include #include -#include #include #include #include @@ -45,8 +45,8 @@ DEFINE_bool(dump_statistics, false, "Dump inference statistics."); DEFINE_string(output_path, "", "Output Execution results to target directory."); using namespace torch::executor; -using torch::executor::util::FileDataLoader; using namespace torch::executor::enn; +using executorch::backends::enn::ExynosFileDataLoader; std::vector split(std::string str, char delimiter = ' ') { std::vector result; @@ -118,27 +118,11 @@ void saveOutput(const exec_aten::Tensor& tensor, int32_t output_index) { fout.close(); } -struct EnnApiDeinit { - void operator()(EnnApi* ptr) const { - if (ptr == nullptr) { - return; - } - - auto ret = ptr->EnnDeinitialize(); - ET_CHECK_MSG(ret == ENN_RET_SUCCESS, "Enn Deinitialize failed."); - } -}; - -std::unique_ptr exynos_npu_init() { - EnnApi* enn_api_inst = EnnApi::getEnnApiInstance(); - auto ret = enn_api_inst->EnnInitialize(); - ET_CHECK_MSG(ret == ENN_RET_SUCCESS, "Enn initialize failed."); - return std::unique_ptr(enn_api_inst); -} - int main(int argc, char** argv) { auto before_init = std::chrono::high_resolution_clock::now(); - std::unique_ptr instance = exynos_npu_init(); + // The EnnApi singleton initializes the NPU on construction and deinitializes + // it on process teardown. + EnnApi::getEnnApiInstance(); auto after_init = std::chrono::high_resolution_clock::now(); double interval_init = std::chrono::duration_cast( after_init - before_init) @@ -159,10 +143,10 @@ int main(int argc, char** argv) { // DataLoaders that use mmap() or point to data that's already in memory, and // users can create their own DataLoaders to load from arbitrary sources. const char* model_path = FLAGS_model.c_str(); - Result loader = FileDataLoader::from(model_path); + Result loader = ExynosFileDataLoader::from(model_path); ET_CHECK_MSG( loader.ok(), - "FileDataLoader::from() failed: 0x%" PRIx32, + "ExynosFileDataLoader::from() failed: 0x%" PRIx32, (uint32_t)loader.error()); // Parse the program file. This is immutable, and can also be reused between From 135a1099ea5ae90ae0a4ba636a579a0426ca47cf Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 2 Sep 2026 22:30:47 -0700 Subject: [PATCH 012/190] Bump MLX pin to v32.2 (#22495) As titled --- .github/workflows/mlx.yml | 31 +++++ backends/mlx/CMakeLists.txt | 9 +- .../mlx/patches/mlx_gather_mm_rhs_lda.patch | 34 ------ .../mlx_metal_remove_addrspace_compat.patch | 39 ++++++ .../mlx/patches/mlx_nax_jit_sdk_gate.patch | 112 ------------------ .../mlx/patches/mlx_qmm_splitk_bk_align.patch | 42 ------- backends/mlx/runtime/MLXInterpreter.h | 11 +- backends/mlx/third-party/mlx | 2 +- 8 files changed, 85 insertions(+), 195 deletions(-) delete mode 100644 backends/mlx/patches/mlx_gather_mm_rhs_lda.patch create mode 100644 backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch delete mode 100644 backends/mlx/patches/mlx_nax_jit_sdk_gate.patch delete mode 100644 backends/mlx/patches/mlx_qmm_splitk_bk_align.patch diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index d25a2b89c43..13a7f9d35b9 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -71,6 +71,37 @@ jobs: ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test mlx_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 )) echo "::endgroup::" + echo "::group::Check MLX artifact sizes" + size_check_failed=0 + check_artifact_size() { + local path="$1" + local max_size="$2" + local baseline_size="$3" + + if [[ ! -f "${path}" ]]; then + echo "::error file=${path}::Expected MLX artifact not found" + size_check_failed=1 + return + fi + + local actual_size + actual_size=$(wc -c < "${path}" | tr -d '[:space:]') + echo "${path}: ${actual_size} bytes (baseline ${baseline_size}, limit ${max_size})" + if (( actual_size > max_size )); then + echo "::error file=${path}::MLX artifact is ${actual_size} bytes, exceeding the ${max_size}-byte limit (baseline ${baseline_size})" + size_check_failed=1 + fi + } + + # Baselines recorded from the MLX v0.32.2 Release build on 2026-09-02. + check_artifact_size cmake-out/backends/mlx/libmlxdelegate.a 5000000 4159552 + check_artifact_size cmake-out/backends/mlx/mlx/libmlx.a 14000000 11458544 + check_artifact_size cmake-out/backends/mlx/mlx/mlx/backend/metal/kernels/mlx.metallib 2000000 1460840 + echo "::endgroup::" + if (( size_check_failed )); then + exit 1 + fi + echo "::group::Run mutable-state (multi-session) unit test" ./cmake-out/backends/mlx/test/mlx_mutable_state_test echo "::endgroup::" diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 58f36465357..75bfc1f6ed6 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -167,8 +167,9 @@ endif() # with add_subdirectory would drop its whole project into ExecuTorch's # target/option namespace, which collides with shared deps MLX fetches (e.g. # nlohmann_json) and leaks MLX's MLX_BUILD_* options into our cache. The -# isolated scope runs MLX's FetchContent in its own namespace, so no collision -# and no submodule patching are needed. +# isolated scope runs MLX's FetchContent in its own namespace, so dependency +# collisions do not require patching. The platform packaging patches below are +# still applied before the external configure step. include(ExternalProject) set(_mlx_binary_dir ${CMAKE_CURRENT_BINARY_DIR}/mlx) @@ -185,11 +186,9 @@ message( # mismatch after an MLX bump fails loudly rather than silently no-op'ing. See # each patch file under patches/ for its rationale. set(_mlx_patches - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_nax_jit_sdk_gate.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_qmm_splitk_bk_align.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_gather_mm_rhs_lda.patch ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_sdk_per_platform.patch ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_swiftpm_metallib_name.patch + ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_remove_addrspace_compat.patch ) # In a framework build the delegate is static, so MLX cannot find a colocated # metallib and instead loads one from a SwiftPM resource bundle. SWIFTPM_BUNDLE diff --git a/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch b/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch deleted file mode 100644 index 8430d83e9cd..00000000000 --- a/backends/mlx/patches/mlx_gather_mm_rhs_lda.patch +++ /dev/null @@ -1,34 +0,0 @@ -Fix the activation row stride in sorted RHS gather_mm. - -The specialized gather_mm_rhs paths flatten all leading activation dimensions -into M, but derive lda from the original second-to-last dimension. For a valid -[N, 1, K] view produced by expand_dims, that singleton dimension can have -stride 1. The kernel then reads row n from A + n instead of A + n * K. - -The activation is made row-contiguous before this calculation, so flattened -rows are K elements apart. Use K as lda in both the Steel and NAX paths. - -Upstream candidate. - -diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp -index 87d2bf52..4fc2c1de 100644 ---- a/mlx/backend/metal/matmul.cpp -+++ b/mlx/backend/metal/matmul.cpp -@@ -1897,7 +1897,7 @@ void gather_mm_rhs( - int K = a.shape(-1); - int M = a.size() / K; - int N = b.shape(-1); -- int lda = a.strides()[a.ndim() - 2]; // should be K -+ int lda = K; - - // Define the dispatch blocks - int bm = 16, bn = 64, bk = 16; -@@ -2030,7 +2030,7 @@ void gather_mm_rhs_nax( - int K = a.shape(-1); - int M = a.size() / K; - int N = b.shape(-1); -- int lda = a.strides()[a.ndim() - 2]; // should be K -+ int lda = K; - int E = b.shape(0); - - // Define the dispatch blocks diff --git a/backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch b/backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch new file mode 100644 index 00000000000..6e1ad402c65 --- /dev/null +++ b/backends/mlx/patches/mlx_metal_remove_addrspace_compat.patch @@ -0,0 +1,39 @@ +Keep Steel integral constants compatible with older Metal frontends. + +MLX supports macOS 14 and selects Metal language 3.1 there. The Metal 4.1 +address-space fix uses metal::remove_addrspace_t, but that trait is unavailable +in the Xcode 15.4 Metal standard library. Because integral_constant.h is embedded +in runtime-JIT sources, this prevents kernels from compiling on that supported +toolchain. + +Use remove_addrspace_t on Metal 4.1+, where explicit thread qualifiers make it +necessary and the trait is available. For older language versions, preserve the +previous decltype behavior with an identity alias. + +Upstream candidate; carried locally until MLX includes the compatibility fix. + +diff --git a/mlx/backend/metal/kernels/steel/utils/integral_constant.h b/mlx/backend/metal/kernels/steel/utils/integral_constant.h +index 2f153f48..97022416 100644 +--- a/mlx/backend/metal/kernels/steel/utils/integral_constant.h ++++ b/mlx/backend/metal/kernels/steel/utils/integral_constant.h +@@ -47,11 +47,19 @@ using Int = integral_constant; + // Binary Operators on Integral constants + /////////////////////////////////////////////////////////////////////////////// + ++#if __METAL_VERSION__ >= 410 ++template ++using remove_addrspace_t = metal::remove_addrspace_t; ++#else ++template ++using remove_addrspace_t = T; ++#endif ++ + #define integral_const_binop(__op__, __operator__) \ + template \ + METAL_FUNC constexpr auto __operator__( \ + integral_constant, integral_constant) { \ + constexpr auto res = tv __op__ uv; \ +- using res_t = metal::remove_addrspace_t; \ ++ using res_t = remove_addrspace_t; \ + return integral_constant{}; \ + } diff --git a/backends/mlx/patches/mlx_nax_jit_sdk_gate.patch b/backends/mlx/patches/mlx_nax_jit_sdk_gate.patch deleted file mode 100644 index b1b1fcfd820..00000000000 --- a/backends/mlx/patches/mlx_nax_jit_sdk_gate.patch +++ /dev/null @@ -1,112 +0,0 @@ -Gate the NAX JIT kernel sources behind the SDK requirement. - -MLX's NAX kernels (GEMM and attention) include -, a framework that only -ships in the macOS 26 / Xcode 26 SDK. With MLX_METAL_JIT=ON (which ExecuTorch -uses) on an older SDK, the JIT preamble generator (make_compiled_preamble.sh) -runs `metal -E` over these headers, which fatals on the missing include and -fails the build. MLX already gates NAX on the non-JIT metallib path -(kernels/CMakeLists.txt), but the JIT path was ungated. - -Instead of guarding the includes with __has_include, gate the NAX -make_jit_source() calls behind the same -MLX_METAL_VERSION/MACOS_SDK_VERSION/CMAKE_OSX_DEPLOYMENT_TARGET check the -metallib path uses, and define MLX_METAL_NO_NAX when the requirement is unmet. -On those SDKs NAX is already runtime-gated via is_nax_available() -(device.cpp), so the get_*_nax_kernel entry points in jit_kernels.cpp are -unreachable; guarded empty preamble getters keep that translation unit linking. -macOS 26+ SDKs still build NAX. - -Upstream candidate; carried locally until an MLX release gates the JIT path. - -diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt ---- a/mlx/backend/metal/CMakeLists.txt -+++ b/mlx/backend/metal/CMakeLists.txt -@@ -84,19 +84,32 @@ if(MLX_METAL_JIT) - - make_jit_source(steel/attn/kernels/steel_attention) - -- make_jit_source( -- steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h -- kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) -- make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) -- make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) -- make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) -- make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) -+ if(MLX_METAL_VERSION GREATER_EQUAL 400 -+ AND MACOS_SDK_VERSION VERSION_GREATER_EQUAL 26.2 -+ AND CMAKE_OSX_DEPLOYMENT_TARGET VERSION_GREATER_EQUAL 26.2) - -- make_jit_source(quantized_nax kernels/quantized_utils.h) -- make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h -- kernels/fp4.h) -+ make_jit_source( -+ steel/gemm/gemm_nax kernels/steel/utils.h kernels/steel/gemm/nax.h -+ kernels/steel/gemm/params.h kernels/steel/gemm/transforms.h) -+ make_jit_source(steel/gemm/kernels/steel_gemm_fused_nax) -+ make_jit_source(steel/gemm/kernels/steel_gemm_gather_nax) -+ make_jit_source(steel/gemm/kernels/steel_gemm_splitk_nax) -+ make_jit_source(steel/gemm/kernels/steel_gemm_segmented_nax) -+ -+ make_jit_source(quantized_nax kernels/quantized_utils.h) -+ make_jit_source(fp_quantized_nax kernels/quantized_utils.h kernels/fp8.h -+ kernels/fp4.h) -+ -+ make_jit_source(steel/attn/kernels/steel_attention_nax) - -- make_jit_source(steel/attn/kernels/steel_attention_nax) -+ else() -+ message( -+ WARNING "NAX kernels require Metal 4, macOS SDK >= 26.2, and " -+ "MACOSX_DEPLOYMENT_TARGET >= 26.2 (SDK ${MACOS_SDK_VERSION}, " -+ "CMAKE_OSX_DEPLOYMENT_TARGET=${CMAKE_OSX_DEPLOYMENT_TARGET}). " -+ "Building without NAX kernels.") -+ target_compile_definitions(mlx PRIVATE MLX_METAL_NO_NAX) -+ endif() - - else() - target_sources(mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/nojit_kernels.cpp) -diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp ---- a/mlx/backend/metal/jit_kernels.cpp -+++ b/mlx/backend/metal/jit_kernels.cpp -@@ -8,6 +8,40 @@ using namespace fmt::literals; - - namespace mlx::core { - -+#ifdef MLX_METAL_NO_NAX -+// NAX JIT preambles are only generated (via make_jit_source) when the SDK -+// requirement is met. On older SDKs they are skipped and MLX_METAL_NO_NAX is -+// defined, so is_nax_available() returns false and the get_*_nax_kernel entry -+// points below are never reached. These empty definitions only exist to satisfy -+// the linker for this translation unit. -+namespace metal { -+const char* gemm_nax() { -+ return ""; -+} -+const char* steel_gemm_fused_nax() { -+ return ""; -+} -+const char* steel_gemm_gather_nax() { -+ return ""; -+} -+const char* steel_gemm_splitk_nax() { -+ return ""; -+} -+const char* steel_gemm_segmented_nax() { -+ return ""; -+} -+const char* quantized_nax() { -+ return ""; -+} -+const char* fp_quantized_nax() { -+ return ""; -+} -+const char* steel_attention_nax() { -+ return ""; -+} -+} // namespace metal -+#endif // MLX_METAL_NO_NAX -+ - MTL::ComputePipelineState* get_arange_kernel( - metal::Device& d, - const std::string& kernel_name, diff --git a/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch b/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch deleted file mode 100644 index f4405b5ae9f..00000000000 --- a/backends/mlx/patches/mlx_qmm_splitk_bk_align.patch +++ /dev/null @@ -1,42 +0,0 @@ -Align split-K partitions to the qmm K-tile (BK=32), fixing nvfp4. - -MLX v0.32.0's qmm_splitk caps split_k only by the quantization group count -(K / group_size), not by the kernel's K-tile width BK (=32). For nvfp4 -(group_size=16) this yields a per-partition K of 16 < BK, and fp_qmm_t_splitk's -tile load reads a full BK-wide K-tile with no K bound -- spilling 16 columns -past the partition into the next group's packed weights and fp8 scales. That -corrupts every partial (non-uniform ~2x error) and reads past the buffer on the -last partition (NaN/inf). Affine (group_size >= 32) is unaffected because its -partitions are already >= BK. - -Fix in the dispatch (MLX's pattern for tile-alignment constraints): require each -K partition to be a whole number of BK-wide tiles as well as whole groups, i.e. -align split_k to max(group_size, BK). Only changes behavior for group_size < 32. - -Upstream candidate. - -diff --git a/mlx/backend/metal/quantized.cpp b/mlx/backend/metal/quantized.cpp -index 62d48714..94c56307 100644 ---- a/mlx/backend/metal/quantized.cpp -+++ b/mlx/backend/metal/quantized.cpp -@@ -884,11 +884,15 @@ void qmm_splitk( - int current_tgs = n_tiles * m_tiles; - int split_k = std::max(1, 512 / current_tgs); - -- // Cap split_k by the number of quantization groups -- split_k = std::min(split_k, K / group_size); -- -- // Ensure K divides evenly by split_k * group_size -- while (split_k > 1 && (K % (split_k * group_size) != 0)) { -+ // Each K partition must be a whole number of BK-wide (32) K-tiles as well as -+ // whole quantization groups. The qmm_t_splitk kernels tile K by BK=32 and do -+ // not bound the K dimension, so a partition smaller than BK (e.g. nvfp4's -+ // group_size=16) would over-read into the next group's weights/scales. -+ int k_align = group_size > 32 ? group_size : 32; -+ split_k = std::min(split_k, K / k_align); -+ -+ // Ensure K divides evenly by split_k * k_align -+ while (split_k > 1 && (K % (split_k * k_align) != 0)) { - split_k--; - } - if (split_k <= 1) { diff --git a/backends/mlx/runtime/MLXInterpreter.h b/backends/mlx/runtime/MLXInterpreter.h index d9b77e8f116..001c3fd4bf2 100644 --- a/backends/mlx/runtime/MLXInterpreter.h +++ b/backends/mlx/runtime/MLXInterpreter.h @@ -293,7 +293,15 @@ inline void exec_sdpa(const SdpaNode& n, ExecutionState& st, StreamOrDevice s) { } array out = fast::scaled_dot_product_attention( - Q, K, V, static_cast(n.scale), mask_mode, mask_arr, sinks, s); + Q, + K, + V, + static_cast(n.scale), + mask_mode, + mask_arr, + sinks, + false, + s); st.set_tensor(n.out, std::move(out)); } @@ -382,6 +390,7 @@ inline void exec_update_and_attend( mask_mode, spec.mask, std::nullopt, + false, s); // Honor the op's output-dtype contract (unset -> SDPA's native output). if (n.out_dtype) { diff --git a/backends/mlx/third-party/mlx b/backends/mlx/third-party/mlx index 7a1d4f5c12a..1f8e74e3f12 160000 --- a/backends/mlx/third-party/mlx +++ b/backends/mlx/third-party/mlx @@ -1 +1 @@ -Subproject commit 7a1d4f5c12ac82f4b4d0a6e71538d89ca0605247 +Subproject commit 1f8e74e3f12f31365464a6867c6579f0e9b29d85 From 8f50849cc00ed5409d703880ad53f8bfe664552e Mon Sep 17 00:00:00 2001 From: zhaoxul-qti Date: Thu, 3 Sep 2026 15:43:04 +0800 Subject: [PATCH 013/190] Qualcomm AI Engine Direct - [doc] QNN ExecuTorch on Windows (#22502) ### Summary - Add documentation for QNN ExecuTorch on Windows, including build, test, and inference instructions. - Add support for the `SC8380XP` SoC model. Refer to the [Supported Snapdragon devices](https://docs.qualcomm.com/doc/80-63442-10/topic/QNN_general_overview.html#supported-snapdragon-devices) for details. ### Test plan ```powershell python -m examples.qualcomm.scripts.deeplab_v3 --build_folder build-x86_64-windows --soc_model SC8380XP --compile_only --download ``` --- .../serialization/qc_compiler_spec.fbs | 1 + backends/qualcomm/serialization/qc_schema.py | 2 + backends/qualcomm/utils/utils.py | 2 + docs/source/backends-qualcomm.md | 236 +++++++++++++----- 4 files changed, 184 insertions(+), 57 deletions(-) diff --git a/backends/qualcomm/serialization/qc_compiler_spec.fbs b/backends/qualcomm/serialization/qc_compiler_spec.fbs index 57708c959e9..4d3d59c4fbe 100644 --- a/backends/qualcomm/serialization/qc_compiler_spec.fbs +++ b/backends/qualcomm/serialization/qc_compiler_spec.fbs @@ -47,6 +47,7 @@ enum QcomChipset: int { UNKNOWN_SM = 0, SA8295 = 39, SA8797 = 72, + SC8380XP = 60, SM8350 = 30, SM8450 = 36, SM8475 = 42, diff --git a/backends/qualcomm/serialization/qc_schema.py b/backends/qualcomm/serialization/qc_schema.py index aeffbc069b6..b3992bc97f1 100644 --- a/backends/qualcomm/serialization/qc_schema.py +++ b/backends/qualcomm/serialization/qc_schema.py @@ -54,6 +54,7 @@ class QcomChipset(IntEnum): UNKNOWN_SM = 0 SA8295 = 39 # v68 SA8797 = 72 # v81 + SC8380XP = 60 # v73 SM8350 = 30 # v68 SM8450 = 36 # v69 SM8475 = 42 # v69 @@ -84,6 +85,7 @@ class SocInfo: _soc_info_table = { QcomChipset.SA8295: SocInfo(QcomChipset.SA8295, HtpInfo(HtpArch.V68, 8)), QcomChipset.SA8797: SocInfo(QcomChipset.SA8797, HtpInfo(HtpArch.V81, 16)), + QcomChipset.SC8380XP: SocInfo(QcomChipset.SC8380XP, HtpInfo(HtpArch.V73, 8)), QcomChipset.SM8350: SocInfo(QcomChipset.SM8350, HtpInfo(HtpArch.V68, 4)), QcomChipset.SM8450: SocInfo(QcomChipset.SM8450, HtpInfo(HtpArch.V69, 8)), QcomChipset.SM8475: SocInfo(QcomChipset.SM8475, HtpInfo(HtpArch.V69, 8)), diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 4a1715e8c03..75ef9c1f128 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -1322,6 +1322,7 @@ def get_soc_to_htp_arch_map(): return { "SA8295": HtpArch.V68, "SA8797": HtpArch.V81, + "SC8380XP": HtpArch.V73, "SM8350": HtpArch.V68, "SM8450": HtpArch.V69, "SM8475": HtpArch.V69, @@ -1355,6 +1356,7 @@ def get_soc_to_chipset_map(): return { "SA8295": QcomChipset.SA8295, "SA8797": QcomChipset.SA8797, + "SC8380XP": QcomChipset.SC8380XP, "SM8350": QcomChipset.SM8350, "SM8450": QcomChipset.SM8450, "SM8475": QcomChipset.SM8475, diff --git a/docs/source/backends-qualcomm.md b/docs/source/backends-qualcomm.md index caf2426ed16..cf9369cb214 100644 --- a/docs/source/backends-qualcomm.md +++ b/docs/source/backends-qualcomm.md @@ -36,15 +36,24 @@ Currently, this ExecuTorch Backend can delegate AI computations to Hexagon proce ### Host OS -The QNN Backend is currently verified on the following Linux host operating systems: +The QNN Backend is verified on the following host operating systems: - **Ubuntu 22.04 LTS (x64)** - **CentOS Stream 9** +- **Windows 10 / 11 (x64)** +- **Windows 10 / 11 (ARM64)** with Qualcomm NPU - **Windows Subsystem for Linux (WSL)** with Ubuntu 22.04 In general, we verify the backend on the same OS versions that the QNN SDK is officially validated against. The exact supported versions are documented in the QNN SDK. +#### Windows (x64 / ARM64) Setup + +To build on native Windows platforms, the MSVC toolchain must be installed. +The required MSVC Build Tools can be installed through **Visual Studio Installer**. + +For installation instructions, refer to the official [Microsoft Visual Studio Downloads page](https://visualstudio.microsoft.com/downloads/). + #### Windows (WSL) Setup To install Ubuntu 22.04 on WSL, run the following command in PowerShell or Windows Terminal: @@ -55,24 +64,23 @@ wsl --install -d ubuntu 22.04 This command will install WSL and set up Ubuntu 22.04 as the default Linux distribution. -For more details and troubleshooting, refer to the official Microsoft WSL installation guide: -👉 [Install WSL | Microsoft Learn](https://learn.microsoft.com/en-us/windows/wsl/install) +For more details and troubleshooting, refer to the official Microsoft WSL installation guide: [Install WSL | Microsoft Learn](https://learn.microsoft.com/en-us/windows/wsl/install). ### Hardware: -You will need an Android / Linux device with adb-connected running on one of Qualcomm SoCs listed in `QcomChipset`. Please navigate to [qc_schema.py](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/serialization/qc_schema.py). -This example is verified with SM8550 and SM8450. +The QNN backend runs on Qualcomm SoCs (Systems on Chips) across two device families: + +- **Android / Linux devices** — connected over `adb`. This example is verified with SM8550 and SM8450. +- **Windows on ARM64 (WoA) devices** — This example is verified with SC8380XP (Qualcomm Snapdragon X Elite). + +The target SoC must be one of those listed in the `QcomChipset` enum; see [qc_schema.py](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/serialization/qc_schema.py). ### Software: - - Follow ExecuTorch recommended Python version. - - A compiler to compile AOT parts, e.g., the GCC compiler comes with Ubuntu LTS. g++ version need to be 13 or higher. - - [Android NDK](https://developer.android.com/ndk). This example is verified with NDK 26c. - - (Optional) Target toolchain for linux embedded platform. - - [Qualcomm AI Engine Direct SDK](https://developer.qualcomm.com/software/qualcomm-ai-engine-direct-sdk) - - Click the "Get Software" button to download the latest version of the QNN SDK. - - Although newer versions are available, we have verified and recommend using QNN 2.37.0 for stability. - - You can download it directly from the following link: [QNN 2.37.0](https://softwarecenter.qualcomm.com/api/download/software/sdks/Qualcomm_AI_Runtime_Community/All/2.37.0.250724/v2.37.0.250724.zip) +[Qualcomm AI Engine Direct SDK](https://developer.qualcomm.com/software/qualcomm-ai-engine-direct-sdk) + - Click the "Get Software" button to download the latest version of the QNN SDK. + - Although newer versions are available, we have verified and recommend using QNN 2.37.0 for stability. + - You can download it directly from the following link: [QNN 2.37.0](https://softwarecenter.qualcomm.com/api/download/software/sdks/Qualcomm_AI_Runtime_Community/All/2.37.0.250724/v2.37.0.250724.zip) The directory with installed Qualcomm AI Engine Direct SDK looks like: ``` @@ -94,6 +102,18 @@ The directory with installed Qualcomm AI Engine Direct SDK looks like: └── share ``` +On Android / Linux devices: + + - Follow ExecuTorch recommended Python version. + - A compiler to compile AOT parts, e.g., the GCC compiler comes with Ubuntu LTS. g++ version need to be 13 or higher. + - [Android NDK](https://developer.android.com/ndk). This example is verified with NDK 26c. + - (Optional) Target toolchain for linux embedded platform. + +On Windows on ARM64 (WoA) devices: + + - Install the **AMD64 version of Python** to run AOT compilation under x64 emulation. This is required because certain Python modules used in the AOT workflow do not currently provide ARM64 prebuilt wheels. + - MSVC Build Tools. + ## Setting up your developer environment @@ -106,9 +126,9 @@ i.e., the directory containing `QNN_README.txt`. `$EXECUTORCH_ROOT` refers to the root of executorch git repository. -### Setup environment variables +### Setup QNN SDK paths and environment variables -Source the QNN SDK environment setup script to configure paths and environment variables: +For Linux platform: ```bash source $QNN_SDK_ROOT/bin/envsetup.sh @@ -116,35 +136,66 @@ source $QNN_SDK_ROOT/bin/envsetup.sh This sets up `LD_LIBRARY_PATH` and other required variables for the QNN SDK tools and libraries. -Additionally, set `PYTHONPATH` for ExecuTorch Python APIs: +For Windows platform: + +```powershell +& "$env:QNN_SDK_ROOT\bin\envsetup.ps1" +``` + +### Setup `PYTHONPATH` for ExecuTorch Python APIs + +For Linux platform: ```bash export PYTHONPATH=$EXECUTORCH_ROOT/..:$PYTHONPATH ``` +For Windows platform: + +```powershell +$env:PYTHONPATH="$env:EXECUTORCH_ROOT\..;$env:PYTHONPATH" +``` + ## Build -An example script for the below building instructions is [here](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/scripts/build.sh). +**On Linux platform**, an example script for the below building instructions is [`build.sh`](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/scripts/build.sh). We recommend to use the script because the ExecuTorch build-command can change from time to time. The above script is actively used. It is updated more frequently than this tutorial. An example usage is ```bash cd $EXECUTORCH_ROOT -# android target +# Android target ./backends/qualcomm/scripts/build.sh -# (optional) linux embedded target +# (Optional) Linux embedded target ./backends/qualcomm/scripts/build.sh --enable_linux_embedded -# for release build +# Android target for release build ./backends/qualcomm/scripts/build.sh --release ``` +**On Windows platform**, use the PowerShell script [`build.ps1`](https://github.com/pytorch/executorch/blob/main/backends/qualcomm/scripts/build.ps1) for the building instructions. Both Windows x64 and ARM64 architectures are supported. +Here's the example usage +```powershell +cd $env:EXECUTORCH_ROOT +# Generate both Windows x64 and ARM64 target libraries +.\backends\qualcomm\scripts\build.ps1 -Release +# Generate only Windows x64 target libraries +.\backends\qualcomm\scripts\build.ps1 -SkipArm64Windows -Release +# Generate only Windows ARM64 target libraries +.\backends\qualcomm\scripts\build.ps1 -SkipX86Windows -Release +``` + +> **Notes** +> +> The script supports building both x64 and cross-compiling ARM64 target artifacts on Windows x64 host. After the build completes, the ARM64 libraries and executables can be copied to a Windows on Snapdragon (WoS) device using `scp`. +> This allows a `.pte` generated on Windows x64 host to be executed on WoS device. ## Deploying and running on device ### AOT compile a model Refer to [this script](https://github.com/pytorch/executorch/blob/main/examples/qualcomm/scripts/deeplab_v3.py) for the exact flow. -We use deeplab-v3-resnet101 as an example in this tutorial. Run below commands to compile: +We use deeplab-v3-resnet101 as an example in this tutorial. +Run below commands to compile on Linux platform: ```bash cd $EXECUTORCH_ROOT @@ -152,14 +203,26 @@ cd $EXECUTORCH_ROOT python -m examples.qualcomm.scripts.deeplab_v3 --build_folder build-android --soc_model SM8550 --compile_only --download ``` -You might see something like below: +For Windows x64 and ARM64 platforms, run the following commands for the AOT compilation: +```powershell +cd $env:EXECUTORCH_ROOT +python -m examples.qualcomm.scripts.deeplab_v3 --build_folder build-x86_64-windows --soc_model SC8380XP --compile_only --download ``` -[INFO][Qnn ExecuTorch] Destroy Qnn context -[INFO][Qnn ExecuTorch] Destroy Qnn device -[INFO][Qnn ExecuTorch] Destroy Qnn backend -Finish compile_only and save to ./deeplab_v3/dlv3_qnn.pte +> **Notes** +> +> AOT compilation on Windows on ARM64 (WoA) device currently relies on an AMD64 Python environment running under x64 emulation, since some AOT dependencies are not yet distributed as ARM64 prebuilt wheels. + +You might see something like below: + +``` +Completed stage: Finalizing Graph Sequence (8966 us) +Starting stage: Completion +Completed stage: Completion (1388 us) +[INFO] [Qnn ExecuTorch]: Destroy Qnn context +[INFO] [Qnn ExecuTorch]: Destroy Qnn device +[INFO] [Qnn ExecuTorch]: Destroy Qnn backend ``` The compiled model is `./deeplab_v3/dlv3_qnn.pte`. @@ -167,26 +230,9 @@ The compiled model is `./deeplab_v3/dlv3_qnn.pte`. Note that the model is compiled for specific backend (e.g., HTP), so you can specify the target backend via `--backend gpu` or `--backend lpai`. If not specified, it will be default to HTP. -### Test model inference on QNN HTP emulator / QNN LPAI emulator - -We can test model inferences before deploying it to a device by HTP emulator. - -Let's build `qnn_executor_runner` for a x64 host: -```bash -# assuming the AOT component is built. -cd $EXECUTORCH_ROOT/build-x86 -cmake ../examples/qualcomm \ - -DCMAKE_PREFIX_PATH="$PWD/lib/cmake/ExecuTorch;$PWD/third-party/gflags;" \ - -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ - -DPYTHON_EXECUTABLE=python3 \ - -Bexamples/qualcomm - -cmake --build examples/qualcomm -j$(nproc) +### Test model inference on Linux x64 host with QNN HTP emulator / QNN LPAI emulator -# qnn_executor_runner can be found under examples/qualcomm/executor_runner -# The full path is $EXECUTORCH_ROOT/build-x86/examples/qualcomm/executor_runner/qnn_executor_runner -ls examples/qualcomm/executor_runner -``` +Before deploying a model to a physical device, inference execution can be tested and validated on a Linux x64 host using the HTP / LPAI emulator. To run the HTP emulator / LPAI emulator, the dynamic linker needs to access QNN libraries and `libqnn_executorch_backend.so`. We set the below two paths to `LD_LIBRARY_PATH` environment variable: @@ -199,27 +245,55 @@ The second path is for `libqnn_executorch_backend.so`. So, we can run `./deeplab_v3/dlv3_qnn.pte` by: ```bash -cd $EXECUTORCH_ROOT/build-x86 +cd $EXECUTORCH_ROOT export LD_LIBRARY_PATH=$EXECUTORCH_ROOT/build-x86/lib/:$LD_LIBRARY_PATH -examples/qualcomm/executor_runner/qnn_executor_runner --model_path ../deeplab_v3/dlv3_qnn.pte +build-x86/examples/qualcomm/executor_runner/qnn_executor_runner --model_path ./deeplab_v3/dlv3_qnn.pte ``` We should see some outputs like the below. Note that the emulator can take some time to finish. ```bash -I 00:00:00.354662 executorch:qnn_executor_runner.cpp:213] Method loaded. -I 00:00:00.356460 executorch:qnn_executor_runner.cpp:261] ignoring error from set_output_data_ptr(): 0x2 -I 00:00:00.357991 executorch:qnn_executor_runner.cpp:261] ignoring error from set_output_data_ptr(): 0x2 -I 00:00:00.357996 executorch:qnn_executor_runner.cpp:265] Inputs prepared. - -I 00:01:09.328144 executorch:qnn_executor_runner.cpp:414] Model executed successfully. -I 00:01:09.328159 executorch:qnn_executor_runner.cpp:421] Write etdump to etdump.etdp, Size = 424 -[INFO] [Qnn ExecuTorch]: Destroy Qnn backend parameters +I 00:00:00.174364 executorch:qnn_executor_runner.cpp:416] Method loaded. +E 00:00:00.179250 executorch:method.cpp:1373] Output 0 is memory planned, or is a constant. Cannot override the existing data pointer. +I 00:00:00.179264 executorch:qnn_executor_runner.cpp:473] ignoring error from set_output_data_ptr(): 0x2 +E 00:00:00.183296 executorch:method.cpp:1373] Output 1 is memory planned, or is a constant. Cannot override the existing data pointer. +I 00:00:00.183305 executorch:qnn_executor_runner.cpp:473] ignoring error from set_output_data_ptr(): 0x2 +I 00:00:00.183310 executorch:qnn_executor_runner.cpp:479] Inputs prepared. +I 00:00:00.184008 executorch:qnn_executor_runner.cpp:684] Input list not provided. Inputs prepared with default values set. +I 00:01:19.663283 executorch:qnn_executor_runner.cpp:695] Model executed successfully. +I 00:01:19.663299 executorch:qnn_executor_runner.cpp:698] Perform 0 inferences for warming up +I 00:01:53.881349 executorch:qnn_executor_runner.cpp:715] 1 inferences took 34218.046000 ms, avg 34218.046000 ms +I 00:01:53.881426 executorch:qnn_executor_runner.cpp:727] Write etdump to etdump.etdp, Size = 576 [INFO] [Qnn ExecuTorch]: Destroy Qnn context [INFO] [Qnn ExecuTorch]: Destroy Qnn device [INFO] [Qnn ExecuTorch]: Destroy Qnn backend ``` -### Run model inference on an Android smartphone with Qualcomm SoCs +### Test model inference on Windows x64 host with QNN HTP emulator / QNN LPAI emulator + +Unlike Linux, which set `LD_LIBRARY_PATH` to access shared libraries, Windows uses the `$env:PATH` environment variable. To enable runtime loading of `qnn_executorch_backend.dll`, ensure that it is discoverable by the Windows DLL loader. + +This can be achieved by either: +- Placing `qnn_executorch_backend.dll` in the same directory as `qnn_executor_runner.exe`; or +- Adding the directory containing `qnn_executorch_backend.dll` to `$env:PATH` environment variable. + +The generated artifacts can be found at: +- `$env:EXECUTORCH_ROOT\build-x86_64-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe` +- `$env:EXECUTORCH_ROOT\build-x86_64-windows\backends\qualcomm\Release\qnn_executorch_backend.dll` + +To add the directory containing `qnn_executorch_backend.dll` to the `$env:PATH` environment variable: +```powershell +$env:PATH="$env:EXECUTORCH_ROOT\build-x86_64-windows\backends\qualcomm\Release;$env:PATH" +``` + +Once configured, `qnn_executorch_backend.dll` will be accessed by `qnn_executor_runner.exe` at runtime. + +To test the model inference on Windows x64 host with QNN HTP emulator / QNN LPAI emulator: +```powershell +cd $env:EXECUTORCH_ROOT\build-x86_64-windows\examples\qualcomm\executor_runner\Release +.\qnn_executor_runner.exe --model_path $env:EXECUTORCH_ROOT\deeplab_v3\dlv3_qnn.pte +``` + +### Run model inference on Android smartphone with Qualcomm SoCs ***Step 1***. We need to push required QNN libraries to the device. @@ -251,7 +325,7 @@ adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnGpu.so ${DEVICE_DIR} adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnLpai.so ${DEVICE_DIR} adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnLpaiStub.so ${DEVICE_DIR} adb push ${QNN_SDK_ROOT}/lib/aarch64-android/libQnnSystem.so ${DEVICE_DIR} -# make sure the skel lib is signed for LPAI backend. +# Make sure the skel lib is signed for LPAI backend. adb push ${QNN_SDK_ROOT}/lib/lpai-v6/signed/libQnnLpaiSkel.so ${DEVICE_DIR} ``` @@ -298,6 +372,54 @@ After the above command, pre-processed inputs and outputs are put in `$EXECUTORC The command-line arguments are written in [utils.py](https://github.com/pytorch/executorch/blob/main/examples/qualcomm/utils.py#L139). The model, inputs, and output location are passed to `qnn_executorch_runner` by `--model_path`, `--input_list_path`, and `--output_folder_path`. +### Run model inference on Windows on Snapdragon (WoS) with Qualcomm SoCs + +Before running inference on Windows on Snapdragon (WoS) with Qualcomm SoCs, ensure that `qnn_executorch_backend.dll` and all required QNN libraries are discoverable by the Windows loader. This can be achieved by either: +- Copying `qnn_executorch_backend.dll` and the required QNN libraries into the same directory as `qnn_executor_runner.exe`; or +- Adding the directories containing these libraries to the `$env:PATH` environment variable. + +The generated artifacts can be found at: +- `$env:EXECUTORCH_ROOT\build-arm64-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe` +- `$env:EXECUTORCH_ROOT\build-arm64-windows\backends\qualcomm\Release\qnn_executorch_backend.dll` + +Depending on the selected QNN backend, the corresponding QNN libraries can be found under: + +```powershell +# For HTP +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtp.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnSystem.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV69Stub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV73Stub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV75Stub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnHtpV79Stub.dll +$env:QNN_SDK_ROOT\lib\hexagon-v69\unsigned\libQnnHtpV69Skel.so +$env:QNN_SDK_ROOT\lib\hexagon-v73\unsigned\libQnnHtpV73Skel.so +$env:QNN_SDK_ROOT\lib\hexagon-v75\unsigned\libQnnHtpV75Skel.so +$env:QNN_SDK_ROOT\lib\hexagon-v79\unsigned\libQnnHtpV79Skel.so +``` + +```powershell +# For GPU +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnGpu.dll +``` + +```powershell +# For LPAI +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnLpai.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnLpaiStub.dll +$env:QNN_SDK_ROOT\lib\aarch64-windows-msvc\QnnSystem.dll +# Make sure the skel lib is signed for LPAI backend. +$env:QNN_SDK_ROOT\lib\lpai-v6\signed\libQnnLpaiSkel.so +``` + +Once configured, `qnn_executorch_backend.dll` and the required QNN libraries can be accessed by `qnn_executor_runner.exe` at runtime. + +To test the model inference on Windows on Snapdragon (WoS) with Qualcomm SoCs: +```powershell +cd $env:EXECUTORCH_ROOT +.\qnn_executor_runner.exe --model_path .\deeplab_v3\dlv3_qnn.pte +``` + ### Run [Android LlamaDemo](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android/LlamaDemo) with QNN backend `$DEMO_APP` refers to the root of the executorch android demo, i.e., the directory containing `build.gradle.kts`. From ddaabb9615440d33b089c5bf415e88be879742be Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Thu, 3 Sep 2026 10:23:01 +0200 Subject: [PATCH 014/190] Add QAT support and pass hooks to QuantizationRecipe. (#21935) ### Summary Extend `QuantizationRecipe` with QAT support and four ordered pass hooks, then wire them into `QuantizeStage`. `QuantizationRecipe` gains `is_qat`, `train_fn`, `calibration_inputs_fn`, and four pass-list fields (`pre_prepare_passes`, `post_prepare_passes`, `pre_convert_passes`, `post_convert_passes`). When `is_qat=True`, `QuantizeStage` calls `prepare_qat_pt2e` instead of `prepare_pt2e`, invokes `train_fn` on the prepared model, and skips calibration. When `is_qat=False` (default), PTQ proceeds as before, but callers may now supply a `calibration_inputs_fn` to supply their own calibration data. If omitted, the existing example-inputs fallback is used. The four pass hooks are applied at the appropriate points in both flows and compose cleanly with `_combine_recipes`. `_combine_recipes` is overhauled to handle all previously dropped fields: the four pass lists are concatenated, `is_qat` must agree across all combined recipes, at most one `train_fn` is allowed, and multiple `calibration_inputs_fn` values are chained into a single factory. The `strict`, `mode`, `pipeline_stages`, and `source_transform_in_place` scalar fields are now also validated for agreement and propagated to the combined recipe. `edge_manager_transform_passes` from `LoweringRecipe` is similarly merged. ### Test Plan ``` pytest -q export/tests/test_export_recipe.py export/tests/test_export_stages.py ``` --- export/recipe.py | 399 +++++++++++++++----- export/stages.py | 102 ++++- export/tests/test_export_recipe.py | 453 ++++++++++++++++++++++- export/tests/test_export_stages.py | 573 ++++++++++++++++++++++++++++- 4 files changed, 1423 insertions(+), 104 deletions(-) diff --git a/export/recipe.py b/export/recipe.py index 1609b989273..8ea2256b7c3 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -1,15 +1,18 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import copy import dataclasses +import logging from abc import ABCMeta, abstractmethod -from dataclasses import dataclass +from dataclasses import dataclass, field from enum import Enum, EnumMeta -from typing import Callable, Dict, List, Optional, Union +from typing import Callable, Dict, Iterable, List, Optional, Union import torch from executorch.exir import EdgeProgramManager, ExportedProgram @@ -129,16 +132,54 @@ class QuantizationRecipe: """ Configuration recipe for quantization. - This class holds the configuration parameters for quantizing a model. + This class holds the configuration parameters for quantizing a model, supporting + both post-training quantization (PTQ) and quantization-aware training (QAT). Attributes: - quantizers: Optional list of quantizers for model quantization + quantizers: Optional list of quantizers for model quantization. ao_quantization_configs: Optional list of AOQuantizationConfig objects that pair - AOBaseConfig with optional filter functions + AOBaseConfig with optional filter functions. + is_qat: If True, use the QAT flow (prepare_qat_pt2e -> train_fn -> convert_pt2e). + If False (default), use the PTQ flow (prepare_pt2e -> calibrate -> convert_pt2e). + dynamic_batch_size: If True, dimension 0 (batch) of the calibration/QAT + training inputs may vary. Otherwise it is fixed to the + example inputs' batch size. + calibration_inputs_fn: Optional callable returning an iterable of input tuples used for + PTQ calibration. When None (default), the example inputs are used. + Ignored when is_qat=True. + train_fn: Callable that receives the prepared GraphModule and trains it. + Required when is_qat=True; ignored otherwise. + pre_prepare_passes: Optional list of callables applied to the captured GraphModule + before prepare_pt2e / prepare_qat_pt2e. + Each callable receives a GraphModule and must return a GraphModule. + post_prepare_passes: Optional list of callables applied to the prepared GraphModule + after prepare_pt2e / prepare_qat_pt2e and before calibration / training. + Each callable receives a GraphModule and must return a GraphModule. + pre_convert_passes: Optional list of callables applied to the GraphModule after + calibration (PTQ) or training (QAT) and before convert_pt2e. + Each callable receives a GraphModule and must return a GraphModule. + post_convert_passes: Optional list of callables applied to the GraphModule after convert_pt2e. + Each callable receives a GraphModule and must return a GraphModule. """ quantizers: Optional[List[Quantizer]] = None ao_quantization_configs: Optional[List[AOQuantizationConfig]] = None + is_qat: bool = False + dynamic_batch_size: bool = False + calibration_inputs_fn: Optional[Callable[[], Iterable[tuple]]] = None + train_fn: Optional[Callable[["torch.fx.GraphModule"], None]] = None + pre_prepare_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None + post_prepare_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None + pre_convert_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None + post_convert_passes: Optional[ + List[Callable[["torch.fx.GraphModule"], "torch.fx.GraphModule"]] + ] = None def get_quantizers(self) -> Optional[List[Quantizer]]: """ @@ -184,6 +225,32 @@ class LoweringRecipe: edge_compile_config: Optional[EdgeCompileConfig] = None +@dataclass +class _CombineAccumulator: + """Private accumulator used by ExportRecipe._collect_recipe_fields.""" + + partitioners: list = field(default_factory=list) + partitioners_by_method: dict = field(default_factory=dict) + quantizers: list = field(default_factory=list) + ao_quantization_configs: list = field(default_factory=list) + pre_edge_passes: list = field(default_factory=list) + edge_transform_passes: list = field(default_factory=list) + edge_manager_transform_passes: list = field(default_factory=list) + pre_prepare_passes: list = field(default_factory=list) + post_prepare_passes: list = field(default_factory=list) + pre_convert_passes: list = field(default_factory=list) + post_convert_passes: list = field(default_factory=list) + is_qat_values: list = field(default_factory=list) + dynamic_batch_size_values: list = field(default_factory=list) + train_fn_values: list = field(default_factory=list) + calibration_inputs_fn_values: list = field(default_factory=list) + strict_values: list = field(default_factory=list) + mode_values: list = field(default_factory=list) + pipeline_stages_values: list = field(default_factory=list) + source_transform_in_place_values: list = field(default_factory=list) + backend_config: object = None + + @experimental( "This API and all of its related functionality such as ExportSession and ExportRecipe are experimental." ) @@ -281,89 +348,102 @@ def combine( return cls._combine_recipes(recipes, recipe_name) + @staticmethod + def _assert_scalar_fields_agree(field_name: str, values: list) -> None: + """Raise ValueError when a scalar field has conflicting values across recipes.""" + unique = set(values) + if len(unique) > 1: + raise ValueError( + f"Cannot combine recipes with conflicting '{field_name}' values: {unique}" + ) + @classmethod - def _combine_recipes( # noqa: C901 - cls, backend_recipes: List["ExportRecipe"], recipe_name: Optional[str] = None - ) -> "ExportRecipe": + def _combine_quantization_recipe( + cls, + is_qat_values: list, + dynamic_batch_size_values: list, + train_fn_values: list, + calibration_inputs_fn_values: list, + all_quantizers: list, + all_ao_quantization_configs: list, + all_pre_prepare_passes: list, + all_post_prepare_passes: list, + all_pre_convert_passes: list, + all_post_convert_passes: list, + ) -> "Optional[QuantizationRecipe]": """ - Util to combine multiple backend recipes into a single multi-backend recipe. - - Args: - backend_recipes: List of ExportRecipe objects to combine - recipe_name: Optional name for the combined recipe + Build the combined QuantizationRecipe from per-recipe collected lists. - Returns: - Combined ExportRecipe for multi-backend deployment + Returns None when no recipe contributed any quantization fields, and logs + an INFO message so callers know quantization is absent from the combination. """ - overriding = [ - r.name or f"recipes[{i}]" - for i, r in enumerate(backend_recipes) - if r.pipeline_stages - ] - if overriding: + # is_qat must agree: the two flows (QAT vs PTQ) are incompatible. + cls._assert_scalar_fields_agree("is_qat", is_qat_values) + + # At most one recipe may supply a train_fn. + non_none_train_fns = [f for f in train_fn_values if f is not None] + if len(non_none_train_fns) > 1: raise ValueError( - "Cannot combine recipes that override pipeline_stages, there is no " - f"correct way to merge the orderings: {overriding}" + "Cannot combine recipes: more than one recipe provides a train_fn." ) - # Extract components from individual recipes - all_partitioners = [] - all_partitioners_by_method = {} - all_quantizers = [] - all_ao_quantization_configs = [] - all_pre_edge_passes = [] - all_transform_passes = [] - combined_backend_config = None + # Multiple calibration_inputs_fn values are chained into a single factory. + non_none_calib_fns = [f for f in calibration_inputs_fn_values if f is not None] + if len(non_none_calib_fns) > 1: + _fns = non_none_calib_fns - for recipe in backend_recipes: - # Collect pre-edge transform passes - if recipe.aten_transform_passes: - all_pre_edge_passes.extend(recipe.aten_transform_passes) - - # Collect partitioners from lowering recipes - if recipe.lowering_recipe and recipe.lowering_recipe.partitioners: - partitioners = recipe.lowering_recipe.partitioners - if isinstance(partitioners, dict): - for method_name, method_partitioners in partitioners.items(): - all_partitioners_by_method.setdefault(method_name, []).extend( - method_partitioners - ) - else: - all_partitioners.extend(partitioners) - - # Collect transform passes from lowering recipes - if recipe.lowering_recipe and recipe.lowering_recipe.edge_transform_passes: - all_transform_passes.extend( - recipe.lowering_recipe.edge_transform_passes - ) - - # Collect for quantize stage - if quantization_recipe := recipe.quantization_recipe: - # Collect PT2E quantizers - if quantization_recipe.quantizers: - all_quantizers.extend(quantization_recipe.quantizers) - - # Collect source transform configs - if quantization_recipe.ao_quantization_configs: - all_ao_quantization_configs.extend( - quantization_recipe.ao_quantization_configs - ) + def _combined_calib_fn(): + for _fn in _fns: + yield from _fn() - # Use the first backend config as base - if combined_backend_config is None and recipe.executorch_backend_config: - combined_backend_config = copy.deepcopy( - recipe.executorch_backend_config - ) - - # Create combined quantization recipe - combined_quantization_recipe = None - if all_quantizers or all_ao_quantization_configs: - combined_quantization_recipe = QuantizationRecipe( - quantizers=all_quantizers if all_quantizers else None, - ao_quantization_configs=( - all_ao_quantization_configs if all_ao_quantization_configs else None - ), + combined_calib_fn: "Optional[Callable[[], Iterable[tuple]]]" = ( + _combined_calib_fn + ) + else: + combined_calib_fn = non_none_calib_fns[0] if non_none_calib_fns else None + + if not ( + all_quantizers + or all_ao_quantization_configs + or all_pre_prepare_passes + or all_post_prepare_passes + or all_pre_convert_passes + or all_post_convert_passes + ): + logging.info( + "Combined recipe has no quantizers, quantization configs, or " + "quantization passes; quantization_recipe will be None." ) + return None + + return QuantizationRecipe( + quantizers=all_quantizers or None, + ao_quantization_configs=all_ao_quantization_configs or None, + is_qat=is_qat_values[0] if is_qat_values else False, + dynamic_batch_size=any(dynamic_batch_size_values), + train_fn=non_none_train_fns[0] if non_none_train_fns else None, + calibration_inputs_fn=combined_calib_fn, + pre_prepare_passes=all_pre_prepare_passes or None, + post_prepare_passes=all_post_prepare_passes or None, + pre_convert_passes=all_pre_convert_passes or None, + post_convert_passes=all_post_convert_passes or None, + ) + + @classmethod + def _combine_lowering_recipe( + cls, + backend_recipes: "List[ExportRecipe]", + all_partitioners: list, + all_partitioners_by_method: dict, + all_edge_transform_passes: list, + all_edge_manager_transform_passes: list, + ) -> "Optional[LoweringRecipe]": + """ + Build the combined LoweringRecipe from per-recipe collected lists. + + Returns None when no recipe contributed any lowering fields, and logs + an INFO message so callers know lowering is absent from the combination. + """ if all_partitioners and all_partitioners_by_method: raise ValueError( @@ -372,7 +452,7 @@ def _combine_recipes( # noqa: C901 ) combined_partitioners = all_partitioners_by_method or all_partitioners - # By value, not identity: every provider builds a fresh config object, + # Compare edge_compile_confgs by value, not identity: every provider builds a fresh config object, # so asking for the same thing twice is not a conflict. distinct: List[tuple[str, EdgeCompileConfig]] = [] for i, recipe in enumerate(backend_recipes): @@ -394,23 +474,164 @@ def _combine_recipes( # noqa: C901 ) edge_compile_config = copy.deepcopy(distinct[0][1]) if distinct else None - combined_lowering_recipe = None - if combined_partitioners or all_transform_passes or edge_compile_config: - combined_lowering_recipe = LoweringRecipe( - partitioners=combined_partitioners if combined_partitioners else None, - edge_transform_passes=( - all_transform_passes if all_transform_passes else None - ), - edge_compile_config=edge_compile_config or EdgeCompileConfig(), + if not ( + combined_partitioners + or all_edge_transform_passes + or all_edge_manager_transform_passes + or edge_compile_config + ): + logging.info( + "Combined recipe has no lowering fields; lowering_recipe will be None." ) + return None + + return LoweringRecipe( + partitioners=combined_partitioners or None, + edge_transform_passes=all_edge_transform_passes or None, + edge_manager_transform_passes=all_edge_manager_transform_passes or None, + edge_compile_config=edge_compile_config or EdgeCompileConfig(), + ) + + @staticmethod + def _collect_lowering_fields( + acc: "_CombineAccumulator", lr: "LoweringRecipe" + ) -> None: + """Accumulate fields from a single LoweringRecipe into acc.""" + if lr.partitioners: + if isinstance(lr.partitioners, dict): + for method_name, method_partitioners in lr.partitioners.items(): + acc.partitioners_by_method.setdefault(method_name, []).extend( + method_partitioners + ) + else: + acc.partitioners.extend(lr.partitioners) + if lr.edge_transform_passes: + acc.edge_transform_passes.extend(lr.edge_transform_passes) + if lr.edge_manager_transform_passes: + acc.edge_manager_transform_passes.extend(lr.edge_manager_transform_passes) + + @staticmethod + def _collect_quantization_fields( + acc: "_CombineAccumulator", qr: "QuantizationRecipe" + ) -> None: + """Accumulate fields from a single QuantizationRecipe into acc.""" + if qr.quantizers: + acc.quantizers.extend(qr.quantizers) + if qr.ao_quantization_configs: + acc.ao_quantization_configs.extend(qr.ao_quantization_configs) + acc.is_qat_values.append(qr.is_qat) + acc.dynamic_batch_size_values.append(qr.dynamic_batch_size) + acc.train_fn_values.append(qr.train_fn) + acc.calibration_inputs_fn_values.append(qr.calibration_inputs_fn) + if qr.pre_prepare_passes: + acc.pre_prepare_passes.extend(qr.pre_prepare_passes) + if qr.post_prepare_passes: + acc.post_prepare_passes.extend(qr.post_prepare_passes) + if qr.pre_convert_passes: + acc.pre_convert_passes.extend(qr.pre_convert_passes) + if qr.post_convert_passes: + acc.post_convert_passes.extend(qr.post_convert_passes) + + @classmethod + def _collect_recipe_fields( + cls, + backend_recipes: "List[ExportRecipe]", + ) -> "_CombineAccumulator": + """ + Iterate over all recipes and accumulate their fields into a single + _CombineAccumulator for later merging. + """ + acc = _CombineAccumulator() + + for recipe in backend_recipes: + if recipe.aten_transform_passes: + acc.pre_edge_passes.extend(recipe.aten_transform_passes) + + if lr := recipe.lowering_recipe: + cls._collect_lowering_fields(acc, lr) + + if qr := recipe.quantization_recipe: + cls._collect_quantization_fields(acc, qr) + + acc.strict_values.append(recipe.strict) + acc.mode_values.append(recipe.mode) + acc.pipeline_stages_values.append( + tuple(recipe.pipeline_stages) if recipe.pipeline_stages else None + ) + acc.source_transform_in_place_values.append( + recipe.source_transform_in_place + ) + + # Use the executorch_backend_config from the first recipe that supplies one. + if acc.backend_config is None and recipe.executorch_backend_config: + acc.backend_config = copy.deepcopy(recipe.executorch_backend_config) + + return acc + + @classmethod + def _combine_recipes( + cls, backend_recipes: "List[ExportRecipe]", recipe_name: "Optional[str]" = None + ) -> "ExportRecipe": + """ + Util to combine multiple backend recipes into a single multi-backend recipe. + + Args: + backend_recipes: List of ExportRecipe objects to combine + recipe_name: Optional name for the combined recipe + + Returns: + Combined ExportRecipe for multi-backend deployment + """ + acc = cls._collect_recipe_fields(backend_recipes) + + # Validate scalar fields that must agree across all recipes. + cls._assert_scalar_fields_agree("strict", acc.strict_values) + cls._assert_scalar_fields_agree("mode", acc.mode_values) + cls._assert_scalar_fields_agree("pipeline_stages", acc.pipeline_stages_values) + cls._assert_scalar_fields_agree( + "source_transform_in_place", acc.source_transform_in_place_values + ) + + combined_quantization_recipe = cls._combine_quantization_recipe( + is_qat_values=acc.is_qat_values, + dynamic_batch_size_values=acc.dynamic_batch_size_values, + train_fn_values=acc.train_fn_values, + calibration_inputs_fn_values=acc.calibration_inputs_fn_values, + all_quantizers=acc.quantizers, + all_ao_quantization_configs=acc.ao_quantization_configs, + all_pre_prepare_passes=acc.pre_prepare_passes, + all_post_prepare_passes=acc.post_prepare_passes, + all_pre_convert_passes=acc.pre_convert_passes, + all_post_convert_passes=acc.post_convert_passes, + ) + + combined_lowering_recipe = cls._combine_lowering_recipe( + backend_recipes=backend_recipes, + all_partitioners=acc.partitioners, + all_partitioners_by_method=acc.partitioners_by_method, + all_edge_transform_passes=acc.edge_transform_passes, + all_edge_manager_transform_passes=acc.edge_manager_transform_passes, + ) recipe_name = recipe_name or "_".join( [r.name for r in backend_recipes if r.name is not None] ) + # All pipeline_stages values are equal (enforced above); use the first non-None one. + shared_pipeline_stages = next( + (r.pipeline_stages for r in backend_recipes if r.pipeline_stages), None + ) return cls( name=recipe_name, quantization_recipe=combined_quantization_recipe, - aten_transform_passes=all_pre_edge_passes, + aten_transform_passes=acc.pre_edge_passes or None, lowering_recipe=combined_lowering_recipe, - executorch_backend_config=combined_backend_config, + executorch_backend_config=acc.backend_config, + pipeline_stages=shared_pipeline_stages, + strict=acc.strict_values[0] if acc.strict_values else True, + mode=acc.mode_values[0] if acc.mode_values else Mode.RELEASE, + source_transform_in_place=( + acc.source_transform_in_place_values[0] + if acc.source_transform_in_place_values + else False + ), ) diff --git a/export/stages.py b/export/stages.py index 96f46d750c3..3d2c8c86a4d 100644 --- a/export/stages.py +++ b/export/stages.py @@ -1,6 +1,7 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -23,7 +24,16 @@ from torch._export.pass_base import PassType from torch.fx.passes.infra.pass_manager import PassManager as GraphModulePassManager from torchao.quantization import quantize_ -from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +from torchao.quantization.pt2e import ( + allow_exported_model_train_eval, + move_exported_model_to_eval, + move_exported_model_to_train, +) +from torchao.quantization.pt2e.quantize_pt2e import ( + convert_pt2e, + prepare_pt2e, + prepare_qat_pt2e, +) from torchao.quantization.pt2e.quantizer import ( ComposableQuantizer, Quantizer as TorchAOPT2EQuantizer, @@ -465,16 +475,31 @@ def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]): else: raise ValueError("No quantizers detected") + @staticmethod + def _apply_passes( + model: "torch.fx.GraphModule", + passes: Optional[List[Callable]], + ) -> "torch.fx.GraphModule": + for pass_fn in passes or []: + try: + model = pass_fn(model) + except Exception as exc: + raise RuntimeError( + f"QuantizeStage: Pass '{pass_fn!r}' raised an error: {exc}" + ) from exc + return model + def run(self, artifact: PipelineArtifact) -> None: if not self._quantization_recipe or not self._quantization_recipe.quantizers: logging.info( - "Quantization recipe is invalid to run QunatizeStage, returning original model" + "Quantization recipe is invalid to run QuantizeStage, returning original model" ) self._artifact = artifact return assert isinstance(artifact.data, dict) + recipe = self._quantization_recipe models = artifact.data example_inputs = artifact.get_context("example_inputs") @@ -487,17 +512,78 @@ def run(self, artifact: PipelineArtifact) -> None: ) inputs = example_inputs[method_name][0] - captured_graph = torch.export.export(model, inputs, strict=True).module() - quantizer = self._get_quantizer_for_prepare_pt2e( - self._quantization_recipe.quantizers # pyre-ignore + # When dynamic_batch_size is requested, mark dimension 0 of every + # tensor input as dynamic so that a QAT training loop can feed + # mini-batches of arbitrary size through the prepared graph. + export_dynamic_shapes = None + if recipe.dynamic_batch_size: + from torch.export import Dim + + batch = Dim("batch", min=1) + export_dynamic_shapes = tuple( + {0: batch} if isinstance(t, torch.Tensor) else None for t in inputs + ) + + # QAT requires the model to be in training mode at capture time so + # that batch_norm and dropout decompose with training-mode semantics. + if recipe.is_qat: + model.train() + + captured_graph = torch.export.export( + model, inputs, dynamic_shapes=export_dynamic_shapes, strict=True + ).module() + + # Pass 1: pre-prepare passes. + captured_graph = self._apply_passes( + captured_graph, recipe.pre_prepare_passes ) - prepared_model = prepare_pt2e(captured_graph, quantizer) - for calibration_input in example_inputs[method_name]: - prepared_model(*calibration_input) + quantizer = self._get_quantizer_for_prepare_pt2e(recipe.quantizers) + + if recipe.is_qat: + if recipe.train_fn is None: + raise ValueError("train_fn must be provided when is_qat=True") + prepared_model = prepare_qat_pt2e(captured_graph, quantizer) + + # Pass 2: post-prepare passes. + prepared_model = self._apply_passes( + prepared_model, recipe.post_prepare_passes + ) + + allow_exported_model_train_eval(prepared_model) + move_exported_model_to_train(prepared_model) + recipe.train_fn(prepared_model) + move_exported_model_to_eval(prepared_model) + else: + prepared_model = prepare_pt2e(captured_graph, quantizer) + + # Pass 2: post-prepare passes. + prepared_model = self._apply_passes( + prepared_model, recipe.post_prepare_passes + ) + + # Use custom calibration inputs when provided; fall back to example inputs. + if recipe.calibration_inputs_fn is not None: + calibration_inputs = recipe.calibration_inputs_fn() + else: + calibration_inputs = example_inputs[method_name] + + for calibration_input in calibration_inputs: + prepared_model(*calibration_input) + + # Pass 3: pre-convert passes. + prepared_model = self._apply_passes( + prepared_model, recipe.pre_convert_passes + ) quantized_model = convert_pt2e(prepared_model) + + # Pass 4: post-convert passes. + quantized_model = self._apply_passes( + quantized_model, recipe.post_convert_passes + ) + quantized_models[method_name] = quantized_model self._artifact = artifact.copy_with_new_data(quantized_models) diff --git a/export/tests/test_export_recipe.py b/export/tests/test_export_recipe.py index 5fc4701a5f6..e1488ffe58c 100644 --- a/export/tests/test_export_recipe.py +++ b/export/tests/test_export_recipe.py @@ -7,12 +7,18 @@ # pyre-strict import unittest -from typing import Any, Dict, Optional, Sequence +from typing import Any, Dict, List, Optional, Sequence from unittest.mock import Mock import torch -from executorch.export.recipe import ExportRecipe, RecipeType +from executorch.export.recipe import ( + ExportRecipe, + LoweringRecipe, + Mode, + QuantizationRecipe, + RecipeType, +) from executorch.export.recipe_provider import BackendRecipeProvider from executorch.export.recipe_registry import recipe_registry @@ -293,8 +299,8 @@ def test_combine_keeps_partitioners(self) -> None: def test_combine_rejects_pipeline_stages(self) -> None: from executorch.export.types import StageType - # Unnamed recipes fall back to their position; named ones are named. - with self.assertRaisesRegex(ValueError, r"pipeline_stages.*recipes\[0\]"): + # Recipes with different pipeline_stages (including None vs. a list) must be rejected. + with self.assertRaisesRegex(ValueError, r"pipeline_stages"): ExportRecipe.combine( [ ExportRecipe( @@ -303,7 +309,7 @@ def test_combine_rejects_pipeline_stages(self) -> None: ExportRecipe(name="b"), ] ) - with self.assertRaisesRegex(ValueError, r"pipeline_stages.*'stagey'"): + with self.assertRaisesRegex(ValueError, r"pipeline_stages"): ExportRecipe.combine( [ ExportRecipe( @@ -334,3 +340,440 @@ def test_combine_keeps_edge_transform_passes(self) -> None: self.assertEqual( combined.lowering_recipe.edge_transform_passes, [first, second] ) + + +# --------------------------------------------------------------------------- +# Helpers shared by combine-recipe tests +# --------------------------------------------------------------------------- + + +def _make_pass(name: str, call_log: List[str]): + """Return a graph-module pass that appends *name* to *call_log*.""" + + def pass_fn(m): + call_log.append(name) + return m + + return pass_fn + + +class TestCombineRecipesEmpty(unittest.TestCase): + def test_empty_recipes_raises(self) -> None: + with self.assertRaises(ValueError): + ExportRecipe.combine([]) + + +class TestCombineRecipesSingleRecipe(unittest.TestCase): + def test_single_recipe_returned_unchanged(self) -> None: + recipe = ExportRecipe(name="solo") + result = ExportRecipe.combine([recipe]) + self.assertIs(result, recipe) + + +class TestCombineRecipesScalarFields(unittest.TestCase): + """Fields that must be identical across all combined recipes.""" + + def test_conflicting_strict_raises(self) -> None: + r1 = ExportRecipe(name="a", strict=True) + r2 = ExportRecipe(name="b", strict=False) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("strict", str(cm.exception)) + + def test_conflicting_mode_raises(self) -> None: + r1 = ExportRecipe(name="a", mode=Mode.DEBUG) + r2 = ExportRecipe(name="b", mode=Mode.RELEASE) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("mode", str(cm.exception)) + + def test_conflicting_source_transform_in_place_raises(self) -> None: + r1 = ExportRecipe(name="a", source_transform_in_place=True) + r2 = ExportRecipe(name="b", source_transform_in_place=False) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("source_transform_in_place", str(cm.exception)) + + def test_agreeing_scalar_fields_are_preserved(self) -> None: + r1 = ExportRecipe( + name="a", strict=False, mode=Mode.DEBUG, source_transform_in_place=True + ) + r2 = ExportRecipe( + name="b", strict=False, mode=Mode.DEBUG, source_transform_in_place=True + ) + result = ExportRecipe.combine([r1, r2]) + self.assertFalse(result.strict) + self.assertEqual(result.mode, Mode.DEBUG) + self.assertTrue(result.source_transform_in_place) + + def test_name_is_joined_from_input_recipe_names(self) -> None: + r1 = ExportRecipe(name="backend_a") + r2 = ExportRecipe(name="backend_b") + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.name, "backend_a_backend_b") + + def test_custom_recipe_name_is_used(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2], recipe_name="custom_name") + self.assertEqual(result.name, "custom_name") + + +class TestCombineRecipesAtenTransformPasses(unittest.TestCase): + def test_aten_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe(name="a", aten_transform_passes=[pass1]) + r2 = ExportRecipe(name="b", aten_transform_passes=[pass2]) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.aten_transform_passes, [pass1, pass2]) + + def test_aten_transform_passes_none_when_both_empty(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.aten_transform_passes) + + def test_aten_transform_passes_one_side_none(self) -> None: + pass1 = Mock() + r1 = ExportRecipe(name="a", aten_transform_passes=[pass1]) + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.aten_transform_passes, [pass1]) + + +class TestCombineRecipesLowering(unittest.TestCase): + def test_partitioners_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe(name="a", lowering_recipe=LoweringRecipe(partitioners=[p1])) + r2 = ExportRecipe(name="b", lowering_recipe=LoweringRecipe(partitioners=[p2])) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNotNone(result.lowering_recipe) + self.assertEqual(result.lowering_recipe.partitioners, [p1, p2]) + + def test_edge_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe( + name="a", lowering_recipe=LoweringRecipe(edge_transform_passes=[pass1]) + ) + r2 = ExportRecipe( + name="b", lowering_recipe=LoweringRecipe(edge_transform_passes=[pass2]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.lowering_recipe.edge_transform_passes, [pass1, pass2]) + + def test_edge_manager_transform_passes_merged(self) -> None: + pass1 = Mock() + pass2 = Mock() + r1 = ExportRecipe( + name="a", + lowering_recipe=LoweringRecipe(edge_manager_transform_passes=[pass1]), + ) + r2 = ExportRecipe( + name="b", + lowering_recipe=LoweringRecipe(edge_manager_transform_passes=[pass2]), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual( + result.lowering_recipe.edge_manager_transform_passes, [pass1, pass2] + ) + + def test_lowering_recipe_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.lowering_recipe) + + def test_edge_compile_config_taken_from_first_recipe_with_one(self) -> None: + from executorch.exir.capture import EdgeCompileConfig + + config = EdgeCompileConfig() + r1 = ExportRecipe(name="a") + r2 = ExportRecipe( + name="b", + lowering_recipe=LoweringRecipe( + partitioners=[Mock()], edge_compile_config=config + ), + ) + result = ExportRecipe.combine([r1, r2]) + # combine() deepcopies the config so the combined recipe cannot mutate + # the provider's shared object; assert value-equality, not identity. + self.assertIsNotNone(result.lowering_recipe) + self.assertEqual(result.lowering_recipe.edge_compile_config, config) + self.assertIsNot(result.lowering_recipe.edge_compile_config, config) + + +class TestCombineRecipesQuantization(unittest.TestCase): + def test_quantizers_merged(self) -> None: + q1 = Mock() + q2 = Mock() + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[q1]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[q2]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNotNone(result.quantization_recipe) + self.assertEqual(result.quantization_recipe.quantizers, [q1, q2]) + + def test_ao_quantization_configs_merged(self) -> None: + from executorch.export.recipe import AOQuantizationConfig + from torchao.core.config import AOBaseConfig + + cfg1 = AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig)) + cfg2 = AOQuantizationConfig(ao_base_config=Mock(spec=AOBaseConfig)) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(ao_quantization_configs=[cfg1]), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(ao_quantization_configs=[cfg2]), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual( + result.quantization_recipe.ao_quantization_configs, [cfg1, cfg2] + ) + + def test_quantization_recipe_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe(name="a") + r2 = ExportRecipe(name="b") + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.quantization_recipe) + + def test_conflicting_is_qat_raises(self) -> None: + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=False), + ) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("is_qat", str(cm.exception)) + + def test_agreeing_is_qat_preserved(self) -> None: + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertTrue(result.quantization_recipe.is_qat) + + def test_two_train_fns_raises(self) -> None: + fn1 = Mock() + fn2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn2 + ), + ) + with self.assertRaises(ValueError) as cm: + ExportRecipe.combine([r1, r2]) + self.assertIn("train_fn", str(cm.exception)) + + def test_single_train_fn_preserved(self) -> None: + fn = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], is_qat=True, train_fn=fn + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe(quantizers=[Mock()], is_qat=True), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.quantization_recipe.train_fn, fn) + + def test_single_calibration_inputs_fn_preserved(self) -> None: + fn = Mock(return_value=[(1,), (2,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn + ), + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIs(result.quantization_recipe.calibration_inputs_fn, fn) + + def test_two_calibration_inputs_fns_chained(self) -> None: + fn1 = Mock(return_value=[(1,), (2,)]) + fn2 = Mock(return_value=[(3,), (4,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn2 + ), + ) + result = ExportRecipe.combine([r1, r2]) + combined_fn = result.quantization_recipe.calibration_inputs_fn + self.assertIsNotNone(combined_fn) + # Each factory must be called exactly once when the combined factory is consumed. + all_inputs = list(combined_fn()) + fn1.assert_called_once_with() + fn2.assert_called_once_with() + self.assertEqual(all_inputs, [(1,), (2,), (3,), (4,)]) + + def test_three_calibration_inputs_fns_chained_in_order(self) -> None: + fn1 = Mock(return_value=[(1,)]) + fn2 = Mock(return_value=[(2,)]) + fn3 = Mock(return_value=[(3,)]) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn1 + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn2 + ), + ) + r3 = ExportRecipe( + name="c", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], calibration_inputs_fn=fn3 + ), + ) + result = ExportRecipe.combine([r1, r2, r3]) + combined_fn = result.quantization_recipe.calibration_inputs_fn + self.assertEqual(list(combined_fn()), [(1,), (2,), (3,)]) + + def test_no_calibration_inputs_fn_stays_none(self) -> None: + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertIsNone(result.quantization_recipe.calibration_inputs_fn) + + def test_pre_prepare_passes_merged(self) -> None: + log: List[str] = [] + p1 = _make_pass("pre_a", log) + p2 = _make_pass("pre_b", log) + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_prepare_passes, [p1, p2]) + + def test_post_prepare_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_prepare_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_prepare_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.post_prepare_passes, [p1, p2]) + + def test_pre_convert_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_convert_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_convert_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_convert_passes, [p1, p2]) + + def test_post_convert_passes_merged(self) -> None: + p1 = Mock() + p2 = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_convert_passes=[p1] + ), + ) + r2 = ExportRecipe( + name="b", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], post_convert_passes=[p2] + ), + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.post_convert_passes, [p1, p2]) + + def test_all_pass_lists_none_when_nothing_contributed(self) -> None: + r1 = ExportRecipe( + name="a", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + qr = result.quantization_recipe + self.assertIsNone(qr.pre_prepare_passes) + self.assertIsNone(qr.post_prepare_passes) + self.assertIsNone(qr.pre_convert_passes) + self.assertIsNone(qr.post_convert_passes) + + def test_pass_lists_preserved_when_only_one_recipe_contributes(self) -> None: + p = Mock() + r1 = ExportRecipe( + name="a", + quantization_recipe=QuantizationRecipe( + quantizers=[Mock()], pre_prepare_passes=[p] + ), + ) + r2 = ExportRecipe( + name="b", quantization_recipe=QuantizationRecipe(quantizers=[Mock()]) + ) + result = ExportRecipe.combine([r1, r2]) + self.assertEqual(result.quantization_recipe.pre_prepare_passes, [p]) diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index 935a796591a..46eb000b743 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 NXP # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -409,6 +410,7 @@ def test_run_no_quantizers(self) -> None: result_artifact = stage.get_artifacts() self.assertEqual(result_artifact, artifact) + @patch("executorch.export.stages.move_exported_model_to_eval") @patch("executorch.export.stages.convert_pt2e") @patch("executorch.export.stages.prepare_pt2e") @patch("executorch.export.stages.ComposableQuantizer") @@ -419,11 +421,19 @@ def test_run_with_quantizers( mock_composable_quantizer: Mock, mock_prepare_pt2e: Mock, mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, ) -> None: """Test execution with quantizers""" mock_quantizer = self.create_dummy_quantizer() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None stage = QuantizeStage(mock_recipe) # Mock the torch.export.export chain @@ -443,9 +453,9 @@ def test_run_with_quantizers( artifact = PipelineArtifact(data=self.models_dict, context=self.context) stage.run(artifact) - # Verify torch.export.export was called + # Verify torch.export.export was called with dynamic_shapes=None (no dynamic batch) mock_torch_export.assert_called_once_with( - self.model, self.example_inputs[0], strict=True + self.model, self.example_inputs[0], dynamic_shapes=None, strict=True ) # Verify ComposableQuantizer was created with the quantizers @@ -471,11 +481,401 @@ def test_run_with_quantizers( self.assertEqual(artifact.data["forward"], self.model) self.assertIsNot(result_artifact.data["forward"], self.model) + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_calls_prepare_qat_pt2e( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """QAT flow: prepare_qat_pt2e is called and train_fn is invoked with the prepared model. + allow_exported_model_train_eval must be called after preparation. + move_exported_model_to_train must be called before train_fn, and + move_exported_model_to_eval must be called after train_fn.""" + mock_quantizer = self.create_dummy_quantizer() + call_order = [] + mock_allow_train_eval.side_effect = lambda m: call_order.append( + "allow_train_eval" + ) + mock_move_to_train.side_effect = lambda m: call_order.append("to_train") + mock_move_to_eval.side_effect = lambda m: call_order.append("to_eval") + + train_fn = Mock(side_effect=lambda m: call_order.append("train_fn")) + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = train_fn + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_captured_graph = Mock() + mock_exported_program.module.return_value = mock_captured_graph + mock_torch_export.return_value = mock_exported_program + + mock_composed_quantizer = Mock() + mock_composable_quantizer.return_value = mock_composed_quantizer + mock_prepared_model = Mock() + mock_prepare_qat_pt2e.return_value = mock_prepared_model + mock_quantized_model = Mock() + mock_convert_pt2e.return_value = mock_quantized_model + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # prepare_qat_pt2e must be called, not prepare_pt2e + mock_prepare_qat_pt2e.assert_called_once_with( + mock_captured_graph, mock_composed_quantizer + ) + # allow_exported_model_train_eval before move_to_train, then train_fn, then to_eval + self.assertEqual( + call_order, ["allow_train_eval", "to_train", "train_fn", "to_eval"] + ) + mock_allow_train_eval.assert_called_once_with(mock_prepared_model) + mock_move_to_train.assert_called_once_with(mock_prepared_model) + mock_move_to_eval.assert_called_once_with(mock_prepared_model) + # train_fn must be called with the prepared model + train_fn.assert_called_once_with(mock_prepared_model) + # convert_pt2e must still be called after training + mock_convert_pt2e.assert_called_once_with(mock_prepared_model) + + result_artifact = stage.get_artifacts() + self.assertEqual(result_artifact.data["forward"], mock_quantized_model) + + @patch("torch.export.export") + def test_run_qat_missing_train_fn_raises(self, mock_torch_export: Mock) -> None: + """QAT flow with train_fn=None must raise ValueError.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + + with self.assertRaises(ValueError) as cm: + stage.run(artifact) + self.assertIn("train_fn must be provided when is_qat=True", str(cm.exception)) + + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_model_put_in_train_mode_before_export( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """QAT: model.train() must be called before torch.export.export so that + batch_norm and dropout decompose with training-mode semantics.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = Mock() + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + # Start model in eval mode; the stage must switch it to train. + self.model.eval() + self.assertFalse(self.model.training) + + training_at_export_time = [] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + + def capture_training_flag(model, *args, **kwargs): + training_at_export_time.append(model.training) + return mock_exported_program + + mock_torch_export.side_effect = capture_training_flag + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + # The model must have been in training mode when export was called. + self.assertEqual(training_at_export_time, [True]) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_does_not_call_prepare_qat_pt2e( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """PTQ flow must not call prepare_qat_pt2e (regression guard).""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + mock_prepare_pt2e.assert_called_once() + mock_prepare_qat_pt2e.assert_not_called() + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_four_passes_called_in_order( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """All four pass hooks are called at the correct points in the PTQ flow.""" + call_order = [] + + def make_pass(name): + def pass_fn(m): + call_order.append(name) + return m + + return pass_fn + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")] + mock_recipe.post_prepare_passes = [make_pass("post_prepare")] + mock_recipe.pre_convert_passes = [make_pass("pre_convert")] + mock_recipe.post_convert_passes = [make_pass("post_convert")] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_graph = Mock() + mock_exported_program.module.return_value = mock_graph + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + self.assertEqual( + call_order, + ["pre_prepare", "post_prepare", "pre_convert", "post_convert"], + ) + + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_qat_four_passes_called_in_order( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """All four pass hooks are called at the correct points in the QAT flow.""" + call_order = [] + + def make_pass(name): + def pass_fn(m): + call_order.append(name) + return m + + return pass_fn + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = True + mock_recipe.dynamic_batch_size = False + mock_recipe.train_fn = Mock() + mock_recipe.pre_prepare_passes = [make_pass("pre_prepare")] + mock_recipe.post_prepare_passes = [make_pass("post_prepare")] + mock_recipe.pre_convert_passes = [make_pass("pre_convert")] + mock_recipe.post_convert_passes = [make_pass("post_convert")] + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + self.assertEqual( + call_order, + ["pre_prepare", "post_prepare", "pre_convert", "post_convert"], + ) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_uses_calibration_inputs_fn_when_provided( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When calibration_inputs_fn is set, it is called and its output is used for calibration.""" + custom_input = (torch.randn(2, 10),) + calibration_inputs_fn = Mock(return_value=[custom_input]) + + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = calibration_inputs_fn + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepared_model = Mock() + mock_prepare_pt2e.return_value = mock_prepared_model + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # calibration_inputs_fn must be called with no arguments + calibration_inputs_fn.assert_called_once_with() + # prepared model must be called with the custom calibration input + mock_prepared_model.assert_called_once_with(*custom_input) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_run_ptq_falls_back_to_example_inputs_when_no_calibration_fn( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When calibration_inputs_fn is None, example inputs are used for calibration.""" + mock_quantizer = self.create_dummy_quantizer() + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.dynamic_batch_size = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + + mock_exported_program = Mock(spec=ExportedProgram) + mock_exported_program.module.return_value = Mock() + mock_torch_export.return_value = mock_exported_program + mock_composable_quantizer.return_value = Mock() + mock_prepared_model = Mock() + mock_prepare_pt2e.return_value = mock_prepared_model + mock_convert_pt2e.return_value = Mock() + + stage = QuantizeStage(mock_recipe) + artifact = PipelineArtifact(data=self.models_dict, context=self.context) + stage.run(artifact) + + # The prepared model must be called with the example inputs (one tuple) + mock_prepared_model.assert_called_once_with(*self.example_inputs[0]) + def test_run_empty_example_inputs(self) -> None: """Test error when example inputs list is empty.""" mock_quantizer = Mock() mock_recipe = Mock(spec=QuantizationRecipe) mock_recipe.quantizers = [mock_quantizer] + mock_recipe.is_qat = False + mock_recipe.calibration_inputs_fn = None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None stage = QuantizeStage(mock_recipe) context = {"example_inputs": {"forward": []}} artifact = PipelineArtifact(data=self.models_dict, context=context) @@ -650,6 +1050,175 @@ def test_run_edge_manager_none(self) -> None: self.assertIn("Edge program manager is not set", str(cm.exception)) +class TestQuantizeStageExportDynamicShapes(unittest.TestCase): + """Tests for the dynamic_batch_size export behavior in QuantizeStage.""" + + def setUp(self) -> None: + self.model = torch.nn.Linear(10, 5) + self.models_dict = {"forward": self.model} + self.example_inputs = [(torch.randn(1, 10),)] + self.context = {"example_inputs": {"forward": self.example_inputs}} + + @staticmethod + def _make_recipe(is_qat: bool, dynamic_batch_size: bool) -> Mock: + mock_recipe = Mock(spec=QuantizationRecipe) + mock_recipe.quantizers = [Mock(spec=TorchAOPT2EQuantizer)] + mock_recipe.is_qat = is_qat + mock_recipe.dynamic_batch_size = dynamic_batch_size + mock_recipe.calibration_inputs_fn = None + mock_recipe.train_fn = Mock() if is_qat else None + mock_recipe.pre_prepare_passes = None + mock_recipe.post_prepare_passes = None + mock_recipe.pre_convert_passes = None + mock_recipe.post_convert_passes = None + return mock_recipe + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_dynamic_batch_size_false_exports_without_dynamic_shapes( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When dynamic_batch_size=False, torch.export.export is called with dynamic_shapes=None.""" + mock_ep = Mock(spec=ExportedProgram) + mock_ep.module.return_value = Mock() + mock_torch_export.return_value = mock_ep + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + recipe = self._make_recipe(is_qat=False, dynamic_batch_size=False) + stage = QuantizeStage(recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + mock_torch_export.assert_called_once_with( + self.model, + self.example_inputs[0], + dynamic_shapes=None, + strict=True, + ) + + @patch("executorch.export.stages.allow_exported_model_train_eval") + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.move_exported_model_to_train") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_qat_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_dynamic_batch_size_true_exports_with_dynamic_batch_dim( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_qat_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_train: Mock, + mock_move_to_eval: Mock, + mock_allow_train_eval: Mock, + ) -> None: + """When dynamic_batch_size=True, torch.export.export is called with a + dynamic_shapes tuple where dimension 0 of every tensor is dynamic.""" + mock_ep = Mock(spec=ExportedProgram) + mock_ep.module.return_value = Mock() + mock_torch_export.return_value = mock_ep + mock_composable_quantizer.return_value = Mock() + mock_prepare_qat_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + recipe = self._make_recipe(is_qat=True, dynamic_batch_size=True) + stage = QuantizeStage(recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + call_kwargs = mock_torch_export.call_args + dynamic_shapes_arg = call_kwargs.kwargs.get( + "dynamic_shapes", call_kwargs.args[2] if len(call_kwargs.args) > 2 else None + ) + # dynamic_shapes must be a non-None tuple with one entry per input tensor. + self.assertIsNotNone(dynamic_shapes_arg) + self.assertIsInstance(dynamic_shapes_arg, tuple) + self.assertEqual(len(dynamic_shapes_arg), len(self.example_inputs[0])) + # The entry for the single tensor input must map dim 0 to a Dim. + first_entry = dynamic_shapes_arg[0] + self.assertIsInstance(first_entry, dict) + self.assertIn(0, first_entry) + + @patch("executorch.export.stages.move_exported_model_to_eval") + @patch("executorch.export.stages.convert_pt2e") + @patch("executorch.export.stages.prepare_pt2e") + @patch("executorch.export.stages.ComposableQuantizer") + @patch("torch.export.export") + def test_dynamic_batch_size_true_ptq_exports_with_dynamic_shapes( + self, + mock_torch_export: Mock, + mock_composable_quantizer: Mock, + mock_prepare_pt2e: Mock, + mock_convert_pt2e: Mock, + mock_move_to_eval: Mock, + ) -> None: + """When dynamic_batch_size=True and is_qat=False, export is called with dynamic shapes.""" + mock_ep = Mock(spec=ExportedProgram) + mock_ep.module.return_value = Mock() + mock_torch_export.return_value = mock_ep + mock_composable_quantizer.return_value = Mock() + mock_prepare_pt2e.return_value = Mock() + mock_convert_pt2e.return_value = Mock() + + recipe = self._make_recipe(is_qat=False, dynamic_batch_size=True) + stage = QuantizeStage(recipe) + stage.run(PipelineArtifact(data=self.models_dict, context=self.context)) + + call_kwargs = mock_torch_export.call_args + dynamic_shapes_arg = call_kwargs.kwargs.get( + "dynamic_shapes", call_kwargs.args[2] if len(call_kwargs.args) > 2 else None + ) + self.assertIsNotNone(dynamic_shapes_arg) + self.assertIsInstance(dynamic_shapes_arg, tuple) + self.assertEqual(len(dynamic_shapes_arg), len(self.example_inputs[0])) + first_entry = dynamic_shapes_arg[0] + self.assertIsInstance(first_entry, dict) + self.assertIn(0, first_entry) + + def test_dynamic_batch_size_true_ptq_calibration_with_variable_batch_sizes( + self, + ) -> None: + """PTQ calibration runs without error when batch sizes vary across calibration inputs.""" + from executorch.export.recipe import QuantizationRecipe + + class PassthroughQuantizer(TorchAOPT2EQuantizer): + def annotate(self, model): + return model + + def validate(self, model): + pass + + def calibration_inputs_fn(): + for batch_size in (2, 4, 8): + yield (torch.randn(batch_size, 10),) + + recipe = QuantizationRecipe( + quantizers=[PassthroughQuantizer()], + is_qat=False, + dynamic_batch_size=True, + calibration_inputs_fn=calibration_inputs_fn, + ) + stage = QuantizeStage(recipe) + # Use batch size 2 for the example input so torch.export does not + # specialize dim 0 as the constant 1. + context = {"example_inputs": {"forward": [(torch.randn(2, 10),)]}} + artifact = PipelineArtifact( + data={"forward": torch.nn.Linear(10, 5)}, + context=context, + ) + stage.run(artifact) + self.assertIn("forward", stage.get_artifacts().data) + + class TestEmptyPassDictIsNotApplied(unittest.TestCase): """`EdgeProgramManager.transform` deep-copies the graph and weights of every method the pass dict does not name, so handing it an empty dict copies From 7490bd48c599f1c5eabfb84dfc67e6eef3bda743 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Thu, 3 Sep 2026 11:52:58 +0200 Subject: [PATCH 015/190] Add pass to replace input dim order clones with permutations. (#21241) ### Summary This pass replaces `dim_order_clone` and `to_dim_order_copy` nodes that consume model inputs by a sequence of permute nodes according to the proposal in https://github.com/pytorch/executorch/issues/19299. Fixes #20095. ### Test plan `pytest backends/transforms/test/test_replace_channels_last_input_clones.py` --- .../replace_channels_last_input_clones.py | 133 ++++++ backends/transforms/targets.bzl | 29 ++ ...test_replace_channels_last_input_clones.py | 402 ++++++++++++++++++ 3 files changed, 564 insertions(+) create mode 100644 backends/transforms/replace_channels_last_input_clones.py create mode 100644 backends/transforms/test/test_replace_channels_last_input_clones.py diff --git a/backends/transforms/replace_channels_last_input_clones.py b/backends/transforms/replace_channels_last_input_clones.py new file mode 100644 index 00000000000..6d620de8ca4 --- /dev/null +++ b/backends/transforms/replace_channels_last_input_clones.py @@ -0,0 +1,133 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Sequence + +import executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops + +from executorch.exir.pass_base import ExportPass, NodeMetadata, ProxyValue +from torch.fx.passes.infra.pass_base import PassResult + +_DIM_ORDER_CHANGING_OPS: frozenset = frozenset( + { + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + } +) + +# kwargs that carry no tensor semantics and require no value inspection. +_PASS_THROUGH_KWARGS: frozenset = frozenset({"dim_order", "non_blocking"}) + +# TensorOptions that the exporter may annotate explicitly on `_to_dim_order_copy` nodes even +# when no actual change occurs. These are allowed only when their value is identical to the +# source tensor's property (i.e. truly a no-op). Any value that would actually change the +# tensor (e.g. a different dtype) blocks replacement. Any kwarg not in either set is unknown +# and is also rejected. +_TENSOR_OPTION_KWARGS: frozenset = frozenset( + {"dtype", "layout", "device", "pin_memory"} +) + +_ALLOWED_KWARGS: frozenset = _PASS_THROUGH_KWARGS | _TENSOR_OPTION_KWARGS + +_TO_NHWC_PERMUTATION: list[int] = [0, 2, 3, 1] +_TO_NCHW_PERMUTATION: list[int] = [0, 3, 1, 2] + + +def _is_4d_contiguous(dim_order: Sequence[int]) -> bool: + return list(dim_order) == [0, 1, 2, 3] + + +def _is_4d_channels_last(dim_order: Sequence[int]) -> bool: + return list(dim_order) == [0, 2, 3, 1] + + +def _is_replaceable_input_boundary_clone(op, args, kwargs) -> bool: + """Return True if the op/args/kwargs describe a 4D channels-last-to-contiguous clone of a model input. + + These are the input boundary clones inserted by `EnforceContiguousDimOrder`: they consume a + channels-last placeholder and produce a contiguous tensor. Replacing them with a permute pair + allows them to be optimized out by other passes, leaving only the no-op `aten.permute_copy`. + """ + if op not in _DIM_ORDER_CHANGING_OPS: + return False + if not args or not hasattr(args[0], "node"): + return False + src = args[0].node + if not isinstance(src, torch.fx.Node) or src.op != "placeholder": + return False + val = src.meta.get("val") + if not isinstance(val, torch.Tensor): + return False + # Primary guard: reject any kwarg outside the known set. Unknown kwargs may carry + # semantics we cannot reason about, so we conservatively block replacement. + if not set(kwargs.keys()) <= _ALLOWED_KWARGS: + return False + # Secondary guard: each TensorOption that is present must be a no-op relative to the + # source tensor. The exporter annotates these even when no actual change occurs, but a + # differing value (e.g. a dtype cast) must not be silently dropped by the replacement. + if kwargs.get("dtype") not in (None, val.dtype): + return False + if kwargs.get("layout") not in (None, val.layout): + return False + if kwargs.get("device") not in (None, val.device): + return False + if kwargs.get("pin_memory") not in (None, False): + return False + return _is_4d_channels_last(val.dim_order()) and _is_4d_contiguous( + kwargs.get("dim_order", []) + ) + + +class ReplaceChannelsLastInputClones(ExportPass): + """Replace `_to_dim_order_copy` and `_clone_dim_order` with an equivalent sequence in the following pattern. This + approach allows the `channels_last.permute_copy` to be optimized out if there are subsequent channels last + operators in the model, leaving only the `aten.permute_copy`, which is effectively a no-op. As a result, the + input data doesn't have to be permuted in memory. + + │ [N, C, H, W] shape, (0, 2, 3, 1) dim order + │ data is stored channels last + │ [N, C, H, W] shape, (0, 2, 3, 1) dim order ┌─────────▼─────────┐ + │ data is stored channels last │ aten.permute_copy ◄──── [0, 2, 3, 1] permutation + ┌───────────▼────────────┐ └─────────┬─────────┘ + │ │ ────────────────► │ [N, H, W, C] shape, (0, 1, 2, 3) dim order + └───────────┬────────────┘ │ data is stored channels last + │ [N, C, H, W] shape, (0, 1, 2, 3) dim order ┌──────────────▼─────────────┐ + ▼ data is stored channels first │ channels_last.permute_copy ◄──── [0, 3, 1, 2] permutation + └──────────────┬─────────────┘ + │ [N, C, H, W] shape, (0, 1, 2, 3) dim order + ▼ data is stored channels first + """ + + _modified: bool + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self._modified = False + result = super().call(graph_module) + return PassResult(result.graph_module, self._modified) + + def call_operator(self, op, args, kwargs, meta: NodeMetadata) -> ProxyValue: + if not _is_replaceable_input_boundary_clone(op, args, kwargs): + return super().call_operator(op, args, kwargs, meta) + + x = super().call_operator( + exir_ops.edge.aten.permute_copy.default, + (args[0], _TO_NHWC_PERMUTATION), + {}, + meta, + ) + x = super().call_operator( + exir_ops.edge.channels_last.permute_copy.default, + (x, _TO_NCHW_PERMUTATION), + {}, + meta, + ) + + self._modified = True + + return x diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 1e456249d3f..d34767b84ae 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -716,3 +716,32 @@ def define_common_targets(): ":enforce_contiguous_dim_order", ], ) + + runtime.python_library( + name = "replace_channels_last_input_clones", + srcs = ["replace_channels_last_input_clones.py"], + visibility = [ + "//executorch/backends/...", + ], + deps = [ + "//caffe2:torch", + ":channels_last_ops", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_replace_channels_last_input_clones", + srcs = [ + "test/test_replace_channels_last_input_clones.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + ":channels_last_ops", + ":enforce_contiguous_dim_order", + ":replace_channels_last_input_clones", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) diff --git a/backends/transforms/test/test_replace_channels_last_input_clones.py b/backends/transforms/test/test_replace_channels_last_input_clones.py new file mode 100644 index 00000000000..e6b30df204f --- /dev/null +++ b/backends/transforms/test/test_replace_channels_last_input_clones.py @@ -0,0 +1,402 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import executorch.backends.transforms.channels_last_ops # noqa: F401 +import pytest +import torch + +from executorch.backends.transforms.enforce_contiguous_dim_order import ( + EnforceContiguousDimOrder, +) +from executorch.backends.transforms.replace_channels_last_input_clones import ( + _is_4d_channels_last, + _is_4d_contiguous, + ReplaceChannelsLastInputClones, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export import ExportedProgram +from torch.fx import GraphModule +from torch.fx.node import Target + +_CLONE_DIM_ORDER = exir_ops.edge.dim_order_ops._clone_dim_order.default +_TO_DIM_ORDER_COPY = exir_ops.edge.dim_order_ops._to_dim_order_copy.default +_ATEN_PERMUTE_COPY = exir_ops.edge.aten.permute_copy.default +_CHANNELS_LAST_PERMUTE_COPY = exir_ops.edge.channels_last.permute_copy.default + + +class SingleInputToDimOrderCopyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + # Expecting `x` to use the channels last memory format. + x = x.to(memory_format=torch.contiguous_format) + x = self.avg_pool(x) + return x + + +class MultiInputToDimOrderCopyModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, *inputs): + contiguous_inputs = [ + input_.to(memory_format=torch.contiguous_format) for input_ in inputs + ] + x = torch.concatenate(contiguous_inputs) + x = self.avg_pool(x) + return x + + +class ToDimOrderCopyAfterAddModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + x = ( + x + x + ) # Make sure the `_to_dim_order_copy` is not consuming the model input. + x = x.to(memory_format=torch.contiguous_format) + x = self.avg_pool(x) + return x + + +class SingleInputModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + x = self.avg_pool(x) + x = torch.relu(x) + return x + + +class MultiInputModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.single_input_module = SingleInputModule() + + def forward(self, *inputs): + x = torch.concatenate(inputs) + x = self.single_input_module(x) + return x + + +class IncompatibleDimOrderModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + x = x.to(memory_format=torch.channels_last) # Incompatible dim order. + x = self.avg_pool(x) + return x + + +class DtypeCastAndLayoutChangeModule(torch.nn.Module): + """Module whose forward casts dtype AND changes memory format via `_to_dim_order_copy`. + The pass must not replace such a node, since doing so would silently drop the cast. + """ + + def __init__(self): + super().__init__() + self.avg_pool = torch.nn.AvgPool2d(3) + + def forward(self, x): + # Both a memory-format change and a dtype cast happen here. + x = x.to(dtype=torch.float64, memory_format=torch.contiguous_format) + x = x.to(dtype=torch.float32) # Cast back so avg_pool accepts it. + x = self.avg_pool(x) + return x + + +def _export_to_edge(module: torch.nn.Module, inputs: tuple) -> ExportedProgram: + ep = torch.export.export(module.eval(), inputs) + return to_edge(ep).exported_program() + + +def _find_nodes(gm: GraphModule, target: Target) -> list[torch.fx.Node]: + return [n for n in gm.graph.nodes if n.op == "call_function" and n.target == target] + + +def _count(gm: GraphModule, target: Target) -> int: + return len(_find_nodes(gm, target)) + + +def _run_pass(ep_or_result) -> tuple[GraphModule, bool]: + gm = ep_or_result.graph_module + result = ReplaceChannelsLastInputClones()(gm) + return result.graph_module, result.modified + + +def _assert_expected_result_pattern( + input_: torch.fx.Node, + aten_permute: torch.fx.Node, + channels_last_permute: torch.fx.Node, +): + assert input_.op == "placeholder" + assert input_.meta["val"].dim_order() == (0, 2, 3, 1) + assert aten_permute.args[0] == input_ + assert aten_permute.target == _ATEN_PERMUTE_COPY + assert aten_permute.args[1] == [0, 2, 3, 1] + assert aten_permute.meta["val"].dim_order() == (0, 1, 2, 3) + assert channels_last_permute.args[0] == aten_permute + assert channels_last_permute.target == _CHANNELS_LAST_PERMUTE_COPY + assert channels_last_permute.args[1] == [0, 3, 1, 2] + assert channels_last_permute.meta["val"].dim_order() == (0, 1, 2, 3) + + +@pytest.fixture(autouse=True) +def _reseed(): + torch.manual_seed(42) + yield + + +class TestReplaceChannelsLastInputDimOrderCopies: + """These tests use models with an explicit dim order change in their `forward()` method, which results in a + `_to_dim_order_copy` operator in edge dialect. + """ + + def test_single_input(self): + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(SingleInputToDimOrderCopyModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == 1 + output_before = ep.module()(*example_inputs) + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 0 + assert _count(gm, _CLONE_DIM_ORDER) == 0 + + nodes = list(gm.graph.nodes) + _assert_expected_result_pattern(*nodes[:3]) + + outputs_after = gm(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test_multiple_inputs(self): + num_inputs = 3 + example_inputs = tuple( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + for _ in range(num_inputs) + ) + + ep = _export_to_edge(MultiInputToDimOrderCopyModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == num_inputs + output_before = ep.module()(*example_inputs) + + gm, modified = _run_pass(ep) + + assert modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 0 + assert _count(gm, _CLONE_DIM_ORDER) == 0 + + nodes = list(gm.graph.nodes) + for i in range(num_inputs): + # Each input should have the expected pattern + input_ = nodes[i] + start_idx = num_inputs + i * (num_inputs - 1) + end_idx = start_idx + 2 + pattern = nodes[start_idx:end_idx] + _assert_expected_result_pattern(input_, *pattern) + + outputs_after = gm(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test__not_applied__not_consuming_model_input(self): + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(ToDimOrderCopyAfterAddModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == 1 + + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 1 + + def test__not_applied__incompatible_dim_order(self): + # This test uses a `to_dim_order_copy` which goes from contiguous to channels_last, which is not what the pass + # was made for. + example_inputs = (torch.randn(1, 3, 8, 8),) + + ep = _export_to_edge(IncompatibleDimOrderModule(), example_inputs) + assert _count(ep.graph_module, _TO_DIM_ORDER_COPY) == 1 + + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 1 + + def test__not_applied__dtype_cast(self): + """A `_to_dim_order_copy` that also changes dtype must not be replaced. + Replacing it with permutes would silently drop the cast, changing the graph's + semantics. + """ + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(DtypeCastAndLayoutChangeModule(), example_inputs) + # The first `_to_dim_order_copy` carries a dtype change; the pass must leave it. + to_dim_order_copies_before = _count(ep.graph_module, _TO_DIM_ORDER_COPY) + + gm, modified = _run_pass(ep) + + assert not modified + assert _count(gm, _TO_DIM_ORDER_COPY) == to_dim_order_copies_before + + @pytest.mark.parametrize( + "extra_kwarg", + [ + # dtype differs from the source tensor's dtype -> cast must not be silently dropped. + {"dtype": torch.float64}, + # layout, device, and pin_memory are in _ALLOWED_KWARGS and are allowed through + # when their value matches the source tensor (which is always the case on CPU + # hardware). The whitelist guards against UNKNOWN kwargs; value checks guard + # against CHANGED TensorOptions. These cases verify that a differing dtype is + # caught; layout/device/pin_memory rejection can only be tested by providing a + # value that actually differs from the source tensor (not possible on CPU-only). + ], + ) + def test__not_applied__dtype_cast_kwarg(self, extra_kwarg): + # Verify that a `_to_dim_order_copy` carrying a dtype change is not replaced. + # The guard must not drop the cast silently, so replacement is blocked. + g = torch.fx.Graph() + ph = g.placeholder("x") + ph.meta["val"] = torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + kwargs = {"dim_order": [0, 1, 2, 3], **extra_kwarg} + clone = g.call_function(_TO_DIM_ORDER_COPY, args=(ph,), kwargs=kwargs) + clone.meta["val"] = torch.randn(1, 3, 8, 8) + g.output((clone,)) + gm = torch.fx.GraphModule({}, g) + + result = ReplaceChannelsLastInputClones()(gm) + + assert not result.modified + assert _count(result.graph_module, _TO_DIM_ORDER_COPY) == 1 + + +class TestReplaceChannelsLastInputCloneDimOrders: + """These tests use channels last example inputs for export and apply the `EnforceContiguousDimOrder` pass which + inserts a `clone_dim_order` operator right after the model inputs to make the dim order contiguous. This is + precisely the intended use-case for the `ReplaceChannelsLastInputClones`. + """ + + def test_single_input(self): + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(SingleInputModule(), example_inputs) + assert _count(ep.graph_module, _CLONE_DIM_ORDER) == 0 + + # Turn the model contiguous and create the input `clone_dim_order` operator. + # SingleInputModule preserves channels-last format (avg_pool + relu), so ECDO inserts + # both an input boundary clone and an output boundary clone. + res1 = EnforceContiguousDimOrder()(ep.graph_module) + assert res1.modified + assert ( + _count(res1.graph_module, _CLONE_DIM_ORDER) == 2 + ) # 1 input + 1 output boundary + + output_before = ep.module()(*example_inputs) + gm, modified = _run_pass(res1) + + assert modified + assert _count(gm, _TO_DIM_ORDER_COPY) == 0 + # Input boundary clone is replaced by permutes; the output boundary clone remains. + assert _count(gm, _CLONE_DIM_ORDER) == 1 + + nodes = list(gm.graph.nodes) + _assert_expected_result_pattern(*nodes[:3]) + + outputs_after = gm(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test_multi_input(self): + num_inputs = 3 + example_inputs = tuple( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + for _ in range(num_inputs) + ) + + ep = _export_to_edge(MultiInputModule(), example_inputs) + assert _count(ep.graph_module, _CLONE_DIM_ORDER) == 0 + + # Turn the model contiguous and create the input `clone_dim_order` operator. + # MultiInputModule (avg_pool + relu) preserves channels-last, so ECDO inserts one + # input boundary clone per input PLUS one output boundary clone. + res1 = EnforceContiguousDimOrder()(ep.graph_module) + assert res1.modified + assert _count(res1.graph_module, _CLONE_DIM_ORDER) == num_inputs + 1 + + output_before = ep.module()(*example_inputs) + res2 = ReplaceChannelsLastInputClones()(res1.graph_module) + + assert res2.modified + assert _count(res2.graph_module, _TO_DIM_ORDER_COPY) == 0 + # Input boundary clones are replaced by permutes; the output boundary clone remains. + assert _count(res2.graph_module, _CLONE_DIM_ORDER) == 1 + + nodes = list(res2.graph_module.graph.nodes) + for i in range(num_inputs): + # Each input should have the expected pattern. + _assert_expected_result_pattern(*nodes[i * num_inputs : i * num_inputs + 3]) + + outputs_after = res2.graph_module(*example_inputs)[0] + assert torch.allclose(output_before, outputs_after) + + def test_idempotency(self): + """Running the pass twice must not alter the graph on the second run.""" + example_inputs = ( + torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last), + ) + + ep = _export_to_edge(SingleInputModule(), example_inputs) + res1 = EnforceContiguousDimOrder()(ep.graph_module) + assert res1.modified + + res2 = ReplaceChannelsLastInputClones()(res1.graph_module) + assert res2.modified + + # Second pass: input boundary clones are gone; only the output boundary clone remains. + res3 = ReplaceChannelsLastInputClones()(res2.graph_module) + assert not res3.modified + # Input boundary clones have been replaced; the output boundary clone is untouched. + assert _count(res3.graph_module, _CLONE_DIM_ORDER) == 1 + assert _count(res3.graph_module, _TO_DIM_ORDER_COPY) == 0 + assert _count(res3.graph_module, _ATEN_PERMUTE_COPY) == _count( + res2.graph_module, _ATEN_PERMUTE_COPY + ) + assert _count(res3.graph_module, _CHANNELS_LAST_PERMUTE_COPY) == _count( + res2.graph_module, _CHANNELS_LAST_PERMUTE_COPY + ) + + +class TestGuardPredicates: + """Unit tests for the `call_operator` guard conditions using hand-built graphs.""" + + def test_is_4d_channels_last(self): + assert _is_4d_channels_last([0, 2, 3, 1]) + assert not _is_4d_channels_last([0, 1, 2, 3]) + assert not _is_4d_channels_last([0, 2, 1]) + assert not _is_4d_channels_last([0, 2, 3, 4, 1]) + + def test_is_4d_contiguous(self): + assert _is_4d_contiguous([0, 1, 2, 3]) + assert not _is_4d_contiguous([0, 2, 3, 1]) + assert not _is_4d_contiguous([0, 1, 2]) + assert not _is_4d_contiguous([0, 1, 2, 3, 4]) From c7bbe0be82d6759b4476cf2d7e02c7109fc873f6 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson <38941629+Sebastian-Larsson@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:37:54 +0200 Subject: [PATCH 016/190] Arm backend: Cast integer comparisons in FP profile (#22503) TOSA section 2.13 permits PRO-FP casts from signed integer types to fp16 or fp32. Cast int8 comparison inputs to fp16 and int16 inputs to fp32, where every value is exactly representable. This allows the comparisons to be lowered without changing their results. Signed-off-by: Sebastian Larsson --- backends/arm/_passes/__init__.py | 1 + backends/arm/_passes/arm_pass_manager.py | 2 + .../cast_int_comparison_inputs_pass.py | 61 +++++++ .../test_cast_int_comparison_inputs_pass.py | 152 ++++++++++++++++++ 4 files changed, 216 insertions(+) create mode 100644 backends/arm/_passes/cast_int_comparison_inputs_pass.py create mode 100644 backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 2b9fcc2e7eb..cd99b928aea 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -17,6 +17,7 @@ from .canonicalize_gather_pass import CanonicalizeGatherPass # noqa from .canonicalize_view_copy_permute_pass import CanonicalizeViewCopyPermutePass # noqa from .cast_int64_pass import CastInt64BuffersToInt32Pass # noqa +from .cast_int_comparison_inputs_pass import CastIntComparisonInputsPass # noqa from .cast_to_int32_pass import CastToInt32Pass # noqa from .constant_folding_pass import ConstantFoldingPass # noqa from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index ef30bec9f00..59d4989fb90 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -17,6 +17,7 @@ CanonicalizeGatherPass, CanonicalizeViewCopyPermutePass, CastInt64BuffersToInt32Pass, + CastIntComparisonInputsPass, CastToInt32Pass, ComputeConstantOpsAOTPass, ConstantFoldingPass, @@ -603,6 +604,7 @@ def _tosa_pipeline( self.add_passes( [ ReplaceScalarWithTensorByProfilePass(), + CastIntComparisonInputsPass(), RewriteLeLtToGeGtPass(), DecomposeLeakyReLUPass(), # Emits full_like so before ConvertFullLikeToFullPass DecomposePReLUPass(), diff --git a/backends/arm/_passes/cast_int_comparison_inputs_pass.py b/backends/arm/_passes/cast_int_comparison_inputs_pass.py new file mode 100644 index 00000000000..6a0b19b0703 --- /dev/null +++ b/backends/arm/_passes/cast_int_comparison_inputs_pass.py @@ -0,0 +1,61 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Set, Type + +import torch + +from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.backends.arm.tosa.specification import get_context_spec +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +class CastIntComparisonInputsPass(ArmOpTargetedPass): + """Cast integer comparison inputs to a lossless floating-point type.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + + target_ops = { + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.lt.Tensor, + } + castable_dtypes = {torch.int8, torch.int16} + + def should_run_pass(self, graph_module: torch.fx.GraphModule) -> bool: + tosa_spec = get_context_spec() + return ( + tosa_spec.support_float() + and not tosa_spec.support_integer() + and super().should_run_pass(graph_module) + ) + + def call_operator(self, op, args, kwargs, meta): + if op not in self.target_ops: + return super().call_operator(op, args, kwargs, meta) + + if not all(arg.data.dtype in self.castable_dtypes for arg in args): + return super().call_operator(op, args, kwargs, meta) + + cast_dtype = ( + torch.float16 + if all(arg.data.dtype == torch.int8 for arg in args) + else torch.float32 + ) + casted_args = [] + for arg in args: + casted_args.append( + super().call_operator( + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + (arg,), + {"dtype": cast_dtype}, + meta, + ) + ) + return super().call_operator(op, tuple(casted_args), kwargs, meta) diff --git a/backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py b/backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py new file mode 100644 index 00000000000..ca3f8cb5cd9 --- /dev/null +++ b/backends/arm/test/passes/test_cast_int_comparison_inputs_pass.py @@ -0,0 +1,152 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import operator +from collections.abc import Callable + +import pytest +import torch +from executorch.backends.arm._passes import CastIntComparisonInputsPass +from executorch.backends.arm.test.tester.test_pipeline import ( + PassPipeline, + TosaPipelineFP, +) +from executorch.backends.test.harness.stages import StageType +from executorch.exir.dialects._ops import ops as edge_ops + + +class Comparison(torch.nn.Module): + def __init__(self, op: Callable[[torch.Tensor, torch.Tensor], torch.Tensor]): + super().__init__() + self.op = op + + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return self.op(x, y) + + +class ScalarComparison(torch.nn.Module): + def __init__(self, op: Callable[[torch.Tensor, int], torch.Tensor]): + super().__init__() + self.op = op + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.op(x, 0) + + +comparison_ops = { + "eq": operator.eq, + "ne": operator.ne, + "ge": operator.ge, + "gt": operator.gt, + "le": operator.le, + "lt": operator.lt, +} +aten_ops = { + "eq": "torch.ops.aten.eq.Tensor", + "ne": "torch.ops.aten.ne.Tensor", + "ge": "torch.ops.aten.ge.Tensor", + "gt": "torch.ops.aten.gt.Tensor", + "le": "torch.ops.aten.le.Tensor", + "lt": "torch.ops.aten.lt.Tensor", +} +aten_scalar_ops = { + name: target.replace("Tensor", "Scalar") for name, target in aten_ops.items() +} +exir_ops = { + name: f"executorch_exir_dialects_edge__ops_aten_{name}_Tensor" + for name in comparison_ops +} +exir_scalar_ops = { + name: target.replace("Tensor", "Scalar") for name, target in exir_ops.items() +} + + +def comparison_inputs(dtype: torch.dtype) -> tuple[torch.Tensor, torch.Tensor]: + limits = torch.iinfo(dtype) + return ( + torch.tensor( + [limits.min, limits.min + 1, limits.max - 1, limits.max], dtype=dtype + ), + torch.tensor( + [limits.min + 1, limits.min, limits.max, limits.max - 1], dtype=dtype + ), + ) + + +@pytest.mark.parametrize("op", comparison_ops.values(), ids=comparison_ops.keys()) +@pytest.mark.parametrize( + ("dtypes", "expected_dtype"), + ( + ((torch.int8, torch.int8), torch.float16), + ((torch.int16, torch.int16), torch.float32), + ((torch.int8, torch.int16), torch.float32), + ), +) +def test_cast_int_comparison_inputs(op, dtypes, expected_dtype) -> None: + inputs = ( + comparison_inputs(dtypes[0])[0], + comparison_inputs(dtypes[1])[1], + ) + pipeline = PassPipeline( + Comparison(op), + inputs, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2 + }, + pass_list=[CastIntComparisonInputsPass], + ) + pipeline.run() + + graph_module = ( + pipeline.tester.get_artifact(StageType.RUN_PASSES) + .exported_program() + .graph_module + ) + cast_op = edge_ops.edge.dim_order_ops._to_dim_order_copy.default + cast_nodes = [node for node in graph_module.graph.nodes if node.target == cast_op] + assert len(cast_nodes) == 2 + assert all(node.kwargs["dtype"] == expected_dtype for node in cast_nodes) + + +def test_cast_int_comparison_inputs_keeps_int32() -> None: + inputs = ( + torch.tensor([2**24, 2**24 + 1], dtype=torch.int32), + torch.tensor([2**24 + 1, 2**24], dtype=torch.int32), + ) + pipeline = PassPipeline( + Comparison(operator.eq), + inputs, + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default" + ], + pass_list=[CastIntComparisonInputsPass], + ) + pipeline.run() + + +@pytest.mark.parametrize("name", comparison_ops) +@pytest.mark.parametrize("dtype", (torch.int8, torch.int16)) +def test_int_comparison_tosa_fp(name, dtype) -> None: + inputs = comparison_inputs(dtype) + pipeline = TosaPipelineFP( + Comparison(comparison_ops[name]), + inputs, + aten_ops[name], + exir_ops[name], + ) + pipeline.run() + + +@pytest.mark.parametrize("name", comparison_ops) +@pytest.mark.parametrize("dtype", (torch.int8, torch.int16)) +def test_int_scalar_comparison_tosa_fp(name, dtype) -> None: + inputs = (comparison_inputs(dtype)[0],) + pipeline = TosaPipelineFP( + ScalarComparison(comparison_ops[name]), + inputs, + aten_scalar_ops[name], + exir_scalar_ops[name], + ) + pipeline.run() From 1e3b7fb263a1b4e90a51170bd2f3f2e685669dc3 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Thu, 3 Sep 2026 13:04:23 +0100 Subject: [PATCH 017/190] Arm backend: Support rank-3 max_pool2d inputs (#22441) Normalize unbatched max_pool2d inputs from [C, H, W] to [1, C, H, W] before decomposition, then restore the rank-3 output. Enable rank-3 support checks and add FP, INT, and VGF test coverage, including quantized VGF execution. Change-Id: Ib8852053b5248601d59685669523079fa19c0955 Signed-off-by: Yufeng Shi --- backends/arm/_passes/__init__.py | 3 + backends/arm/_passes/arm_pass_manager.py | 2 + .../normalize_max_pool2d_input_rank_pass.py | 97 +++++++++++++++++++ .../arm/operator_support/pool_2d_support.py | 8 ++ backends/arm/test/ops/test_max_pool.py | 24 ++++- ...st_normalize_max_pool2d_input_rank_pass.py | 87 +++++++++++++++++ 6 files changed, 217 insertions(+), 4 deletions(-) create mode 100644 backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py create mode 100644 backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index cd99b928aea..d287d9d306c 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -170,6 +170,9 @@ from .normalize_index_put_none_indices_pass import ( # noqa NormalizeIndexPutNoneIndicesPass, ) +from .normalize_max_pool2d_input_rank_pass import ( # noqa + NormalizeMaxPool2dInputRankPass, +) from .normalize_while_initial_args_pass import NormalizeWhileInitialArgsPass # noqa from .promote_bool_operands_pass import PromoteBoolOperandsPass # noqa from .propagate_view_copy_permute_pass import ( # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 59d4989fb90..b141deb698f 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -138,6 +138,7 @@ NormalizeDelegateIOLayoutPass, NormalizeIndexPutBoolIndexTensorPass, NormalizeIndexPutNoneIndicesPass, + NormalizeMaxPool2dInputRankPass, NormalizeTransformInputPlaceholdersPass, NormalizeWhileInitialArgsPass, PromoteBoolOperandsPass, @@ -637,6 +638,7 @@ def _tosa_pipeline( UnsqueezeBeforeRepeatPass(), DecomposeCumsumPass(exported_program), DecomposeAsStridedCopyPass(), + NormalizeMaxPool2dInputRankPass(), DecomposeMaxPool2dPass(), DecomposeLargeStrideMaxPool2dForU55Pass(), SizeAdjustInputPass(), diff --git a/backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py b/backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py new file mode 100644 index 00000000000..98d417f6f56 --- /dev/null +++ b/backends/arm/_passes/normalize_max_pool2d_input_rank_pass.py @@ -0,0 +1,97 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Set, Type + +import torch +from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.backends.arm._passes.arm_pass_utils import ( + create_node, + get_first_fake_tensor, +) +from executorch.backends.arm._passes.convert_squeezes_to_view import ( + ConvertSqueezesToViewPass, +) +from executorch.backends.arm._passes.decompose_maxpool2d_with_dilation_pass import ( + DecomposeMaxPool2dPass, +) +from executorch.backends.arm._passes.rewrite_max_pool2d_pass import RewriteMaxPool2dPass +from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import Node + + +class NormalizeMaxPool2dInputRankPass(ArmOpTargetedPass): + """Normalize unbatched rank-3 max_pool2d inputs to rank 4. + + Unsqueeze inputs from ``[C, H, W]`` to ``[1, C, H, W]``. + + Squeeze the leading dimension after pooling to restore the rank-3 output. + + The complete shape transformation is:: + + [C, H, W] + -> unsqueeze(0) -> [1, C, H, W] + -> max_pool2d -> [1, C, H_out, W_out] + -> squeeze(0) -> [C, H_out, W_out] + + """ + + target_ops = (exir_ops.edge.aten.max_pool2d.default,) + _passes_required_after: Set[Type[ExportPass]] = { + ConvertSqueezesToViewPass, + DecomposeMaxPool2dPass, + RewriteMaxPool2dPass, + SizeAdjustInputPass, + } + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + pool_nodes = graph.find_nodes(op="call_function", target=self.target_ops[0]) + + for pool_node in pool_nodes: + input_node = pool_node.args[0] + if not isinstance(input_node, Node): + raise RuntimeError("Expected max_pool2d input to be a node") + + input_fake = get_first_fake_tensor(input_node) + if input_fake.dim() != 3: + continue + + output_fake = get_first_fake_tensor(pool_node) + with graph.inserting_before(pool_node): + unsqueeze = create_node( + graph, + exir_ops.edge.aten.unsqueeze_copy.default, + args=(input_node, 0), + from_node=pool_node, + inherit_qparams=False, + ) + unsqueeze.meta["val"] = input_fake.unsqueeze(0) + pool_node.replace_input_with(input_node, unsqueeze) + + pool_node.meta["val"] = output_fake.unsqueeze(0) + original_users = list(pool_node.users) + with graph.inserting_after(pool_node): + squeeze = create_node( + graph, + exir_ops.edge.aten.squeeze_copy.dims, + args=(pool_node, [0]), + from_node=pool_node, + inherit_qparams=False, + ) + squeeze.meta["val"] = output_fake + for user in original_users: + user.replace_input_with(pool_node, squeeze) + + modified = True + + if modified: + graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/arm/operator_support/pool_2d_support.py b/backends/arm/operator_support/pool_2d_support.py index a022ed942fd..03b52fbb85f 100644 --- a/backends/arm/operator_support/pool_2d_support.py +++ b/backends/arm/operator_support/pool_2d_support.py @@ -228,6 +228,14 @@ def is_node_tosa_supported(self, node: fx.Node, tosa_spec: TosaSpecification): """ shape = cast(torch.Tensor, node.all_input_nodes[0].meta["val"]).shape + if len(shape) == 3: + shape = torch.Size((1, *shape)) + elif len(shape) != 4: + self.reporter.report_reject( + node, f"Maxpool2d needs rank 3 or 4 input, got shape {list(shape)}" + ) + return False + kernel = cast(tuple[int, int], node.args[1]) stride = cast(tuple[int, int], node.args[2]) padding = cast(tuple[int, int], node.args[3]) if len(node.args) >= 4 else (0, 0) diff --git a/backends/arm/test/ops/test_max_pool.py b/backends/arm/test/ops/test_max_pool.py index 22dfe09b070..1dbe5536c31 100644 --- a/backends/arm/test/ops/test_max_pool.py +++ b/backends/arm/test/ops/test_max_pool.py @@ -82,6 +82,12 @@ [3, 2, 1], ), } + +test_data_suite_rank3 = { + "bev_class_slice": lambda: (torch.rand(1, 8, 8), [1, 1, 0]), + "spatial_pool": lambda: (torch.rand(3, 9, 11), [3, 2, 1]), +} + test_data_suite_fp8 = { "rand_fp8e4m3": lambda: ( torch.rand(1, 8, 20, 20).to(torch.float8_e4m3fn), @@ -168,7 +174,11 @@ def forward(self, x): @common.parametrize( - "test_data", test_data_suite | test_data_suite_fp16 | test_data_suite_bf16 + "test_data", + test_data_suite + | test_data_suite_rank3 + | test_data_suite_fp16 + | test_data_suite_bf16, ) def test_max_pool2d_tosa_FP(test_data: torch.Tensor): test_data, model_params = test_data() @@ -197,7 +207,7 @@ def test_max_pool2d_tosa_FP_fp8(test_data: torch.Tensor): pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("test_data", test_data_suite | test_data_suite_rank3) def test_max_pool2d_tosa_INT(test_data: torch.Tensor): test_data, model_params = test_data() pipeline = TosaPipelineINT[input_t1]( @@ -374,22 +384,28 @@ def test_max_pool2d_tosa_INT_dilation(test_data): # VGF tests @common.parametrize( - "test_data", test_data_suite | test_data_suite_bf16 | test_data_suite_fp16 + "test_data", + test_data_suite + | test_data_suite_rank3 + | test_data_suite_bf16 + | test_data_suite_fp16, ) @common.SkipIfNoModelConverter def test_max_pool2d_vgf_no_quant(test_data: torch.Tensor): test_data, model_params = test_data() + run_on_vulkan_runtime = test_data.dim() == 4 pipeline = VgfPipeline[input_t1]( MaxPool2d(*model_params), (test_data,), aten_op, exir_op, quantize=False, + run_on_vulkan_runtime=run_on_vulkan_runtime, ) pipeline.run() -@common.parametrize("test_data", test_data_suite) +@common.parametrize("test_data", test_data_suite | test_data_suite_rank3) @common.SkipIfNoModelConverter def test_max_pool2d_vgf_quant(test_data: torch.Tensor): test_data, model_params = test_data() diff --git a/backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py b/backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py new file mode 100644 index 00000000000..11aee89778c --- /dev/null +++ b/backends/arm/test/passes/test_normalize_max_pool2d_input_rank_pass.py @@ -0,0 +1,87 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch +from executorch.backends.arm._passes import ( + NormalizeMaxPool2dInputRankPass, + RemoveGetItemPass, +) +from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.test.harness.stages import StageType +from executorch.exir.dialects._ops import ops as exir_ops + + +input_t = Tuple[torch.Tensor] + + +class MaxPool2d(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.nn.functional.max_pool2d( + x, + kernel_size=(3, 2), + stride=(2, 1), + padding=(1, 0), + dilation=(1, 1), + ceil_mode=True, + ) + + +def test_normalize_rank3_max_pool2d_input() -> None: + pipeline = PassPipeline[input_t]( + MaxPool2d(), + (torch.rand(3, 9, 11),), + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1, + "executorch_exir_dialects_edge__ops_aten_max_pool2d_default": 1, + "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims": 1, + }, + pass_list=[RemoveGetItemPass, NormalizeMaxPool2dInputRankPass], + ) + pipeline.run() + + exported_program = pipeline.tester.get_artifact( + StageType.RUN_PASSES + ).exported_program() + pool_node = next( + node + for node in exported_program.graph.nodes + if node.target == exir_ops.edge.aten.max_pool2d.default + ) + unsqueeze_node = pool_node.args[0] + assert isinstance(unsqueeze_node, torch.fx.Node) + assert unsqueeze_node.target == exir_ops.edge.aten.unsqueeze_copy.default + assert unsqueeze_node.args[1] == 0 + assert tuple(pool_node.args[1]) == (3, 2) + assert tuple(pool_node.args[2]) == (2, 1) + assert tuple(pool_node.args[3]) == (1, 0) + assert tuple(pool_node.args[4]) == (1, 1) + assert pool_node.args[5] is True + + squeeze_node = next(iter(pool_node.users)) + assert squeeze_node.target == exir_ops.edge.aten.squeeze_copy.dims + assert squeeze_node.args == (pool_node, [0]) + + +def test_normalize_rank4_max_pool2d_input_is_noop() -> None: + PassPipeline[input_t]( + MaxPool2d(), + (torch.rand(1, 3, 9, 11),), + ops_before_pass={ + "executorch_exir_dialects_edge__ops_aten_max_pool2d_with_indices_default": 1, + }, + ops_after_pass={ + "executorch_exir_dialects_edge__ops_aten_max_pool2d_default": 1, + }, + ops_not_after_pass=[ + "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default", + "executorch_exir_dialects_edge__ops_aten_squeeze_copy_dims", + ], + pass_list=[RemoveGetItemPass, NormalizeMaxPool2dInputRankPass], + ).run() From 7e9ac9bc8e07db1d38fd70115634066ea2a8434e Mon Sep 17 00:00:00 2001 From: SaoirseARM <44364573+SaoirseARM@users.noreply.github.com> Date: Thu, 3 Sep 2026 13:54:41 +0100 Subject: [PATCH 018/190] Arm backend: Clean up of TOSA dialect operators (#22515) - move test_tosa_* to tosa_dialect folder - move comparison operators out of binary operators - add aten->tosa for comparison operators Signed-off-by: Saoirse Stewart --- .../arm/_passes/aten_to_tosa_comparison.py | 27 ++++++ .../_passes/aten_to_tosa_tensor_operators.py | 6 -- backends/arm/_passes/exir_to_tosa_pass.py | 17 +++- .../test_tosa_data_layout_visitors.py | 0 .../test_tosa_dialect_activation.py | 0 .../test_tosa_dialect_argmax.py | 0 .../test_tosa_dialect_binary_ops.py | 52 ---------- .../test_tosa_dialect_comparison.py | 97 +++++++++++++++++++ .../test_tosa_dialect_fft.py | 0 .../test_tosa_dialect_max_pool2d_adaptive.py | 0 .../test_tosa_dialect_scatter.py | 0 .../test_tosa_dialect_unary_ops.py | 0 .../test_tosa_shape_node_visitors.py | 0 .../test_tosa_shape_support.py | 0 .../misc/{ => tosa_dialect}/test_tosa_spec.py | 0 backends/arm/test/targets.bzl | 2 +- backends/arm/tosa/dialect/__init__.py | 1 + backends/arm/tosa/dialect/ops/_common.py | 12 +++ .../tosa/dialect/ops/binary_elementwise.py | 74 ++++---------- backends/arm/tosa/dialect/ops/comparison.py | 78 +++++++++++++++ 20 files changed, 247 insertions(+), 119 deletions(-) create mode 100644 backends/arm/_passes/aten_to_tosa_comparison.py rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_data_layout_visitors.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_dialect_activation.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_dialect_argmax.py (100%) create mode 100644 backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_dialect_fft.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_dialect_max_pool2d_adaptive.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_dialect_scatter.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_dialect_unary_ops.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_shape_node_visitors.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_shape_support.py (100%) rename backends/arm/test/misc/{ => tosa_dialect}/test_tosa_spec.py (100%) create mode 100644 backends/arm/tosa/dialect/ops/comparison.py diff --git a/backends/arm/_passes/aten_to_tosa_comparison.py b/backends/arm/_passes/aten_to_tosa_comparison.py new file mode 100644 index 00000000000..d444ce47a96 --- /dev/null +++ b/backends/arm/_passes/aten_to_tosa_comparison.py @@ -0,0 +1,27 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.transforms.aten_to_dialect_pass import ( + AtenToDialectPass, + DialectNodeSpec, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Node + + +def rewrite_comparison_operator( + node: Node, pass_: AtenToDialectPass +) -> DialectNodeSpec | None: + match node.target: + case exir_ops.edge.aten.eq.Tensor: + target = exir_ops.backend.tosa.EQUAL.default + case exir_ops.edge.aten.ge.Tensor: + target = exir_ops.backend.tosa.GREATER_EQUAL.default + case exir_ops.edge.aten.gt.Tensor: + target = exir_ops.backend.tosa.GREATER.default + case _: + return None + + return DialectNodeSpec(target, node.args, dict(node.kwargs)) diff --git a/backends/arm/_passes/aten_to_tosa_tensor_operators.py b/backends/arm/_passes/aten_to_tosa_tensor_operators.py index 6705514c606..84afb096c4e 100644 --- a/backends/arm/_passes/aten_to_tosa_tensor_operators.py +++ b/backends/arm/_passes/aten_to_tosa_tensor_operators.py @@ -56,12 +56,6 @@ def rewrite_binary_operator( target = exir_ops.backend.tosa.ARITHMETIC_RIGHT_SHIFT.default case exir_ops.edge.aten.bitwise_xor.Tensor: target = exir_ops.backend.tosa.BITWISE_XOR.default - case exir_ops.edge.aten.eq.Tensor: - target = exir_ops.backend.tosa.EQUAL.default - case exir_ops.edge.aten.ge.Tensor: - target = exir_ops.backend.tosa.GREATER_EQUAL.default - case exir_ops.edge.aten.gt.Tensor: - target = exir_ops.backend.tosa.GREATER.default case exir_ops.edge.aten.logical_and.default: target = exir_ops.backend.tosa.LOGICAL_AND.default case exir_ops.edge.aten.logical_or.default: diff --git a/backends/arm/_passes/exir_to_tosa_pass.py b/backends/arm/_passes/exir_to_tosa_pass.py index e03b880dfd5..b0d6b11f857 100644 --- a/backends/arm/_passes/exir_to_tosa_pass.py +++ b/backends/arm/_passes/exir_to_tosa_pass.py @@ -9,6 +9,9 @@ from executorch.backends.arm._passes.aten_to_tosa_activation_functions import ( get_activation_replacement, ) +from executorch.backends.arm._passes.aten_to_tosa_comparison import ( + rewrite_comparison_operator, +) from executorch.backends.arm._passes.aten_to_tosa_data_layout import ( rewrite_data_layout_operator, ) @@ -79,9 +82,6 @@ def _get_fft_replacement( exir_ops.edge.aten.bitwise_or.Tensor, exir_ops.edge.aten.bitwise_right_shift.Tensor, exir_ops.edge.aten.bitwise_xor.Tensor, - exir_ops.edge.aten.eq.Tensor, - exir_ops.edge.aten.ge.Tensor, - exir_ops.edge.aten.gt.Tensor, exir_ops.edge.aten.logical_and.default, exir_ops.edge.aten.logical_or.default, exir_ops.edge.aten.logical_xor.default, @@ -97,6 +97,17 @@ def _get_binary_operator_replacement( return rewrite_binary_operator(node, pass_) +@register_dialect_substitutions( + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, +) +def _get_comparison_operator_replacement( + node: Node, pass_: AtenToDialectPass +) -> DialectNodeSpec | None: + return rewrite_comparison_operator(node, pass_) + + @register_dialect_substitutions( exir_ops.edge.aten.abs.default, exir_ops.edge.aten.bitwise_not.default, diff --git a/backends/arm/test/misc/test_tosa_data_layout_visitors.py b/backends/arm/test/misc/tosa_dialect/test_tosa_data_layout_visitors.py similarity index 100% rename from backends/arm/test/misc/test_tosa_data_layout_visitors.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_data_layout_visitors.py diff --git a/backends/arm/test/misc/test_tosa_dialect_activation.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_activation.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_activation.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_activation.py diff --git a/backends/arm/test/misc/test_tosa_dialect_argmax.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_argmax.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_argmax.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_argmax.py diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py index 920454f5a9b..c4a811f3a2c 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_binary_ops.py @@ -81,33 +81,6 @@ def _to_fake(mode: FakeTensorMode, *values): (2, 3), torch.int8, ), - pytest.param( - "EQUAL", - "TOSA-1.1+INT", - torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), - torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), - {}, - (2, 4, 3), - torch.bool, - ), - pytest.param( - "GREATER", - "TOSA-1.1+FP", - torch.randn((2, 1, 3), dtype=torch.float32), - torch.randn((1, 4, 3), dtype=torch.float32), - {}, - (2, 4, 3), - torch.bool, - ), - pytest.param( - "GREATER_EQUAL", - "TOSA-1.1+INT", - torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), - torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), - {}, - (2, 4, 3), - torch.bool, - ), pytest.param( "INTDIV", "TOSA-1.1+INT", @@ -400,31 +373,6 @@ def test_intdiv_supports_int32_on_fp_profile() -> None: assert tuple(output.shape) == tuple(input1.shape) -def test_equal_rejects_int8() -> None: - input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) - input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) - - with TosaLoweringContext( - TosaSpecification.create_from_string("TOSA-1.1+INT") - ), FakeTensorMode() as mode: - with pytest.raises(TosaValueError, match="Unsupported dtype"): - exir_ops.backend.tosa.EQUAL.default(*_to_fake(mode, input1, input2)) - - -@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"]) -def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None: - input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32) - input2 = torch.randint(1, 8, (2, 3), dtype=torch.int32) - - with TosaLoweringContext( - TosaSpecification.create_from_string("TOSA-1.1+FP") - ), FakeTensorMode() as mode: - with pytest.raises(TosaValueError, match="doesn't support int32"): - getattr(exir_ops.backend.tosa, op_name).default( - *_to_fake(mode, input1, input2) - ) - - @pytest.mark.parametrize("op_name", ["MAXIMUM", "MINIMUM"]) def test_extrema_ops_reject_int32_on_fp_profile(op_name: str) -> None: input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32) diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py new file mode 100644 index 00000000000..6fbf25ed380 --- /dev/null +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_comparison.py @@ -0,0 +1,97 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import executorch.backends.arm.tosa.dialect # noqa: F401 +import pytest +import torch +from executorch.backends.arm.tosa.dialect.lib import TosaValueError +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode + + +def _to_fake(mode: FakeTensorMode, *values): + return [ + mode.from_tensor(value) if isinstance(value, torch.Tensor) else value + for value in values + ] + + +@pytest.mark.parametrize( + ( + "op_name", + "spec", + "input1", + "input2", + "expected_shape", + ), + [ + pytest.param( + "EQUAL", + "TOSA-1.1+INT", + torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), + torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), + (2, 4, 3), + ), + pytest.param( + "GREATER", + "TOSA-1.1+FP", + torch.randn((2, 1, 3), dtype=torch.float32), + torch.randn((1, 4, 3), dtype=torch.float32), + (2, 4, 3), + ), + pytest.param( + "GREATER_EQUAL", + "TOSA-1.1+INT", + torch.randint(1, 16, (2, 1, 3), dtype=torch.int32), + torch.randint(1, 8, (1, 4, 3), dtype=torch.int32), + (2, 4, 3), + ), + ], +) +def test_tosa_comparison_ops( + op_name: str, + spec: str, + input1: torch.Tensor, + input2: torch.Tensor, + expected_shape: tuple[int, ...], +) -> None: + with TosaLoweringContext( + TosaSpecification.create_from_string(spec) + ), FakeTensorMode() as mode: + output = getattr(exir_ops.backend.tosa, op_name).default( + *_to_fake(mode, input1, input2) + ) + + assert output.dtype == torch.bool + assert tuple(output.shape) == expected_shape + + +def test_equal_rejects_int8() -> None: + input1 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) + input2 = torch.randint(-8, 8, (2, 3), dtype=torch.int8) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+INT") + ), FakeTensorMode() as mode: + with pytest.raises(TosaValueError, match="Unsupported dtype"): + exir_ops.backend.tosa.EQUAL.default(*_to_fake(mode, input1, input2)) + + +@pytest.mark.parametrize("op_name", ["EQUAL", "GREATER", "GREATER_EQUAL"]) +def test_compare_ops_reject_int32_on_fp_profile(op_name: str) -> None: + input1 = torch.randint(1, 16, (2, 3), dtype=torch.int32) + input2 = torch.randint(1, 8, (2, 3), dtype=torch.int32) + + with TosaLoweringContext( + TosaSpecification.create_from_string("TOSA-1.1+FP") + ), FakeTensorMode() as mode: + with pytest.raises(TosaValueError, match="doesn't support int32"): + getattr(exir_ops.backend.tosa, op_name).default( + *_to_fake(mode, input1, input2) + ) diff --git a/backends/arm/test/misc/test_tosa_dialect_fft.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_fft.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_fft.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_fft.py diff --git a/backends/arm/test/misc/test_tosa_dialect_max_pool2d_adaptive.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_max_pool2d_adaptive.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_max_pool2d_adaptive.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_max_pool2d_adaptive.py diff --git a/backends/arm/test/misc/test_tosa_dialect_scatter.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_scatter.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_scatter.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_scatter.py diff --git a/backends/arm/test/misc/test_tosa_dialect_unary_ops.py b/backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py similarity index 100% rename from backends/arm/test/misc/test_tosa_dialect_unary_ops.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_dialect_unary_ops.py diff --git a/backends/arm/test/misc/test_tosa_shape_node_visitors.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py similarity index 100% rename from backends/arm/test/misc/test_tosa_shape_node_visitors.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py diff --git a/backends/arm/test/misc/test_tosa_shape_support.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py similarity index 100% rename from backends/arm/test/misc/test_tosa_shape_support.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py diff --git a/backends/arm/test/misc/test_tosa_spec.py b/backends/arm/test/misc/tosa_dialect/test_tosa_spec.py similarity index 100% rename from backends/arm/test/misc/test_tosa_spec.py rename to backends/arm/test/misc/tosa_dialect/test_tosa_spec.py diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 1b08a0ec4ef..aeb783045a9 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -65,7 +65,7 @@ def define_arm_tests(): "misc/tosa_dialect/test_tosa_dialect_mxfp_conv2d.py", "misc/tosa_dialect/test_tosa_dialect_mxfp_linear.py", "misc/tosa_dialect/test_tosa_resize.py", - "misc/test_tosa_spec.py", + "misc/tosa_dialect/test_tosa_spec.py", "misc/test_bn_relu_folding_qat.py", "misc/test_custom_partition.py", "misc/test_debug_hook.py", diff --git a/backends/arm/tosa/dialect/__init__.py b/backends/arm/tosa/dialect/__init__.py index 504323365ac..ee6e56b39fb 100644 --- a/backends/arm/tosa/dialect/__init__.py +++ b/backends/arm/tosa/dialect/__init__.py @@ -9,6 +9,7 @@ avg_pool2d, avg_pool2d_adaptive, binary_elementwise, + comparison, conv2d, conv2d_block_scaled, conv3d, diff --git a/backends/arm/tosa/dialect/ops/_common.py b/backends/arm/tosa/dialect/ops/_common.py index daeef30b097..28ac9ac84ef 100644 --- a/backends/arm/tosa/dialect/ops/_common.py +++ b/backends/arm/tosa/dialect/ops/_common.py @@ -29,6 +29,18 @@ def require_same_dtype(input1: torch.Tensor, input2: torch.Tensor, op: str) -> N ) +def binary_meta( + input1: torch.Tensor, + input2: torch.Tensor, + op: str, + *, + output_dtype: torch.dtype | None = None, +) -> torch.Tensor: + require_same_dtype(input1, input2, op) + output_shape = broadcast_shape(input1, input2, op) + return torch.empty(output_shape, dtype=output_dtype or input1.dtype) + + def validate_nan_mode(nan_mode: str, op: str) -> None: if nan_mode not in _VALID_NAN_MODES: raise TosaValueError( diff --git a/backends/arm/tosa/dialect/ops/binary_elementwise.py b/backends/arm/tosa/dialect/ops/binary_elementwise.py index 1a3f7222419..5c181ae0a80 100644 --- a/backends/arm/tosa/dialect/ops/binary_elementwise.py +++ b/backends/arm/tosa/dialect/ops/binary_elementwise.py @@ -6,8 +6,7 @@ import torch from executorch.backends.arm.tosa.dialect.lib import TosaValueError from executorch.backends.arm.tosa.dialect.ops._common import ( - broadcast_shape, - require_same_dtype, + binary_meta, validate_nan_mode, ) from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op @@ -37,18 +36,6 @@ def _raise_unsupported_profile(dtype: torch.dtype, op: str) -> None: ) -def _binary_meta( - input1: torch.Tensor, - input2: torch.Tensor, - op: str, - *, - output_dtype: torch.dtype | None = None, -) -> torch.Tensor: - require_same_dtype(input1, input2, op) - output_shape = broadcast_shape(input1, input2, op) - return torch.empty(output_shape, dtype=output_dtype or input1.dtype) - - def _require_int_profile_support(dtype: torch.dtype, op: str) -> None: if not get_context_spec().support_integer(): _raise_unsupported_profile(dtype, op) @@ -140,7 +127,7 @@ def _validate_and_infer_mul_output_dtype(dtype: torch.dtype) -> torch.dtype: # ) def ADD(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_add_sub_dtype(input1.dtype, "ADD") - return _binary_meta(input1, input2, "ADD") + return binary_meta(input1, input2, "ADD") @register_fake_tosa_op( @@ -154,7 +141,7 @@ def ARITHMETIC_RIGHT_SHIFT( round: bool = False, ) -> torch.Tensor: _validate_any_profile_int_dtype(input1.dtype, "ARITHMETIC_RIGHT_SHIFT") - return _binary_meta(input1, input2, "ARITHMETIC_RIGHT_SHIFT") + return binary_meta(input1, input2, "ARITHMETIC_RIGHT_SHIFT") @register_fake_tosa_op( @@ -163,7 +150,7 @@ def ARITHMETIC_RIGHT_SHIFT( ) def BITWISE_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bitwise_and_dtype(input1.dtype) - return _binary_meta(input1, input2, "BITWISE_AND") + return binary_meta(input1, input2, "BITWISE_AND") @register_fake_tosa_op( @@ -172,7 +159,7 @@ def BITWISE_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def BITWISE_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_int_dtype(input1.dtype, "BITWISE_OR") - return _binary_meta(input1, input2, "BITWISE_OR") + return binary_meta(input1, input2, "BITWISE_OR") @register_fake_tosa_op( @@ -181,34 +168,7 @@ def BITWISE_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def BITWISE_XOR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_int_dtype(input1.dtype, "BITWISE_XOR") - return _binary_meta(input1, input2, "BITWISE_XOR") - - -@register_fake_tosa_op( - "EQUAL(Tensor input1, Tensor input2) -> Tensor", - TosaSpecification.all_versions_and_profiles(), -) -def EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: - _validate_profile_int32_or_fp_dtype(input1.dtype, "EQUAL") - return _binary_meta(input1, input2, "EQUAL", output_dtype=torch.bool) - - -@register_fake_tosa_op( - "GREATER(Tensor input1, Tensor input2) -> Tensor", - TosaSpecification.all_versions_and_profiles(), -) -def GREATER(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: - _validate_profile_int32_or_fp_dtype(input1.dtype, "GREATER") - return _binary_meta(input1, input2, "GREATER", output_dtype=torch.bool) - - -@register_fake_tosa_op( - "GREATER_EQUAL(Tensor input1, Tensor input2) -> Tensor", - TosaSpecification.all_versions_and_profiles(), -) -def GREATER_EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: - _validate_profile_int32_or_fp_dtype(input1.dtype, "GREATER_EQUAL") - return _binary_meta(input1, input2, "GREATER_EQUAL", output_dtype=torch.bool) + return binary_meta(input1, input2, "BITWISE_XOR") @register_fake_tosa_op( @@ -217,7 +177,7 @@ def GREATER_EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def INTDIV(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_int32_dtype(input1.dtype, "INTDIV") - return _binary_meta(input1, input2, "INTDIV") + return binary_meta(input1, input2, "INTDIV") @register_fake_tosa_op( @@ -226,7 +186,7 @@ def INTDIV(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def LOGICAL_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bool_dtype(input1.dtype, "LOGICAL_AND") - return _binary_meta(input1, input2, "LOGICAL_AND") + return binary_meta(input1, input2, "LOGICAL_AND") @register_fake_tosa_op( @@ -235,7 +195,7 @@ def LOGICAL_AND(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def LOGICAL_LEFT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_any_profile_int_dtype(input1.dtype, "LOGICAL_LEFT_SHIFT") - return _binary_meta(input1, input2, "LOGICAL_LEFT_SHIFT") + return binary_meta(input1, input2, "LOGICAL_LEFT_SHIFT") @register_fake_tosa_op( @@ -244,7 +204,7 @@ def LOGICAL_LEFT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tens ) def LOGICAL_RIGHT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_any_profile_int_dtype(input1.dtype, "LOGICAL_RIGHT_SHIFT") - return _binary_meta(input1, input2, "LOGICAL_RIGHT_SHIFT") + return binary_meta(input1, input2, "LOGICAL_RIGHT_SHIFT") @register_fake_tosa_op( @@ -253,7 +213,7 @@ def LOGICAL_RIGHT_SHIFT(input1: torch.Tensor, input2: torch.Tensor) -> torch.Ten ) def LOGICAL_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bool_dtype(input1.dtype, "LOGICAL_OR") - return _binary_meta(input1, input2, "LOGICAL_OR") + return binary_meta(input1, input2, "LOGICAL_OR") @register_fake_tosa_op( @@ -262,7 +222,7 @@ def LOGICAL_OR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def LOGICAL_XOR(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_bool_dtype(input1.dtype, "LOGICAL_XOR") - return _binary_meta(input1, input2, "LOGICAL_XOR") + return binary_meta(input1, input2, "LOGICAL_XOR") @register_fake_tosa_op( @@ -277,7 +237,7 @@ def MAXIMUM( ) -> torch.Tensor: validate_nan_mode(nan_mode, "MAXIMUM") _validate_profile_int32_or_fp_dtype(input1.dtype, "MAXIMUM") - return _binary_meta(input1, input2, "MAXIMUM") + return binary_meta(input1, input2, "MAXIMUM") @register_fake_tosa_op( @@ -292,7 +252,7 @@ def MINIMUM( ) -> torch.Tensor: validate_nan_mode(nan_mode, "MINIMUM") _validate_profile_int32_or_fp_dtype(input1.dtype, "MINIMUM") - return _binary_meta(input1, input2, "MINIMUM") + return binary_meta(input1, input2, "MINIMUM") @register_fake_tosa_op( @@ -315,7 +275,7 @@ def MUL( op="MUL", ) - return _binary_meta(input1, input2, "MUL", output_dtype=output_dtype) + return binary_meta(input1, input2, "MUL", output_dtype=output_dtype) @register_fake_tosa_op( @@ -324,7 +284,7 @@ def MUL( ) def POW(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_fp_dtype(input1.dtype, "POW") - return _binary_meta(input1, input2, "POW") + return binary_meta(input1, input2, "POW") @register_fake_tosa_op( @@ -333,4 +293,4 @@ def POW(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: ) def SUB(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: _validate_add_sub_dtype(input1.dtype, "SUB") - return _binary_meta(input1, input2, "SUB") + return binary_meta(input1, input2, "SUB") diff --git a/backends/arm/tosa/dialect/ops/comparison.py b/backends/arm/tosa/dialect/ops/comparison.py new file mode 100644 index 00000000000..f01b9e0180b --- /dev/null +++ b/backends/arm/tosa/dialect/ops/comparison.py @@ -0,0 +1,78 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.backends.arm.tosa.dialect.lib import TosaValueError +from executorch.backends.arm.tosa.dialect.ops._common import binary_meta +from executorch.backends.arm.tosa.dialect.ops_registration import register_fake_tosa_op +from executorch.backends.arm.tosa.specification import ( + get_context_spec, + TosaSpecification, +) + +FP_DTYPES = (torch.float16, torch.float32) + + +def _dtype_name(dtype: torch.dtype) -> str: + return str(dtype).removeprefix("torch.") + + +def _raise_unsupported_dtype(dtype: torch.dtype, op: str) -> None: + raise TosaValueError(f"Unsupported dtype {dtype} for {op}", op=op) + + +def _raise_unsupported_profile(dtype: torch.dtype, op: str) -> None: + raise TosaValueError( + f"TOSA spec {get_context_spec()} doesn't support {_dtype_name(dtype)} for {op}", + op=op, + ) + + +def _validate_comparison_dtype(dtype: torch.dtype, op: str) -> None: + tosa_spec = get_context_spec() + + if dtype == torch.int32: + if not tosa_spec.support_integer(): + _raise_unsupported_profile(dtype, op) + return + + if dtype in FP_DTYPES: + if not tosa_spec.support_float(): + _raise_unsupported_profile(dtype, op) + return + + if dtype == torch.bfloat16: + if not (tosa_spec.support_float() and tosa_spec.support_extension("bf16")): + _raise_unsupported_profile(dtype, op) + return + + _raise_unsupported_dtype(dtype, op) + + +@register_fake_tosa_op( + "EQUAL(Tensor input1, Tensor input2) -> Tensor", + TosaSpecification.all_versions_and_profiles(), +) +def EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + _validate_comparison_dtype(input1.dtype, "EQUAL") + return binary_meta(input1, input2, "EQUAL", output_dtype=torch.bool) + + +@register_fake_tosa_op( + "GREATER(Tensor input1, Tensor input2) -> Tensor", + TosaSpecification.all_versions_and_profiles(), +) +def GREATER(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + _validate_comparison_dtype(input1.dtype, "GREATER") + return binary_meta(input1, input2, "GREATER", output_dtype=torch.bool) + + +@register_fake_tosa_op( + "GREATER_EQUAL(Tensor input1, Tensor input2) -> Tensor", + TosaSpecification.all_versions_and_profiles(), +) +def GREATER_EQUAL(input1: torch.Tensor, input2: torch.Tensor) -> torch.Tensor: + _validate_comparison_dtype(input1.dtype, "GREATER_EQUAL") + return binary_meta(input1, input2, "GREATER_EQUAL", output_dtype=torch.bool) From 2c1da324c9ecc84885b38ce3e2be4ec4a37c7786 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 2 Sep 2026 18:33:48 -0700 Subject: [PATCH 019/190] [ET-VK][ops] Extend arange, clamp, and index.Tensor support Pull Request resolved: https://github.com/pytorch/executorch/pull/22481 Support runtime-varying `arange` bounds, fractional and negative steps, and symbolic `clamp` bounds for integer and floating-point tensors. Extend buffer `index.Tensor` to gather on any axis while preserving multidimensional index shapes. Authored with Codex. ghstack-source-id: 423931774 @exported-using-ghexport Differential Revision: [D118507314](https://our.internmc.facebook.com/intern/diff/D118507314/) --- backends/vulkan/op_registry.py | 87 +++++++++-------- .../runtime/graph/ops/glsl/arange_buffer.glsl | 16 +++- .../graph/ops/glsl/arange_texture.glsl | 16 +++- .../graph/ops/glsl/index_tensor_buffer.glsl | 88 ++++++++++++++---- .../graph/ops/glsl/index_tensor_buffer.yaml | 4 + .../runtime/graph/ops/glsl/unary_op.glsl | 33 ++++++- .../runtime/graph/ops/glsl/unary_op.yaml | 8 ++ .../vulkan/runtime/graph/ops/impl/Arange.cpp | 56 +++++------ .../runtime/graph/ops/impl/IndexTensor.cpp | 72 ++++++++++---- .../vulkan/runtime/graph/ops/impl/UnaryOp.cpp | 44 +++++++++ backends/vulkan/test/test_vulkan_delegate.py | 93 +++++++++++++++++++ 11 files changed, 398 insertions(+), 119 deletions(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index eed287b3a08..91da9134850 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1435,58 +1435,56 @@ def register_where(): # ============================================================================= -@update_features(exir_ops.edge.aten.index.Tensor) -def register_index_tensor(): - def _index_tensor_shapes(node: torch.fx.Node): - """(self_val, index_val) for the supported single-index form, else None.""" - self_arg = node.args[0] - indices = node.args[1] - - if not isinstance(self_arg, torch.fx.Node): - return None - self_val = self_arg.meta.get("val", None) - if self_val is None: - return None - - # Only support exactly one non-None index tensor, applied to dim 0. - if not isinstance(indices, (list, tuple)): - return None - non_none = [idx for idx in indices if idx is not None] - if len(non_none) != 1 or indices[0] is None: - return None - index_arg = non_none[0] - if not isinstance(index_arg, torch.fx.Node): - return None - index_val = index_arg.meta.get("val", None) - if index_val is None: - return None +def _index_tensor_shapes(node: torch.fx.Node): + """Return self, index, and axis for the supported form, else None.""" + self_arg = node.args[0] + indices = node.args[1] + + if not isinstance(self_arg, torch.fx.Node): + return None + self_val = self_arg.meta.get("val", None) + if self_val is None or not isinstance(indices, (list, tuple)): + return None + + non_none = [(dim, index) for dim, index in enumerate(indices) if index is not None] + if len(non_none) != 1: + return None + index_dim, index_arg = non_none[0] + if index_dim >= len(self_val.size()) or not isinstance(index_arg, torch.fx.Node): + return None + index_val = index_arg.meta.get("val", None) + if index_val is None: + return None + + return self_val, index_val, index_dim + + +def _check_index_tensor_node(node: torch.fx.Node) -> bool: + shapes = _index_tensor_shapes(node) + if shapes is None: + return False + _, index_val, _ = shapes + # The gather is expressed as "one index position per output slice", so + # the index must be 1-D. `self` may be any rank. + return len(index_val.size()) == 1 - return self_val, index_val - def check_index_tensor_node(node: torch.fx.Node) -> bool: - shapes = _index_tensor_shapes(node) - if shapes is None: - return False - _, index_val = shapes - # The gather is expressed as "one index position per output slice", so - # the index must be 1-D. `self` may be any rank: the buffer shader - # copies self's trailing dims through unchanged. - return len(index_val.size()) == 1 +def _pick_index_tensor_storage(node: torch.fx.Node): + shapes = _index_tensor_shapes(node) + # Only the buffer shader handles a higher-rank `self`. + if shapes is not None and len(shapes[0].size()) > 1: + return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER + return utils.ANY_STORAGE, utils.ANY_STORAGE - def pick_index_tensor_storage(node: torch.fx.Node): - shapes = _index_tensor_shapes(node) - # Only the buffer shader handles a higher-rank `self`; the texture - # variant still assumes the 1-D form (it reads self[idx, 0, 0, 0]). - if shapes is not None and len(shapes[0].size()) > 1: - return utils.CONTIGUOUS_BUFFER, utils.CONTIGUOUS_BUFFER - return utils.ANY_STORAGE, utils.ANY_STORAGE +@update_features(exir_ops.edge.aten.index.Tensor) +def register_index_tensor(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, supports_resize=True, - are_node_inputs_supported_fn=check_index_tensor_node, - pick_io_storage_fn=pick_index_tensor_storage, + are_node_inputs_supported_fn=_check_index_tensor_node, + pick_io_storage_fn=_pick_index_tensor_storage, ) @@ -1500,6 +1498,7 @@ def register_arange(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_T, + supports_resize=True, ) diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl index 2e9377533c8..9bffc1c4132 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_buffer.glsl @@ -23,18 +23,28 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_ubo(B, "BufferMetadata", "outp")} -${layout_declare_ubo(B, "float", "start")} -${layout_declare_ubo(B, "float", "step")} +${layout_declare_ubo(B, "uint", "start")} +${layout_declare_ubo(B, "uint", "step")} + +layout(push_constant) uniform restrict Block { + ivec2 params_are_int; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" +float decode_param(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); +} + void main() { const uint out_bufi = linear_idx_from_gid(); if (out_of_bounds(out_bufi, outp)) { return; } - t_out[out_bufi] = T(start + out_bufi * step); + const float start_val = decode_param(start, params_are_int.x); + const float step_val = decode_param(step, params_are_int.y); + t_out[out_bufi] = T(start_val + out_bufi * step_val); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl index 0a5636b300f..73c2b5e5dd6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/arange_texture.glsl @@ -23,14 +23,22 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "texture3d")} ${layout_declare_ubo(B, "TextureMetadata", "outp")} -${layout_declare_ubo(B, "float", "start")} -${layout_declare_ubo(B, "float", "step")} +${layout_declare_ubo(B, "uint", "start")} +${layout_declare_ubo(B, "uint", "step")} + +layout(push_constant) uniform restrict Block { + ivec2 params_are_int; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; ${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} const int packed_dim = get_packed_dim(out_layout); +float decode_param(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); +} + void main() { const ivec3 out_pos = ivec3(gl_GlobalInvocationID); @@ -44,11 +52,13 @@ void main() { // arange output is 1D, so the W dimension holds the element index. // Compute the value for each element in the texel along the packed dim. VEC4_T outtex = VEC4_T(0); + const float start_val = decode_param(start, params_are_int.x); + const float step_val = decode_param(step, params_are_int.y); int limit = min( 4, safe_idx(outp.sizes, packed_dim) - out_tidx.data[packed_dim]); for (int comp = 0; comp < limit; comp++) { int elem_idx = out_tidx.data[0]; // W index is the linear element index - outtex[comp] = VEC4_T(start + elem_idx * step).x; + outtex[comp] = VEC4_T(start_val + elem_idx * step_val).x; out_tidx.data[packed_dim]++; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl index db61e0859f2..b2497f98ccc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.glsl @@ -9,6 +9,7 @@ #version 450 core ${define_required_extensions("buffer", DTYPE)} +${define_required_extensions(INDEX_STORAGE, "int")} #define PRECISION ${PRECISION} @@ -22,19 +23,68 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_out", DTYPE, "buffer")} ${layout_declare_tensor(B, "r", "t_self", DTYPE, "buffer")} -${layout_declare_tensor(B, "r", "t_index", "int", "buffer")} +${layout_declare_tensor(B, "r", "t_index", "int", INDEX_STORAGE)} ${layout_declare_ubo(B, "BufferMetadata", "outp")} ${layout_declare_ubo(B, "BufferMetadata", "inp")} -${layout_declare_ubo(B, "BufferMetadata", "index")} +$if INDEX_STORAGE == "buffer": + ${layout_declare_ubo(B, "BufferMetadata", "index")} +$else: + ${layout_declare_ubo(B, "TextureMetadata", "index")} + +layout(push_constant) uniform restrict Block { + ivec2 index_params; +}; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" -// Implements aten.index.Tensor for the case where self is 1D and there is -// exactly one index tensor. Each output element is: +// Implements aten.index.Tensor with exactly one index tensor. Each output +// element is: // output[...] = self[index[...]] +${layout_declare_spec_const(C, "int", "out_layout", "CONTIG_LAYOUT_INT")} +${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} +${layout_declare_spec_const(C, "int", "index_layout", "CONTIG_LAYOUT_INT")} + +int load_index(const TensorIndex out_tidx) { +$if INDEX_STORAGE == "buffer": + uint index_bufi = 0; + for (int d = 0; d < index_params.y; ++d) { + index_bufi += + stride_at(index, d) * idx_at(out_tidx, index_params.x + d); + } + return t_index[index_bufi]; +$else: + TensorIndex4D index_tidx = zero_tensor4d_idx(); + index_tidx.data.x = int(idx_at(out_tidx, index_params.x)); + if (index_params.y > 1) { + index_tidx.data.y = int(idx_at(out_tidx, index_params.x + 1)); + } + if (index_params.y > 2) { + index_tidx.data.z = int(idx_at(out_tidx, index_params.x + 2)); + } + if (index_params.y > 3) { + index_tidx.data.w = int(idx_at(out_tidx, index_params.x + 3)); + } + const TextureElementIndex index_elem = + tensor4d_idx_to_texture_element_idx_simple( + index, index_tidx, index_layout); + return texelFetch(t_index, index_elem.pos, 0)[index_elem.comp]; +} + +uint self_idx_at( + const TensorIndex out_tidx, + const int self_axis, + const uint index_value) { + if (self_axis == index_params.x) { + return index_value; + } + const int out_axis = self_axis < index_params.x + ? self_axis + : self_axis + index_params.y - 1; + return idx_at(out_tidx, out_axis); +} void main() { const uint out_bufi = linear_idx_from_gid(); @@ -45,22 +95,20 @@ void main() { // Convert output buffer index to tensor index TensorIndex out_tidx = linear_idx_to_tensor_idx(outp, out_bufi); - const uint self_rank = ndim(inp); - const uint index_rank = ndim(index); - // WHCN order places self's trailing axes before the index axes. - const uint index_axis_offset = self_rank - 1; - - uint index_bufi = 0; - for (uint d = 0; d < index_rank; ++d) { - index_bufi += - stride_at(index, d) * idx_at(out_tidx, index_axis_offset + d); - } - const int idx = t_index[index_bufi]; - - uint self_bufi = stride_at(inp, self_rank - 1) * uint(idx); - for (uint d = 0; d + 1 < self_rank; ++d) { - self_bufi += stride_at(inp, d) * idx_at(out_tidx, d); - } + const int idx = load_index(out_tidx); + + TensorIndex self_tidx; + initialize(self_tidx); + const int self_rank = int_ndim(inp); + if (self_rank > 0) self_tidx.data[0].x = self_idx_at(out_tidx, 0, uint(idx)); + if (self_rank > 1) self_tidx.data[0].y = self_idx_at(out_tidx, 1, uint(idx)); + if (self_rank > 2) self_tidx.data[0].z = self_idx_at(out_tidx, 2, uint(idx)); + if (self_rank > 3) self_tidx.data[0].w = self_idx_at(out_tidx, 3, uint(idx)); + if (self_rank > 4) self_tidx.data[1].x = self_idx_at(out_tidx, 4, uint(idx)); + if (self_rank > 5) self_tidx.data[1].y = self_idx_at(out_tidx, 5, uint(idx)); + if (self_rank > 6) self_tidx.data[1].z = self_idx_at(out_tidx, 6, uint(idx)); + if (self_rank > 7) self_tidx.data[1].w = self_idx_at(out_tidx, 7, uint(idx)); + const uint self_bufi = tensor_idx_to_linear_idx(inp, self_tidx); t_out[out_bufi] = t_self[self_bufi]; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml index ef79704203f..f4f168dfb37 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/index_tensor_buffer.yaml @@ -8,7 +8,11 @@ index_tensor_buffer: parameter_names_with_default_values: DTYPE: float STORAGE: buffer + INDEX_STORAGE: buffer generate_variant_forall: + INDEX_STORAGE: + - VALUE: buffer + - VALUE: texture3d DTYPE: - VALUE: half - VALUE: float diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl b/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl index 45aa3ed7133..5ef891e7439 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.glsl @@ -23,16 +23,23 @@ ${define_active_storage_type(STORAGE)} layout(std430) buffer; -${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} -${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} +${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)} +${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)} + +$if DYNAMIC_PARAMS: + ${layout_declare_ubo(B, "uint", "minimum")} + ${layout_declare_ubo(B, "uint", "maximum")} layout(push_constant) uniform restrict Block { $if STORAGE == "buffer": int numel; $else: ivec4 out_limits; -float minimum; -float maximum; +$if DYNAMIC_PARAMS: + ivec2 bounds_are_int; +$else: + float minimum; + float maximum; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -40,6 +47,11 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" #include "activations.h" +$if DYNAMIC_PARAMS: + float decode_bound(const uint value, const int is_int) { + return is_int != 0 ? float(int(value)) : uintBitsToFloat(value); + } + #ifdef USING_BUFFER void main() { @@ -48,7 +60,13 @@ void main() { return; } - float in_val = float(t_in[i]); +$if DYNAMIC_PARAMS: + const T in_val = T(t_in[i]); + const T minimum_val = T(decode_bound(minimum, bounds_are_int.x)); + const T maximum_val = T(decode_bound(maximum, bounds_are_int.y)); + t_out[i] = T(op(in_val, minimum_val, maximum_val)); +$else: + const float in_val = float(t_in[i]); t_out[i] = T(op(in_val, minimum, maximum)); } @@ -62,6 +80,11 @@ void main() { } VEC4_T in_texel = texelFetch(t_in, pos, 0); +$if DYNAMIC_PARAMS: + const VEC4_T minimum_val = VEC4_T(decode_bound(minimum, bounds_are_int.x)); + const VEC4_T maximum_val = VEC4_T(decode_bound(maximum, bounds_are_int.y)); + imageStore(t_out, pos, op(in_texel, minimum_val, maximum_val)); +$else: imageStore(t_out, pos, VEC4_T(op(in_texel, minimum, maximum))); } diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml index 46d12806149..0331a15fde6 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml @@ -3,6 +3,7 @@ unary_op: OPERATOR: clamp(X, A, B) DTYPE: float STORAGE: texture3d + DYNAMIC_PARAMS: false generate_variant_forall: DTYPE: - VALUE: half @@ -18,6 +19,13 @@ unary_op: - NAME: clamp_int32 OPERATOR: clamp(X, A, B) DTYPE: int32 + - NAME: clamp_dynamic_int32 + OPERATOR: clamp(X, A, B) + DTYPE: int32 + DYNAMIC_PARAMS: true + - NAME: clamp_dynamic + OPERATOR: clamp(X, A, B) + DYNAMIC_PARAMS: true - NAME: cos OPERATOR: cos(X) - NAME: exp diff --git a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp index f635c9282f2..839b94f5e75 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Arange.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Arange.cpp @@ -15,6 +15,8 @@ #include +#include + namespace vkcompute { void resize_arange_node( @@ -23,18 +25,22 @@ void resize_arange_node( const std::vector& extra_args) { const ValueRef out = args.at(0).refs.at(0); - int start_val = 0; - int step_val = 1; + double start_val = 0.0; + double step_val = 1.0; if (!graph->val_is_none(extra_args.at(0))) { - start_val = graph->extract_scalar(extra_args.at(0)); + start_val = graph->extract_scalar(extra_args.at(0)); } - const int end_val = graph->extract_scalar(extra_args.at(1)); + const double end_val = graph->extract_scalar(extra_args.at(1)); if (!graph->val_is_none(extra_args.at(2))) { - step_val = graph->extract_scalar(extra_args.at(2)); + step_val = graph->extract_scalar(extra_args.at(2)); } + VK_CHECK_COND(step_val != 0.0, "arange: step must be nonzero"); + const double range_size = (end_val - start_val) / step_val; + VK_CHECK_COND( + range_size >= 0.0, "arange: bounds are inconsistent with step sign"); const std::vector out_sizes = { - utils::div_up(end_val - start_val, step_val)}; + static_cast(std::ceil(range_size))}; graph->virtual_resize(out, out_sizes); } @@ -55,39 +61,35 @@ void check_arange_input( } } +vkapi::BufferBindInfo get_arange_param_buffer( + ComputeGraph& graph, + const ValueRef value, + const float default_value) { + if (graph.val_is_symint(value)) { + return graph.get_or_create_int_param_buffer(value); + } + return graph.create_params_buffer( + graph.extract_scalar_or(value, default_value)); +} + void add_arange_node( ComputeGraph& graph, const ValueRef start, const ValueRef end, const ValueRef step, const ValueRef out) { - float start_val = 0.0f; - float step_val = 1.0f; - if (graph.val_is_none(end)) { VK_THROW("arange: end must be specified!"); } - if (!graph.val_is_none(start)) { - if (graph.val_is_int(start)) { - start_val = static_cast(graph.extract_scalar(start)); - } else { - start_val = graph.extract_scalar(start); - } - } - if (!graph.val_is_none(step)) { - if (graph.val_is_int(step)) { - step_val = static_cast(graph.extract_scalar(step)); - } else { - step_val = graph.extract_scalar(step); - } - } - std::string kernel_name("arange"); kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); add_dtype_suffix(kernel_name, graph.dtype_of(out)); + const utils::ivec2 params_are_int = { + graph.val_is_symint(start) ? 1 : 0, graph.val_is_symint(step) ? 1 : 0}; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), @@ -97,10 +99,10 @@ void add_arange_node( {{out, vkapi::kWrite}}, // Shader params buffers {graph.meta_ubo(out), - graph.create_params_buffer(start_val), - graph.create_params_buffer(step_val)}, + get_arange_param_buffer(graph, start, 0.0f), + get_arange_param_buffer(graph, step, 1.0f)}, // Push Constants - {}, + {PushConstantDataInfo(¶ms_are_int, sizeof(params_are_int))}, // Specialization Constants {graph.hashed_layout_of(out)}, // Resize Args diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp index ddd8e8994b1..f490f60f75c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexTensor.cpp @@ -9,6 +9,7 @@ #include #include +#include #include @@ -18,19 +19,31 @@ void resize_index_tensor_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { - (void)resize_args; const ValueRef out = args.at(0).refs.at(0); const ValueRef self = args.at(1).refs.at(0); const ValueRef index = args.at(1).refs.at(1); - // aten.index.Tensor with a single index tensor gathers along dim 0, so - // out.sizes = index.sizes ++ self.sizes[1:] - // Using the index's sizes alone is only correct when self is 1-D; for any - // higher-rank self it also changes the tensor's RANK, which virtual_resize - // rejects outright ("new sizes cannot modify the dimensionality"). + int64_t index_dim = -1; + { + const ValueListPtr indices = graph->get_value_list(resize_args.at(0)); + for (size_t dim = 0; dim < indices->size(); ++dim) { + if (!graph->val_is_none(indices->at(dim))) { + index_dim = utils::safe_downcast(dim); + break; + } + } + } + VK_CHECK_COND(index_dim >= 0, "index.Tensor: an index tensor is required"); + const std::vector self_sizes = graph->sizes_of(self); - std::vector out_sizes = graph->sizes_of(index); - out_sizes.insert(out_sizes.end(), self_sizes.begin() + 1, self_sizes.end()); + const std::vector index_sizes = graph->sizes_of(index); + std::vector out_sizes; + out_sizes.reserve(self_sizes.size() + index_sizes.size() - 1); + out_sizes.insert( + out_sizes.end(), self_sizes.begin(), self_sizes.begin() + index_dim); + out_sizes.insert(out_sizes.end(), index_sizes.begin(), index_sizes.end()); + out_sizes.insert( + out_sizes.end(), self_sizes.begin() + index_dim + 1, self_sizes.end()); graph->virtual_resize(out, out_sizes); } @@ -39,14 +52,24 @@ void add_index_tensor_node( ComputeGraph& graph, const ValueRef self, const ValueRef index, + const int64_t index_dim, + const ValueRef indices_list_ref, const ValueRef out) { std::string kernel_name = "index_tensor"; kernel_name.reserve(kShaderNameReserve); add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + if (graph.is_buffer_storage(out)) { + add_storage_type_suffix(kernel_name, graph.storage_type_of(index)); + } add_dtype_suffix(kernel_name, graph.dtype_of(out)); vkapi::ParamsBindList param_ubos = { graph.meta_ubo(out), graph.meta_ubo(self), graph.meta_ubo(index)}; + const utils::ivec2 index_params = { + utils::safe_downcast(graph.dim_of(self) - 1 - index_dim), + utils::safe_downcast(graph.dim_of(index))}; + std::vector push_constants = { + PushConstantDataInfo(&index_params, sizeof(index_params))}; graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -58,11 +81,13 @@ void add_index_tensor_node( // Shader params buffers param_ubos, // Push Constants - {}, + push_constants, // Specialization Constants - {graph.hashed_layout_of(out), graph.hashed_layout_of(self)}, + {graph.hashed_layout_of(out), + graph.hashed_layout_of(self), + graph.hashed_layout_of(index)}, // Resize Args - {}, + {indices_list_ref}, // Resizing Logic resize_index_tensor_node)); } @@ -72,14 +97,27 @@ void index_tensor(ComputeGraph& graph, const std::vector& args) { ValueRef indices_list_ref = args[1]; ValueRef out = args[2]; - ValueListPtr indices_list = graph.get_value_list(indices_list_ref); + ValueRef index = -1; + int64_t index_dim = -1; + { + const ValueListPtr indices_list = graph.get_value_list(indices_list_ref); + for (size_t dim = 0; dim < indices_list->size(); ++dim) { + const ValueRef candidate = indices_list->at(dim); + if (graph.val_is_none(candidate)) { + continue; + } + VK_CHECK_COND( + index_dim < 0, "index.Tensor: only one index tensor is supported"); + index = candidate; + index_dim = utils::safe_downcast(dim); + } + } + VK_CHECK_COND(index_dim >= 0, "index.Tensor: an index tensor is required"); VK_CHECK_COND( - indices_list->size() == 1, - "index.Tensor: only one index tensor is supported"); - - ValueRef index = indices_list->at(0); + index_dim < graph.dim_of(self), + "index.Tensor: index dimension is invalid"); - add_index_tensor_node(graph, self, index, out); + add_index_tensor_node(graph, self, index, index_dim, indices_list_ref, out); } REGISTER_OPERATORS { diff --git a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp index 6a50cb2f6a9..d17f57774f7 100644 --- a/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp @@ -69,6 +69,46 @@ void add_unary_op_node( resize_unary_op_node)); } +void add_dynamic_clamp_node( + ComputeGraph& graph, + const ValueRef in, + const ValueRef min, + const ValueRef max, + const ValueRef out) { + std::string kernel_name("clamp_dynamic"); + add_dtype_suffix(kernel_name, graph.dtype_of(out)); + add_storage_type_suffix(kernel_name, graph.storage_type_of(out)); + + const bool output_is_int = graph.dtype_of(out) == vkapi::kInt; + const utils::ivec2 bounds_are_int = { + output_is_int || graph.val_is_symint(min) ? 1 : 0, + output_is_int || graph.val_is_symint(max) ? 1 : 0}; + const vkapi::BufferBindInfo min_param = bounds_are_int[0] + ? graph.get_or_create_int_param_buffer( + min, std::numeric_limits::min()) + : graph.create_params_buffer(graph.extract_scalar_or( + min, -std::numeric_limits::infinity())); + const vkapi::BufferBindInfo max_param = bounds_are_int[1] + ? graph.get_or_create_int_param_buffer( + max, std::numeric_limits::max()) + : graph.create_params_buffer(graph.extract_scalar_or( + max, std::numeric_limits::infinity())); + + graph.execute_nodes().emplace_back(new DynamicDispatchNode( + graph, + VK_KERNEL_FROM_STR(kernel_name), + default_pick_gwg, + default_pick_lwg, + {{out, vkapi::kWrite}, {in, vkapi::kRead}}, + {min_param, max_param}, + {graph.is_buffer_storage(out) ? graph.numel_pc_of(out) + : graph.logical_limits_pc_of(out), + PushConstantDataInfo(&bounds_are_int, sizeof(bounds_are_int))}, + {}, + {}, + resize_unary_op_node)); +} + float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) { if (!graph.val_is_none(val)) { return graph.extract_scalar(val); @@ -85,6 +125,10 @@ float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) { #define DEFINE_CLAMP_FN(op_name) \ void op_name(ComputeGraph& graph, const std::vector& args) { \ + if (graph.val_is_symint(args[1]) || graph.val_is_symint(args[2])) { \ + return add_dynamic_clamp_node( \ + graph, args[0], args[1], args[2], args[3]); \ + } \ return add_unary_op_node( \ graph, \ args[0], \ diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index c6915d37684..34a4f12f62c 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -13,6 +13,7 @@ import executorch.backends.vulkan.test.utils as test_utils import torch +import torch.nn.functional as F from executorch.backends.transforms.convert_dtype_pass import I64toI32 from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.vulkan.vulkan_preprocess import VulkanBackend @@ -538,6 +539,20 @@ def forward(self, x): self.lower_module_and_test_output(ClampModule(), sample_inputs) + def test_vulkan_backend_dynamic_float_clamp(self): + class ClampModule(torch.nn.Module): + def forward(self, x): + return torch.clamp(x, max=x.shape[0]) + + sample_inputs = (torch.arange(32).reshape(8, 4).float(),) + length = Dim("length", min=2, max=16) + self.lower_module_and_test_output( + ClampModule(), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.arange(12).reshape(3, 4).float(),)], + ) + def test_vulkan_backend_cos(self): class CosModule(torch.nn.Module): def __init__(self): @@ -1625,6 +1640,84 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_index_tensor_nonzero_axis(self): + class IndexTensorModule(torch.nn.Module): + def __init__(self, dim): + super().__init__() + self.dim = dim + self.index = torch.tensor([0, 2]) + + def forward(self, x): + indices = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + sample_inputs = (torch.arange(24).reshape(1, 3, 8).float(),) + for dim in (1, 2): + self.lower_module_and_test_output( + IndexTensorModule(dim), + sample_inputs, + ) + + def test_vulkan_backend_dynamic_replicate_pad_time_reduction(self): + class TimeReductionModule(torch.nn.Module): + def forward(self, x): + padded_frames = 8 * ((x.shape[1] + 7) // 8) + x = F.pad( + x, + (0, 0, 0, padded_frames - x.shape[1]), + mode="replicate", + ) + return x.view(x.shape[0], -1, 640) + + sample_inputs = (torch.randn(1, 24, 80),) + frames = Dim("frames", min=1, max=24) + self.lower_module_and_test_output( + TimeReductionModule(), + sample_inputs, + dynamic_shapes={"x": {1: frames}}, + test_inputs=[ + (torch.randn(1, 8, 80),), + (torch.randn(1, 9, 80),), + (torch.randn(1, 17, 80),), + ], + ) + + def test_vulkan_backend_dynamic_arange_float_step(self): + class ArangeModule(torch.nn.Module): + def __init__(self, end_scale, step): + super().__init__() + self.end_scale = end_scale + self.step = step + + def forward(self, x): + return torch.arange(0, self.end_scale * x.shape[0], self.step) + + sample_inputs = (torch.randn(8),) + length = Dim("length", min=2, max=16) + for end_scale, step in ((1, 0.5), (-1, -0.5)): + with self.subTest(end_scale=end_scale, step=step): + self.lower_module_and_test_output( + ArangeModule(end_scale, step), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.randn(3),), (torch.randn(7),)], + ) + + def test_vulkan_backend_dynamic_arange_start(self): + class ArangeModule(torch.nn.Module): + def forward(self, x): + return torch.arange(x.shape[0], 32, 2) + + sample_inputs = (torch.randn(8),) + length = Dim("length", min=2, max=16) + self.lower_module_and_test_output( + ArangeModule(), + sample_inputs, + dynamic_shapes={"x": {0: length}}, + test_inputs=[(torch.randn(3),), (torch.randn(15),)], + ) + def test_vulkan_backend_arange_int(self): class ArangeModule(torch.nn.Module): def __init__(self, input): From 8ea353dbb5cf4ab434fc439e2a27116bcb8c13d8 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Thu, 3 Sep 2026 15:18:29 +0100 Subject: [PATCH 020/190] Arm backend: Fix ConvTranspose2d batch norm fusion (#22517) Pass `transpose=True` to PyTorch's batch norm fusion helper for transposed convolutions. Fuse grouped transposed convolutions one group at a time and restore the grouped parameter layout before decomposition. Add tests for unequal input and output channel counts, grouped fusion with and without bias, and non-affine batch norm. Authored with Codex. Change-Id: I7fbd523fa1627ce5d198aea5adbd083c91c0635f Signed-off-by: Yufeng Shi --- .../arm/_passes/fuse_batch_norm2d_pass.py | 119 ++++++++++++++++-- backends/arm/test/ops/test_batch_norm.py | 57 +++++++++ .../test/passes/test_fuse_batchnorm_pass.py | 55 ++++++++ .../source/backends/arm-vgf/VGF_op_support.md | 2 +- 4 files changed, 220 insertions(+), 13 deletions(-) diff --git a/backends/arm/_passes/fuse_batch_norm2d_pass.py b/backends/arm/_passes/fuse_batch_norm2d_pass.py index a13ed9da922..69cd52c40b9 100644 --- a/backends/arm/_passes/fuse_batch_norm2d_pass.py +++ b/backends/arm/_passes/fuse_batch_norm2d_pass.py @@ -12,6 +12,9 @@ create_node, get_first_fake_tensor, ) +from executorch.backends.arm._passes.decompose_grouped_conv_pass import ( + DecomposeGroupedConvPass, +) from executorch.backends.arm.common.debug import get_node_debug_info from executorch.backends.transforms.utils import ( create_constant_placeholder, @@ -27,11 +30,13 @@ class FuseBatchNorm2dPass(ArmPass): - """Fuses the pattern convolution -> batchnorm by updating the weights and - bias of the convolution and removing the batchnorm. + """Fuse convolution followed by BatchNorm. + + Update the convolution weights and bias and remove the BatchNorm operation. + """ - _passes_required_after: Set[Type[ExportPass]] = set() + _passes_required_after: Set[Type[ExportPass]] = {DecomposeGroupedConvPass} def __init__(self, exported_program: ExportedProgram, *args, **kwargs): super().__init__(*args, **kwargs) @@ -45,6 +50,79 @@ def get_bias_name(self, weight_node: Node, bias_node: Node | None) -> str: else: return weight_node.name + "_bias_fused_bn" + @staticmethod + def _fuse_grouped_transposed_conv_bn_weights( + conv_weight: torch.Tensor, + conv_bias: torch.Tensor | None, + bn_mean: torch.Tensor, + bn_var: torch.Tensor, + bn_epsilon: float, + bn_weight: torch.Tensor | None, + bn_bias: torch.Tensor | None, + groups: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Fuse BatchNorm into grouped transposed-convolution parameters. + + This helper runs before ``DecomposeGroupedConvPass`` and transforms:: + + grouped ConvTranspose -> BatchNorm + + into a grouped ConvTranspose with fused weights and bias. A transposed + convolution weight has layout ``[Cin, Cout/groups, ...]``. The weight + is split on its input-channel dimension, while the bias and BatchNorm + parameters are split on their output-channel dimension. Each group is + fused independently before the original grouped layout is restored. + + Args: + conv_weight (torch.Tensor): Grouped transposed-convolution weight. + conv_bias (torch.Tensor | None): Convolution bias. + bn_mean (torch.Tensor): BatchNorm running mean. + bn_var (torch.Tensor): BatchNorm running variance. + bn_epsilon (float): BatchNorm numerical-stability constant. + bn_weight (torch.Tensor | None): BatchNorm weight. + bn_bias (torch.Tensor | None): BatchNorm bias. + groups (int): Number of convolution groups. + + Returns: + tuple[torch.Tensor, torch.Tensor]: Fused weight and bias in the + original grouped layout. + + Raises: + RuntimeError: If the grouped channel dimensions are inconsistent. + + """ + if conv_weight.size(0) % groups != 0 or bn_mean.numel() % groups != 0: + raise RuntimeError("Grouped transposed convolution has invalid channels") + + input_channels_per_group = conv_weight.size(0) // groups + output_channels_per_group = bn_mean.numel() // groups + if conv_weight.size(1) != output_channels_per_group: + raise RuntimeError("BatchNorm channels do not match convolution output") + + fused_weights: list[torch.Tensor] = [] + fused_biases: list[torch.Tensor] = [] + for group in range(groups): + input_start = group * input_channels_per_group + input_end = input_start + input_channels_per_group + output_start = group * output_channels_per_group + output_end = output_start + output_channels_per_group + output_slice = slice(output_start, output_end) + + fused_weight, fused_bias = fuse_conv_bn_weights( + conv_weight[input_start:input_end], + conv_bias[output_slice] if conv_bias is not None else None, + bn_mean[output_slice], + bn_var[output_slice], + bn_epsilon, + bn_weight[output_slice] if bn_weight is not None else None, + bn_bias[output_slice] if bn_bias is not None else None, + transpose=True, + ) + fused_weights.append(fused_weight) + fused_biases.append(fused_bias) + + return torch.cat(fused_weights, dim=0), torch.cat(fused_biases, dim=0) + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 modified = False constant_placeholders_to_delete = set() @@ -176,15 +254,32 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 ) # Fuse bn weights/bias with input weights/bias - fused_weight, fused_bias = fuse_conv_bn_weights( - input_weight_tensor, - input_bias_tensor, - bn_mean_tensor, - bn_var_tensor, - epsilon, - bn_weight_tensor, - bn_bias_tensor, - ) + transposed = bool(input_node.args[6]) + groups = int(input_node.args[8]) + if transposed and groups > 1: + fused_weight, fused_bias = ( + self._fuse_grouped_transposed_conv_bn_weights( + input_weight_tensor, + input_bias_tensor, + bn_mean_tensor, + bn_var_tensor, + epsilon, + bn_weight_tensor, + bn_bias_tensor, + groups, + ) + ) + else: + fused_weight, fused_bias = fuse_conv_bn_weights( + input_weight_tensor, + input_bias_tensor, + bn_mean_tensor, + bn_var_tensor, + epsilon, + bn_weight_tensor, + bn_bias_tensor, + transpose=transposed, + ) # Create fused weights and bias to conv and replace conv args with graph_module.graph.inserting_before(input_weight_node): diff --git a/backends/arm/test/ops/test_batch_norm.py b/backends/arm/test/ops/test_batch_norm.py index 4fb458ee918..b3df52c8f64 100644 --- a/backends/arm/test/ops/test_batch_norm.py +++ b/backends/arm/test/ops/test_batch_norm.py @@ -22,6 +22,7 @@ Input = Tuple[torch.Tensor] ATEN_BATCH_NORM = "torch.ops.aten.batch_norm.default" ATEN_CONV2D = "torch.ops.aten.conv2d.default" +ATEN_CONV_TRANSPOSE2D = "torch.ops.aten.conv_transpose2d.input" @dataclass(frozen=True) @@ -110,6 +111,30 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.batch_norm(self.conv2d(x)) +class BatchNorm2dConvTranspose(torch.nn.Module): + aten_ops = [ATEN_CONV_TRANSPOSE2D, ATEN_BATCH_NORM] + + def __init__(self, groups: int) -> None: + super().__init__() + self.conv_transpose2d = torch.nn.ConvTranspose2d( + in_channels=4, + out_channels=6, + kernel_size=3, + padding=1, + groups=groups, + ) + self.batch_norm = _make_batch_norm( + 6, + affine=True, + weight=torch.rand(6), + bias=torch.rand(6), + track_running_stats=True, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.batch_norm(self.conv_transpose2d(x)) + + class BatchNorm2dNoStats(torch.nn.Module): def __init__( self, @@ -205,6 +230,27 @@ def test_native_batch_norm_legit_no_training_tosa_FP_conv_fuses_before_decompose pipeline.run() +@common.parametrize("groups", {"groups=1": 1, "groups=2": 2}) +def test_conv_transpose_batch_norm_fuses_before_decompose_tosa_FP( + groups: int, +) -> None: + model = BatchNorm2dConvTranspose(groups) + pipeline = TosaPipelineFP[Input]( + model, + (torch.rand(1, 4, 5, 6),), + aten_op=model.aten_ops, + ) + pipeline.count_tosa_ops( + { + "TRANSPOSE_CONV2D": groups, + "CONCAT": int(groups > 1), + "RSQRT": 0, + "SUB": 0, + } + ) + pipeline.run() + + @common.parametrize("case", batch_norm_cases) def test_native_batch_norm_legit_no_training_tosa_INT_conv(case: BatchNormCase) -> None: test_data, model_params = case.make_input_and_parameters() @@ -254,6 +300,17 @@ def test_native_batch_norm_legit_no_training_vgf_no_quant_conv( ).run() +@common.SkipIfNoModelConverter +def test_grouped_conv_transpose_batch_norm_vgf_no_quant() -> None: + model = BatchNorm2dConvTranspose(groups=2) + VgfPipeline[Input]( + model, + (torch.rand(1, 4, 5, 6),), + aten_op=model.aten_ops, + quantize=False, + ).run() + + @common.parametrize("case", batch_norm_cases) @common.SkipIfNoModelConverter def test_native_batch_norm_legit_no_training_vgf_quant_conv( diff --git a/backends/arm/test/passes/test_fuse_batchnorm_pass.py b/backends/arm/test/passes/test_fuse_batchnorm_pass.py index 1c4d862d356..4b09cdaa647 100644 --- a/backends/arm/test/passes/test_fuse_batchnorm_pass.py +++ b/backends/arm/test/passes/test_fuse_batchnorm_pass.py @@ -104,6 +104,47 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x +class MergeConvTransposeBN(torch.nn.Module): + ops_before_pass: ClassVar[Dict[str, int]] = { + "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 1, + "executorch_exir_dialects_edge__ops_aten_convolution_default": 1, + } + ops_after_pass: ClassVar[Dict[str, int]] = { + "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 0, + "executorch_exir_dialects_edge__ops_aten_convolution_default": 1, + } + + def __init__( + self, + groups: int = 1, + bias: bool = False, + affine: bool = True, + in_channels: int = 4, + out_channels: int = 6, + ) -> None: + super().__init__() + self.conv_transpose2d = torch.nn.ConvTranspose2d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2, + stride=2, + groups=groups, + bias=bias, + ) + self.batch_norm2d = torch.nn.BatchNorm2d(out_channels, affine=affine) + self.batch_norm2d.running_mean = torch.rand(out_channels) + self.batch_norm2d.running_var = torch.rand(out_channels) + if affine: + self.batch_norm2d.weight = torch.nn.Parameter(torch.rand(out_channels)) + self.batch_norm2d.bias = torch.nn.Parameter(torch.rand(out_channels)) + + def get_inputs(self) -> input_t: + return (torch.randn(1, self.conv_transpose2d.in_channels, 8, 8),) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.batch_norm2d(self.conv_transpose2d(x)) + + class MergeMultipleUsersBN(torch.nn.Module): ops_before_pass: ClassVar[Dict[str, int]] = { "executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default": 2, @@ -154,6 +195,20 @@ def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: "merge_two_of_two_bn_affine": cast( ModuleWithBatchNormAttrs, MergeTwosOfTwoBN(True) ), + "merge_conv_transpose_bn": cast(ModuleWithBatchNormAttrs, MergeConvTransposeBN()), + "merge_grouped_conv_transpose_bn": cast( + ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2) + ), + "merge_grouped_conv_transpose_bn_bias": cast( + ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2, bias=True) + ), + "merge_grouped_conv_transpose_bn_no_affine": cast( + ModuleWithBatchNormAttrs, MergeConvTransposeBN(groups=2, affine=False) + ), + "merge_grouped_conv_transpose_bn_equal_channels": cast( + ModuleWithBatchNormAttrs, + MergeConvTransposeBN(groups=2, in_channels=4, out_channels=4), + ), "merge_multiple_users_bn_affine": cast( ModuleWithBatchNormAttrs, MergeMultipleUsersBN(True) ), diff --git a/docs/source/backends/arm-vgf/VGF_op_support.md b/docs/source/backends/arm-vgf/VGF_op_support.md index fd27dcde42e..fc77f589bbc 100644 --- a/docs/source/backends/arm-vgf/VGF_op_support.md +++ b/docs/source/backends/arm-vgf/VGF_op_support.md @@ -43,7 +43,7 @@ Total supported PyTorch APIs: **154**. | `torch.conv1d` | FP, INT | `FP32`, `INT8`, `INT4` | 8x8, 8x4 | | `torch.conv2d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | | `torch.conv3d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | -| `torch.conv_transpose2d` | FP, INT | `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | +| `torch.conv_transpose2d` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `INT16`, `INT4` | 8x8, 8x4, 16x8 | | `torch.cos` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 | | `torch.cosh` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.cumsum` | FP, INT | `FP32`, `INT8` | 8x8 | From c8d518916b5e0f2f0ab0441566254d2ec1b3a3a3 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Thu, 3 Sep 2026 16:13:08 +0100 Subject: [PATCH 021/190] Arm backend: Add pre-decomposition partitioner pipeline (#22514) Route the EXIR pre-decomposition hook through an Arm pass pipeline. Store compile specs on concrete partitioners so the pipeline can be configured consistently for TOSA, Ethos-U, and VGF. Keep the pipeline a no-op until backend-specific passes are registered, and update generated partitioner docs and public API manifests. Assisted by Codex. Change-Id: I1963bad577714907e5ba22eed765333d9816308c Signed-off-by: Yufeng Shi --- backends/arm/_passes/arm_pass_manager.py | 6 +++++ backends/arm/ethosu/partitioner.py | 1 + .../api_manifest_running.toml | 8 ++++++ backends/arm/tosa/partitioner.py | 26 +++++++++++++++++++ backends/arm/vgf/partitioner.py | 1 + .../arm-ethos-u/arm-ethos-u-partitioner.md | 16 ++++++++++++ .../backends/arm-vgf/arm-vgf-partitioner.md | 16 ++++++++++++ 7 files changed, 74 insertions(+) diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index b141deb698f..b94b40e05c4 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -458,6 +458,12 @@ def _tosa_context(self, graph_module: GraphModule) -> TosaLoweringContext: shape_env = _get_shape_env_from_gm(graph_module) return TosaLoweringContext(self.tosa_spec, shape_env) + def transform_for_pre_decomposition_pipeline( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Apply Arm passes before default ATen decompositions.""" + return exported_program + def _transform_graph_module(self, graph_module: GraphModule): # TFA and control-flow submodule paths operate on bare GraphModules # without a standalone ExportedProgram to keep in sync. diff --git a/backends/arm/ethosu/partitioner.py b/backends/arm/ethosu/partitioner.py index a36584a975a..eb04e69c8bf 100644 --- a/backends/arm/ethosu/partitioner.py +++ b/backends/arm/ethosu/partitioner.py @@ -34,6 +34,7 @@ def __init__( self.delegation_spec = DelegationSpec( EthosUBackend.__name__, compile_spec._to_list() ) + self.compile_spec = compile_spec self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) diff --git a/backends/arm/public_api_manifests/api_manifest_running.toml b/backends/arm/public_api_manifests/api_manifest_running.toml index 9ff031c41bf..dc87ac9ee96 100644 --- a/backends/arm/public_api_manifests/api_manifest_running.toml +++ b/backends/arm/public_api_manifests/api_manifest_running.toml @@ -60,6 +60,10 @@ signature = "EthosUPartitioner.partition(self, exported_program: torch.export.ex kind = "function" signature = "EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" +[python.EthosUPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + [python.EthosUQuantizer] kind = "class" signature = "EthosUQuantizer(compile_spec: 'EthosUCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" @@ -180,6 +184,10 @@ signature = "VgfPartitioner.partition(self, exported_program: torch.export.expor kind = "function" signature = "VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" +[python.VgfPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + [python.VgfQuantizer] kind = "class" signature = "VgfQuantizer(compile_spec: 'VgfCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 5e5ea1d8423..43449ace3c3 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -21,6 +21,7 @@ from typing import Callable, cast, List, Mapping, Optional, Sequence, Tuple import torch +from executorch.backends.arm._passes.arm_pass_manager import ArmPassManager from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( calculate_multiples, @@ -32,6 +33,7 @@ is_exact_tosa_boundary_bilinear_downscale, ) +from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec from executorch.backends.arm.common.type import ensure_type from executorch.backends.arm.constants import DQ_OPS, Q_OPS from executorch.backends.arm.operator_support.tosa_supported_operators import ( @@ -361,6 +363,8 @@ class TOSAPartitioner(Partitioner): """ + compile_spec: ArmCompileSpec + def __init__( self, compile_spec: TosaCompileSpec, @@ -381,12 +385,34 @@ def __init__( self.delegation_spec = DelegationSpec( TOSABackend.__name__, compile_spec._to_list() ) + self.compile_spec = compile_spec self.tosa_spec = compile_spec.tosa_spec self.additional_checks = additional_checks self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) self._custom_partition_ops: set[torch._ops.OpOverload] = set() self.intermediate_path = compile_spec._get_intermediate_path() + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Apply required Arm passes before default ATen decompositions. + + EXIR invokes this backend extension hook automatically through + ``to_edge_transform_and_lower``. Model export users should not call it + directly. + + Args: + exported_program (ExportedProgram): The ATen-dialect program to + transform. + + Returns: + ExportedProgram: The transformed ATen-dialect program. + + """ + return ArmPassManager( + self.compile_spec + ).transform_for_pre_decomposition_pipeline(exported_program) + def register_custom_partition_op(self, op: torch._ops.OpOverload) -> None: """Register a custom op to be considered supported.""" self._custom_partition_ops.add(op) diff --git a/backends/arm/vgf/partitioner.py b/backends/arm/vgf/partitioner.py index 8ed8d8941e3..be7de275af4 100644 --- a/backends/arm/vgf/partitioner.py +++ b/backends/arm/vgf/partitioner.py @@ -35,6 +35,7 @@ def __init__( self.delegation_spec = DelegationSpec( VgfBackend.__name__, compile_spec._to_list() ) + self.compile_spec = compile_spec self.additional_checks = additional_checks self.tosa_spec = compile_spec.tosa_spec self._decomposable_resize_support = DecomposableResizeSupported(self.tosa_spec) diff --git a/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md b/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md index 26df600045b..4b674b681dd 100644 --- a/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md +++ b/docs/source/backends/arm-ethos-u/arm-ethos-u-partitioner.md @@ -50,3 +50,19 @@ Returns: def EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None: ``` Register a custom op to be considered supported. + +```python +def EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram: +``` +Apply required Arm passes before default ATen decompositions. + +EXIR invokes this backend extension hook automatically through +``to_edge_transform_and_lower``. Model export users should not call it +directly. + +Args: +- **exported_program (ExportedProgram)**: The ATen-dialect program to + transform. + +Returns: +- **ExportedProgram**: The transformed ATen-dialect program. diff --git a/docs/source/backends/arm-vgf/arm-vgf-partitioner.md b/docs/source/backends/arm-vgf/arm-vgf-partitioner.md index 66701acebaa..620088956bd 100644 --- a/docs/source/backends/arm-vgf/arm-vgf-partitioner.md +++ b/docs/source/backends/arm-vgf/arm-vgf-partitioner.md @@ -50,3 +50,19 @@ Returns: def VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None: ``` Register a custom op to be considered supported. + +```python +def VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram: +``` +Apply required Arm passes before default ATen decompositions. + +EXIR invokes this backend extension hook automatically through +``to_edge_transform_and_lower``. Model export users should not call it +directly. + +Args: +- **exported_program (ExportedProgram)**: The ATen-dialect program to + transform. + +Returns: +- **ExportedProgram**: The transformed ATen-dialect program. From 843f77ee2ca451d77991c1570220c10cb7e1a953 Mon Sep 17 00:00:00 2001 From: Per Held Date: Mon, 17 Aug 2026 09:06:20 +0200 Subject: [PATCH 022/190] Arm backend: Delegate reflection padding on U55 Enable legal static ReflectionPad1d, ReflectionPad2d, and ReflectionPad3d cases proven through Vela and Corstone-300. Reflection padding mirrors edge values when extending a tensor: Original: [A B C D] Pad 2 each side: [C B | A B C D | C B] Cover ranks, padding shapes, batching, and quantization modes, and leave dynamic shapes undelegated when padding validity cannot be proven at compile time. Authored with Codex. Change-Id: Ib0b833c2c866b13e504920671303ee22ceac01b6 Signed-off-by: Per Held --- .../arm/operator_support/ethos_u55_support.py | 32 ++++++- .../tosa_supported_operators.py | 9 ++ .../arm/test/ops/test_reflection_pad1d.py | 75 ++++++++++++++++ .../arm/test/ops/test_reflection_pad2d.py | 87 +++++++++++++++++++ .../arm/test/ops/test_reflection_pad3d.py | 75 ++++++++++++++++ 5 files changed, 275 insertions(+), 3 deletions(-) create mode 100644 backends/arm/test/ops/test_reflection_pad1d.py create mode 100644 backends/arm/test/ops/test_reflection_pad2d.py create mode 100644 backends/arm/test/ops/test_reflection_pad3d.py diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py index 35ce0e51df9..8cf8e56560a 100644 --- a/backends/arm/operator_support/ethos_u55_support.py +++ b/backends/arm/operator_support/ethos_u55_support.py @@ -208,9 +208,6 @@ class EthosU55NotSupported(OperatorSupportBase): exir_ops.edge.aten.scatter_reduce.two, exir_ops.edge.aten.scatter_add.default, exir_ops.edge.aten.upsample_bilinear2d.vec, # RESIZE - exir_ops.edge.aten.reflection_pad1d.default, # REVERSE - exir_ops.edge.aten.reflection_pad2d.default, # REVERSE - exir_ops.edge.aten.reflection_pad3d.default, # REVERSE exir_ops.edge.aten.where.self, # SELECT ] @@ -355,6 +352,35 @@ def is_node_supported( ) return False + reflection_pad_constraints = { + exir_ops.edge.aten.reflection_pad1d.default: ((2, 3), (2,)), + exir_ops.edge.aten.reflection_pad2d.default: ((3, 4), (2, 4)), + exir_ops.edge.aten.reflection_pad3d.default: ((4, 5), (6,)), + } + if node.target in reflection_pad_constraints: + input_shape = get_first_fake_tensor(node.all_input_nodes[0]).shape + padding = typing.cast(typing.Sequence[int], node.args[1]) + supported_ranks, supported_padding_lengths = reflection_pad_constraints[ + node.target + ] + if ( + len(input_shape) in supported_ranks + and len(padding) in supported_padding_lengths + ): + spatial_sizes = tuple(reversed(input_shape[-(len(padding) // 2) :])) + pad_pairs = tuple(zip(padding[::2], padding[1::2])) + if all( + isinstance(size, int) and 0 <= before < size and 0 <= after < size + for (before, after), size in zip(pad_pairs, spatial_sizes) + ): + return True + self.reporter.report_reject( + node, + "U55 reflection padding requires a supported static input rank " + "and nonnegative padding smaller than its spatial dimension.", + ) + return False + return True diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index c1fa7015623..c1cf33764b9 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -762,6 +762,15 @@ def is_node_supported( input_node = node.all_input_nodes[0] input_quantized = FuseQuantizedActivationPass._is_fuseable_input(input_node) + if any( + isinstance(input_node.meta["val"], torch.SymInt) + for input_node in node.all_input_nodes + ): + self.reporter.report_reject( + node, "Symbolic scalar inputs cannot be delegated." + ) + return False + input_quantized = input_quantized or all( (input_node.target in DQ_OPS) or _is_integer_dtype(get_first_fake_tensor(input_node).dtype) diff --git a/backends/arm/test/ops/test_reflection_pad1d.py b/backends/arm/test/ops/test_reflection_pad1d.py new file mode 100644 index 00000000000..45a8f0ee9bd --- /dev/null +++ b/backends/arm/test/ops/test_reflection_pad1d.py @@ -0,0 +1,75 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester +from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite_u55 = { + "rank2_symmetric": lambda: (torch.rand(2, 5), (1, 1)), + "rank3_symmetric": lambda: (torch.rand(1, 2, 5), (1, 1)), + "asymmetric": lambda: (torch.rand(1, 2, 5), (1, 3)), + "maximum_legal": lambda: (torch.rand(1, 2, 5), (4, 4)), + "batched": lambda: (torch.rand(2, 2, 5), (1, 1)), +} + + +class ReflectionPad1d(torch.nn.Module): + def __init__(self, padding): + super().__init__() + self.padding = padding + + def forward(self, x): + return torch.nn.functional.pad(x, self.padding, mode="reflect") + + +@common.parametrize("test_data", test_data_suite_u55) +@common.XfailIfNoCorstone300 +def test_reflection_pad1d_u55_INT(test_data): + data, padding = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad1d(padding), + (data,), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad1d_u55_INT_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad1d((1, 1)), + (torch.rand(1, 2, 5),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_reflection_pad1d_u55_INT_symbolic_width_not_delegated(): + width = torch.export.Dim("width", min=3, max=8) + tester = ArmTester( + ReflectionPad1d((1, 1)), + (torch.rand(1, 2, 5),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {2: width}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.aten.scalar_tensor.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/ops/test_reflection_pad2d.py b/backends/arm/test/ops/test_reflection_pad2d.py new file mode 100644 index 00000000000..a9b91bc64a8 --- /dev/null +++ b/backends/arm/test/ops/test_reflection_pad2d.py @@ -0,0 +1,87 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester +from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite_u55 = { + "both_axes": lambda: (torch.rand(1, 2, 4, 3), (1, 1, 1, 1)), + "width_only": lambda: (torch.rand(1, 2, 4, 3), (1, 1, 0, 0)), + "height_only": lambda: (torch.rand(1, 2, 4, 3), (0, 0, 1, 1)), + "asymmetric": lambda: (torch.rand(1, 2, 5, 4), (1, 2, 3, 1)), + "maximum_legal": lambda: (torch.rand(1, 2, 4, 3), (2, 2, 3, 3)), + "batched": lambda: (torch.rand(2, 2, 4, 3), (1, 1, 1, 1)), +} + + +class ReflectionPad2d(torch.nn.Module): + def __init__(self, padding): + super().__init__() + self.padding = padding + + def forward(self, x): + return torch.nn.functional.pad(x, self.padding, mode="reflect") + + +@common.parametrize("test_data", test_data_suite_u55) +@common.XfailIfNoCorstone300 +def test_reflection_pad2d_u55_INT(test_data): + data, padding = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad2d(padding), + (data,), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad2d_u55_INT_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad2d((1, 1, 1, 1)), + (torch.rand(1, 2, 4, 3),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad2d_u55_INT_rank3(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad2d((1, 1, 1, 1)), + (torch.rand(2, 4, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +def test_reflection_pad2d_u55_INT_symbolic_width_not_delegated(): + width = torch.export.Dim("width", min=3, max=8) + tester = ArmTester( + ReflectionPad2d((1, 1, 1, 1)), + (torch.rand(1, 2, 4, 5),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {3: width}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.aten.scalar_tensor.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/ops/test_reflection_pad3d.py b/backends/arm/test/ops/test_reflection_pad3d.py new file mode 100644 index 00000000000..4bd343eafe9 --- /dev/null +++ b/backends/arm/test/ops/test_reflection_pad3d.py @@ -0,0 +1,75 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester +from executorch.backends.arm.test.tester.test_pipeline import EthosU55PipelineINT + + +input_t1 = Tuple[torch.Tensor] + +test_data_suite_u55 = { + "rank4_symmetric": lambda: (torch.rand(2, 4, 4, 4), (1, 1, 1, 1, 1, 1)), + "rank5_symmetric": lambda: (torch.rand(1, 2, 4, 4, 4), (1, 1, 1, 1, 1, 1)), + "asymmetric": lambda: (torch.rand(1, 2, 5, 5, 5), (1, 2, 2, 1, 3, 1)), + "maximum_legal": lambda: (torch.rand(1, 2, 4, 4, 4), (3, 3, 3, 3, 3, 3)), + "batched": lambda: (torch.rand(2, 2, 4, 4, 4), (1, 1, 1, 1, 1, 1)), +} + + +class ReflectionPad3d(torch.nn.Module): + def __init__(self, padding): + super().__init__() + self.padding = padding + + def forward(self, x): + return torch.nn.functional.pad(x, self.padding, mode="reflect") + + +@common.parametrize("test_data", test_data_suite_u55) +@common.XfailIfNoCorstone300 +def test_reflection_pad3d_u55_INT(test_data): + data, padding = test_data() + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad3d(padding), + (data,), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_reflection_pad3d_u55_INT_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ReflectionPad3d((1, 1, 1, 1, 1, 1)), + (torch.rand(1, 2, 4, 4, 4),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_reflection_pad3d_u55_INT_symbolic_width_not_delegated(): + width = torch.export.Dim("width", min=3, max=8) + tester = ArmTester( + ReflectionPad3d((1, 1, 1, 1, 1, 1)), + (torch.rand(1, 2, 4, 4, 5),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {4: width}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.aten.scalar_tensor.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets From d3e5f1b4f539f63cbde54f89a5e8a3e9ea4c0452 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 2 Sep 2026 18:35:10 -0700 Subject: [PATCH 023/190] [ET-VK][runtime] Replace resize update set with generation stamps Pull Request resolved: https://github.com/pytorch/executorch/pull/22496 Dense generation stamps replace allocation-heavy `std::unordered_set` update tracking while preserving recursive `ValueList` semantics. Authored with Codex. ghstack-source-id: 423933048 @exported-using-ghexport Differential Revision: [D118543865](https://our.internmc.facebook.com/intern/diff/D118543865/) --- .../vulkan/runtime/graph/ComputeGraph.cpp | 41 ++++++-- backends/vulkan/runtime/graph/ComputeGraph.h | 8 +- .../vulkan/test/vulkan_compute_api_test.cpp | 96 +++++++++++++++++++ 3 files changed, 133 insertions(+), 12 deletions(-) diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index f23d1f19c66..e337f1b9791 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -11,6 +11,8 @@ #include +#include + #include #include @@ -245,12 +247,12 @@ bool ComputeGraph::was_value_updated(const ValueRef idx) const noexcept { return false; } - // Check if this ValueRef itself was updated - if (updated_values_.find(idx) != updated_values_.end()) { + const size_t value_idx = static_cast(idx); + if (value_idx < value_update_generations_.size() && + value_update_generations_[value_idx] == current_update_generation_) { return true; } - // If this is a ValueList, check each ValueRef in the list if (val_is_value_list(idx)) { const auto& value_list = values_.at(idx).toConstValueList(); for (const auto& nested_idx : value_list) { @@ -263,6 +265,26 @@ bool ComputeGraph::was_value_updated(const ValueRef idx) const noexcept { return false; } +void ComputeGraph::mark_value_updated(const ValueRef idx) { + if (!is_valid_value_idx(idx)) { + return; + } + if (value_update_generations_.size() < values_.size()) { + value_update_generations_.resize(values_.size()); + } + value_update_generations_[static_cast(idx)] = + current_update_generation_; +} + +void ComputeGraph::advance_update_generation() noexcept { + current_update_generation_++; + if (current_update_generation_ == 0) { + std::fill( + value_update_generations_.begin(), value_update_generations_.end(), 0); + current_update_generation_ = 1; + } +} + utils::GPUMemoryLayout ComputeGraph::suggested_memory_layout( const std::vector& sizes) { if (config_.enable_memory_layout_override) { @@ -775,8 +797,7 @@ void ComputeGraph::set_symint(const ValueRef idx, const int32_t val) { int32_t cur_val = read_symint(idx); if (cur_val != val) { get_symint(idx)->set(val); - // Track that this ValueRef was updated - updated_values_.insert(idx); + mark_value_updated(idx); } } @@ -1047,6 +1068,8 @@ void ComputeGraph::maybe_cast_and_copy_from_staging( } void ComputeGraph::prepare() { + value_update_generations_.resize(values_.size()); + #define MERGE_FIELD(field) \ static_cast(std::ceil( \ std::max( \ @@ -1258,8 +1281,7 @@ void ComputeGraph::execute() { execute_count_++; - // Clear the set of updated values at the end of inference - updated_values_.clear(); + advance_update_generation(); // Reset the re-encoding flag at the end of inference requires_reencode_ = false; @@ -1281,7 +1303,7 @@ void ComputeGraph::resize_input( const std::vector& new_sizes) { IOValueRef io_val = inputs_.at(idx); virtual_resize(io_val.value, new_sizes); - updated_values_.insert(io_val.staging); + mark_value_updated(io_val.staging); } void ComputeGraph::virtual_resize( @@ -1290,8 +1312,7 @@ void ComputeGraph::virtual_resize( std::vector cur_sizes = sizes_of(idx); if (cur_sizes != new_sizes) { get_tensor(idx)->virtual_resize(new_sizes); - // Track that this ValueRef was updated - updated_values_.insert(idx); + mark_value_updated(idx); } } diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 1de890efb38..22cd3d9e692 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -204,8 +204,8 @@ class ComputeGraph final { // List of command buffers deferred for submission std::vector deferred_cmd_list_; - // Set to track which ValueRefs were updated during inference - std::unordered_set updated_values_; + std::vector value_update_generations_; + uint32_t current_update_generation_ = 1; // Cache to prevent duplicate prepacking of the same weight tensor with the // same kernel. Key is (inputValueRef, kernel_name). @@ -1222,6 +1222,10 @@ class ComputeGraph final { void print_readable(); + private: + void mark_value_updated(const ValueRef idx); + void advance_update_generation() noexcept; + // // Friend classes // diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 95776e42304..c64a330e767 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -2152,6 +2152,102 @@ TEST(VulkanComputeGraphTest, test_simple_graph_with_symint) { } } +TEST(VulkanComputeGraphTest, was_value_updated_tracks_tensor_changes) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef tensor = graph.add_tensor({2, 4}, vkapi::kFloat); + + EXPECT_FALSE(graph.was_value_updated(kDummyValueRef)); + EXPECT_FALSE(graph.was_value_updated(tensor)); + + graph.virtual_resize(tensor, {2, 4}); + EXPECT_FALSE(graph.was_value_updated(tensor)); + + graph.virtual_resize(tensor, {1, 4}); + EXPECT_TRUE(graph.was_value_updated(tensor)); +} + +TEST(VulkanComputeGraphTest, was_value_updated_tracks_symint_changes) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef symint = graph.add_symint(3); + + EXPECT_FALSE(graph.was_value_updated(symint)); + + graph.set_symint(symint, 3); + EXPECT_FALSE(graph.was_value_updated(symint)); + + graph.set_symint(symint, 5); + EXPECT_TRUE(graph.was_value_updated(symint)); +} + +TEST(VulkanComputeGraphTest, was_value_updated_checks_nested_value_lists) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef unchanged = graph.add_symint(1); + const ValueRef changed = graph.add_symint(2); + const ValueRef inner_list = graph.add_value_list({unchanged, changed}); + const ValueRef outer_list = + graph.add_value_list({kDummyValueRef, inner_list}); + + EXPECT_FALSE(graph.was_value_updated(inner_list)); + EXPECT_FALSE(graph.was_value_updated(outer_list)); + + graph.set_symint(changed, 3); + + EXPECT_FALSE(graph.was_value_updated(unchanged)); + EXPECT_TRUE(graph.was_value_updated(changed)); + EXPECT_TRUE(graph.was_value_updated(inner_list)); + EXPECT_TRUE(graph.was_value_updated(outer_list)); +} + +TEST(VulkanComputeGraphTest, resize_input_marks_staging_value_updated) { + GraphConfig config; + ComputeGraph graph(config); + + const IOValueRef input = graph.add_input_tensor({2, 4}, vkapi::kFloat); + + EXPECT_FALSE(graph.was_value_updated(input.value)); + EXPECT_FALSE(graph.was_value_updated(input.staging)); + + graph.resize_input(0, {2, 4}); + + EXPECT_FALSE(graph.was_value_updated(input.value)); + EXPECT_TRUE(graph.was_value_updated(input.staging)); +} + +TEST(VulkanComputeGraphTest, execute_advances_value_update_generation) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef symint = graph.add_symint(1); + const ValueRef values = graph.add_value_list({symint}); + + graph.prepare(); + graph.set_symint(symint, 2); + + EXPECT_TRUE(graph.was_value_updated(symint)); + EXPECT_TRUE(graph.was_value_updated(values)); + + graph.execute(); + + EXPECT_FALSE(graph.was_value_updated(symint)); + EXPECT_FALSE(graph.was_value_updated(values)); + + graph.set_symint(symint, 3); + + EXPECT_TRUE(graph.was_value_updated(symint)); + EXPECT_TRUE(graph.was_value_updated(values)); + + graph.execute(); + + EXPECT_FALSE(graph.was_value_updated(symint)); + EXPECT_FALSE(graph.was_value_updated(values)); +} + #define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val) \ std::vector data_##name(utils::multiply_integers(sizes)); \ std::fill(data_##name.begin(), data_##name.end(), val); \ From cb0db9eba2c4b32bde76a775c48fb8999c42ac1b Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 2 Sep 2026 18:35:11 -0700 Subject: [PATCH 024/190] [ET-VK][runtime] Check resize inputs before outputs Pull Request resolved: https://github.com/pytorch/executorch/pull/22497 Prioritize read dependencies before write-only outputs so active nodes short-circuit without changing resize semantics. Authored with Codex. ghstack-source-id: 423933065 @exported-using-ghexport Differential Revision: [D118543866](https://our.internmc.facebook.com/intern/diff/D118543866/) --- .../vulkan/runtime/graph/ops/ExecuteNode.cpp | 19 ++- .../vulkan/test/vulkan_compute_api_test.cpp | 131 ++++++++++++++++++ 2 files changed, 148 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp index a1a089c88e5..a2bf4268b01 100644 --- a/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp +++ b/backends/vulkan/runtime/graph/ops/ExecuteNode.cpp @@ -39,8 +39,11 @@ bool ExecuteNode::trigger_resize(ComputeGraph* graph) { } bool ExecuteNode::was_any_arg_updated(const ComputeGraph* const graph) const { - // Check all ValueRefs in ArgGroups + // Check input args. for (const auto& arg_group : args_) { + if (!(arg_group.access & vkapi::kRead)) { + continue; + } for (const auto& value_ref : arg_group.refs) { if (graph->was_value_updated(value_ref)) { return true; @@ -48,13 +51,25 @@ bool ExecuteNode::was_any_arg_updated(const ComputeGraph* const graph) const { } } - // Check all ValueRefs in resize_args + // Check resize args. for (const auto& value_ref : resize_args_) { if (graph->was_value_updated(value_ref)) { return true; } } + // Check output args. + for (const auto& arg_group : args_) { + if (arg_group.access & vkapi::kRead) { + continue; + } + for (const auto& value_ref : arg_group.refs) { + if (graph->was_value_updated(value_ref)) { + return true; + } + } + } + return false; } diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index c64a330e767..5eb1b62d8ff 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -2204,6 +2204,137 @@ TEST(VulkanComputeGraphTest, was_value_updated_checks_nested_value_lists) { EXPECT_TRUE(graph.was_value_updated(outer_list)); } +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_read_arg_updates) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef output = graph.add_symint(1); + const ValueRef input = graph.add_symint(2); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{output, vkapi::kWrite}, {input, vkapi::kRead}}); + + graph.set_symint(input, 3); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_write_arg_updates) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef output = graph.add_symint(1); + const ValueRef input = graph.add_symint(2); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{output, vkapi::kWrite}, {input, vkapi::kRead}}); + + graph.set_symint(output, 3); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_read_write_updates) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef value = graph.add_symint(1); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{value, vkapi::kReadWrite}}); + + graph.set_symint(value, 2); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_tracks_nested_resize_args) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef value = graph.add_symint(1); + const ValueRef inner_list = graph.add_value_list({value}); + const ValueRef outer_list = graph.add_value_list({inner_list}); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {outer_list}); + + graph.set_symint(value, 2); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST(VulkanComputeGraphTest, execute_node_resize_skips_unchanged_args) { + GraphConfig config; + ComputeGraph graph(config); + + const ValueRef output = graph.add_symint(1); + const ValueRef input = graph.add_symint(2); + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {{output, vkapi::kWrite}, {input, vkapi::kRead}}); + + EXPECT_FALSE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 0); +} + +TEST(VulkanComputeGraphTest, execute_node_force_resize_ignores_arg_updates) { + GraphConfig config; + config.force_resize = true; + ComputeGraph graph(config); + + size_t resize_count = 0; + ExecuteNode node([&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + +TEST( + VulkanComputeGraphTest, + execute_node_data_dependent_resize_is_unconditional) { + GraphConfig config; + ComputeGraph graph(config); + + size_t resize_count = 0; + ExecuteNode node( + [&resize_count](ComputeGraph*, const auto&, const auto&) { + ++resize_count; + }, + {}, + {}, + "data_dependent_node", + true); + + EXPECT_TRUE(node.trigger_resize(&graph)); + EXPECT_EQ(resize_count, 1); +} + TEST(VulkanComputeGraphTest, resize_input_marks_staging_value_updated) { GraphConfig config; ComputeGraph graph(config); From 8d33e935c5cd61b44d79c1454836c5f0e3af5c27 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 2 Sep 2026 18:35:11 -0700 Subject: [PATCH 025/190] [ET-VK][runtime] Inline resize update checks Pull Request resolved: https://github.com/pytorch/executorch/pull/22498 Inline direct generation-stamp lookup while keeping rare nested ValueList traversal out of line. Authored with Codex. ghstack-source-id: 423933067 @exported-using-ghexport Differential Revision: [D118543873](https://our.internmc.facebook.com/intern/diff/D118543873/) --- .../vulkan/runtime/graph/ComputeGraph.cpp | 23 ++++--------------- backends/vulkan/runtime/graph/ComputeGraph.h | 18 ++++++++++++++- 2 files changed, 22 insertions(+), 19 deletions(-) diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index e337f1b9791..74407cb00bc 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -242,26 +242,13 @@ utils::StorageType ComputeGraph::suggested_storage_type() { return utils::kTexture3D; } -bool ComputeGraph::was_value_updated(const ValueRef idx) const noexcept { - if (!is_valid_value_idx(idx)) { - return false; - } - - const size_t value_idx = static_cast(idx); - if (value_idx < value_update_generations_.size() && - value_update_generations_[value_idx] == current_update_generation_) { - return true; - } - - if (val_is_value_list(idx)) { - const auto& value_list = values_.at(idx).toConstValueList(); - for (const auto& nested_idx : value_list) { - if (was_value_updated(nested_idx)) { - return true; - } +bool ComputeGraph::was_value_list_updated(const ValueRef idx) const noexcept { + const auto& value_list = values_[static_cast(idx)].toConstValueList(); + for (const auto nested_idx : value_list) { + if (was_value_updated(nested_idx)) { + return true; } } - return false; } diff --git a/backends/vulkan/runtime/graph/ComputeGraph.h b/backends/vulkan/runtime/graph/ComputeGraph.h index 22cd3d9e692..eb01e3abf5e 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.h +++ b/backends/vulkan/runtime/graph/ComputeGraph.h @@ -712,6 +712,7 @@ class ComputeGraph final { private: void check_no_active_value_ptrs(); + bool was_value_list_updated(const ValueRef idx) const noexcept; public: /* @@ -1174,7 +1175,22 @@ class ComputeGraph final { // Check if a specific ValueRef (or ValueList) was updated, with recursive // handling - bool was_value_updated(const ValueRef idx) const noexcept; + inline bool was_value_updated(const ValueRef idx) const noexcept { + if (idx < 0) { + return false; + } + + const size_t value_idx = static_cast(idx); + if (value_idx >= values_.size()) { + return false; + } + if (value_idx < value_update_generations_.size() && + value_update_generations_[value_idx] == current_update_generation_) { + return true; + } + + return values_[value_idx].isValueList() && was_value_list_updated(idx); + } // Set the flag to indicate that re-encoding is required inline void set_requires_reencode() noexcept { From fe5d8d62ddd549cbae3f77c3a0aa6ebd68ad26db Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 2 Sep 2026 18:35:12 -0700 Subject: [PATCH 026/190] [ET-VK][runtime] Resolve invariant quantization shaders once Pull Request resolved: https://github.com/pytorch/executorch/pull/22499 Resolve resize-invariant shader names during graph construction while preserving dynamic workgroups and the M==1 quantize gate. Authored with Codex. ghstack-source-id: 423933078 @exported-using-ghexport Differential Revision: [D118543874](https://our.internmc.facebook.com/intern/diff/D118543874/) --- .../runtime/graph/ops/impl/ChooseQParams.cpp | 23 +- .../graph/ops/impl/QuantizeDequantize.cpp | 40 +--- .../vulkan/test/vulkan_compute_api_test.cpp | 211 ++++++++++++++++++ 3 files changed, 228 insertions(+), 46 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp index cd1f9510bad..1b27a53628e 100644 --- a/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/ChooseQParams.cpp @@ -34,23 +34,6 @@ void resize_choose_qparams_per_row( graph->virtual_resize(input_zeros, new_sizes); } -vkapi::ShaderInfo pick_choose_qparams_per_row_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - (void)resize_args; - - const ValueRef input = args.at(1).refs.at(0); - const ValueRef input_zps = args.at(0).refs.at(1); - - std::string kernel_name = "choose_qparams_per_row"; - add_storage_type_suffix(kernel_name, graph->storage_type_of(input)); - add_dtype_suffix(kernel_name, graph->dtype_of(input)); - add_zp_dtype_mode_suffix(kernel_name, graph->dtype_of(input_zps)); - - return VK_KERNEL_FROM_STR(kernel_name); -} - GlobalWorkGrid pick_choose_qparams_per_row_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -99,9 +82,13 @@ void add_choose_qparams_per_row_node( PushConstantDataInfo(&quant_max_val, sizeof(int32_t)), }; + std::string kernel_name = "choose_qparams_per_row"; + add_storage_type_suffix(kernel_name, graph.storage_type_of(input)); + add_dtype_suffix(kernel_name, graph.dtype_of(input)); + add_zp_dtype_mode_suffix(kernel_name, graph.dtype_of(input_zps)); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - pick_choose_qparams_per_row_shader, + VK_KERNEL_FROM_STR(kernel_name), pick_choose_qparams_per_row_gwg, pick_required_lwg, // Inputs and Outputs diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp index 97c939dcabf..7cf03d98a94 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizeDequantize.cpp @@ -61,33 +61,6 @@ GlobalWorkGrid quantize_and_pack_4h4w_gwg( kTiledWorkGrid); } -vkapi::ShaderInfo pick_quantize_and_pack_4h4w_with_group_sums_shader( - ComputeGraph* graph, - const std::vector& args, - const std::vector& resize_args) { - const ValueRef packed_int_input = args.at(0).refs.at(0); - const ValueRef fp_input = args.at(1).refs.at(0); - const ValueRef packed_input_zps = args.at(1).refs.at(2); - const ValueRef group_size = resize_args.at(0); - - const int64_t group_size_val = graph->extract_scalar(group_size); - - std::string shader_name = "quantize_and_pack_4h4w_with_group_sums"; - if (group_size_val >= 128) { - shader_name += "_o2w32"; - } else { - shader_name += "_o4w16"; - } - - add_storage_type_suffix( - shader_name, graph->storage_type_of(packed_int_input)); - add_storage_type_suffix(shader_name, graph->storage_type_of(fp_input)); - add_dtype_suffix(shader_name, graph->dtype_of(fp_input)); - add_zp_dtype_mode_suffix(shader_name, graph->dtype_of(packed_input_zps)); - - return VK_KERNEL_FROM_STR(shader_name); -} - GlobalWorkGrid pick_quantize_and_pack_4h4w_with_group_sums_gwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -199,9 +172,20 @@ void add_quantize_and_pack_4h4w_with_group_sums_node( const int32_t group_size_val = graph.extract_scalar(group_size); const int32_t blocks_per_group = utils::div_up(group_size_val, int32_t(4)); + std::string shader_name = "quantize_and_pack_4h4w_with_group_sums"; + if (group_size_val >= 128) { + shader_name += "_o2w32"; + } else { + shader_name += "_o4w16"; + } + add_storage_type_suffix(shader_name, graph.storage_type_of(packed_int_input)); + add_storage_type_suffix(shader_name, graph.storage_type_of(fp_input)); + add_dtype_suffix(shader_name, graph.dtype_of(fp_input)); + add_zp_dtype_mode_suffix(shader_name, graph.dtype_of(packed_input_zps)); + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, - pick_quantize_and_pack_4h4w_with_group_sums_shader, + VK_KERNEL_FROM_STR(shader_name), pick_quantize_and_pack_4h4w_with_group_sums_gwg, pick_required_lwg, // Inputs and Outputs diff --git a/backends/vulkan/test/vulkan_compute_api_test.cpp b/backends/vulkan/test/vulkan_compute_api_test.cpp index 5eb1b62d8ff..a0ca49cff6b 100644 --- a/backends/vulkan/test/vulkan_compute_api_test.cpp +++ b/backends/vulkan/test/vulkan_compute_api_test.cpp @@ -32,6 +32,7 @@ #include #include +#include #include @@ -2379,6 +2380,216 @@ TEST(VulkanComputeGraphTest, execute_advances_value_update_generation) { EXPECT_FALSE(graph.was_value_updated(values)); } +TEST(VulkanComputeGraphTest, choose_qparams_handles_dynamic_row_counts) { + constexpr int64_t kMaxM = 8; + constexpr int64_t kK = 128; + + GraphConfig config; + config.enable_querypool = true; + config.expect_dynamic_shapes = true; + ComputeGraph graph(config); + + const IOValueRef input = + graph.add_input_tensor({kMaxM, kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef quant_min = graph.add_scalar(-128); + const ValueRef quant_max = graph.add_scalar(127); + const ValueRef scales = graph.add_tensor( + {kMaxM}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {kMaxM}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + + VK_GET_OP_FN("etvk.choose_qparams_per_row.default") + (graph, {input.value, quant_min, quant_max, scales, zero_points}); + + const ValueRef scales_staging = graph.set_output_tensor(scales); + const ValueRef zero_points_staging = graph.set_output_tensor(zero_points); + + graph.prepare(); + graph.prepack(); + + for (const int64_t M : std::vector{kMaxM, 4, 1, kMaxM}) { + graph.resize_input(0, {M, kK}); + graph.propagate_resize(); + + EXPECT_EQ(graph.sizes_of(scales), std::vector({M})); + EXPECT_EQ(graph.sizes_of(zero_points), std::vector({M})); + + std::vector input_data(M * kK); + for (int64_t m = 0; m < M; ++m) { + std::fill_n(input_data.begin() + m * kK, kK, float(m + 1)); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, input_data.data(), input_data.size(), vkapi::kFloat); + + graph.execute(); + + std::vector scale_data(M); + std::vector zero_point_data(M); + graph.maybe_cast_and_copy_from_staging( + scales_staging, scale_data.data(), scale_data.size(), vkapi::kFloat); + graph.maybe_cast_and_copy_from_staging( + zero_points_staging, + zero_point_data.data(), + zero_point_data.size(), + vkapi::kChar); + + for (int64_t m = 0; m < M; ++m) { + EXPECT_NEAR(scale_data[m], float(m + 1) / 255.0f, 1e-6f); + EXPECT_EQ(zero_point_data[m], -128); + } + + graph.context()->querypool().extract_results(); + const auto shader_results = + graph.context()->querypool().get_shader_timestamp_data(); + const auto choose_result = std::find_if( + shader_results.begin(), shader_results.end(), [](const auto& result) { + return result.kernel_name.find("choose_qparams_per_row") != + std::string::npos; + }); + ASSERT_NE(choose_result, shader_results.end()); + EXPECT_EQ(choose_result->metadata.gwg[0], 1u); + EXPECT_EQ( + choose_result->metadata.gwg[1], + utils::div_up_4(utils::safe_downcast(M))); + EXPECT_EQ(choose_result->metadata.gwg[2], 1u); + EXPECT_EQ(choose_result->metadata.lwg[0], 64u); + EXPECT_EQ(choose_result->metadata.lwg[1], 1u); + EXPECT_EQ(choose_result->metadata.lwg[2], 1u); + } +} + +void test_quantize_and_pack_handles_dynamic_row_counts( + const int64_t group_size_value, + const utils::uvec3& expected_local_wg_size) { + if (!api::context()->adapter_ptr()->supports_int8_dot_product()) { + GTEST_SKIP() << "Quantize and pack requires integer dot product support"; + } + + constexpr int64_t kMaxM = 8; + constexpr int64_t kK = 128; + const int64_t num_groups = kK / group_size_value; + const int64_t max_m4 = utils::div_up(kMaxM, int64_t(4)); + + GraphConfig config; + config.enable_querypool = true; + config.expect_dynamic_shapes = true; + ComputeGraph graph(config); + + const IOValueRef input = + graph.add_input_tensor({kMaxM, kK}, vkapi::kFloat, utils::kBuffer); + const ValueRef quant_min = graph.add_scalar(-128); + const ValueRef quant_max = graph.add_scalar(127); + const ValueRef scales = graph.add_tensor( + {kMaxM}, vkapi::kFloat, utils::kTexture3D, utils::kWidthPacked); + const ValueRef zero_points = graph.add_tensor( + {kMaxM}, vkapi::kChar, utils::kTexture3D, utils::kWidthPacked); + + VK_GET_OP_FN("etvk.choose_qparams_per_row.default") + (graph, {input.value, quant_min, quant_max, scales, zero_points}); + + const ValueRef packed_input = graph.add_tensor( + {kMaxM, kK}, vkapi::kInt8x4, utils::kBuffer, utils::kPackedInt8_4H4W); + const ValueRef input_sums = graph.add_tensor( + {num_groups * max_m4 * 4}, + vkapi::kInt, + utils::kBuffer, + utils::kWidthPacked); + const ValueRef group_size = graph.add_scalar(group_size_value); + const QuantizationConfig input_quant_config( + 8, kPerChannel, {1, kK}, false, true); + + add_quantize_and_pack_4h4w_with_group_sums_node( + graph, + input_quant_config, + input.value, + input_sums, + scales, + zero_points, + packed_input, + group_size); + + const ValueRef packed_input_staging = graph.set_output_tensor(packed_input); + const ValueRef input_sums_staging = graph.set_output_tensor(input_sums); + + graph.prepare(); + graph.prepack(); + + for (const int64_t M : std::vector{kMaxM, 4, 1, kMaxM}) { + graph.resize_input(0, {M, kK}); + graph.propagate_resize(); + + std::vector input_data(M * kK); + for (int64_t m = 0; m < M; ++m) { + std::fill_n(input_data.begin() + m * kK, kK, float(m + 1)); + } + graph.maybe_cast_and_copy_into_staging( + input.staging, input_data.data(), input_data.size(), vkapi::kFloat); + + graph.execute(); + + graph.context()->querypool().extract_results(); + const auto shader_results = + graph.context()->querypool().get_shader_timestamp_data(); + const auto quantize_result = std::find_if( + shader_results.begin(), shader_results.end(), [](const auto& result) { + return result.kernel_name.find( + "quantize_and_pack_4h4w_with_group_sums") != + std::string::npos; + }); + + if (M == 1) { + EXPECT_EQ(quantize_result, shader_results.end()); + continue; + } + + ASSERT_NE(quantize_result, shader_results.end()); + EXPECT_EQ( + quantize_result->metadata.gwg[0], + utils::safe_downcast(num_groups)); + EXPECT_EQ( + quantize_result->metadata.gwg[1], + utils::div_up_4(utils::safe_downcast(M))); + EXPECT_EQ(quantize_result->metadata.gwg[2], 1u); + EXPECT_EQ(quantize_result->metadata.lwg[0], expected_local_wg_size[0]); + EXPECT_EQ(quantize_result->metadata.lwg[1], expected_local_wg_size[1]); + EXPECT_EQ(quantize_result->metadata.lwg[2], expected_local_wg_size[2]); + + const size_t packed_numel = graph.staging_buffer_numel_of(packed_input); + std::vector packed_data(packed_numel); + graph.maybe_cast_and_copy_from_staging( + packed_input_staging, + packed_data.data(), + packed_data.size(), + vkapi::kInt8x4); + for (int64_t i = 0; i < M * kK / 4; ++i) { + EXPECT_EQ(packed_data[i], 0x7f7f7f7f); + } + + std::vector sums_data(num_groups * max_m4 * 4); + graph.maybe_cast_and_copy_from_staging( + input_sums_staging, sums_data.data(), sums_data.size(), vkapi::kInt); + const int64_t current_m4 = utils::div_up(M, int64_t(4)); + for (int64_t group = 0; group < num_groups; ++group) { + for (int64_t m = 0; m < M; ++m) { + EXPECT_EQ( + sums_data[group * current_m4 * 4 + m], 127 * group_size_value); + } + } + } +} + +TEST( + VulkanComputeGraphTest, + quantize_and_pack_handles_dynamic_row_counts_with_small_groups) { + test_quantize_and_pack_handles_dynamic_row_counts(32, {4u, 1u, 16u}); +} + +TEST( + VulkanComputeGraphTest, + quantize_and_pack_handles_dynamic_row_counts_with_large_groups) { + test_quantize_and_pack_handles_dynamic_row_counts(128, {2u, 1u, 32u}); +} + #define CREATE_WEIGHT_TENSOR(name, sizes, dtype, val) \ std::vector data_##name(utils::multiply_integers(sizes)); \ std::fill(data_##name.begin(), data_##name.end(), val); \ From 02ee1cadd3017f56622daf91914d4de5f26bb25b Mon Sep 17 00:00:00 2001 From: telgamal-1 Date: Thu, 3 Sep 2026 15:46:41 -0700 Subject: [PATCH 027/190] Honor range-learned scales in the tied embedding quantization path (#22476) Differential Revision: D118365125 Pull Request resolved: https://github.com/pytorch/executorch/pull/22476 --- examples/models/llama/source_transformation/quantize.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/examples/models/llama/source_transformation/quantize.py b/examples/models/llama/source_transformation/quantize.py index 6bcf35b2a69..65981c57f8c 100644 --- a/examples/models/llama/source_transformation/quantize.py +++ b/examples/models/llama/source_transformation/quantize.py @@ -776,6 +776,7 @@ def get_quant_embedding_transform( embedding_quantize: str, use_shared_embedding: bool = False, quantize_with_hqq: bool = True, + range_learning: bool = False, ): if embedding_quantize.startswith("torchao:"): from torchao.prototype.quantization.embedding.api import ( @@ -819,6 +820,7 @@ def _torchao_embedding_quantizer(model): weight_dtype=weight_dtype, granularity=granularity, mapping_type=mapping_type, + range_learning=range_learning, ).quantize(model) return model From c65b23a2f6b6bd6270240d529d5607c430f83b3d Mon Sep 17 00:00:00 2001 From: mcremon-meta <134334895+mcremon-meta@users.noreply.github.com> Date: Thu, 3 Sep 2026 18:17:46 -0700 Subject: [PATCH 028/190] Lower Conv1d atomically through TOSA Conv2d (#22282) Differential Revision: D116686805 Pull Request resolved: https://github.com/pytorch/executorch/pull/22282 --- backends/arm/_passes/BUCK | 1 - backends/arm/_passes/__init__.py | 1 - backends/arm/_passes/arm_pass_manager.py | 2 - backends/arm/_passes/conv1d_unsqueeze_pass.py | 41 ----- .../_passes/decompose_grouped_conv_pass.py | 4 +- backends/arm/_passes/rewrite_conv_pass.py | 111 +++++++++++- ...est_canonicalize_view_copy_permute_pass.py | 32 ++++ ...ve_permutes_around_elementwise_tosa_ops.py | 36 ++++ .../arm/test/passes/test_rewrite_conv_pass.py | 171 ++++++++++++++++++ backends/arm/tosa/partitioner.py | 4 +- .../canonicalize_view_copy_permute_pass.py | 116 ++++++++---- 11 files changed, 428 insertions(+), 91 deletions(-) delete mode 100644 backends/arm/_passes/conv1d_unsqueeze_pass.py diff --git a/backends/arm/_passes/BUCK b/backends/arm/_passes/BUCK index 401cbd8b18e..49b6ba93287 100644 --- a/backends/arm/_passes/BUCK +++ b/backends/arm/_passes/BUCK @@ -67,7 +67,6 @@ fbcode_target( "//executorch/backends/transforms:fuse_duplicate_users_pass", "//executorch/backends/transforms:propagate_view_copy_permute_pass", "//executorch/backends/transforms:fuse_identical_input_transforms_pass", - "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass", "//executorch/backends/transforms:fuse_view_copy", "//executorch/backends/transforms:remove_getitem_op", "//executorch/backends/transforms:replace_scalar_with_tensor", diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index d287d9d306c..fc48569a3da 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -20,7 +20,6 @@ from .cast_int_comparison_inputs_pass import CastIntComparisonInputsPass # noqa from .cast_to_int32_pass import CastToInt32Pass # noqa from .constant_folding_pass import ConstantFoldingPass # noqa -from .conv1d_unsqueeze_pass import Conv1dUnsqueezePass # noqa from .convert_elu_params import ConvertELUParamsPass # noqa from .convert_expand_copy_to_repeat import ConvertExpandCopyToRepeatPass # noqa from .convert_full_like_to_full_pass import ConvertFullLikeToFullPass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index b94b40e05c4..041296ec36f 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -22,7 +22,6 @@ ComputeConstantOpsAOTPass, ConstantFoldingPass, ControlFlowConstInlinePass, - Conv1dUnsqueezePass, ConvertEluFamilyToEluPass, ConvertELUParamsPass, ConvertExpandCopyToRepeatPass, @@ -603,7 +602,6 @@ def _tosa_pipeline( DecomposeAdaptiveAvgPool2dPass(), DecomposeDynamicAdaptiveAvgPool2dPass(), DecomposeAvgPool2dPass(), - Conv1dUnsqueezePass(exported_program), ] ) diff --git a/backends/arm/_passes/conv1d_unsqueeze_pass.py b/backends/arm/_passes/conv1d_unsqueeze_pass.py deleted file mode 100644 index c01fbdb9f60..00000000000 --- a/backends/arm/_passes/conv1d_unsqueeze_pass.py +++ /dev/null @@ -1,41 +0,0 @@ -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# Copyright 2024-2026 Arm Limited and/or its affiliates. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -from typing import Set, Type - -from executorch.backends.arm._passes import ArmOpTargetedPass -from executorch.backends.arm._passes.convert_squeezes_to_view import ( - ConvertSqueezesToViewPass, -) -from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass -from executorch.backends.arm._passes.size_adjust_input_pass import SizeAdjustInputPass -from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import ( - ConvertConv1dToConv2dPass, -) -from executorch.exir import ExportedProgram -from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.pass_base import ExportPass - - -class Conv1dUnsqueezePass(ConvertConv1dToConv2dPass, ArmOpTargetedPass): - """Arm wrapper for the shared Conv1d-to-Conv2d transform.""" - - _passes_required_after: Set[Type[ExportPass]] = { - ConvertSqueezesToViewPass, - RewriteConvPass, - SizeAdjustInputPass, - } - target_ops = (exir_ops.edge.aten.convolution.default,) - - def __init__(self, exported_program: ExportedProgram) -> None: - # Grouped-convolution decomposition creates one graph-local weight - # slice per group. Allow the shared pass to add the unit-height - # dimension after these producers. - super().__init__( - exported_program, - graph_local_weight_targets={exir_ops.edge.aten.slice_copy.Tensor}, - ) diff --git a/backends/arm/_passes/decompose_grouped_conv_pass.py b/backends/arm/_passes/decompose_grouped_conv_pass.py index 7a8b744d9e3..5ed65e4cd7e 100644 --- a/backends/arm/_passes/decompose_grouped_conv_pass.py +++ b/backends/arm/_passes/decompose_grouped_conv_pass.py @@ -8,8 +8,8 @@ import torch from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass -from executorch.backends.arm._passes.conv1d_unsqueeze_pass import Conv1dUnsqueezePass from executorch.backends.arm._passes.quant_args import QuantArgs +from executorch.backends.arm._passes.rewrite_conv_pass import RewriteConvPass from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -46,7 +46,7 @@ class DecomposeGroupedConvPass(ArmOpTargetedPass): """ - _passes_required_after: Set[Type[ExportPass]] = {Conv1dUnsqueezePass} + _passes_required_after: Set[Type[ExportPass]] = {RewriteConvPass} target_ops = ( exir_ops.edge.aten.convolution.default, torch.ops.aten.conv_transpose2d.input, diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index 6fb686e7ac3..dc00baebaac 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -94,7 +94,7 @@ def _adjust_pad_if_needed( pass instead. """ - mod_remainder = ( + mod_remainder: int | torch.SymInt = ( input_len + 2 * pad - dilation * (input_weight - 1) - 1 ) % stride @@ -121,14 +121,14 @@ def _adjust_pad_if_needed( return pad - mod_remainder - def _is_depthwise_conv2d(self, node: torch.fx.Node) -> bool: + def _is_depthwise_conv(self, node: torch.fx.Node) -> bool: if ( node.op != "call_function" or node.target != exir_ops.edge.aten.convolution.default ): return False input_tensor = get_first_fake_tensor(node.all_input_nodes[0]) - if len(input_tensor.shape) != 4: + if len(input_tensor.shape) not in (3, 4): return False groups = node.args[-1] in_channels = input_tensor.shape[1] @@ -524,7 +524,7 @@ def _combine_rescale_scales( @staticmethod def _is_direct_int32_rescale(node: torch.fx.Node) -> bool: """Return whether a node directly rescales its input to INT32.""" - return ( + return bool( node.op == "call_function" and node.target == exir_ops.backend.tosa.RESCALE.default and len(node.args) > 1 @@ -911,7 +911,70 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 dilation = tuple(dilation_list) pad = pad_attr - if self._is_conv3d(len(input_shape), group): + if spatial_rank == 1: + target_op = ( + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default + if self._is_depthwise_conv(node) + else exir_ops.backend.tosa.CONV2D.default + ) + pre_permute_dims = (0, 2, 1) + post_permute_dims = (0, 2, 1) + with graph_module.graph.inserting_before(node): + x = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.permute_copy.default, + args=(x, list(pre_permute_dims)), + from_node=node, + ) + permuted_input_fake = permute_fake_tensor_metadata( + input_fake_tensor, pre_permute_dims + ) + x.meta["val"] = permuted_input_fake + input_tensor_for_tosa_fake = permuted_input_fake.unsqueeze(1) + x = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(x, list(input_tensor_for_tosa_fake.shape)), + from_node=node, + ) + x.meta["val"] = input_tensor_for_tosa_fake + + kernel_width = weight_shape[2] + if target_op == exir_ops.backend.tosa.DEPTHWISE_CONV2D.default: + in_channels = input_fake_tensor.shape[1] + channel_multiplier = weight_shape[0] // in_channels + weight = self._rewrite_weight( + graph_module, + weight, + node, + permute_dims=(1, 2, 0), + name_suffix="hwicm", + reshape_dims=( + 1, + kernel_width, + in_channels, + channel_multiplier, + ), + ) + else: + weight = self._rewrite_weight( + graph_module, + weight, + node, + permute_dims=(0, 2, 1), + name_suffix="ohwi", + reshape_dims=( + weight_shape[0], + 1, + kernel_width, + weight_shape[1], + ), + ) + weight_fake_tensor = get_first_fake_tensor(weight) + stride = (1, stride[0]) + dilation = (1, dilation[0]) + pad = [0, 0, pad[0], pad[1]] + elif self._is_conv3d(len(input_shape), group): target_op = exir_ops.backend.tosa.CONV3D.default pre_permute_dims = ODHWI_ORDER post_permute_dims = ODHWI_INVERSE_ORDER @@ -934,7 +997,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 name_suffix="odhwi", ) weight_fake_tensor = get_first_fake_tensor(weight) - elif self._is_depthwise_conv2d(node): + elif self._is_depthwise_conv(node): target_op = exir_ops.backend.tosa.DEPTHWISE_CONV2D.default pre_permute_dims = NHWC_ORDER post_permute_dims = NHWC_INVERSE_ORDER @@ -1039,7 +1102,28 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 if post_permute_dims is None: raise RuntimeError("Expected post permute dims for explicit layout") + output_conversion_node = node_replacement post_permute_input = node_replacement + squeeze_view: torch.fx.Node | None = None + if spatial_rank == 1: + squeezed_output_fake = cast( + FakeTensor, node_replacement_fake_tensor.squeeze(1) + ) + special_dtype = node_replacement.meta.get(TosaSpecialDtype.meta_key()) + with graph_module.graph.inserting_after(node_replacement): + node_replacement = create_node( + graph=graph_module.graph, + op_target=exir_ops.edge.aten.view_copy.default, + args=(node_replacement, list(squeezed_output_fake.shape)), + from_node=node, + ) + node_replacement.meta["val"] = squeezed_output_fake + if special_dtype: + node_replacement.meta[TosaSpecialDtype.meta_key()] = special_dtype + squeeze_view = node_replacement + post_permute_input = node_replacement + node_replacement_fake_tensor = squeezed_output_fake + with graph_module.graph.inserting_after(node_replacement): node_replacement = create_node( graph=graph_module.graph, @@ -1059,16 +1143,23 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 tosa_node_fake_tensor.dtype == torch.int32 and input_fake_tensor.dtype == torch.int16 ) - if is_a16w8_conv: + if is_a16w8_conv and spatial_rank != 1: # Keep values in INT32 whenever a consumer supports it, even # though the declared output is INT16, by branching from the # accumulator before narrowing. + # + # Rank-three convolutions are excluded. The legacy Conv1d + # expansion placed a rank-changing view between the convolution + # and its INT32 consumers, so the convolution narrowed to its + # exported output domain instead of forking. Forking here would + # give each branch its own boundary rescale and permute, which + # Vela materialises as a second full transpose of the output. self._insert_a16w8_output_branches( graph_module, node, tosa_op, tosa_node_fake_tensor, - post_permute_input, + output_conversion_node, post_permute_dims, ) # Only users not moved to widened branches remain on the @@ -1078,7 +1169,9 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 node.replace_all_uses_with(node_replacement) else: graph_module.graph.erase_node(node_replacement) - graph_module.graph.erase_node(post_permute_input) + if squeeze_view is not None: + graph_module.graph.erase_node(squeeze_view) + graph_module.graph.erase_node(output_conversion_node) else: node.replace_all_uses_with(node_replacement) diff --git a/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py b/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py index 38a55c8ba10..1ed9961d4f0 100644 --- a/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py +++ b/backends/arm/test/passes/test_canonicalize_view_copy_permute_pass.py @@ -298,6 +298,38 @@ def test_canonicalize_moves_permute_before_view() -> None: _validate_numerics(gm_before, result.graph_module, (x_data,)) +def test_canonicalize_sinks_singleton_view_below_permute() -> None: + builder = GraphBuilder() + x_data = torch.randn(2, 8, 64) + x = builder.placeholder("x", x_data) + v1 = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(x, [2, 8, 1, 64]), + ) + p1 = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(v1, [0, 2, 3, 1]), + ) + builder.output([p1]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + pass_instance = CanonicalizeViewCopyPermutePass() + result = cast(PassResult, pass_instance.call(original)) + + assert result.modified + compute_nodes = [ + node for node in result.graph_module.graph.nodes if node.op == "call_function" + ] + assert [node.target for node in compute_nodes] == [ + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, + ] + assert compute_nodes[0].args[1] == [0, 2, 1] + assert compute_nodes[1].args[1] == [2, 1, 64, 8] + _validate_numerics(gm_before, result.graph_module, (x_data,)) + + def test_canonicalize_follows_interleaved_chain_users() -> None: builder = GraphBuilder() x_data = torch.randn(4, 2, 4) diff --git a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py index 864b6c669f9..f7af7cc41e2 100644 --- a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py @@ -26,6 +26,8 @@ RESCALE_TARGET = exir_ops.backend.tosa.RESCALE.default MUL_TARGET = exir_ops.edge.aten.mul.Tensor ADD_TARGET = exir_ops.edge.aten.add.Tensor +SUB_TARGET = exir_ops.edge.aten.sub.Tensor +VIEW_TARGET = exir_ops.edge.aten.view_copy.default ERF_TARGET = exir_ops.edge.aten.erf.default @@ -150,6 +152,40 @@ def test_remove_permutes_around_rescale_tosa_INT() -> None: assert _count_nodes(result.graph_module, RESCALE_TARGET) == 1 +def test_sink_view_preserves_layout_through_rescale_to_broadcast_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 4, 1, 1) + direct = graph.placeholder("direct") + direct.meta["val"] = torch.randn(1, 8, 4) + + permute = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 3, 1])) + permute.meta["val"] = torch.randn(1, 1, 1, 4) + mul = graph.create_node("call_function", MUL_TARGET, args=(permute, permute)) + mul.meta["val"] = torch.randn(1, 1, 1, 4) + sink = graph.create_node("call_function", VIEW_TARGET, args=(mul, [1, 1, 4])) + sink.meta["val"] = torch.randn(1, 1, 4) + rescale = graph.create_node( + "call_function", + RESCALE_TARGET, + args=(sink, torch.int8, [1.0], 0, 0), + ) + rescale.meta["val"] = torch.randn(1, 1, 4) + sub = graph.create_node("call_function", SUB_TARGET, args=(direct, rescale)) + sub.meta["val"] = torch.randn(1, 8, 4) + graph.output(sub) + + graph_module = torch.fx.GraphModule({}, graph) + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert not result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 1 + assert sub.args == (direct, rescale) + + def test_remove_permutes_around_gelu_with_folded_scalar_constants_tosa_FP() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/backends/arm/test/passes/test_rewrite_conv_pass.py b/backends/arm/test/passes/test_rewrite_conv_pass.py index 31c4205f16f..e1605faefa6 100644 --- a/backends/arm/test/passes/test_rewrite_conv_pass.py +++ b/backends/arm/test/passes/test_rewrite_conv_pass.py @@ -90,6 +90,37 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) + x +class A16W8Conv1dInt32Consumer(nn.Module): + """Exercise a rank-three A16W8 convolution consumed only by an INT32 add.""" + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv1d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Feed the convolution output directly to a residual addition.""" + return self.conv(x) + x + + +class A16W8Conv1dSharedConsumers(nn.Module): + """Exercise a rank-three A16W8 convolution read by two INT32 consumers. + + Mirrors the attention tail of an ECAPA-style model, where ``Softmax(dim=2)`` + decomposes into an ``amax`` reduction and a subtraction that both read the + convolution output. + + """ + + def __init__(self) -> None: + super().__init__() + self.conv = nn.Conv1d(4, 4, 1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Share the convolution output between reduction and subtraction.""" + y = self.conv(x) + return y - y.amax(dim=2, keepdim=True) + + class A16W8MixedConsumer(nn.Module): """Exercise a shared A16W8 convolution output with mixed consumers.""" @@ -296,6 +327,27 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) +class Conv1dBiasModule(torch.nn.Module): + def __init__(self, depthwise: bool = False) -> None: + super().__init__() + groups = 4 if depthwise else 1 + out_channels = 8 if depthwise else 6 + self.conv = torch.nn.Conv1d( + 4, + out_channels, + kernel_size=3, + padding=1, + groups=groups, + bias=True, + ) + + def get_inputs(self) -> tuple[torch.Tensor]: + return (torch.randn(1, 4, 8),) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(x) + + class Conv3dBiasModule(torch.nn.Module): def __init__(self) -> None: super().__init__() @@ -491,6 +543,69 @@ def test_rewrite_conv_a16w8_preserves_int32_for_int32_consumers() -> None: assert direct_int32_rescales[0].args[2] == pytest.approx(expected_int32_scales[0]) +def test_rewrite_conv1d_a16w8_narrows_instead_of_forking_int32() -> None: + """Test that a rank-three A16W8 convolution narrows instead of forking. + + Each widened INT32 branch carries its own boundary rescale and layout + permutation. Vela materialises the rank-three permutation as a full + transpose of the convolution output, so a second branch doubles that cost. A + rank-three convolution therefore narrows to its exported INT16 domain and + keeps a single layout boundary. + + """ + inputs = (torch.randn(1, 4, 8),) + gm, _ = _rewrite_a16w8_convs(A16W8Conv1dInt32Consumer(), inputs) + + conv = _get_call_function_node(gm, exir_ops.backend.tosa.CONV2D.default) + forked_int32_rescales = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.backend.tosa.RESCALE.default + and node.args[1] == torch.int32 + and node.all_input_nodes[0] is conv + ] + assert forked_int32_rescales == [] + + (boundary_rescale,) = tuple(conv.users) + assert boundary_rescale.target == exir_ops.backend.tosa.RESCALE.default + (squeeze_view,) = tuple(boundary_rescale.users) + assert squeeze_view.target == exir_ops.edge.aten.view_copy.default + assert squeeze_view.meta["val"].shape == torch.Size((1, 8, 4)) + (boundary_permute,) = tuple(squeeze_view.users) + assert boundary_permute.target == exir_ops.edge.aten.permute_copy.default + assert boundary_permute.meta["val"].shape == torch.Size((1, 4, 8)) + + +def test_rewrite_conv1d_a16w8_shares_one_layout_boundary() -> None: + """Test that consumers of a rank-three A16W8 convolution share one boundary. + + Forking a widened INT32 branch gives every consumer its own boundary rescale + and layout permutation. Vela materialises the rank-three permutation as a + full transpose of the convolution output, so a second branch doubles it. + + """ + inputs = (torch.randn(1, 4, 8),) + gm, _ = _rewrite_a16w8_convs(A16W8Conv1dSharedConsumers(), inputs) + + conv = _get_call_function_node(gm, exir_ops.backend.tosa.CONV2D.default) + boundary_rescales = [ + node + for node in conv.users + if node.target == exir_ops.backend.tosa.RESCALE.default + ] + assert len(boundary_rescales) == 1 + + output_permutes = [ + node + for node in gm.graph.nodes + if node.op == "call_function" + and node.target == exir_ops.edge.aten.permute_copy.default + and node.meta["val"].shape == torch.Size((1, 4, 8)) + ] + assert len(output_permutes) == 1 + + def test_rewrite_conv_a16w8_preserves_int32_after_permute() -> None: r"""Test that an indirect INT32 consumer keeps a widened branch. @@ -533,6 +648,62 @@ def test_rewrite_conv_a16w8_preserves_int32_after_permute() -> None: assert len(widened_paths) == 1 +@pytest.mark.parametrize( + "depthwise,target_op,expected_weight_shape,expected_output_shape", + [ + ( + False, + exir_ops.backend.tosa.CONV2D.default, + (6, 1, 3, 4), + (1, 6, 8), + ), + ( + True, + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default, + (1, 3, 4, 2), + (1, 8, 8), + ), + ], +) +def test_rewrite_conv1d_emits_atomic_rank3_layout_boundaries( + depthwise: bool, + target_op, + expected_weight_shape: tuple[int, ...], + expected_output_shape: tuple[int, ...], +) -> None: + module = Conv1dBiasModule(depthwise).eval() + edge_program = to_edge(export(module, module.get_inputs())).exported_program() + + with TosaLoweringContext(_compile_spec().tosa_spec): + result = RewriteConvPass(edge_program)(edge_program.graph_module) + assert result is not None + graph_module = result.graph_module + + conv = _get_call_function_node(graph_module, target_op) + input_view = conv.args[0] + assert isinstance(input_view, torch.fx.Node) + assert input_view.target == exir_ops.edge.aten.view_copy.default + input_permute = input_view.args[0] + assert isinstance(input_permute, torch.fx.Node) + assert input_permute.target == exir_ops.edge.aten.permute_copy.default + assert input_permute.args[1] == [0, 2, 1] + assert input_view.meta["val"].shape == torch.Size((1, 1, 8, 4)) + + weight = conv.args[1] + assert isinstance(weight, torch.fx.Node) + assert weight.meta["val"].shape == torch.Size(expected_weight_shape) + + output_view = next( + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.view_copy.default and node.args[0] is conv + ) + output_permute = next(iter(output_view.users)) + assert output_permute.target == exir_ops.edge.aten.permute_copy.default + assert output_permute.args[1] == [0, 2, 1] + assert output_permute.meta["val"].shape == torch.Size(expected_output_shape) + + @pytest.mark.skipif(not _VGF_ENABLED, reason="VGF not enabled") def test_fold_and_annotate_q_params_vgf_quant_tracks_fused_relu_qparams() -> None: exported_program = _export_quantized(TinyConvReluCat()) diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 43449ace3c3..a3262af8bc0 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -154,7 +154,7 @@ def _is_noop_as_strided_copy(node: torch.fx.Node) -> bool: else: input_tensor = get_first_fake_tensor(ensure_type(torch.fx.Node, node.args[0])) output_tensor = get_first_fake_tensor(node) - return ( + return bool( len(input_tensor.shape) == len(output_tensor.shape) and all( statically_known_true(input_dim == output_dim) @@ -196,7 +196,7 @@ def _is_noop_squeeze(node: torch.fx.Node) -> bool: else: input_tensor = get_first_fake_tensor(ensure_type(torch.fx.Node, node.args[0])) output_tensor = get_first_fake_tensor(node) - return input_tensor.shape == output_tensor.shape + return bool(input_tensor.shape == output_tensor.shape) def _is_noop_flip(node: torch.fx.node.Node) -> bool: diff --git a/backends/transforms/canonicalize_view_copy_permute_pass.py b/backends/transforms/canonicalize_view_copy_permute_pass.py index 0a76f10011d..c29a246b04a 100644 --- a/backends/transforms/canonicalize_view_copy_permute_pass.py +++ b/backends/transforms/canonicalize_view_copy_permute_pass.py @@ -197,45 +197,83 @@ def _fuse_sequential_ops( any_changed = True continue - if index + 1 < len(updated_chain): - next_node = updated_chain[index + 1] - if ( - node.target == self._VIEW_TARGET - and next_node.target == self._VIEW_TARGET - ): - # Fuse conscutive views - self._set_node_op( - node, self._VIEW_TARGET, input_node, self._shape(next_node) - ) - self._remove_node( - graph_module, updated_chain, index + 1, replacement=node - ) - changed = True - any_changed = True - continue - - if self._is_permute(node) and self._is_permute(next_node): - # Fuse consecutive permutes - dims = self._permute_dims(node) - next_dims = self._permute_dims(next_node) - self._set_node_op( - node, - self._PERMUTE_TARGET, - input_node, - [dims[dim] for dim in next_dims], - ) - self._remove_node( - graph_module, updated_chain, index + 1, replacement=node - ) - changed = True - any_changed = True - continue + if self._fuse_pair(graph_module, updated_chain, index): + changed = True + any_changed = True + continue index += 1 if not changed: return updated_chain, any_changed + def _fuse_pair( + self, graph_module: GraphModule, chain: list[Node], index: int + ) -> bool: + """Fuse or reorder the adjacent pair at ``index``, if possible.""" + if index + 1 >= len(chain): + return False + + node, next_node = chain[index], chain[index + 1] + input_node = cast(Node, node.args[0]) + + if self._sink_singleton_view(chain, index): + return True + + if node.target == self._VIEW_TARGET and next_node.target == self._VIEW_TARGET: + # Fuse conscutive views + self._set_node_op( + node, self._VIEW_TARGET, input_node, self._shape(next_node) + ) + self._remove_node(graph_module, chain, index + 1, replacement=node) + return True + + if self._is_permute(node) and self._is_permute(next_node): + # Fuse consecutive permutes + dims = self._permute_dims(node) + next_dims = self._permute_dims(next_node) + self._set_node_op( + node, + self._PERMUTE_TARGET, + input_node, + [dims[dim] for dim in next_dims], + ) + self._remove_node(graph_module, chain, index + 1, replacement=node) + return True + + return False + + def _sink_singleton_view(self, chain: list[Node], index: int) -> bool: + """Rewrite ``view(S).permute(P)`` to ``permute(P').view(S')``. + + Only applies to a lone pair whose view just inserts unit dimensions. + Longer chains are reordered by the swap loop in ``call()``; a pair never + is. Permuting at the lower rank lets layout boundaries that reach the + same tensor from different ranks converge on one permute. + + """ + if len(chain) != 2: + return False + + view_node, permute_node = chain[index], chain[index + 1] + input_node = cast(Node, view_node.args[0]) + if ( + view_node.target != self._VIEW_TARGET + or not self._is_permute(permute_node) + or not self._only_inserts_singletons( + self._shape(input_node), self._shape(view_node) + ) + ): + return False + + swapped_args = self._view_permute_swap(view_node, permute_node) + if swapped_args is None: + return False + + self._set_node_op(view_node, self._PERMUTE_TARGET, input_node, swapped_args[0]) + self._set_node_op(permute_node, self._VIEW_TARGET, view_node, swapped_args[1]) + return True + def _maybe_swap_args( self, op1: Node, op2: Node ) -> tuple[Sequence[_Dim], Sequence[_Dim]] | None: @@ -319,6 +357,18 @@ def _inverse_permutation(permutation: Sequence[int]) -> list[int]: inverse[dim] = index return inverse + @classmethod + def _only_inserts_singletons( + cls, input_shape: Sequence[_Dim], output_shape: Sequence[_Dim] + ) -> bool: + """Whether a view only adds singleton dimensions to its input.""" + if len(output_shape) <= len(input_shape): + return False + kept = [dim for dim in output_shape if not _dim_equals(dim, 1)] + return cls._shapes_equal( + kept, [dim for dim in input_shape if not _dim_equals(dim, 1)] + ) + @classmethod def _is_singleton_permutation( cls, shape: Sequence[_Dim], permutation: Sequence[int] From 89e108c7425d8dc284008b492c64f2499a3d7707 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 3 Sep 2026 23:59:07 -0700 Subject: [PATCH 029/190] Fix the Samsung MobileBert test setup so it can run (#22550) ### Summary The Samsung model test job runs every test under `backends/samsung/test/models`. One of those tests, the MobileBert fine-tuning test, always errors out before it does any work: ``` ERROR: setUpClass (test_mobilebert_finetuning.Test_Milestone_MobileBertFinetune) AttributeError: type object 'Test_Milestone_MobileBertFinetune' has no attribute 'model_cache_dir' ``` One error fails the whole job, so the other ten Samsung model tests passing does not help. The job only runs the tests when a Samsung device is available, which is why this looks like it comes and goes. There were three separate problems stacked in the test setup, and each one was hidden behind the one before it: 1. `setUpClass` read `cls.model_cache_dir`, which is not defined anywhere. 2. It passed that value to `patch_mobilebert_finetuning()`, which takes no arguments. Python only checks the number of arguments when the call happens, so this could not be seen until the first problem was gone. 3. The setup replaced `load_tokenizer` with a copy of itself that left out the model name, so `AutoTokenizer.from_pretrained()` had nothing to load. The replacement tokenizer loader was the same as the real one except for the missing model name, so it could only ever do less than the real one. Removing the setup that installed it fixes all three problems at once and lets the test use the real loader. The test now has the same shape as the other ten tests in the same directory, which build a model and check it with no extra setup. ### Test plan Ran the real test collection and setup for this file before and after the change, with the heavy third party dependencies stubbed out so only the setup path was under test. - Before: reproduces the same `AttributeError` seen in CI, and never reaches the test body. - After: setup succeeds and the test body is reached. Also checked that `model_cache_dir`, `setUpClass` and the now unused `AutoTokenizer` import are all gone, and that the file still compiles. --- .../test/models/test_mobilebert_finetuning.py | 26 ------------------- 1 file changed, 26 deletions(-) diff --git a/backends/samsung/test/models/test_mobilebert_finetuning.py b/backends/samsung/test/models/test_mobilebert_finetuning.py index ffb60219c39..9ed7440d5ba 100644 --- a/backends/samsung/test/models/test_mobilebert_finetuning.py +++ b/backends/samsung/test/models/test_mobilebert_finetuning.py @@ -12,36 +12,10 @@ ) from executorch.backends.samsung.test.tester import SamsungTester from executorch.backends.samsung.test.utils.utils import TestConfig - from executorch.examples.samsung.scripts.mobilebert_finetune import MobileBertFinetune -from transformers import AutoTokenizer - - -def patch_mobilebert_finetuning(): - def _monkeypatch_load_tokenizer(self): - tokenizer = AutoTokenizer.from_pretrained( - do_lower_case=True, - ) - return tokenizer - - old_func = MobileBertFinetune.load_tokenizer - MobileBertFinetune.load_tokenizer = _monkeypatch_load_tokenizer - return old_func - - -def recover_mobilebert_finetuning(old_func): - MobileBertFinetune.load_tokenizer = old_func class Test_Milestone_MobileBertFinetune(unittest.TestCase): - @classmethod - def setUpClass(cls): - cls._old_func = patch_mobilebert_finetuning(cls.model_cache_dir) - - @classmethod - def tearDownClass(cls): - recover_mobilebert_finetuning(cls._old_func) - def test_mobilebert_finetuning_fp16(self): mobilebert_finetune = MobileBertFinetune() model, _ = mobilebert_finetune.get_finetune_mobilebert(None) From 5410b1a5676238d7114819858136ebd4a26007fc Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 3 Sep 2026 23:59:16 -0700 Subject: [PATCH 030/190] Stop shipping unusable MKL search paths in the Linux wheel (#22541) Fixes #21611 ## The problem, in plain terms Five shared libraries in the Linux wheel search three directories that exist on nobody's machine: ``` /lib/intel64 /lib/intel64_win /lib/win-x64 ``` Two of them name a Windows layout, in a Linux wheel. They come from PyTorch. Its exported CMake package creates a `caffe2::mkl` imported target with a hardcoded list of link directories, and linking torch brings them in even though this project never asks for MKL: ``` ${MKL_ROOT}/lib ${MKL_ROOT}/lib/intel64 ${MKL_ROOT}/lib/intel64_win ${MKL_ROOT}/lib/win-x64 ``` `MKL_ROOT` resolves to nothing here, so what the linker records is left anchored at the filesystem root. Packaging copies the built libraries out of the build tree rather than installing them, so whatever the linker recorded ships as is. The bare `/lib` form does not survive, because CMake filters its own implicit link directories out of the link line, which is why three entries appear rather than four. This only affects Linux x86_64. The aarch64 nightly records no absolute entries at all, since MKL is not found there, and the macOS and Windows wheels record none either. ## Why it is worth fixing Nothing needs those directories. No shipped library names an MKL or OpenMP runtime among its dependencies, so nothing resolves through them. They are not merely untidy either. They sit ahead of the relative entries packaging appends, and the loader searches in order, so a user who happens to have a matching directory resolves a library from there instead of from the one the wheel installed. That is the same shadowing the release check already rejects a CUDA toolkit prefix for. ## What changed Two small pieces. Packaging now drops these entries, in the same place the other unusable ones are already dropped. The match is deliberately narrow: only the exact `/lib/` form that an empty prefix produces. A real MKL installation spells the same arch directory below a prefix, as `/opt/intel/mkl/lib/intel64`, and that is a directory the environment genuinely provides, so it is kept. The release check that rejects absolute search paths listed these three as allowed. That is why they shipped while a check whose whole purpose is rejecting absolute paths reported the wheel clean. It now rejects the empty prefix form specifically, and still accepts a real installation's prefixed directory, so the two halves agree rather than contradict each other. ## Test plan New unit tests in `.ci/scripts/tests/test_runtime_path_filter.py`. They read the functions out of `setup.py`, so they exercise the code that ships rather than a copy, and they run on every pull request through the existing unit test job. They need no wheel build, which is what the previous check could not manage: it could only see this after a full build, and only on the platform that built one. ``` pytest .ci/scripts/tests/test_runtime_path_filter.py 27 passed ``` They cover both directions, because a filter that satisfies either one alone is wrong: - the three unresolved entries are dropped - a real MKL installation and ordinary system directories are kept - every relative entry on a shipped library survives, so a library can still find its siblings - the absolute torch directory is kept when it is a library's only route to torch, which is what stops this becoming a blanket rule that breaks importing I checked by mutation that each part of the fix is load bearing: | what I broke | tests that failed | | --- | ---: | | removed the packaging filter | 4 | | removed the release check's maths rejection | 3 | | removed its build directory branch | 2 | | removed its catch-all rejection | 1 | | severed the call that applies it to a shipped library | 1 | | reordered its guards so a build directory is accepted | 1 | | narrowed what it accepts, refusing a real install | 3 | | made the packaging filter reject an ordinary system directory | 1 | | added an arch to the shared constant, untaught to the check | 2 | The release check's per-entry decision is a small module level function, so the unit test calls the same code the wheel check runs and compares the reason it returns. Asserting only that a path was rejected was not enough: the check rejects every absolute path it does not recognise, so an unknown arch passed for the wrong reason and two of its three branches could each be deleted on their own. I also measured the published artifacts directly, reading the recorded search paths out of every shared library in each wheel: | wheel | libraries with a search path | unusable entries | | --- | ---: | ---: | | 1.4.1 release, CPython 3.12 | 6 | 0 | | current nightly, CPU, CPython 3.12 | 18 | 15 | | current nightly, CUDA, CPython 3.13 | 21 | 15 | | current nightly, CPU, aarch64 | 16 | 0 | The 1.4.1 release predates the code changed here. Both nightly rows show the same 15 entries across the same five libraries, which is what this removes. Reading those same libraries' declared dependencies shows no MKL or OpenMP runtime in any of them. ## Scope This covers the first of the two things the issue suggests: not shipping a runtime search path a user cannot use. It does not do the second, which is to resolve MKL through something the wheel or its declared dependencies own. That one belongs upstream, since the directory list is set in PyTorch's own CMake package, and its own comment there marks it as a hack. Nothing in this wheel links MKL, so there is no dependency here left to redirect. ## What I did not verify No wheel was built for this. The change is to the filter those builds already run, and the tests exercise it directly. The link line that emits the paths is unchanged; the entries are dropped at packaging, the same way the CUDA toolkit and torch directories already are. Co-authored-by: PyTorch Bot --- .ci/scripts/tests/test_runtime_path_filter.py | 331 ++++++++++++++++++ .ci/scripts/wheel/test_shared_libraries.py | 144 ++++---- setup.py | 38 ++ 3 files changed, 445 insertions(+), 68 deletions(-) create mode 100644 .ci/scripts/tests/test_runtime_path_filter.py diff --git a/.ci/scripts/tests/test_runtime_path_filter.py b/.ci/scripts/tests/test_runtime_path_filter.py new file mode 100644 index 00000000000..74af1b03066 --- /dev/null +++ b/.ci/scripts/tests/test_runtime_path_filter.py @@ -0,0 +1,331 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the runtime search path filter packaging applies to shipped libraries. + +Here rather than in the wheel checks, because the decision under test is a pure function of one +string. The wheel checks can only see it after a full wheel build, and only on the platform that +built one, so a filter that dropped the wrong entry reached the published artifact before anything +ran that would notice. + +Two properties are covered, and they pull in opposite directions, which is why both are needed. +The filter must drop the MKL arch directories torch's exported link interface leaves anchored at +the filesystem root, and it must keep the absolute torch directory that is a library's only route +to torch when no relative one was recorded. A filter that satisfies either alone is wrong: the +first way ships unusable paths, the second way stops the extensions importing. + +The functions are read out of setup.py rather than restated, so the test exercises what ships. +setup.py calls setup() at module scope, so it is loaded by compiling the definitions this needs +instead of importing it, which would exit during setuptools argument parsing. +""" + +import ast +import importlib.util +import re +from pathlib import Path, PurePosixPath + +import pytest + +REPO_ROOT = Path(__file__).resolve().parents[3] + +# Compiled from setup.py, so a change to the filter is exercised here rather than duplicated. +_WANTED = ( + "_MKL_ARCH_DIRECTORIES", + "_is_cuda_toolkit_directory", + "_is_unresolved_math_library_directory", + "_is_usable_runtime_path", +) + + +def _setup_source() -> str: + """setup.py's text, decoded as UTF-8. + + The encoding is named because `read_text()` defaults to the locale's, and both files this test + parses contain non-ASCII characters. On a Windows runner that resolves to a code page, which + mangles them, and the mangled text is what gets parsed. + """ + return (REPO_ROOT / "setup.py").read_text(encoding="utf-8") + + +def _setup_namespace() -> dict: + """The runtime path helpers from setup.py, compiled without running setup().""" + tree = ast.parse(_setup_source()) + wanted = [] + for node in tree.body: + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.name in _WANTED: + wanted.append(node) + elif isinstance(node, ast.Assign): + names = [t.id for t in node.targets if isinstance(t, ast.Name)] + if any(name in _WANTED for name in names): + wanted.append(node) + found = set() + for node in wanted: + found.add( + node.name + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)) + else next(t.id for t in node.targets if isinstance(t, ast.Name)) + ) + missing = sorted(set(_WANTED) - found) + assert not missing, ( + f"setup.py no longer defines {missing} at module scope, so this test would silently " + "check nothing. Update the names here to match." + ) + namespace = {"re": re, "PurePosixPath": PurePosixPath} + exec( + compile(ast.Module(body=wanted, type_ignores=[]), "", "exec"), + namespace, + ) + return namespace + + +@pytest.fixture(scope="module") +def setup_helpers() -> dict: + return _setup_namespace() + + +# setup.py's own arch names, read at import so the parametrized cases below are driven by the +# shipped constant rather than a second copy of it. Adding an arch to setup.py then extends both +# the reject and the accept cases, which is what keeps packaging and the release check in step. +_MKL_ARCH_DIRECTORIES = _setup_namespace()["_MKL_ARCH_DIRECTORIES"] + +# What the linker records when MKL's prefix resolves empty, leaving its arch subdirectory +# concatenated onto nothing. Two of the three name a Windows layout, in a Linux wheel. +UNRESOLVED_MATH_DIRECTORIES = tuple(f"/lib/{arch}" for arch in _MKL_ARCH_DIRECTORIES) + + +@pytest.mark.parametrize("entry", UNRESOLVED_MATH_DIRECTORIES) +def test_drops_math_directories_with_an_empty_prefix(entry, setup_helpers): + assert setup_helpers["_is_unresolved_math_library_directory"](entry) is True + # A trailing separator is the same directory, which the path type normalises on its own. Asserted + # because patchelf prints entries as recorded, so that spelling can genuinely arrive. + assert setup_helpers["_is_unresolved_math_library_directory"](entry + "/") is True + + +@pytest.mark.parametrize( + "entry", + [ + # A real MKL installation spells the same arch directory below a prefix, and that + # directory genuinely exists, so dropping it would break a library resolving through it. + "/opt/intel/mkl/lib/intel64", + "/opt/intel/oneapi/mkl/latest/lib/intel64", + "/usr/lib/intel64", + "/home/user/lib/intel64", + # The arch name as a parent rather than as the entry itself. + "/lib/intel64/extra", + # Ordinary system directories, which differ from the bad entries only in the last part. + "/lib", + "/lib64", + "/lib/x86_64-linux-gnu", + # Relative entries are decided before this predicate is reached, but it must not claim + # one, or a wheel's own hop into a directory named intel64 would be dropped. + "$ORIGIN/../../lib", + "$ORIGIN/lib/intel64", + ], +) +def test_keeps_directories_that_name_a_real_prefix(entry, setup_helpers): + assert setup_helpers["_is_unresolved_math_library_directory"](entry) is False + # Also through the production predicate, not just the narrow one. Asserting only the narrow + # predicate let _is_usable_runtime_path start rejecting an ordinary system directory with the + # whole suite green, since nothing checked that these entries actually survive the filter. + assert setup_helpers["_is_usable_runtime_path"](entry, True, True) is True + + +# Read off the pybindings extension in the published x86_64 CPU nightly, in the recorded order. The +# MKL block sits after five relative hops and before one. On this wheel the entries are simply dead, +# since nothing resolves through them, so dropping them costs a few wasted lookups at load time and +# removes paths a user cannot have. On the CUDA wheel the same block precedes the hop into the CUDA +# runtime, which is where the ordering matters. +SHIPPED_RUNTIME_PATH = [ + "$ORIGIN/../../../torch/lib", + "$ORIGIN/../../src/executorch/lib", + "$ORIGIN/../../backends/qualcomm", + "$ORIGIN/../../lib", + "$ORIGIN/../../../lib64", + "/lib/intel64", + "/lib/intel64_win", + "/lib/win-x64", + "$ORIGIN/../../backends/cuda", +] + + +def _relative_torch_route_predicate(): + """setup.py's own has_relative_torch_route expression, lifted out of its function. + + Read rather than restated, because it decides the third argument to the filter under test. + A copy here would let the test keep passing after that expression changed, which is the one + failure a regression guard must not have. + """ + tree = ast.parse(_setup_source()) + for node in ast.walk(tree): + if not ( + isinstance(node, ast.Assign) + and any( + isinstance(target, ast.Name) and target.id == "has_relative_torch_route" + for target in node.targets + ) + ): + continue + # The expression iterates a name bound in setup.py's own scope, so it is rebound to this + # function's argument by compiling it as the body of a one-argument lambda. + source = ast.unparse(node.value) + iterated = "for entry in found" + assert iterated in source, ( + f"setup.py computes has_relative_torch_route as {source!r}, which no longer iterates " + "the name this test rebinds. Update the rebinding rather than leaving it a no-op." + ) + return eval( + f"lambda entries: {source.replace(iterated, 'for entry in entries')}" + ) + raise AssertionError( + "setup.py no longer computes has_relative_torch_route, so this test would pass the " + "filter an argument the shipped code never produces." + ) + + +def _filtered(entries, setup_helpers): + is_usable = setup_helpers["_is_usable_runtime_path"] + has_relative_torch_route = _relative_torch_route_predicate()(entries) + return [ + entry + for entry in entries + # True is safe_to_drop_toolkit_paths: this wheel ships no CUDA, so an absolute toolkit path + # in it names only the build machine. Packaging derives the same value from the built tree. + if is_usable(entry, True, has_relative_torch_route) + ] + + +def test_shipped_library_keeps_no_absolute_entry(setup_helpers): + kept = _filtered(SHIPPED_RUNTIME_PATH, setup_helpers) + assert [entry for entry in kept if entry.startswith("/")] == [] + + +def test_shipped_library_keeps_every_relative_hop(setup_helpers): + # Asserted separately from the absence of absolute entries, because a filter that dropped + # everything would satisfy that one while leaving the library unable to find its siblings. + kept = _filtered(SHIPPED_RUNTIME_PATH, setup_helpers) + assert kept == [ + entry for entry in SHIPPED_RUNTIME_PATH if not entry.startswith("/") + ] + + +def test_keeps_the_absolute_torch_directory_when_it_is_the_only_route(setup_helpers): + # The case that stops this being a blanket "drop everything absolute": an extension links + # torch and, with no relative route recorded, reaches it only through the directory the + # linker found it in. Dropping that would stop the extension importing. + entries = [ + "/opt/conda/envs/py_3.12/lib/python3.12/site-packages/torch/lib", + "/lib/intel64", + ] + assert _filtered(entries, setup_helpers) == [entries[0]] + + +def test_drops_the_absolute_torch_directory_when_a_relative_route_exists(setup_helpers): + # The other side of the same rule, and the only case that exercises the route expression at all. + # Without a relative entry in the list the route is False whatever that expression says, so the + # test above cannot tell a correct expression from an inverted one. + entries = [ + "$ORIGIN/../../../torch/lib", + "/opt/conda/envs/py_3.12/lib/python3.12/site-packages/torch/lib", + "/lib/intel64", + ] + assert _filtered(entries, setup_helpers) == [entries[0]] + + +_MATH_DIRECTORY_REASON = "a maths library directory whose prefix resolved empty" +_BUILD_DIRECTORY_REASON = "inside a build of this project" +_UNREACHABLE_REASON = "an absolute directory the wheel has a relative route to" + + +def _release_check_decision(): + """The release check's own per-entry decision, imported rather than replayed. + + Loaded as a module so the unit test exercises the function the wheel check calls, rather than + that function's source text. Reading the text let the rejection be deleted outright. + """ + path = REPO_ROOT / ".ci" / "scripts" / "wheel" / "test_shared_libraries.py" + spec = importlib.util.spec_from_file_location("_release_check_under_test", path) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module._unusable_runtime_path_kind + + +@pytest.mark.parametrize("arch", _MKL_ARCH_DIRECTORIES) +def test_release_check_rejects_the_math_directories(arch): + # The check and packaging are separate code, so a wheel built without patchelf keeps these + # entries and only the check would catch it. Driven off setup.py's own constant, and asserting + # the REASON rather than merely a rejection: the check rejects every absolute path it does not + # recognise, so an arch it was never taught about would otherwise pass through the catch-all. + decide = _release_check_decision() + assert ( + decide(f"/lib/{arch}", "_C.cpython-312-x86_64-linux-gnu.so") + == _MATH_DIRECTORY_REASON + ) + + +def test_release_check_rejects_each_kind_for_its_own_reason(): + # One assertion per rejecting branch, by reason, so deleting any single branch fails here. + # Asserting only "not None" let the build-directory branch and the catch-all each be removed + # on their own with every test green. + decide = _release_check_decision() + assert decide("/home/u/pip-out/lib", "_C.so") == _BUILD_DIRECTORY_REASON + assert decide("/opt/rocm/lib", "_C.so") == _UNREACHABLE_REASON + assert decide("", "_C.so") is not None + + +def test_release_check_rejects_a_build_directory_before_the_allowlist(): + # Order is load bearing and nothing else holds it. A torch directory inside a CI worker tree + # must be rejected, which only happens because the build-directory branch runs before the + # suffix allowlist gets to accept the /torch/lib ending. + decide = _release_check_decision() + entry = "/home/ec2-user/actions-runner/_work/executorch/pytorch/torch/lib" + assert decide(entry, "_C.so") == _BUILD_DIRECTORY_REASON + + +@pytest.mark.parametrize("arch", _MKL_ARCH_DIRECTORIES) +def test_release_check_still_accepts_a_real_mkl_installation(arch, setup_helpers): + # The two must agree on every arch in the shared constant. Packaging KEEPS a prefixed one, + # because the environment provides it, so a check that rejected it would fail a wheel packaging + # deliberately allowed and the builder could not satisfy both. Parametrized off the constant, so + # adding an arch to setup.py without teaching the check about it fails here. + entry = f"/opt/intel/mkl/lib/{arch}" + assert setup_helpers["_is_usable_runtime_path"](entry, True, True) is True + assert ( + _release_check_decision()(entry, "_C.cpython-312-x86_64-linux-gnu.so") is None + ) + + +def test_the_wheel_scan_consults_the_classifier(): + # The unit tests above call the classifier directly, and the wheel scan is the only thing that + # applies it to a real library. That scan needs an installed wheel, so it cannot run here; + # what is checkable is the wiring, and severing it left all tests green. Asserted on the AST so + # a rename or an accidental deletion fails rather than silently disabling the enforcement. + path = REPO_ROOT / ".ci" / "scripts" / "wheel" / "test_shared_libraries.py" + tree = ast.parse(path.read_text(encoding="utf-8")) + scan = next( + ( + node + for node in ast.walk(tree) + if isinstance(node, ast.FunctionDef) + and node.name == "test_no_absolute_runtime_paths" + ), + None, + ) + assert scan is not None, "the wheel scan this check enforces no longer exists" + called = { + node.func.id + for node in ast.walk(scan) + if isinstance(node, ast.Call) and isinstance(node.func, ast.Name) + } + assert "_unusable_runtime_path_kind" in called, ( + "test_no_absolute_runtime_paths no longer calls _unusable_runtime_path_kind, so the wheel " + "scan would report every shipped library clean while these unit tests still pass." + ) + + +def test_release_check_accepts_a_relative_entry(): + # A relative hop is the normal case and must never be rejected, whatever the absolute rules do. + assert _release_check_decision()("$ORIGIN/../../lib", "_C.so") is None diff --git a/.ci/scripts/wheel/test_shared_libraries.py b/.ci/scripts/wheel/test_shared_libraries.py index 2ed6464105b..04646c7549c 100644 --- a/.ci/scripts/wheel/test_shared_libraries.py +++ b/.ci/scripts/wheel/test_shared_libraries.py @@ -2017,6 +2017,79 @@ def _names_a_build_directory(entry: str) -> bool: ) +# Absolute directories a shipped library may name. PyTorch's own is allowed because the wheel +# neither declares nor bundles PyTorch, so an absolute path is the only way to reach it. The maths +# library arch directories are allowed because a real installation spells them below a prefix, as +# /opt/intel/mkl/lib/intel64, which the environment genuinely provides. +# +# Matched as a suffix. A substring test exempted any path merely CONTAINING one of these, so a +# directory such as /home/user/torch/lib.backup/stage passed without reaching the build-directory +# classifier. Both "_win" spellings are listed explicitly now that the match is anchored. +_ALLOWED_ABSOLUTE_SUFFIXES = ( + "/torch/lib", + "/lib/intel64", + "/lib/intel64_win", + "/lib/win-x64", +) + +# The same arch directories with an EMPTY prefix, which is what PyTorch's exported link interface +# records when its MKL_ROOT resolves to nothing. They point nowhere, and they sit ahead of the +# relative hops packaging appends, which is the shadowing a CUDA toolkit prefix is rejected for. +# Matched exactly rather than folded into the allowlist above, because the two differ only in the +# prefix and a suffix match cannot tell them apart. +_UNRESOLVED_MATH_DIRECTORIES = ( + "/lib/intel64", + "/lib/intel64_win", + "/lib/win-x64", +) + +# Held for a wheel that bundles PyTorch's libraries rather than declaring them: such a copy records +# the CUDA toolkit directory of the machine that built IT, which is not this project's to fix. +# +# No wheel ships one today, so this clause never fires. It stays as a guard for a future +# bundling change; if that never comes, delete it rather than leaving an unexercised exemption. +_VENDORED_PREFIXES = ( + "libtorch", + "libc10", + "libshm", + "libcaffe2", + "libgomp", + "libiomp", +) + + +def _unusable_runtime_path_kind(entry: str, library_name: str) -> str | None: + """Why a recorded runtime search path is one a user cannot use, or None if it is fine. + + A library must not name an absolute directory the wheel has a relative route to. The one that + shipped was a CUDA toolkit prefix recorded on the build machine: it sat ahead of the relative + hop, so a user with a toolkit at the same prefix resolved the runtime from there instead of from + the declared dependency, and the builder always has one, so nothing exercised the hop. Stated as + a property rather than a list of known-bad directories, because a list only catches what someone + already thought of and that prefix was not on one. + + Order matters. The build-directory branch runs before the suffix allowlist so that a torch + directory inside a CI worker tree is rejected rather than accepted for its /torch/lib ending. + + Module scope so a unit test can call this directly and compare the reason it returns. Inline in + the caller's loop it could only be reached by building a wheel. + """ + if not entry: + # The loader reads an empty entry as the process working directory. + return "the process working directory" + if not entry.startswith("/") or library_name.startswith(_VENDORED_PREFIXES): + return None + if _names_a_build_directory(entry): + return "inside a build of this project" + if entry.rstrip("/") in _UNRESOLVED_MATH_DIRECTORIES: + return "a maths library directory whose prefix resolved empty" + if any( + entry.rstrip("/").endswith(allowed) for allowed in _ALLOWED_ABSOLUTE_SUFFIXES + ): + return None + return "an absolute directory the wheel has a relative route to" + + def test_no_absolute_runtime_paths() -> None: """No shipped library may search a directory a user does not have. @@ -2059,53 +2132,6 @@ def test_no_absolute_runtime_paths() -> None: package_dir = _installed_package_dir() - # This project's libraries must not name an absolute directory the wheel has a relative route to. The - # one that shipped was a CUDA toolkit prefix recorded on the build machine: it sat ahead of the relative - # hop, so a user with a toolkit at the same prefix resolved the CUDA runtime from there instead of from - # the declared dependency, and the builder always has one, so nothing exercised the hop. - # - # Stated as a property rather than a list of known-bad directories, because a list only catches what - # someone already thought of and that prefix was not on one. - # - # PyTorch's own directory is allowed: the wheel neither declares nor bundles PyTorch, so an absolute - # path is the only way to reach it. The maths library directories are allowed too. They arrive as - # -L flags in PyTorch's exported link interface, which CMake mirrors into the runtime path, so every - # library here that links PyTorch carries them. They point nowhere on any machine: measured on the - # link line as -L/lib/intel64 -L/lib/intel64_win -L/lib/win-x64, which is a prefix variable that - # resolved empty leaving the concatenation at the filesystem root. - # - # Matched as a suffix, the same way packaging decides what to strip at setup.py:1300. A substring - # test exempted any path merely CONTAINING one of these, so a directory such as - # /home/user/torch/lib.backup/stage passed without ever reaching the build-directory classifier. - # Both "_win" spellings are listed explicitly now that the match is anchored. - # - # A torch directory inside a CI worker tree, such as - # /home/ec2-user/actions-runner/_work/.../pytorch/torch/lib, is rejected rather than allowed: the - # build-directory classifier sees the worker components and the allowlist never gets to accept the - # /torch/lib suffix. Packaging strips the same entry, because every extension that names Torch now - # records a relative route to it. - allowed_absolute = ( - "/torch/lib", - "/lib/intel64", - "/lib/intel64_win", - "/lib/win-x64", - ) - # Held for a wheel that bundles PyTorch's libraries rather than declaring them: such a copy records - # the CUDA toolkit directory of the machine that built IT, which is not this project's to fix. - # - # No wheel ships one today. Six wheels across manylinux and macOS contain zero files with these - # prefixes, because the wheel declares torch as a dependency and there is no auditwheel step, so - # this clause is currently never false. It stays as a guard for a future bundling change; if that - # never comes, delete it rather than leaving an unexercised exemption in the check. - vendored_prefixes = ( - "libtorch", - "libc10", - "libshm", - "libcaffe2", - "libgomp", - "libiomp", - ) - offenders = {} inspected = 0 with_a_runtime_path = 0 @@ -2122,27 +2148,9 @@ def test_no_absolute_runtime_paths() -> None: with_a_runtime_path += 1 bad = [] for entry in entries: - if not entry: - bad.append("") - elif ( - entry.startswith("/") - and not library.name.startswith(vendored_prefixes) - and ( - _names_a_build_directory(entry) - or not any( - entry.rstrip("/").endswith(allowed) - for allowed in allowed_absolute - ) - ) - ): - # Named separately so the message says which kind it is: a build directory and a - # toolkit prefix are the same defect with different causes. - kind = ( - "inside a build of this project" - if _names_a_build_directory(entry) - else "an absolute directory the wheel has a relative route to" - ) - bad.append(f"{entry} ({kind})") + kind = _unusable_runtime_path_kind(entry, library.name) + if kind is not None: + bad.append(f"{entry or ''} ({kind})") if bad: offenders[str(library.relative_to(package_dir))] = bad diff --git a/setup.py b/setup.py index c80131657f5..8980d1bd94a 100644 --- a/setup.py +++ b/setup.py @@ -238,6 +238,9 @@ def _minimal_packages() -> List[str]: "13": ("nvidia/cu13/lib",), } +# Arch subdirectory names MKL's exported link interface appends to its prefix. +_MKL_ARCH_DIRECTORIES = ("intel64", "intel64_win", "win-x64") + def _cmake_args() -> List[str]: """CMAKE_ARGS split into arguments, tolerating an unbalanced quote. @@ -600,6 +603,36 @@ def cuda_named(part: str) -> bool: return len(parts) >= 4 and parts[-3] == "targets" and cuda_named(parts[-4]) +def _is_unresolved_math_library_directory(entry: str) -> bool: + """Whether a runtime search path entry is a maths library directory whose prefix resolved empty. + + Torch's exported CMake package creates a caffe2::mkl imported target, and linking torch brings it + in even though this project never asks for MKL. Its link directories are a hardcoded list of four, + spelled below MKL_ROOT, which resolves to nothing here, so what the linker records is left + anchored at the filesystem root: /lib, /lib/intel64, /lib/intel64_win, /lib/win-x64. The bare + /lib does not survive, because CMake filters its own implicit link directories out of the link + line. + + An empty prefix is the whole signature, so the entry must be exactly /lib/. A real + installation spells the same arch directory below a prefix, as /opt/intel/mkl/lib/intel64, and + that one is a directory the environment genuinely provides. + + Only the three arch directories are handled. The bare prefix/lib is deliberately left, because + with a resolving prefix it is an ordinary library directory this project has no business + dropping, and with an empty one CMake filters it out as an implicit link directory before the + link line is built, so it never reaches a shipped library. + + Dropping these loses nothing: no shipped library names an MKL or OpenMP runtime in DT_NEEDED, so + nothing resolves through them, and they sit ahead of the relative hops appended below. + """ + parts = PurePosixPath(entry).parts + return ( + len(parts) == 3 + and parts[:2] == ("/", "lib") + and parts[2] in _MKL_ARCH_DIRECTORIES + ) + + def _package_relative_depth(library: Path) -> int: """How many directories separate a shipped library from the installed package root. @@ -1446,6 +1479,9 @@ def _is_usable_runtime_path( # directory from the machine that built it. It is kept when no relative route exists, because # then it is the only way this library finds torch. # + # A maths library directory whose prefix resolved empty is dropped for the same reason: nothing in + # the wheel resolves through it, and it sits ahead of the relative hops appended below. + # # Anything else absolute stays, because it is a dependency the environment provides and the wheel # has no relative answer for. # @@ -1465,6 +1501,8 @@ def _is_usable_runtime_path( # CUDA version, which is the one absolute path that has to survive. if safe_to_drop_toolkit_paths and _is_cuda_toolkit_directory(entry): return False + if _is_unresolved_math_library_directory(entry): + return False if entry.rstrip("/").endswith("/torch/lib") and has_relative_torch_route: return False return True From af6454ec2071840aab98729869fb3db17cbfea84 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 3 Sep 2026 23:59:32 -0700 Subject: [PATCH 031/190] Bump the PyTorch pin to 2.14 (#22501) ## Summary PyTorch 2.14.0 is released, so move the pin from 2.13 to it. The tree is mid-cycle at 1.5.0 and the pin normally follows the current stable release. The pin itself is a few version strings plus a re-sync of the vendored c10 headers, which CI requires to match PyTorch's tree byte for byte. Everything else here is a change 2.14 forces, each as its own commit. ## What 2.14 forces **C++20 where ATen headers are compiled.** `intrusive_ptr.h` uses `operator<=>` and `TensorBase.h` uses `requires`, neither guarded, so a target that includes ATen cannot parse at C++17. This follows the pattern already in the tree: six CMake targets set `CXX_STANDARD 20` for exactly this reason. Added the Vulkan op tests to that list, and taught the Buck wrapper to make the same per-target distinction. Non-ATen targets, including everything the bare-metal presets build, still compile at C++17. **Two new AOTI entry points.** 2.14's generated wrapper calls `aoti_torch_is_defined` and `aoti_torch_empty_strided_pinned`, neither of which existed here. The first answers whether a tensor holds storage. The second allocates ordinary host memory, because there is no pinned allocator in this runtime and pinning only lets a copy to the device overlap other work. It refuses a device other than CPU, as PyTorch's own does. **Derived shim spellings.** Some custom operators now reach the fallback-kernel check under the name Inductor derives rather than the name they are registered under. Both spellings are accepted for the Metal ops and the CUDA int4 pack matmul. **Build and CI repairs.** `setup.py bdist_wheel` is gone, so the macOS source build uses the standard frontend. 2.14 is published for ROCm 7.2, not 7.1. PyTorch now compiles at C++20, which makes CMake scan for modules with scanners the images do not have. And a pip `cmake` in the build environment made the image build silently produce a wheel with no BLAS, so the image's own cmake is used and the result is now asserted rather than assumed. Two of these were already broken before this branch: the unpinned `katex` install, and the `cmake` interaction. ## Behaviour changes The re-sync is not cosmetic. `overflows()` answers differently for float to integer casts, in both directions. Filling an int8 tensor with 127.5 was refused and now gives 127. A value at 2^63 cast to int64 was accepted and wrapped, and is now refused. The portable kernels reach this through `check_overflow_cast`, so `full`, `full_like`, `fill`, `scalar_tensor`, `hardtanh`, `leaky_relu`, `scatter` and `constant_pad_nd` inherit it. The header is a faithful copy of upstream, so this is not ours to undo, but it should be visible rather than buried in a header diff. Neither direction had a test; both do now. Two smaller ones come with it. Building a delegate sorts its placeholders, and lifted tensor constants were sorted as if they were user inputs, which left a constant after one and made any later constant insert impossible. They now sort with the parameters and buffers. And dropout was missing from the list of operators that take their input's observer. It is an identity once a model is not training, which is how a quantized model is deployed, so measuring it separately gave the operator after it a different scale. XNNPACK then refused a reshape whose two sides disagreed. Both spellings are listed, as several other operators already are. ## Test plan Each commit carries its own. Covering the whole change: `compare_dirs.sh`, the header check CI runs, passes against a real `release/2.14` checkout and fails without the re-sync. All eight vendored headers are byte-identical to upstream; the one build-file edit follows a file upstream moved. The behaviour change was measured by compiling `overflows()` from both branches, not read off the diff. Both new cases fail on the old header, so neither is vacuous. The regenerated import library was checked member by member: 45 advertised names against 42 before, none lost, and its archive metadata is zeroed so the file is reproducible. Not covered locally: the image build, the CUDA, ROCm, Qualcomm and Metal jobs, the export suites, the Buck build, and anything needing Windows. CI has all of it. ## Known outstanding Three things this change does not repair. None of them is new here, and each needs its own change rather than being folded into a version bump. No job links the checked-in Windows link stub with the Microsoft linker. The cross build uses a GNU linker, the same family that produced the archive, so it cannot answer whether the Microsoft one accepts it. That needs a Windows lowering job, which does not exist yet. There is now a note beside the file describing how it is produced. The two Cortex-M model tests expect three quantize pairs where they used to expect one. Two of those three are round trips: the value is converted back to float and immediately converted again at the same scale, measured as 0.004997437 and 0.0000032501507 on both sides. So a pass that used to absorb them no longer matches. The pad operator those pairs sit around is created by this backend's own passes, after the shared pass that folds such pairs has already run, so absorbing them means changing the order or extending a pass another backend shares. Correct output, more work than needed. Asking for pinned host memory gives ordinary host memory. There is no pinned allocator in this runtime, and adding one needs a matching release path, so the entry point allocates ordinary memory and says so. A caller that wants a copy to the device to overlap other work will not get the overlap. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --------- Co-authored-by: PyTorch Bot --- .ci/docker/ci_commit_pins/pytorch.txt | 2 +- .ci/docker/common/install_docs_reqs.sh | 4 +- .ci/docker/common/install_pytorch.sh | 20 +++- .ci/scripts/test-rocm-aoti.sh | 2 +- .ci/scripts/test-rocm-voxtral.sh | 2 +- .ci/scripts/utils.sh | 33 +++++-- .ci/scripts/wheel/test_cpp_sdk.py | 2 +- .github/workflows/rocm.yml | 6 +- .github/workflows/windows-msvc.yml | 2 +- backends/aoti/common_shims.cpp | 5 + backends/aoti/common_shims.h | 4 + backends/aoti/common_shims_slim.cpp | 8 ++ backends/aoti/common_shims_slim.h | 4 + backends/apple/metal/metal_backend.py | 8 ++ .../arm/test/misc/test_transpose_counts.py | 2 +- .../cortex_m/test/misc/test_portable_int8.py | 5 +- backends/cortex_m/test/models/test_ds_cnn.py | 18 +++- .../cortex_m/test/models/test_mobilenet_v2.py | 12 ++- backends/cortex_m/test/tester.py | 4 +- backends/cuda/cuda_backend.py | 3 + backends/cuda/runtime/aoti_cuda_shims.lib | Bin 12738 -> 43802 bytes backends/cuda/runtime/aoti_cuda_shims.lib.md | 36 +++++++ backends/cuda/runtime/shims/memory.cpp | 25 +++++ backends/cuda/runtime/shims/memory.h | 25 +++++ backends/cuda/tests/test_sort_shim.py | 2 + .../node_converter/test_clone_converter.py | 5 +- backends/vulkan/test/op_tests/CMakeLists.txt | 4 + .../quantizer/xnnpack_quantizer_utils.py | 3 + exir/lowered_backend_module.py | 24 +++-- install_requirements.py | 4 +- kernels/test/ScalarOverflowTestMacros.h | 56 ++++++----- kernels/test/op_full_test.cpp | 13 +++ .../core/portable_type/c10/c10/targets.bzl | 1 - .../core/portable_type/c10/c10/util/complex.h | 17 ++-- .../c10/c10/util/llvmMathExtras.h | 1 + .../portable_type/c10/c10/util/overflows.h | 17 +++- .../c10/torch/headeronly/macros/Macros.h | 9 ++ .../c10/torch/headeronly/util/Half.h | 2 +- .../torch/headeronly/util/TypeSafeSignMath.h | 90 ++++-------------- .../c10/torch/headeronly/util/complex.h | 60 ++++++++++++ .../headeronly}/util/complex_utils.h | 8 +- .../executorch/build/runtime_wrapper.bzl | 71 ++++++++++---- torch_pin.py | 2 +- 43 files changed, 439 insertions(+), 182 deletions(-) create mode 100644 backends/cuda/runtime/aoti_cuda_shims.lib.md rename runtime/core/portable_type/c10/{c10 => torch/headeronly}/util/complex_utils.h (80%) diff --git a/.ci/docker/ci_commit_pins/pytorch.txt b/.ci/docker/ci_commit_pins/pytorch.txt index 401a0594d98..6c3fe42ddf3 100644 --- a/.ci/docker/ci_commit_pins/pytorch.txt +++ b/.ci/docker/ci_commit_pins/pytorch.txt @@ -1 +1 @@ -release/2.13 +release/2.14 diff --git a/.ci/docker/common/install_docs_reqs.sh b/.ci/docker/common/install_docs_reqs.sh index ea54d90523e..2794ffb8fc9 100755 --- a/.ci/docker/common/install_docs_reqs.sh +++ b/.ci/docker/common/install_docs_reqs.sh @@ -20,7 +20,9 @@ if [ -n "$BUILD_DOCS" ]; then apt-get update apt-get install -y --no-install-recommends yarn - yarn global add katex --prefix /usr/local + # katex 0.18.5 requires commander@15 / node >= 22.12; pin to the last + # release compatible with the node 16 installed above + yarn global add katex@0.18.4 --prefix /usr/local sudo apt-get -y install doxygen diff --git a/.ci/docker/common/install_pytorch.sh b/.ci/docker/common/install_pytorch.sh index 0ac5e79cf4a..e51a8886afd 100755 --- a/.ci/docker/common/install_pytorch.sh +++ b/.ci/docker/common/install_pytorch.sh @@ -76,21 +76,31 @@ install_pytorch_and_domains() { # the image compiler cannot satisfy. The venv inherits the image's # site-packages, so PyTorch still builds against the same numpy. # - # Keep the list in sync with pytorch/pyproject.toml [build-system].requires. + # Keep in sync with pytorch/pyproject.toml [build-system].requires. local build_venv=/tmp/pytorch-build-venv rm -rf "${build_venv}" conda_run python -m venv --system-site-packages "${build_venv}" + # No pip cmake: scikit-build-core would prefer it over the image's, and it + # searches site-packages, where MKL and libomp are not. conda_run "${build_venv}/bin/pip" install build "scikit-build-core>=1.0" \ - "setuptools>=77.0.0,<82" "cmake>=3.27,<4" ninja "packaging>=24.2" \ - "typing-extensions>=4.10.0" pyyaml six - conda_run "${build_venv}/bin/python" -m build --wheel --no-isolation + ninja "packaging>=24.2" "typing-extensions>=4.10.0" pyyaml six numpy + # These images have no module scanner, and nothing here uses modules. + conda_run env CMAKE_CXX_SCAN_FOR_MODULES=OFF \ + "${build_venv}/bin/python" -m build --wheel --no-isolation rm -rf "${build_venv}" pip_install "$(echo dist/*.whl)" + # A build with no BLAS succeeds silently. Run from / to import the wheel. + (cd / && conda_run python -c " +import torch +assert torch._C.has_lapack, 'built without LAPACK' +torch.linalg.qr(torch.randn(4, 4)) +") + # Grab the pinned audio and vision commits from PyTorch TORCHAUDIO_VERSION=release/2.11 export TORCHAUDIO_VERSION - TORCHVISION_VERSION=release/0.28 + TORCHVISION_VERSION=release/0.29 export TORCHVISION_VERSION install_domains diff --git a/.ci/scripts/test-rocm-aoti.sh b/.ci/scripts/test-rocm-aoti.sh index 0f4ac9d826f..00592bc4bf6 100644 --- a/.ci/scripts/test-rocm-aoti.sh +++ b/.ci/scripts/test-rocm-aoti.sh @@ -7,7 +7,7 @@ set -euo pipefail -ROCM_VERSION="${ROCM_VERSION:-7.1}" +ROCM_VERSION="${ROCM_VERSION:-7.2}" ROCM_PATH="${ROCM_PATH:-/opt/rocm}" PYTORCH_ROCM_INDEX="${PYTORCH_ROCM_INDEX:-https://download.pytorch.org/whl/test/rocm${ROCM_VERSION}}" TORCHAO_ROCM_WHEEL_BASE="${TORCHAO_ROCM_WHEEL_BASE:-https://download.pytorch.org/whl/nightly/rocm${ROCM_VERSION}}" diff --git a/.ci/scripts/test-rocm-voxtral.sh b/.ci/scripts/test-rocm-voxtral.sh index cd0562c3808..eb00c1efc73 100644 --- a/.ci/scripts/test-rocm-voxtral.sh +++ b/.ci/scripts/test-rocm-voxtral.sh @@ -7,7 +7,7 @@ set -euo pipefail -ROCM_VERSION="${ROCM_VERSION:-7.1}" +ROCM_VERSION="${ROCM_VERSION:-7.2}" ROCM_PATH="${ROCM_PATH:-/opt/rocm}" EXPECTED_ROCM_ARCH="${EXPECTED_ROCM_ARCH:-gfx950}" EXPECTED_WARP_SIZE="${EXPECTED_WARP_SIZE:-64}" diff --git a/.ci/scripts/utils.sh b/.ci/scripts/utils.sh index 234e162e48e..3d573daa395 100644 --- a/.ci/scripts/utils.sh +++ b/.ci/scripts/utils.sh @@ -106,8 +106,8 @@ install_pytorch_and_domains() { local python_version=$(python -c 'import platform; v=platform.python_version_tuple(); print(f"{v[0]}{v[1]}")') local torch_release=$(cat version.txt) # Download key must match the upload key below (basename of dist/*.whl, - # which always carries setup.py's resolved +gitHASH). Branch-ref pins - # like `release/2.13` would otherwise produce `+gitrelease` here and + # which always carries the build's resolved +gitHASH). Branch-ref pins + # like `release/2.14` would otherwise produce `+gitrelease` here and # never hit the cache. local torch_short_hash=$(git rev-parse --short=7 HEAD) local torch_wheel_path="cached_artifacts/pytorch/executorch/pytorch_wheels/${system_name}/${python_version}" @@ -127,18 +127,31 @@ install_pytorch_and_domains() { if [[ "${torch_wheel_not_found}" == "1" ]]; then echo "No cached wheel found, continue with building PyTorch at ${TORCH_VERSION}" - # Install PyTorch's own build-time deps so the source build does not - # silently inherit them from whatever else happens to be in the env - # (e.g. executorch's requirements-ci.txt). - pip install -r requirements-build.txt git submodule update --init --recursive if [[ "$(uname -m)" == "aarch64" ]]; then export BUILD_IGNORE_SVE_UNAVAILABLE=1 fi - USE_DISTRIBUTED=1 python setup.py bdist_wheel + # PyTorch dropped setup.py. Build in a throwaway environment that can see the + # active one, so its build requirements, which pin a cmake that would be + # preferred over the one on PATH, cannot disturb what is installed here. + # + # Keep in sync with pytorch/pyproject.toml [build-system].requires. + local build_venv=/tmp/pytorch-build-venv + rm -rf "${build_venv}" + python -m venv --system-site-packages "${build_venv}" + "${build_venv}/bin/pip" install build "scikit-build-core>=1.0" ninja \ + "packaging>=24.2" "typing-extensions>=4.10.0" pyyaml six numpy + USE_DISTRIBUTED=1 "${build_venv}/bin/python" -m build --wheel --no-isolation + rm -rf "${build_venv}" pip install "$(echo dist/*.whl)" - - # Invariant: the basename setup.py just produced must match the cache + # A build with no BLAS succeeds silently, so check rather than assume. + (cd / && python -c " +import torch +assert torch._C.has_lapack, 'built without LAPACK' +torch.linalg.qr(torch.randn(4, 4)) +") + + # Invariant: the basename the build just produced must match the cache # URL we'd reconstruct on the next run. If they diverge (someone edits # torch_wheel_name above, or PyTorch renames its wheels), the cache # will silently miss and every macOS run will fall back to a ~30-min @@ -178,7 +191,7 @@ install_pytorch_and_domains() { # Grab the pinned audio and vision commits from PyTorch TORCHAUDIO_VERSION=release/2.11 export TORCHAUDIO_VERSION - TORCHVISION_VERSION=release/0.28 + TORCHVISION_VERSION=release/0.29 export TORCHVISION_VERSION install_domains diff --git a/.ci/scripts/wheel/test_cpp_sdk.py b/.ci/scripts/wheel/test_cpp_sdk.py index 29b2021d398..202a400e6c7 100644 --- a/.ci/scripts/wheel/test_cpp_sdk.py +++ b/.ci/scripts/wheel/test_cpp_sdk.py @@ -991,7 +991,7 @@ def test_every_shipped_header_compiles(work_dir: Path) -> None: # These say in their own text that they must not be included directly, and name the header to # include instead. Including one anyway is a use error rather than a packaging defect. "c10/util/complex_math.h", - "c10/util/complex_utils.h", + "torch/headeronly/util/complex_utils.h", ) source = work_dir / "header_probe.cpp" diff --git a/.github/workflows/rocm.yml b/.github/workflows/rocm.yml index 97d180885f7..154454eb145 100644 --- a/.github/workflows/rocm.yml +++ b/.github/workflows/rocm.yml @@ -169,7 +169,7 @@ jobs: strategy: fail-fast: false matrix: - rocm-version: ["7.1"] + rocm-version: ["7.2"] uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write @@ -206,7 +206,7 @@ jobs: strategy: fail-fast: false matrix: - rocm-version: ["7.1"] + rocm-version: ["7.2"] with: timeout: 180 no-sudo: true @@ -248,7 +248,7 @@ jobs: strategy: fail-fast: false matrix: - rocm-version: ["7.1"] + rocm-version: ["7.2"] with: timeout: 180 no-sudo: true diff --git a/.github/workflows/windows-msvc.yml b/.github/workflows/windows-msvc.yml index bde38a8288f..269f2cf2381 100644 --- a/.github/workflows/windows-msvc.yml +++ b/.github/workflows/windows-msvc.yml @@ -91,7 +91,7 @@ jobs: - name: Install build dependencies shell: pwsh - run: python -m pip install pyyaml torch==2.13.0 --extra-index-url https://download.pytorch.org/whl/test/cpu + run: python -m pip install pyyaml torch==2.14.0 --extra-index-url https://download.pytorch.org/whl/test/cpu - name: Build ExecuTorch shell: pwsh diff --git a/backends/aoti/common_shims.cpp b/backends/aoti/common_shims.cpp index f3a34a09987..e83a9576b6c 100644 --- a/backends/aoti/common_shims.cpp +++ b/backends/aoti/common_shims.cpp @@ -159,6 +159,11 @@ AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel) { return Error::Ok; } +AOTITorchError aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined) { + *ret_is_defined = tensor != nullptr; + return Error::Ok; +} + // Device and layout utility functions int32_t aoti_torch_device_type_cpu() { // Let's say cpu is 0 for ET as well diff --git a/backends/aoti/common_shims.h b/backends/aoti/common_shims.h index d057279e22a..0b85b09ee51 100644 --- a/backends/aoti/common_shims.h +++ b/backends/aoti/common_shims.h @@ -62,6 +62,10 @@ aoti_torch_get_dim(Tensor* tensor, int64_t* ret_dim); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel); +// PyTorch has an undefined-tensor state with no equivalent here: null check. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined); + // Utility functions for device and layout information AOTI_SHIM_EXPORT int32_t aoti_torch_device_type_cpu(); AOTI_SHIM_EXPORT int32_t aoti_torch_layout_strided(); diff --git a/backends/aoti/common_shims_slim.cpp b/backends/aoti/common_shims_slim.cpp index c8c7408aa62..d9e255bf220 100644 --- a/backends/aoti/common_shims_slim.cpp +++ b/backends/aoti/common_shims_slim.cpp @@ -68,6 +68,14 @@ AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel) { return Error::Ok; } +AOTITorchError aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined) { + if (ret_is_defined == nullptr) { + return Error::InvalidArgument; + } + *ret_is_defined = tensor != nullptr && tensor->defined(); + return Error::Ok; +} + int32_t aoti_torch_layout_strided() { // Slimtensor only support strided layout, the return value will always be 0, // a.k.a at::Layout::Strided; diff --git a/backends/aoti/common_shims_slim.h b/backends/aoti/common_shims_slim.h index c5a5cab9413..0d60f578605 100644 --- a/backends/aoti/common_shims_slim.h +++ b/backends/aoti/common_shims_slim.h @@ -51,6 +51,10 @@ aoti_torch_get_dim(Tensor* tensor, int64_t* ret_dim); AOTI_SHIM_EXPORT AOTITorchError aoti_torch_get_numel(Tensor* tensor, int64_t* ret_numel); +// Undefined means either a null handle or a tensor whose storage was released. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_is_defined(Tensor* tensor, bool* ret_is_defined); + AOTI_SHIM_EXPORT int32_t aoti_torch_layout_strided(); // ============================================================ diff --git a/backends/apple/metal/metal_backend.py b/backends/apple/metal/metal_backend.py index 57ca0ddf83e..0b9a852343b 100644 --- a/backends/apple/metal/metal_backend.py +++ b/backends/apple/metal/metal_backend.py @@ -32,16 +32,24 @@ def get_device_name(cls) -> str: @classmethod def get_supported_fallback_kernels(cls) -> Dict[str, Any]: return { + # An operator named the way it is registered also needs the name + # Inductor derives for it, so several appear under both. "aoti_torch_mps_addmm_out": None, "aoti_torch_mps_bmm_out": None, "aoti_torch_mps_convolution": None, "aoti_torch_mps_mm_out": None, "at::_ops::_scaled_dot_product_attention_math_for_mps::call": None, + "aoti_torch_mps__scaled_dot_product_attention_math_for_mps": None, "at::_ops::_scaled_dot_product_attention_math_for_mps_v2::call": None, + "aoti_torch_mps__scaled_dot_product_attention_math_for_mps_v2": None, "torchao::_linear_fp_act_4bit_weight": None, + "aoti_torch_mps__linear_fp_act_4bit_weight": None, "at::_ops::topk::call": None, + "aoti_torch_mps_topk": None, "metal::gather_qmv": None, + "aoti_torch_mps_gather_qmv": None, "metal::gated_delta_rule": None, + "aoti_torch_mps_gated_delta_rule": None, } @classmethod diff --git a/backends/arm/test/misc/test_transpose_counts.py b/backends/arm/test/misc/test_transpose_counts.py index bd73ddfe0cb..643a14b3301 100644 --- a/backends/arm/test/misc/test_transpose_counts.py +++ b/backends/arm/test/misc/test_transpose_counts.py @@ -543,7 +543,7 @@ def forward(self, x: torch.Tensor): "groupnorm_channels_last": TransposeCountCase( GroupNormModule(), (torch.randn(1, 4, 4, 4).to(memory_format=torch.channels_last),), - 2, + 1, ), "cumsum_rank4_dim3_channels_last": TransposeCountCase( CumsumModule(), diff --git a/backends/cortex_m/test/misc/test_portable_int8.py b/backends/cortex_m/test/misc/test_portable_int8.py index 6efeec9e5b0..41f7f254863 100644 --- a/backends/cortex_m/test/misc/test_portable_int8.py +++ b/backends/cortex_m/test/misc/test_portable_int8.py @@ -716,7 +716,10 @@ def _quantize_and_export( OP_CASES, xfails=xfails, strict=False, - skips={"while_loop": "Has been observed to hang randomly."}, + skips={ + "while_loop": "Has been observed to hang randomly.", + "dropout": "Not training, so it folds away and no node survives to carry int8.", + }, ) def test_shared_qspec_portable_int8_ops(op_case: OpCase) -> None: tester = CortexMTester(op_case.module, op_case.example_inputs) diff --git a/backends/cortex_m/test/models/test_ds_cnn.py b/backends/cortex_m/test/models/test_ds_cnn.py index 206af19a61e..ba5b2da6b78 100644 --- a/backends/cortex_m/test/models/test_ds_cnn.py +++ b/backends/cortex_m/test/models/test_ds_cnn.py @@ -15,7 +15,6 @@ "executorch_exir_dialects_edge__ops_aten_linear_default": 1, "executorch_exir_dialects_edge__ops_aten_relu_default": 9, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 2, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 18, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 17, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 15, @@ -23,14 +22,13 @@ ops_after_transforms: dict[str, int] = { "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 3, "executorch_exir_dialects_edge__ops_cortex_m_pad_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 3, "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 4, "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 5, "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 2, } test_cases = { @@ -43,11 +41,21 @@ } +ops_absent_after_transforms: list[str] = [ + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", +] + + @parametrize("test_case", test_cases) def test_dialect_ds_cnn(test_case): inputs = test_case.get_example_inputs() tester = CortexMTester(test_case.model, inputs) - tester.test_dialect(ops_before_transforms, ops_after_transforms, qtol=1) + tester.test_dialect( + ops_before_transforms, + ops_after_transforms, + qtol=1, + ops_absent_after_transforms=ops_absent_after_transforms, + ) @parametrize("test_case", test_cases) diff --git a/backends/cortex_m/test/models/test_mobilenet_v2.py b/backends/cortex_m/test/models/test_mobilenet_v2.py index 67f0937a006..9bc99e4bf2c 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v2.py +++ b/backends/cortex_m/test/models/test_mobilenet_v2.py @@ -20,7 +20,6 @@ "executorch_exir_dialects_edge__ops_aten_hardtanh_default": 35, "executorch_exir_dialects_edge__ops_aten_linear_default": 1, "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_channel_default": 104, "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 79, "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 67, @@ -28,14 +27,13 @@ ops_after_transforms: dict[str, int] = { "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, - "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 2, + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 2, "executorch_exir_dialects_edge__ops_cortex_m_quantized_add_default": 10, "executorch_exir_dialects_edge__ops_cortex_m_quantized_avg_pool2d_default": 1, "executorch_exir_dialects_edge__ops_cortex_m_quantized_conv2d_default": 35, "executorch_exir_dialects_edge__ops_cortex_m_quantized_depthwise_conv2d_default": 17, "executorch_exir_dialects_edge__ops_cortex_m_quantized_linear_default": 1, - "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default": 1, } # Use larger sample set for calibration to get better quantization @@ -54,6 +52,11 @@ } +ops_absent_after_transforms: list[str] = [ + "executorch_exir_dialects_edge__ops_dim_order_ops__clone_dim_order_default", +] + + @parametrize("test_case", test_cases) def test_dialect_mv2(test_case): inputs = test_case.get_example_inputs() @@ -63,6 +66,7 @@ def test_dialect_mv2(test_case): ops_after_transforms, qtol=10, calibration_samples=calibration_samples, + ops_absent_after_transforms=ops_absent_after_transforms, ) # assert that top 1 output matches diff --git a/backends/cortex_m/test/tester.py b/backends/cortex_m/test/tester.py index a1b5245b80b..b644db4e6c1 100644 --- a/backends/cortex_m/test/tester.py +++ b/backends/cortex_m/test/tester.py @@ -132,6 +132,7 @@ def test_dialect( qtol=0, atol=1e-03, calibration_samples=None, + ops_absent_after_transforms=None, ): """ Test the python dialect op implementation. @@ -142,13 +143,14 @@ def test_dialect( ) else: quantization_stage = None - self.quantize(quantization_stage) self.export() self.to_edge() self.check_count(ops_before_transforms) self.run_passes() self.check_count(ops_after_transforms) + if ops_absent_after_transforms: + self.check_not(ops_absent_after_transforms) self.run_method_and_compare_outputs( inputs=self.example_inputs, qtol=qtol, atol=atol ) diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 854ebf8f952..89e1b5ad79a 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -669,7 +669,10 @@ def get_supported_fallback_kernels(cls) -> Dict[str, Any]: return {} return { "at::_ops::_weight_int4pack_mm::call": None, + # Also under the shim name Inductor derives for it. + "aoti_torch_cuda__weight_int4pack_mm": None, "at::_ops::sort_stable::call": None, + "aoti_torch_cuda_sort_stable": None, "aoti_torch_cuda_randint_low_out": None, "executorch_cuda::int4_plain_mm": None, "aoti_torch_cuda_int4_plain_mm": None, diff --git a/backends/cuda/runtime/aoti_cuda_shims.lib b/backends/cuda/runtime/aoti_cuda_shims.lib index 8bb03cc1c1ee54d3643d7d91456c5d6c4e216ee2..c0d61c611008c3897350d7b0950a8e0229607ba3 100644 GIT binary patch literal 43802 zcmeI5TaO$^6@bgL&dnqw1n1^lJa+8FBw=^wy5lGcmKCB10Q3#SHP$VRT0Of(7fRLXN4?OS-5YGtl#8bjKRo&C4x~k4>O%_Y7?Ut&0 zrs`C6*Qu}1oUW>=KJijG-rRoq$Vye;R>@cOneMCAMyuITBI`$l5Zgi=eOHJ_-caG# zi$Y*{@|X(8|1JcElOL#X`Yjf8dSK&1%7sJ;dP~jW@5(2}8f2i>7zX*Zh&EKi;-5Ww+xcG(&!4pDY z2q6!K?w?fHdQS)pm%gLIpeh805y)T|-%?@sV<9kH`$&Z!ToMAq+pntdBgli{CvZI& z-T_$*KO3s>K3qG7pC47>7oZ0WH~*}{FaIC}hL1i};n#3I7=H7V3cvl25Ey>{Hx)j9 zH`tECZV-;9aj?1D4TH&cJeUOictgmz=?EwYf^ifNr_p$498H61G@OjaA_(Haj>_m3 zB4vb=NxU^oa!-QM?$j*8B11~p+>SOco0(KJmFtg&kws2M=j8&Z(n3Yd~9+Xya0}X|Zof zNxGC)%9VIC3Z~b0z)0-unt*mAUBHzW^fo6mO3>W$oEyFVD4f>Xmb`kpnR7JiJ9A;0 zkwc!UVKEO~d8#zo$Op)|FHj}g^0?0gc_7i2r`>c_3GzUqEl-EZ(lK-FLP^dnDv_Ke z8thE3CqGn{-wr+5l`q(dhr?)I_Ed_LvIT~7b9X$B;HFo;g8T_GE9?@?mDLTWVX!l` zxDr^vEH$4?-2x5^E^|@|du~L`e%<5_qpc;&jV(VuX+r)y%o$`-PJXyMh%7e{madX= z@{{;2o2!EIDk(1yK3cdH1*2XMe*f%MpXf4<60mqazmza?JDe+d#fj%d^NhhO9gMnB z5Kn?=h(8Y&Vz+Ek)>80Ce-Fh_F;Wg(LpSQh*1LwJ(HzqVyeX?ujAuziz8_u(kJx={ z=h8{eVRSWds)KF#MYjEnnj$3s)o?sC+cSxka{6-t_)g&<5-a%!L-Jf``&esrl1In& zBeqhvmq#nsgvlyyY-){7t+S~OHnqv7w%F7*o4UrPcG%RjTxyNgRgKkEjn!3+)m4qv zRgKkEjn!3+)m4qvRgKkEoz+#H)m5Ft0|hf724qlKP>bs?RjiBoAPI;-D)R4KkG-j??AU9Z0$orJRmb;P#ByfMf{Dk_sPkvhN#o$D!?UPTTU-M8eyqrcP@!Qe|P+w5TOQJX3 z-W^_6F?Yxqs%fm`Qc_bTqxMwG6_*UrFlxzK4$1<0{Kx#}~a}pzUt*#lR zF+Xm2#A^})I!|X?MZ5@~-MyxPtuHn8nd&XgHJZzEr%qVDOJ2iG(BcvK$z=I%h<^dn z`-!IUP!+3tWEt2anwF=q>~3U?uBFMSsV1E43OTRHskzhFxyPGRmC;nfL5+ z?ex4S;lpc*Od1C@zV+#8{PrRAf!j9X``Y@oAq%!TUV|dx9{8k4Qq;g#k?=!o30}sM zFnAVJr;=WgP$#vAXMkLaUc#Z1Rati=L0yr$LLu<%iq)s#6K=?}ig`_bSG$VXfDdyt zV03`l)a1ScR$88vuThJwDE|;DG>`J_q~-B&J1HNN^Kz8`&@N|C{*GO4MtK$L=S%q^ z+Qybql_C><67+#vM)<2LLhdVK9gc)>OYIH7Z8FYrv6YS|-;HEIw2B^y;`^u2=JWWz zNVvG2eD}zCIlg!7at7aT+2v+@ABFbv<-0;#SPcLc+G;-}*`-Xski_7jB9L&xG&>T? zSms(yLwBzKQIacLK>lOIkb1A?X7}E;lyMk=yhmw`1Deo6wB&Z0tm1z$eFno(mB;M2%$NHKd>!p>J?X;g(`VOt7 z|KObog8RX=jqRG(dyc*fI2@Q3bAjt@NE5xyEF&LQ{M~-sjXpzmmk@ zp(4Hyd+H?{wKt8kDdK+at2g}UK5O=C{Av$A2z{=)ozH!a=QN}D4xe3d%h_Awx$gtJ z+>CDgz4fL0E?4_%^RiUWS5bTL&PeT2%Df4dopiY^pg3^ zqA$Yy-Imrirs7I6irZ`bETRot)~9ztMaX?cT!N#?yrBgb@}?j|{8=yJ*H+VBl&p8J zBt9b^t(mWtzh@6^jkCpz{IvO8!gogen@J2FDq;ur^NIJJU(fjN4SjGjr5Xny)n~Pz zjkA~Djcl)+7xbCDnRjLMm^i*=nqTxBIEpDbkVKN)hIUQ zo^x9$!kQLX*mX%a^KH)&&(bSrnP0>!KWbrKUmdI1ynBxRB)sbyaIv0~VmS94p@CT{0!^C3Ae_b=mXb~$6V`H$^#v-w!CXtQ@;$|D-va#_mR3YkT=*y(yE z8P_jM&eg6WV8Ljn<|fs;W*B?D<~fepGVS%d(9ZLiUMDQvPNpxB^Kwj2?Q#awZ` zQ7>YfA2VAtT{5`PEI(#!Bs*I1P1dF#(JUUSEDF$TQ%RR+>A7w#y4W1A-&reEEh-b8~#wbM&yZ$}Go=_!v)_9IIER97L8nO~XEnE?ouzU-Yr}~ttxgMG{%c0IhO4)Fky%Vn8?8#;Az5sut5RE(>GMep z9*lp@NTzA+?ObaT4BItrU~xX%sYGV0%ekkStzLxfQx3M}Ye*#l#T{GpoO?>xmbK^5 zRstN^+8|@ftGP}Aw1LIB1h-O&%r=(ejF8S|BFzj4 zuIYXb(EWeRj6Rd67~f+L?{+<4wpi8L>KP|Vj89zGh7spdT(0B9+Ttpt`ufCQGAN#o z#^F{JjC#FEG%YmEI0l7V=2@`oOybktM6bOa>zoY7w>GFapKo*Cf7zRndCZISW<2NO z8;ujNQDkdfMYv(h_&%>9i$8i^Z zkYaP~IcHkgZB=sz(n%smv)>qDIs1d98Gm_wJ&_kQn#W-`7>v465Kn?=7;f|ej|zH{|00$x%`tw`uUb=JA)nXM!k3lcD%@Rqo0Ob=5w?E z9VfLmkfv!I$A066Wq-A7pZ(-`MdNrH4FdUx|3Y)@*1;H{IZIWdt-H)T+I!f^XVMZej|zH3cs>NT$Z^C zslL8+jWRuqt_H~>huh(>+mDL{%$)a literal 12738 zcmcH#U+P!YR#iWmLe(S!l{QiVr8H_$5=t5si$>rgn)IBxGrRY3W_L#$ zuXH>!_dM^t=bn4cz3%v8BpY3Luwf{q{)VMb@h<+33=fCL#zLXd`^QHG$xH(Pd<9^` zReOR+j?j5SK;Jdgi_o7Jum$xHY(p6YTNVTiECCP3rOZF2*8OZXuq`fo1J`YVTf?5WcNZ z3cBe%Pwjiu+P!Of7jnjv86hv9Mqdz8ts{1Pej#tgQ~9w>B>Ie%OoAh+42>x*Ye(|7 zmA6y5bk>^7rjyoeB0Xyz9#!P|Q&lRl0&!+v?PMl@%*y4n@t7TxC8!9)hM1kOi@i$E ze#4IDWp);CNJ-g8yn$JWq+$tM7GWX{{%9hdvaL+k&e_?+w%d)IEX0R`mt-)S<7H)S z^awkPI`gjSG8cz7b)_@MESW3nG#GU8yp6t+PA6pAVxvJH$>rknDc7o;m0nb4guxmT zjzqF4S()=~P(_8WI*V_9F_KmM#zY$Yu5a-CCeq+9GFmyDg?K7vAC;|U@un217wG?X zBx#lGj6!2blg(BrrI_JMV@R9FZbcWfSvy5figAvCawq*rSXr@1K4N9^SzSP(o#o2K zziBIriFArap5BQfCWqV1^2GATGRmqa9iow^k1Ju=p~fbIHW4|7{?7$A4P$URox~s? znYXR<+#GtCK597?TT%iImmj}_W4UAY;dm7H3~p>BR4HC#RX9T>m<7(@c81-WOQa+D z?aJn)B7FE}9mbKc+*`p|!1!=tKn;(XY@*qphqDoRu-Tu7GY@&N*`H?zYuXM!4n1JI zV?bft)KQ#FE1FqU6lti^yci>VnF(ib&u1esE18bjRy=3fsmN?X-Tqm82`LnaONx;R z3^~ySJNk^=0nuR4kq@U9leV%W78`@-@ZmxZzYUZ4s~Z5ungOyc0GSSeDzHb7!i0_vM0cP;~;4r{+JHYWyfMw);8{co>`!(dv_ao01q(LPM z_*@#mca-(?R+NE|!S5xcsqF!1=mt323-CFLfUU3^8elEl4eQ}{xD8gpI%tA>;11}6e&~c==zwnMg3Zta^>8n2fDp7p z3pB$v*a)rA2KT`LY=OIA6Wj@nD7-6lI%ZJG$&bjhN|KcBu`A5&v0HsG@(R1C z6J1Lauwk;Y4a^DorgYlA3)`s=y4>bdayP?W(_bvJL6xcLD3Eelr=i11CA^f9$CgYq z9RA||W_H3M;Lq++$rtswj@56mgLgPb&adNe;Q4h5qU;cvd zq1#Z0&q)|Q?b3zaX3XIM>QSFZN9~gkI$ev;X3Wp&w~py(FRW{9Ke;o6e=@@}a){{h_^6`yO@i%3}?FkwFB21q8h{;@{hxH%1#wM$^d*+Wz9Hnwpz8 zWlm(l!i2yZf=K|!`~|PY>8#M_(u@%eGC_f6yx2nRSWgvxBpqN^;hS{L;#XBAMp?&f z2&R|wbOe4oA=>9 ziEhP4xsnma6M^zKTR}&9V#IY*9BBYr%(0Vip8fBq!r@lQ8G>(Ss0H1?Yh|JkmWwt# zm^k!0h0hucNYXaqmwsv{1RaI(33C*dfB(V1X~=6;3&i077Zro|8}F3Guh{QL4)vd* z2wBHy^||$0fxfpNIPnMS`zp_7`{*ZM70AZA*w84}d@CcKnGT7m0+P=9#>?Wkw%Fiw4LjW>E5mSjiM z`8A*HL_6-N$c{p_RPBf`GxOoq=vy>1cWRcDtjMah!`p z3F;l(I{DhhwsW*jR;xDnRF>Kx1YW7!I{DT=7BA3jtns(t)~^$0Hh%rT?_8(ZSj*bM zl*rm(1fC6MpV;}m#92Jp$DCW)sigb7;ysV1KJZ+GnKk;rp=bY&7Xx?~;IF?>>uUxC z?mx3KuN?T}HN2d_l0ntaoF`PPo$-KB%_3c`;Qa@uy>Mnxt!5@hzjpTLZ_y1>6Q^0Z zQPq?OpY`qLPSx@9>!0DZ65jSO{gw`2RzmkrPJI2Jq=b`QNnzILyVZuhK8QyqAXF}~LF!3&qj*EXn{ne#ABd0=1TvXT0`zq(A> z$VOFr$wRkl?Fciu@$zWyBfL7r8%xFKoeKi*YT}mcxlEM(Zp*6|Datl8x{cFAt@D9B zfZ1iHj{o!$?J_OC`im!zn)<*#FwW%WxgXy58D2}{EwMTF-Z5ykdc#aDde?v6IMH@f^a&=)|LV*yu2N3fsp exports.txt +# add the new names to exports.txt, then +{ echo 'LIBRARY aoti_cuda_shims.dll'; echo 'EXPORTS'; sed 's/^/ /' exports.txt; } \ + > aoti_cuda_shims.def +x86_64-w64-mingw32-dlltool -d aoti_cuda_shims.def -l aoti_cuda_shims.lib \ + --dllname aoti_cuda_shims.dll +# zero the archive metadata so the file is reproducible +mkdir extract && cd extract && x86_64-w64-mingw32-ar x ../aoti_cuda_shims.lib \ + && rm -f ../aoti_cuda_shims.lib \ + && x86_64-w64-mingw32-ar rcsD ../aoti_cuda_shims.lib $(ls | sort) +``` + +The archive has to be built at a fresh path. Updating it in place keeps the reverse +member order the generator produced, so the same export list would not give the same +bytes twice. + +Then check the result is a superset of what it replaced and that no member carries a +timestamp, since this file ships in the wheel and a stamped one makes two builds of +the same source differ. + +A name only resolves if something in the DLL defines it. The DLL is built from the +CUDA shims plus the SlimTensor common shims, so a shim added to the ETensor common +shims will link and then fail to load. diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 8a81916ab6c..976b282a29c 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -195,6 +195,31 @@ AOTITorchError aoti_torch_empty_strided( return Error::Ok; } +AOTITorchError aoti_torch_empty_strided_pinned( + int64_t ndim, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + int32_t dtype, + int32_t device_type, + int32_t device_index, + SlimTensor** ret_new_tensor) { + ET_CHECK_OR_RETURN_ERROR( + static_cast(device_type) == DeviceType::CPU, + InvalidArgument, + "aoti_torch_empty_strided_pinned: pinned memory is host memory, so the " + "device type must be CPU, got %d", + device_type); + + return aoti_torch_empty_strided( + ndim, + sizes_ptr, + strides_ptr, + dtype, + device_type, + device_index, + ret_new_tensor); +} + AOTITorchError aoti_torch_delete_tensor_object(SlimTensor* tensor) { ET_CHECK_OR_RETURN_ERROR( tensor != nullptr, diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index ca464a9acf5..03e2b3ed18d 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -95,6 +95,31 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch_empty_strided( int32_t device_index, SlimTensor** ret_new_tensor); +/** + * Allocates ordinary host memory where pinned host memory was asked for. + * + * There is no pinned allocator here. Pinning only lets a copy to the device + * overlap other work, so ordinary memory is correct and slower. + * + * @param ndim Number of dimensions + * @param sizes_ptr Pointer to the sizes, ndim of them + * @param strides_ptr Pointer to the strides, ndim of them, or null for + * contiguous + * @param dtype Element type, as a scalar type value + * @param device_type Must be CPU, since pinned memory is host memory + * @param device_index Device index, unused for CPU + * @param ret_new_tensor Receives the new tensor + * @return Error::Ok on success, Error::InvalidArgument if the device is not CPU + */ +AOTI_SHIM_EXPORT AOTITorchError aoti_torch_empty_strided_pinned( + int64_t ndim, + const int64_t* sizes_ptr, + const int64_t* strides_ptr, + int32_t dtype, + int32_t device_type, + int32_t device_index, + SlimTensor** ret_new_tensor); + /** * Deletes a tensor object and frees associated resources. * diff --git a/backends/cuda/tests/test_sort_shim.py b/backends/cuda/tests/test_sort_shim.py index fc5f870fc42..da1ee0c3e2c 100644 --- a/backends/cuda/tests/test_sort_shim.py +++ b/backends/cuda/tests/test_sort_shim.py @@ -37,7 +37,9 @@ _CUDA_FALLBACK_KERNELS = frozenset( { "at::_ops::_weight_int4pack_mm::call", + "aoti_torch_cuda__weight_int4pack_mm", "at::_ops::sort_stable::call", + "aoti_torch_cuda_sort_stable", "aoti_torch_cuda_randint_low_out", "executorch_cuda::int4_plain_mm", "aoti_torch_cuda_int4_plain_mm", diff --git a/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py index 5ee3db6752f..1238e31e246 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py @@ -220,8 +220,9 @@ def test_conv_dropout_no_quant( ], ) - # Clone with inplace=True should not produce clone edge op and vice versa - assert inplace_dropout ^ has_clone + # Neither spelling leaves a clone behind on this PyTorch: the out-of-place + # one used to and no longer does. + assert not has_clone @parameterized.expand([("QAT", True), ("PTQ", False)]) def test_clone_pool_view_copy_quant( diff --git a/backends/vulkan/test/op_tests/CMakeLists.txt b/backends/vulkan/test/op_tests/CMakeLists.txt index 0f8456accf5..5facea5a14b 100644 --- a/backends/vulkan/test/op_tests/CMakeLists.txt +++ b/backends/vulkan/test/op_tests/CMakeLists.txt @@ -73,6 +73,8 @@ function(vulkan_op_test test_name test_src) add_executable(${test_name} ${test_src}) target_include_directories(${test_name} PRIVATE ${COMMON_INCLUDES}) + # ATen headers require C++20. + set_target_properties(${test_name} PROPERTIES CXX_STANDARD 20) target_link_libraries( ${test_name} PRIVATE GTest::gtest_main @@ -90,6 +92,8 @@ endfunction() if(TARGET vulkan_backend AND LIB_TORCH) add_library(test_utils ${CMAKE_CURRENT_SOURCE_DIR}/test_utils.cpp) target_include_directories(test_utils PRIVATE ${COMMON_INCLUDES}) + # ATen headers require C++20. + set_target_properties(test_utils PROPERTIES CXX_STANDARD 20) target_link_libraries( test_utils PRIVATE vulkan_backend ${LIB_TORCH} ${LIB_TORCH_CPU} ) diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index fca0f1b14c6..751388d9221 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1105,6 +1105,9 @@ def _is_share_obs_or_fq_op(op: Callable) -> bool: torch.ops.aten.slice.Tensor, torch.ops.aten.slice_copy.Tensor, torch.ops.aten.flatten.using_ints, + # Identity once not training, which is how a quantized model is deployed. + torch.ops.aten.dropout.default, + torch.ops.aten.dropout_.default, ] diff --git a/exir/lowered_backend_module.py b/exir/lowered_backend_module.py index a4c2f2cfe79..358bcb25e8c 100644 --- a/exir/lowered_backend_module.py +++ b/exir/lowered_backend_module.py @@ -384,7 +384,7 @@ def arrange_graph_placeholders( ) -> torch.fx.GraphModule: """ Modifies the graph of the given graphmodule with one that contains the same nodes as the original, - but with placeholders in order of (Params + Buffers) (User Inputs) + but with placeholders in order of (Params + Buffers + Constants) (User Inputs) This is used by the delegate api which disturbs the placeholder ordering when creating a submodule from partitioned nodes @@ -403,32 +403,30 @@ def arrange_graph_placeholders( graph_sign = owning_program.graph_signature # Add all placeholders into the graph first: - # Cache these properties — each call rebuilds the dict from input_specs. + # Cache these properties to avoid rebuilding the dict on each access. params_map = graph_sign.inputs_to_parameters buffers_map = graph_sign.inputs_to_buffers + constants_map = graph_sign.inputs_to_lifted_tensor_constants param_nodes = [] buffer_nodes = [] + constant_nodes = [] input_nodes = [] for node in gm.graph.nodes: if node.op != "placeholder": continue - if node.name in params_map and node.meta.get("delegation_tag", None) == tag: + is_tagged = node.meta.get("delegation_tag", None) == tag + if node.name in params_map and is_tagged: param_nodes.append(node) - elif node.name in buffers_map and node.meta.get("delegation_tag", None) == tag: + elif node.name in buffers_map and is_tagged: buffer_nodes.append(node) + elif node.name in constants_map and is_tagged: + constant_nodes.append(node) else: input_nodes.append(node) - for param_node in param_nodes: - new_node = new_graph.node_copy(param_node, lambda x: node_map[x]) - node_map[param_node] = new_node - for buffer_node in buffer_nodes: - new_node = new_graph.node_copy(buffer_node, lambda x: node_map[x]) - node_map[buffer_node] = new_node - for input_node in input_nodes: - new_node = new_graph.node_copy(input_node, lambda x: node_map[x]) - node_map[input_node] = new_node + for node in param_nodes + buffer_nodes + constant_nodes + input_nodes: + node_map[node] = new_graph.node_copy(node, lambda x: node_map[x]) # Now add all the other nodes in order for node in gm.graph.nodes: diff --git a/install_requirements.py b/install_requirements.py index b7d220179d3..c4a934a0180 100644 --- a/install_requirements.py +++ b/install_requirements.py @@ -66,7 +66,7 @@ def install_requirements(use_pytorch_nightly): # Setting use_pytorch_nightly to false to test the pinned PyTorch commit. Note # that we don't need to set any version number there because they have already # been installed on CI before this step, so pip won't reinstall them - ("torch==2.13.0" if use_pytorch_nightly else "torch"), + ("torch==2.14.0" if use_pytorch_nightly else "torch"), f"torchao=={TORCHAO_NIGHTLY_VERSION}", ] @@ -134,7 +134,7 @@ def install_optional_example_requirements(use_pytorch_nightly): print("Installing torch domain libraries") DOMAIN_LIBRARIES = [ - ("torchvision==0.28.0" if use_pytorch_nightly else "torchvision"), + ("torchvision==0.29.0" if use_pytorch_nightly else "torchvision"), ("torchaudio==2.11.0" if use_pytorch_nightly else "torchaudio"), ] # Then install domain libraries diff --git a/kernels/test/ScalarOverflowTestMacros.h b/kernels/test/ScalarOverflowTestMacros.h index 46a2425b0fa..6567ad71564 100644 --- a/kernels/test/ScalarOverflowTestMacros.h +++ b/kernels/test/ScalarOverflowTestMacros.h @@ -11,28 +11,36 @@ // Macro to generate scalar overflow test cases for a given test suite. // The test suite must have a method called expect_bad_scalar_value_dies // that takes a template parameter for ScalarType and a Scalar value. -#define GENERATE_SCALAR_OVERFLOW_TESTS(TEST_SUITE_NAME) \ - TEST_F(TEST_SUITE_NAME, ByteTensorTooLargeScalarDies) { \ - /* Cannot be represented by a uint8_t. */ \ - expect_bad_scalar_value_dies(256); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, CharTensorTooSmallScalarDies) { \ - /* Cannot be represented by a int8_t. */ \ - expect_bad_scalar_value_dies(-129); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, ShortTensorTooLargeScalarDies) { \ - /* Cannot be represented by a int16_t. */ \ - expect_bad_scalar_value_dies(32768); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, FloatTensorTooSmallScalarDies) { \ - /* Cannot be represented by a float. */ \ - expect_bad_scalar_value_dies(-3.41e+38); \ - } \ - \ - TEST_F(TEST_SUITE_NAME, FloatTensorTooLargeScalarDies) { \ - /* Cannot be represented by a float. */ \ - expect_bad_scalar_value_dies(3.41e+38); \ +#define GENERATE_SCALAR_OVERFLOW_TESTS(TEST_SUITE_NAME) \ + TEST_F(TEST_SUITE_NAME, ByteTensorTooLargeScalarDies) { \ + /* Cannot be represented by a uint8_t. */ \ + expect_bad_scalar_value_dies(256); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, CharTensorTooSmallScalarDies) { \ + /* Cannot be represented by a int8_t. */ \ + expect_bad_scalar_value_dies(-129); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, ShortTensorTooLargeScalarDies) { \ + /* Cannot be represented by a int16_t. */ \ + expect_bad_scalar_value_dies(32768); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, FloatTensorTooSmallScalarDies) { \ + /* Cannot be represented by a float. */ \ + expect_bad_scalar_value_dies(-3.41e+38); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, FloatTensorTooLargeScalarDies) { \ + /* Cannot be represented by a float. */ \ + expect_bad_scalar_value_dies(3.41e+38); \ + } \ + \ + TEST_F(TEST_SUITE_NAME, LongTensorTooLargeScalarDies) { \ + /* 2^63 is one past the largest int64_t, so converting it is undefined \ + * unless the range check rejects it first. The add suites reject it \ + * earlier, on the alpha type, so there this case only repeats their \ + * existing floating-point alpha coverage. */ \ + expect_bad_scalar_value_dies(9223372036854775808.0); \ } diff --git a/kernels/test/op_full_test.cpp b/kernels/test/op_full_test.cpp index 752c4e710f4..e412b67dfc4 100644 --- a/kernels/test/op_full_test.cpp +++ b/kernels/test/op_full_test.cpp @@ -84,6 +84,19 @@ ET_FORALL_REALHBF16_TYPES(GENERATE_TEST) GENERATE_SCALAR_OVERFLOW_TESTS(OpFullOutTest) +// The other half of the boundary change: 127.5 used to be refused for an int8 +// tensor and now truncates to 127. +TEST_F(OpFullOutTest, CharTensorFractionalScalarTruncates) { + TensorFactory tf; + std::vector sizes = {2, 2}; + std::vector sizes_int64_t(sizes.begin(), sizes.end()); + auto aref = IntArrayRef(sizes_int64_t.data(), sizes_int64_t.size()); + Tensor out = tf.zeros(sizes); + + op_full_out(aref, 127.5, out); + EXPECT_TENSOR_EQ(out, tf.full(sizes, 127)); +} + TEST_F(OpFullOutTest, HalfSupport) { TensorFactory tf; diff --git a/runtime/core/portable_type/c10/c10/targets.bzl b/runtime/core/portable_type/c10/c10/targets.bzl index 675c04f97a1..a39b2b9f566 100644 --- a/runtime/core/portable_type/c10/c10/targets.bzl +++ b/runtime/core/portable_type/c10/c10/targets.bzl @@ -114,7 +114,6 @@ def define_common_targets(): "util/bit_cast.h", "util/complex.h", "util/complex_math.h", - "util/complex_utils.h", "util/floating_point_utils.h", "util/irange.h", "util/llvmMathExtras.h", diff --git a/runtime/core/portable_type/c10/c10/util/complex.h b/runtime/core/portable_type/c10/c10/util/complex.h index 4e699684bc3..f9849a94ced 100644 --- a/runtime/core/portable_type/c10/c10/util/complex.h +++ b/runtime/core/portable_type/c10/c10/util/complex.h @@ -31,19 +31,11 @@ C10_HOST_DEVICE T abs(const c10::complex& z) { #endif } -#if defined(USE_ROCM) -#define ROCm_Bug(x) -#else -#define ROCm_Bug(x) x -#endif - template C10_HOST_DEVICE T arg(const c10::complex& z) { - return ROCm_Bug(std)::atan2(std::imag(z), std::real(z)); + return std::atan2(std::imag(z), std::real(z)); } -#undef ROCm_Bug - template constexpr T norm(const c10::complex& z) { return z.real() * z.real() + z.imag() * z.imag(); @@ -73,6 +65,9 @@ constexpr c10::complex conj(const c10::complex& z) { #define C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H // math functions are included in a separate file #include // IWYU pragma: keep -// utilities for complex types -#include // IWYU pragma: keep #undef C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H + +namespace c10 { +using torch::headeronly::is_complex; +using torch::headeronly::scalar_value_type; +} // namespace c10 diff --git a/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h b/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h index da297241449..8ae5cde4f02 100644 --- a/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h +++ b/runtime/core/portable_type/c10/c10/util/llvmMathExtras.h @@ -400,6 +400,7 @@ constexpr inline bool isShiftedUInt(uint64_t x) { N + S <= 64, "isShiftedUInt with N + S > 64 is too wide."); // Per the two static_asserts above, S must be strictly less than 64. So // 1 << S is not undefined behavior. + // NOLINTNEXTLINE(bugprone-chained-comparison) return isUInt(x) && (x % (UINT64_C(1) << S) == 0); } diff --git a/runtime/core/portable_type/c10/c10/util/overflows.h b/runtime/core/portable_type/c10/c10/util/overflows.h index 183a2f62a32..348ee5ffa40 100644 --- a/runtime/core/portable_type/c10/c10/util/overflows.h +++ b/runtime/core/portable_type/c10/c10/util/overflows.h @@ -61,13 +61,28 @@ template std::enable_if_t, bool> overflows( From f, bool strict_unsigned [[maybe_unused]] = false) { - using limit = std::numeric_limits::type>; + using ToScalar = typename scalar_value_type::type; + using limit = std::numeric_limits; if (limit::has_infinity && std::isinf(static_cast(f))) { return false; } if (!limit::has_quiet_NaN && (f != f)) { return true; } + if constexpr (std::is_integral_v) { + // limit::max() for wide integer types is NOT exactly representable in + // floating point (e.g. int64 max = 2^63-1 rounds up to 2^63), so `f > + // limit::max()` lets a just-out-of-range value like 2^63 slip through and + // then become INT64_MIN via static_cast. Compare against the + // exactly-representable upper bound max()+1 == 2^digits instead. lowest() + // is 0 or a negated power of two, so it stays exact. (digits-1 keeps the + // shift < 64 for the uint64 case; the *2 recovers 2^digits without a 1<<64 + // overflow.) + constexpr int digits = limit::digits; + constexpr From upper = + static_cast(uint64_t{1} << (digits - 1)) * From{2}; + return f < static_cast(limit::lowest()) || f >= upper; + } return f < limit::lowest() || f > limit::max(); } diff --git a/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h b/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h index cef99df3f56..08c4e9f1f84 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h +++ b/runtime/core/portable_type/c10/torch/headeronly/macros/Macros.h @@ -123,6 +123,15 @@ #define C10_HAS_CPP_ATTRIBUTE(x) (0) #endif +/// Bind a returned reference/pointer's lifetime to a parameter (or *this) so +/// Clang can warn when it would dangle. Expands to nothing on compilers that +/// lack the attribute (e.g. non-clang, older nvcc). +#if C10_HAS_CPP_ATTRIBUTE(clang::lifetimebound) +#define C10_LIFETIMEBOUND [[clang::lifetimebound]] +#else +#define C10_LIFETIMEBOUND +#endif + #ifndef FBCODE_CAFFE2 /// DEPRECATED: Warn if a type or return value is discarded. #define C10_NODISCARD [[nodiscard]] diff --git a/runtime/core/portable_type/c10/torch/headeronly/util/Half.h b/runtime/core/portable_type/c10/torch/headeronly/util/Half.h index e5aa622656c..401472357ec 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/util/Half.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/Half.h @@ -213,7 +213,7 @@ C10_HOST_DEVICE inline float fp16_ieee_to_fp32_value(uint16_t h) { * Now, remember that denormalized half-precision numbers are represented as: * FP16 = mantissa * 2**(-24). * The trick is to construct a normalized single-precision number with the - * same mantissa and thehalf-precision input and with an exponent which would + * same mantissa and the half-precision input and with an exponent which would * scale the corresponding mantissa bits to 2**(-24). A normalized * single-precision floating-point number is represented as: FP32 = (1 + * mantissa * 2**(-23)) * 2**(exponent - 127) Therefore, when the biased diff --git a/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h b/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h index c33a286bc5b..8e897957fee 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/TypeSafeSignMath.h @@ -14,20 +14,6 @@ C10_CLANG_DIAGNOSTIC_IGNORE("-Wimplicit-int-float-conversion") namespace c10 { -/// Returns false since we cannot have x < 0 if x is unsigned. -template -inline constexpr bool is_negative( - const T& /*x*/, - std::true_type /*is_unsigned*/) { - return false; -} - -/// Returns true if a signed variable x < 0 -template -inline constexpr bool is_negative(const T& x, std::false_type /*is_unsigned*/) { - return x < T(0); -} - /// Returns true if x < 0 /// NOTE: Will fail on an unsigned custom type /// For the most part it's possible to fix this if @@ -35,19 +21,12 @@ inline constexpr bool is_negative(const T& x, std::false_type /*is_unsigned*/) { /// However, notably, c10::Half does not :-( template inline constexpr bool is_negative(const T& x) { - return is_negative(x, std::is_unsigned()); -} - -/// Returns the sign of an unsigned variable x as 0, 1 -template -inline constexpr int signum(const T& x, std::true_type /*is_unsigned*/) { - return T(0) < x; -} - -/// Returns the sign of a signed variable x as -1, 0, 1 -template -inline constexpr int signum(const T& x, std::false_type /*is_unsigned*/) { - return (T(0) < x) - (x < T(0)); + if constexpr (std::is_unsigned_v) { + // An unsigned value can never be less than zero. + return false; + } else { + return x < T(0); + } } /// Returns the sign of x as -1, 0, 1 @@ -57,7 +36,11 @@ inline constexpr int signum(const T& x, std::false_type /*is_unsigned*/) { /// However, notably, c10::Half does not :-( template inline constexpr int signum(const T& x) { - return signum(x, std::is_unsigned()); + if constexpr (std::is_unsigned_v) { + return T(0) < x; + } else { + return (T(0) < x) - (x < T(0)); + } } /// Returns true if a and b are not both negative @@ -86,53 +69,22 @@ inline constexpr bool greater_than_max(const T& x) { #pragma GCC diagnostic pop #endif -/// Returns true if x < lowest(Limit). Standard comparison -template -inline constexpr bool less_than_lowest( - const T& x, - std::false_type /*limit_is_unsigned*/, - std::false_type /*x_is_unsigned*/) { - return x < std::numeric_limits::lowest(); -} - -/// Returns false since all the limit is signed and therefore includes -/// negative values but x cannot be negative because it is unsigned -template -inline constexpr bool less_than_lowest( - const T& /*x*/, - std::false_type /*limit_is_unsigned*/, - std::true_type /*x_is_unsigned*/) { - return false; -} - -/// Returns true if x < 0, where 0 is constructed from T. -/// Limit is not signed, so its lower value is zero -template -inline constexpr bool less_than_lowest( - const T& x, - std::true_type /*limit_is_unsigned*/, - std::false_type /*x_is_unsigned*/) { - return x < T(0); -} - -/// Returns false sign both types are unsigned -template -inline constexpr bool less_than_lowest( - const T& /*x*/, - std::true_type /*limit_is_unsigned*/, - std::true_type /*x_is_unsigned*/) { - return false; -} - -/// Returns true if x is less than the lowest value of type T +/// Returns true if x is less than the lowest value of type Limit /// NOTE: Will fail on an unsigned custom type /// For the most part it's possible to fix this if /// the custom type has a constexpr constructor. /// However, notably, c10::Half does not : template inline constexpr bool less_than_lowest(const T& x) { - return less_than_lowest( - x, std::is_unsigned(), std::is_unsigned()); + if constexpr (std::is_unsigned_v) { + // x is unsigned, so it can never be below the lowest value of any type. + return false; + } else if constexpr (std::is_unsigned_v) { + // Limit is unsigned, so its lowest value is zero. + return x < T(0); + } else { + return x < std::numeric_limits::lowest(); + } } } // namespace c10 diff --git a/runtime/core/portable_type/c10/torch/headeronly/util/complex.h b/runtime/core/portable_type/c10/torch/headeronly/util/complex.h index 733a22d5dbb..c349602dcf0 100644 --- a/runtime/core/portable_type/c10/torch/headeronly/util/complex.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/complex.h @@ -3,6 +3,7 @@ #include #include +#include #include #if defined(__CUDACC__) || defined(__HIPCC__) @@ -588,6 +589,60 @@ struct alignas(4) complex { } }; +template <> +struct alignas(4) complex { + BFloat16 real_; + BFloat16 imag_; + + // Constructors + complex() = default; + // BFloat16 constructor is not constexpr so the following constructor can't + // be constexpr + C10_HOST_DEVICE explicit inline complex( + const BFloat16& real, + const BFloat16& imag) + : real_(real), imag_(imag) {} + C10_HOST_DEVICE inline complex(const c10::complex& value) + : real_(value.real()), imag_(value.imag()) {} + + // Conversion operator + inline C10_HOST_DEVICE operator c10::complex() const { + return {real_, imag_}; + } + + constexpr C10_HOST_DEVICE BFloat16 real() const { + return real_; + } + constexpr C10_HOST_DEVICE BFloat16 imag() const { + return imag_; + } + + C10_HOST_DEVICE complex& operator+=( + const complex& other) { + real_ = static_cast(real_) + static_cast(other.real_); + imag_ = static_cast(imag_) + static_cast(other.imag_); + return *this; + } + + C10_HOST_DEVICE complex& operator-=( + const complex& other) { + real_ = static_cast(real_) - static_cast(other.real_); + imag_ = static_cast(imag_) - static_cast(other.imag_); + return *this; + } + + C10_HOST_DEVICE complex& operator*=( + const complex& other) { + auto a = static_cast(real_); + auto b = static_cast(imag_); + auto c = static_cast(other.real()); + auto d = static_cast(other.imag()); + real_ = a * c - b * d; + imag_ = a * d + b * c; + return *this; + } +}; + } // namespace c10 HIDDEN_NAMESPACE_BEGIN(torch, headeronly) @@ -614,3 +669,8 @@ using c10::complex_literals::operator""_id; HIDDEN_NAMESPACE_END(torch, headeronly) C10_CLANG_DIAGNOSTIC_POP() + +#define C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H +// utilities for complex types +#include // IWYU pragma: keep +#undef C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H diff --git a/runtime/core/portable_type/c10/c10/util/complex_utils.h b/runtime/core/portable_type/c10/torch/headeronly/util/complex_utils.h similarity index 80% rename from runtime/core/portable_type/c10/c10/util/complex_utils.h rename to runtime/core/portable_type/c10/torch/headeronly/util/complex_utils.h index 1ca105f1d0a..ddc66ffe776 100644 --- a/runtime/core/portable_type/c10/c10/util/complex_utils.h +++ b/runtime/core/portable_type/c10/torch/headeronly/util/complex_utils.h @@ -1,11 +1,13 @@ +#pragma once + #if !defined(C10_INTERNAL_INCLUDE_COMPLEX_REMAINING_H) #error \ - "c10/util/complex_utils.h is not meant to be individually included. Include c10/util/complex.h instead." + "torch/headeronly/util/complex_utils.h is not meant to be individually included. Include torch/headeronly/util/complex.h instead." #endif #include -namespace c10 { +HIDDEN_NAMESPACE_BEGIN(torch, headeronly) template struct is_complex : public std::false_type {}; @@ -31,7 +33,7 @@ struct scalar_value_type> { using type = T; }; -} // namespace c10 +HIDDEN_NAMESPACE_END(torch, headeronly) namespace std { diff --git a/shim_et/xplat/executorch/build/runtime_wrapper.bzl b/shim_et/xplat/executorch/build/runtime_wrapper.bzl index e84af76f3d9..445c1aa3b98 100644 --- a/shim_et/xplat/executorch/build/runtime_wrapper.bzl +++ b/shim_et/xplat/executorch/build/runtime_wrapper.bzl @@ -145,30 +145,45 @@ def _has_pytorch_dep(dep_list): return True return False -def _patch_test_compiler_flags(kwargs): +def _is_aten_target(kwargs): + """Whether a target compiles against ATen. + + Keyed on exact dep names, not a substring: every label contains "torch". + """ + aten_external_deps = [ + "c10", + "gmock_aten", + "gtest_aten", + "libtorch", + "libtorch_python", + "torch-core-cpp", + ] + for key in ["external_deps", "exported_external_deps"]: + for dep in kwargs.get(key) or []: + if dep in aten_external_deps: + return True + for key in ["xplat_deps", "fbcode_deps"]: + if _has_pytorch_dep(kwargs.get(key)): + return True + return False + +def _patch_test_compiler_flags(kwargs, aten_mode = False): if "compiler_flags" not in kwargs: kwargs["compiler_flags"] = [] - # Determine C++ standard based on whether this is an aten test. - # Aten tests require at least C++20 to compile against PyTorch, while - # non-aten tests are pinned to C++17 for embedded. + # A test that compiles against ATen needs C++20, which PyTorch's headers + # require. Every other test stays at C++17, which the embedded builds use. name = kwargs.get("name", "") - external_deps = kwargs.get("external_deps", []) - deps = kwargs.get("deps", []) - xplat_deps = kwargs.get("xplat_deps", []) - fbcode_deps = kwargs.get("fbcode_deps", []) is_aten_test = ( + aten_mode or "_aten" in name or - "aten_" in name or - "libtorch" in external_deps or - "gtest_aten" in external_deps or - "gmock_aten" in external_deps or - _has_pytorch_dep(deps) or - _has_pytorch_dep(xplat_deps) or - _has_pytorch_dep(fbcode_deps) + "aten_" in name ) - - if not is_aten_test: + if is_aten_test: + kwargs["compiler_flags"] += [ + "-std=c++20", + ] + else: kwargs["compiler_flags"] += [ "-std=c++17", ] @@ -267,9 +282,21 @@ def _patch_kwargs_cxx(kwargs): env.remove_platform_specific_args(kwargs) return _patch_kwargs_common(kwargs) +def _patch_aten_mode_std(kwargs, aten_mode): + """Raises an ATen-mode target to C++20, which PyTorch's headers require. + + A plain compiler flag, which the prelude places after the toolchain's. + """ + if aten_mode: + kwargs["compiler_flags"] = kwargs.get("compiler_flags", []) + ["-std=c++20"] + return kwargs + def _cxx_library_common(*args, **kwargs): + # Before _patch_kwargs_cxx, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_cxx(kwargs) _patch_build_mode_flags(kwargs) + _patch_aten_mode_std(kwargs, aten_mode) env.patch_platform_build_mode_flags(kwargs) env.patch_headers(kwargs) @@ -294,8 +321,11 @@ def _cxx_library(*args, **kwargs): _cxx_library_common(*args, **kwargs) def _cxx_binary_helper(*args, **kwargs): + # Before _patch_kwargs_cxx, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_cxx(kwargs) _patch_build_mode_flags(kwargs) + _patch_aten_mode_std(kwargs, aten_mode) env.patch_platform_build_mode_flags(kwargs) env.patch_cxx_compiler_flags(kwargs) @@ -315,20 +345,25 @@ def _cxx_test(*args, **kwargs): kwargs["deps"] = [] kwargs["deps"].append("//executorch/test/utils:utils") + # Before _patch_kwargs_cxx, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_cxx(kwargs) env.patch_headers(kwargs) _patch_build_mode_flags(kwargs) - _patch_test_compiler_flags(kwargs) + _patch_test_compiler_flags(kwargs, aten_mode) env.patch_platform_build_mode_flags(kwargs) env.cxx_test(*args, **kwargs) def _cxx_python_extension(*args, **kwargs): + # Before _patch_kwargs_common, which consumes external_deps. + aten_mode = _is_aten_target(kwargs) _patch_kwargs_common(kwargs) _remove_caffe2_deps(kwargs) kwargs["srcs"] = _patch_executorch_references(kwargs["srcs"]) if "types" in kwargs: kwargs["types"] = _patch_executorch_references(kwargs["types"]) + _patch_aten_mode_std(kwargs, aten_mode) env.cxx_python_extension(*args, **kwargs) def _export_file(*args, **kwargs): diff --git a/torch_pin.py b/torch_pin.py index ca593b1ef05..f46d5b67ec0 100644 --- a/torch_pin.py +++ b/torch_pin.py @@ -1,2 +1,2 @@ -TORCH_VERSION = "2.13.0" +TORCH_VERSION = "2.14.0" # NIGHTLY_VERSION = "dev20260318" Temporarily pinning to stable release candidate. Revert https://github.com/pytorch/executorch/pull/18287 From 9036d840faceb5a6ca563ad3585b65e1e55e1b08 Mon Sep 17 00:00:00 2001 From: qti-horodnic Date: Fri, 4 Sep 2026 00:00:15 -0700 Subject: [PATCH 032/190] Qualcomm AI Engine Direct - Fix a QNN crash on AMD (#22543) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Fixed a crash on AMD hosts when PyTorch's global backend flags are frozen. `QnnQuantizer.__init__` calls `disable_mkldnn_on_amd()`, which assigns `torch.backends.mkldnn.enabled = False`. PyTorch forbids that assignment once `disable_global_flags()` has been called, which happens whenever `torch.testing._internal.common_utils` is imported — as the Qualcomm test suite does. The result is that constructing a quantizer on an AMD host fails with: ``` RuntimeError: not allowed to set torch.backends.mkldnn flags after disable_global_flags; please use flags() context manager instead ``` The assignment is now guarded, so a frozen-flags environment leaves the existing MKLDNN setting in place instead of raising. Disabling MKLDNN is a workaround for AMD-host crashes, not a correctness requirement, so skipping it when the flag is locked is safe. Note that this is a temporary workaround for an issue introduced in a previous [pr](https://github.com/pytorch/executorch/pull/22395), not a root-cause fix, I haven't validated that this change has not broken anything else. ### Test plan Ran the tests mentioned in the test plan section of https://github.com/pytorch/executorch/pull/22542. --- backends/qualcomm/utils/qnn_sdk_setup.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/backends/qualcomm/utils/qnn_sdk_setup.py b/backends/qualcomm/utils/qnn_sdk_setup.py index 9ef1b66392c..bce3b761f9e 100644 --- a/backends/qualcomm/utils/qnn_sdk_setup.py +++ b/backends/qualcomm/utils/qnn_sdk_setup.py @@ -166,7 +166,8 @@ def disable_mkldnn_on_amd() -> None: import torch - torch.backends.mkldnn.enabled = False + if not torch.backends.flags_frozen(): + torch.backends.mkldnn.enabled = False def _host_is_amd() -> bool: From f5d3e54da6f591bc2f82335f000930c3bd13bf6b Mon Sep 17 00:00:00 2001 From: Emma Kujala <47500215+emmakujala@users.noreply.github.com> Date: Fri, 4 Sep 2026 11:20:11 +0200 Subject: [PATCH 033/190] Arm backend: Enable per-delegate profiling in perf_monitor (#22511) Add the option for per-delegate profiling to make it easier to understand model performance that contains multiple delegate calls. Signed-off-by: Emma Kujala --- backends/arm/runtime/EthosUBackend.cpp | 18 ++++ backends/arm/runtime/EthosUBackend_Internal.h | 4 + examples/arm/executor_runner/CMakeLists.txt | 25 +++++ .../arm/executor_runner/arm_perf_monitor.cpp | 99 +++++++++++++++++++ 4 files changed, 146 insertions(+) diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index 1305c5b4995..2e752e99995 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -56,16 +56,30 @@ namespace arm { extern "C" { void __attribute__((weak)) EthosUBackend_execute_begin() {} void __attribute__((weak)) EthosUBackend_execute_end() {} +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +void __attribute__((weak)) EthosUBackend_delegate_begin(const void*) {} +void __attribute__((weak)) EthosUBackend_delegate_end() {} +#endif __attribute__((weak)) unsigned char* ethosu_fast_scratch = nullptr; __attribute__((weak)) size_t ethosu_fast_scratch_size = 0; } class EthosUBackendExecuteCallbacks { public: +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + explicit EthosUBackendExecuteCallbacks(const void* handle) { + EthosUBackend_execute_begin(); + EthosUBackend_delegate_begin(handle); + } +#else EthosUBackendExecuteCallbacks() { EthosUBackend_execute_begin(); } +#endif ~EthosUBackendExecuteCallbacks() { +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + EthosUBackend_delegate_end(); +#endif EthosUBackend_execute_end(); } }; @@ -152,7 +166,11 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { // and EthosUBackend_execute_end() is called while CollectArm_CPU_Cycles is // in scope. e.g. We meassure from now until we exit this metod (in any way // we might do it). +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + EthosUBackendExecuteCallbacks CollectArm_CPU_Cycles(input_handle); +#else EthosUBackendExecuteCallbacks CollectArm_CPU_Cycles; +#endif ExecutionHandle* execution_handle = static_cast(input_handle); diff --git a/backends/arm/runtime/EthosUBackend_Internal.h b/backends/arm/runtime/EthosUBackend_Internal.h index 48fc4aa3a79..01069a9af72 100644 --- a/backends/arm/runtime/EthosUBackend_Internal.h +++ b/backends/arm/runtime/EthosUBackend_Internal.h @@ -74,6 +74,10 @@ struct ExecutionHandle { extern "C" { void EthosUBackend_execute_begin(); void EthosUBackend_execute_end(); +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +void EthosUBackend_delegate_begin(const void* handle); +void EthosUBackend_delegate_end(); +#endif extern unsigned char* ethosu_fast_scratch; extern size_t ethosu_fast_scratch_size; } diff --git a/examples/arm/executor_runner/CMakeLists.txt b/examples/arm/executor_runner/CMakeLists.txt index ba80df19ddf..4280dde15d8 100644 --- a/examples/arm/executor_runner/CMakeLists.txt +++ b/examples/arm/executor_runner/CMakeLists.txt @@ -54,6 +54,13 @@ set(ET_NUM_INFERENCES option(ET_LOG_DUMP_INPUT "Dump input in log" OFF) option(ET_LOG_DUMP_OUTPUT "Dump output in log" ON) +option(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING + "Report Ethos-U PMU statistics per delegate" OFF +) +set(ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES + "16" + CACHE STRING "Maximum delegates tracked by per-delegate profiling" +) option(ET_BUNDLE_IO "Set to compile in BundleIO support" OFF) set(BUNDLED_PROGRAM_LIBRARY_DIR @@ -338,6 +345,24 @@ if(ET_NUM_INFERENCES) ) endif() +if(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + if(NOT ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES MATCHES "^[1-9][0-9]*$") + message( + FATAL_ERROR + "ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES must be a positive integer" + ) + endif() + target_compile_definitions( + arm_executor_runner + PRIVATE + ET_ARM_ETHOSU_PER_DELEGATE_PROFILING + ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES=${ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES} + ) + target_compile_definitions( + executorch_delegate_ethos_u PRIVATE ET_ARM_ETHOSU_PER_DELEGATE_PROFILING + ) +endif() + if(ET_LOG_DUMP_INPUT) target_compile_definitions(arm_executor_runner PUBLIC ET_LOG_DUMP_INPUT) endif() diff --git a/examples/arm/executor_runner/arm_perf_monitor.cpp b/examples/arm/executor_runner/arm_perf_monitor.cpp index 59daede6920..9d92a878a2b 100644 --- a/examples/arm/executor_runner/arm_perf_monitor.cpp +++ b/examples/arm/executor_runner/arm_perf_monitor.cpp @@ -40,6 +40,36 @@ uint64_t ethosu_ArmBackendExecuteCycleCount = 0; uint64_t ethosu_ArmWhenNPURunCycleCountStart = 0; uint64_t ethosu_ArmWhenNPURunCycleCount = 0; uint64_t ethosu_pmuCycleCount = 0; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +struct DelegateStats { + const void* handle = nullptr; + uint64_t backend_invocations = 0; + uint64_t npu_invocations = 0; + uint64_t pmu_cycles = 0; + std::array pmu_events{}; +}; + +std::array + ethosu_delegateStats; +size_t ethosu_delegateCount = 0; +DelegateStats* ethosu_activeDelegate = nullptr; +bool ethosu_delegateCapacityExceeded = false; + +DelegateStats* get_delegate_stats(const void* handle) { + for (size_t i = 0; i < ethosu_delegateCount; ++i) { + if (ethosu_delegateStats[i].handle == handle) { + return ðosu_delegateStats[i]; + } + } + if (ethosu_delegateCount == ethosu_delegateStats.size()) { + ethosu_delegateCapacityExceeded = true; + return nullptr; + } + DelegateStats& stats = ethosu_delegateStats[ethosu_delegateCount++]; + stats.handle = handle; + return &stats; +} +#endif std::array ethosu_pmuEventCounts = {0}; // ethosu_pmuCountersUsed should match numbers of counters setup in @@ -50,6 +80,19 @@ static_assert(ETHOSU_PMU_NCOUNTERS >= ethosu_pmuCountersUsed); extern "C" { +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) +void EthosUBackend_delegate_begin(const void* handle) { + ethosu_activeDelegate = get_delegate_stats(handle); + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->backend_invocations++; + } +} + +void EthosUBackend_delegate_end() { + ethosu_activeDelegate = nullptr; +} +#endif + // Callback invoked at start of NPU execution void ethosu_inference_begin(struct ethosu_driver* drv, void*) { // Enable PMU @@ -100,10 +143,27 @@ void ethosu_inference_begin(struct ethosu_driver* drv, void*) { // Callback invoked at end of NPU execution void ethosu_inference_end(struct ethosu_driver* drv, void*) { ethosu_delegation_count++; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + const uint64_t pmu_cycles = ETHOSU_PMU_Get_CCNTR(drv); + ethosu_pmuCycleCount += pmu_cycles; + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->npu_invocations++; + ethosu_activeDelegate->pmu_cycles += pmu_cycles; + } +#else ethosu_pmuCycleCount += ETHOSU_PMU_Get_CCNTR(drv); +#endif for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) { +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + const uint64_t event_count = ETHOSU_PMU_Get_EVCNTR(drv, i); + ethosu_pmuEventCounts[i] += event_count; + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->pmu_events[i] += event_count; + } +#else ethosu_pmuEventCounts[i] += ETHOSU_PMU_Get_EVCNTR(drv, i); +#endif } ETHOSU_PMU_Disable(drv); // Add Cortex-M cycle clock used during this NPU execution @@ -131,6 +191,12 @@ void StartMeasurements() { ethosu_ArmBackendExecuteCycleCount = 0; ethosu_ArmWhenNPURunCycleCount = 0; ethosu_pmuCycleCount = 0; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + ethosu_delegateStats = {}; + ethosu_delegateCount = 0; + ethosu_activeDelegate = nullptr; + ethosu_delegateCapacityExceeded = false; +#endif for (size_t i = 0; i < ethosu_pmuCountersUsed; i++) { ethosu_pmuEventCounts[i] = 0; @@ -221,6 +287,39 @@ void StopMeasurements(int num_inferences) { ethosu_pmuEventCounts[i], (double)ethosu_pmuEventCounts[i] / num_inferences); } +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + ET_LOG(Info, "Ethos-U per-delegate PMU report:"); + for (size_t delegate_id = 0; delegate_id < ethosu_delegateCount; + ++delegate_id) { + const DelegateStats& stats = ethosu_delegateStats[delegate_id]; + ET_LOG( + Info, + "Ethos-U delegate %zu: %" PRIu64 " backend invocations, %" PRIu64 + " NPU invocations", + delegate_id, + stats.backend_invocations, + stats.npu_invocations); + ET_LOG( + Info, + "Ethos-U delegate %zu PMU cycles: %" PRIu64, + delegate_id, + stats.pmu_cycles); + for (size_t event = 0; event < ethosu_pmuCountersUsed; ++event) { + ET_LOG( + Info, + "Ethos-U delegate %zu PMU counter %zu: %" PRIu64, + delegate_id, + event, + stats.pmu_events[event]); + } + } + if (ethosu_delegateCapacityExceeded) { + ET_LOG( + Error, + "Ethos-U per-delegate profiling exceeded its capacity of %zu delegates", + ethosu_delegateStats.size()); + } +#endif #if defined(ETHOSU55) || defined(ETHOSU65) ET_LOG( Info, From ecd46df329ceb25f381fe13315c19fecf02a96b5 Mon Sep 17 00:00:00 2001 From: Ethan Ng Date: Fri, 4 Sep 2026 05:04:10 -0700 Subject: [PATCH 034/190] Add quantized stacked-halves RoPE operator (#22423) Differential Revision: D118320273 Pull Request resolved: https://github.com/pytorch/executorch/pull/22423 --- backends/cadence/aot/ops_registrations.py | 21 ++++++++++++++ backends/cadence/aot/ref_implementations.py | 32 +++++++++++++++++++++ 2 files changed, 53 insertions(+) diff --git a/backends/cadence/aot/ops_registrations.py b/backends/cadence/aot/ops_registrations.py index da82a1ea3ec..6789d17f16b 100644 --- a/backends/cadence/aot/ops_registrations.py +++ b/backends/cadence/aot/ops_registrations.py @@ -489,6 +489,13 @@ def register_fake( "rope_rotate_stacked_halves.out(Tensor input, Tensor sin_tensor, Tensor cos_tensor, Tensor? pos, *, Tensor(a!) out) -> Tensor(a!)" ) +lib.define( + "quantized_rope_rotate_stacked_halves(Tensor input, Tensor sin_tensor, Tensor cos_tensor, Tensor? pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point) -> (Tensor out)" +) +lib.define( + "quantized_rope_rotate_stacked_halves.out(Tensor input, Tensor sin_tensor, Tensor cos_tensor, Tensor? pos, float in_scale, int in_zero_point, float out_scale, int out_zero_point, *, Tensor(a!) out) -> Tensor(a!)" +) + lib.define( "quantized_softmax(Tensor input, Tensor mask, int dim, int mask_type, Tensor pos, Tensor in_scale, Tensor in_zero_point, Tensor out_scale, Tensor out_zero_point) -> (Tensor out)" ) @@ -3136,6 +3143,20 @@ def rope_rotate_stacked_halves_meta( return input.new_empty(input.shape, dtype=input.dtype) +@register_fake("cadence::quantized_rope_rotate_stacked_halves") +def quantized_rope_rotate_stacked_halves_meta( + input: torch.Tensor, + sin_tensor: torch.Tensor, + cos_tensor: torch.Tensor, + pos: Optional[torch.Tensor], + in_scale: float, + in_zero_point: int, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + return rope_rotate_stacked_halves_meta(input, sin_tensor, cos_tensor, pos) + + @register_fake("cadence::idma_copy") def copy_idma_copy_impl( src: torch.Tensor, diff --git a/backends/cadence/aot/ref_implementations.py b/backends/cadence/aot/ref_implementations.py index d3a5c853a4a..16768a8b68e 100644 --- a/backends/cadence/aot/ref_implementations.py +++ b/backends/cadence/aot/ref_implementations.py @@ -2265,6 +2265,38 @@ def rope_rotate_stacked_halves( return rotated.view(original_shape) +@impl_tracked(m, "quantized_rope_rotate_stacked_halves") +def quantized_rope_rotate_stacked_halves( + input_tensor: torch.Tensor, + sin_tensor: torch.Tensor, + cos_tensor: torch.Tensor, + pos: torch.Tensor | None, + in_scale: float, + in_zero_point: int, + out_scale: float, + out_zero_point: int, +) -> torch.Tensor: + dtype = input_tensor.dtype + dtype_limits = torch.iinfo(dtype) + dequantized = dequantize_per_tensor_common( + input_tensor, + in_scale, + in_zero_point, + dtype_limits.min, + dtype_limits.max, + dtype, + ) + rotated = rope_rotate_stacked_halves(dequantized, sin_tensor, cos_tensor, pos) + return quantize_per_tensor_common( + rotated, + out_scale, + out_zero_point, + dtype_limits.min, + dtype_limits.max, + dtype, + ) + + @impl_tracked(m, "im2row") def im2row( input_tensor: torch.Tensor, From 696a144ffc19b0250a393368da6fd9b39aacb6b4 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson <38941629+Sebastian-Larsson@users.noreply.github.com> Date: Fri, 4 Sep 2026 14:04:18 +0200 Subject: [PATCH 035/190] Arm backend: Reject unsupported comparisons in FP profile (#22558) TOSA section 2.8 limits PRO-FP comparison inputs to fp16 and fp32, with bf16 provided by EXT-BF16. Allow int8 and int16 inputs legalized by the preceding pass. Reject other unsupported input types during partitioning so they remain on CPU instead of failing during TOSA lowering. Signed-off-by: Sebastian Larsson --- .../tosa_supported_operators.py | 36 +++++++++------ .../test/misc/test_tosa_operator_support.py | 45 +++++++++++++++++++ 2 files changed, 68 insertions(+), 13 deletions(-) diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index c1cf33764b9..a1cdd630876 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -222,7 +222,7 @@ def _floating_profile_negative_checks( ) -> list[OperatorSupportBase]: checks: list[OperatorSupportBase] = [CheckMixedFloatingInputs(reporter)] if not tosa_spec.support_integer(): - checks.append(CheckInt32ComparisonInputs(reporter)) + checks.append(CheckFPComparisonInputs(reporter)) return checks @@ -1146,12 +1146,14 @@ def is_node_supported( return True -class CheckInt32ComparisonInputs(OperatorSupportBase): - """Reject int32 comparisons under the FP profile.""" +class CheckFPComparisonInputs(OperatorSupportBase): + """Reject unsupported comparison inputs under the FP profile.""" target_ops = { exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.eq.Scalar, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ne.Scalar, exir_ops.edge.aten.ge.Tensor, exir_ops.edge.aten.ge.Scalar, exir_ops.edge.aten.gt.Tensor, @@ -1161,6 +1163,8 @@ class CheckInt32ComparisonInputs(OperatorSupportBase): exir_ops.edge.aten.lt.Tensor, exir_ops.edge.aten.lt.Scalar, } + supported_dtypes = {torch.float16, torch.float32, torch.bfloat16} + castable_comparison_dtypes = {torch.int8, torch.int16} def __init__(self, reporter: WhyNoPartitionReporter) -> None: self.reporter = reporter @@ -1172,19 +1176,25 @@ def is_node_supported( if node.target not in self.target_ops: return True - for input_node in ( - input_node + input_dtypes = [ + get_first_fake_tensor(input_node).dtype for input_node in node.all_input_nodes if input_node.op != "get_attr" - ): - if get_first_fake_tensor(input_node).dtype == torch.int32: - self.reporter.report_reject( - node, - "FP profile does not support int32 comparison inputs.", - ) - return False + ] + if all(dtype in self.supported_dtypes for dtype in input_dtypes): + return True - return True + if all(dtype in self.castable_comparison_dtypes for dtype in input_dtypes): + return True + + unsupported_dtype = next( + dtype for dtype in input_dtypes if dtype not in self.supported_dtypes + ) + self.reporter.report_reject( + node, + f"FP profile does not support {unsupported_dtype} comparison inputs.", + ) + return False class CheckScalarReductionInputs(OperatorSupportBase): diff --git a/backends/arm/test/misc/test_tosa_operator_support.py b/backends/arm/test/misc/test_tosa_operator_support.py index a3dce64fefc..662b428a21f 100644 --- a/backends/arm/test/misc/test_tosa_operator_support.py +++ b/backends/arm/test/misc/test_tosa_operator_support.py @@ -3,8 +3,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import pytest import torch from executorch.backends.arm.operator_support.tosa_supported_operators import ( + CheckFPComparisonInputs, CheckKnownUnsupportedTOSASemantics, ) from executorch.exir.backend.utils import WhyNoPartitionReporter @@ -27,6 +29,49 @@ def _checker() -> CheckKnownUnsupportedTOSASemantics: return CheckKnownUnsupportedTOSASemantics(WhyNoPartitionReporter()) +def _fp_comparison_checker() -> CheckFPComparisonInputs: + return CheckFPComparisonInputs(WhyNoPartitionReporter()) + + +@pytest.mark.parametrize( + "target", + ( + exir_ops.edge.aten.eq.Tensor, + exir_ops.edge.aten.ne.Tensor, + exir_ops.edge.aten.ge.Tensor, + exir_ops.edge.aten.gt.Tensor, + exir_ops.edge.aten.le.Tensor, + exir_ops.edge.aten.lt.Tensor, + ), +) +@pytest.mark.parametrize( + "dtype", + (torch.bool, torch.uint8, torch.int32, torch.int64), +) +def test_fp_comparison_rejects_unsupported_inputs(target, dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (3, 4), dtype) + y = _placeholder(graph, "y", (3, 4), dtype) + node = graph.call_function(target, (x, y)) + node.meta["val"] = _fake_tensor((3, 4), torch.bool) + + assert not _fp_comparison_checker().is_node_supported({}, node) + + +@pytest.mark.parametrize( + "dtype", + (torch.float16, torch.float32, torch.bfloat16, torch.int8, torch.int16), +) +def test_fp_comparison_accepts_supported_inputs(dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (3, 4), dtype) + y = _placeholder(graph, "y", (3, 4), dtype) + node = graph.call_function(exir_ops.edge.aten.eq.Tensor, (x, y)) + node.meta["val"] = _fake_tensor((3, 4), torch.bool) + + assert _fp_comparison_checker().is_node_supported({}, node) + + def test_rejects_argmax_without_int32_cast_user() -> None: graph = torch.fx.Graph() x = _placeholder(graph, "x", (3, 4)) From 44111b7ad2a133bba270cddaab90d138ca4a3641 Mon Sep 17 00:00:00 2001 From: Erik Lundell Date: Fri, 4 Sep 2026 14:27:44 +0200 Subject: [PATCH 036/190] Arm backend: Reject complex dtypes (#22555) They are not supported by the backend. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Erik Lundell --- backends/arm/operator_support/tosa_supported_operators.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index a1cdd630876..9a6149b0a7f 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -352,7 +352,7 @@ def _positive_checks( def _disallowed_dtypes(tosa_spec: TosaSpecification) -> list[torch.dtype]: - dtypes = [torch.float64] + dtypes = [torch.float64, torch.complex32, torch.complex64, torch.complex128] if not tosa_spec.support_extension("bf16"): dtypes.append(torch.bfloat16) if not ( From 6ac72d7c8401492b0f0e92b3cae0710f5bffe53e Mon Sep 17 00:00:00 2001 From: Oscar Andersson <87121123+oscarandersson8218@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:08:14 +0200 Subject: [PATCH 037/190] Arm backend: Allow FP64 operators to decompose (#22551) The TOSA partitioner preserves certain operators when the backend can lower them directly. Since the Arm backend does not support FP64, preserving FP64 variants prevents framework decompositions that may expose supported downstream operations. Allow FP64 variants of preserved operators to decompose while retaining the existing behavior for explicitly registered custom operators. Add coverage for the general preservation policy and the linspace portable fallback. cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Oscar Andersson --- backends/arm/test/ops/test_arange.py | 76 +++++++++++++++++++++++++++- backends/arm/tosa/partitioner.py | 41 ++++++++------- 2 files changed, 98 insertions(+), 19 deletions(-) diff --git a/backends/arm/test/ops/test_arange.py b/backends/arm/test/ops/test_arange.py index 90ab437b9e7..31165ba48a0 100644 --- a/backends/arm/test/ops/test_arange.py +++ b/backends/arm/test/ops/test_arange.py @@ -1,4 +1,4 @@ -# Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2025-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -17,6 +17,11 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.dialects._ops import ops as exir_ops +from torch._subclasses.fake_tensor import FakeTensorMode input_t = tuple[torch.Tensor] test_data_t = tuple[Callable[[], input_t], tuple[float, float, float, torch.dtype]] @@ -173,6 +178,75 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: } +class _LinspaceToFloatAdd(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + values = torch.linspace(0.0, 1.0, x.shape[0], dtype=torch.float64) + return values.to(torch.float32) + x + + +@pytest.mark.parametrize( + ("dtype", "expected"), + ((torch.float32, True), (torch.float64, False)), +) +def test_linspace_preservation_depends_on_dtype( + dtype: torch.dtype, expected: bool +) -> None: + exported_program = torch.export.export( + LinspaceAdd(0.0, 1.0, 10, dtype), + (torch.randn(10),), + ) + partitioner = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + linspace = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.linspace.default + ) + + assert filter_fn is not None + assert filter_fn(linspace) is expected + + +def test_ops_to_not_decompose_are_not_preserved_for_fp64() -> None: + exported_program = torch.export.export( + LinspaceAdd(0.0, 1.0, 10, torch.float32), + (torch.randn(10),), + ) + partitioner = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")) + ops_to_not_decompose, filter_fn = partitioner.ops_to_not_decompose(exported_program) + graph = torch.fx.Graph() + fake_tensor = FakeTensorMode().from_tensor(torch.empty(1, dtype=torch.float64)) + + assert filter_fn is not None + for op in ops_to_not_decompose: + node = graph.call_function(op) + node.meta["val"] = fake_tensor + assert not filter_fn(node), f"FP64 {op} should be decomposed" + + +def test_linspace_fp64_decomposes_for_portable_fallback() -> None: + inputs = (torch.randn(10),) + exported_program = torch.export.export(_LinspaceToFloatAdd(), inputs) + partitioner = TOSAPartitioner(TosaCompileSpec("TOSA-1.0+FP")) + + edge_manager = to_edge_transform_and_lower( + exported_program, + partitioner=[partitioner], + ) + targets = { + node.target + for node in edge_manager.exported_program().graph.nodes + if node.op == "call_function" + } + + assert exir_ops.edge.aten.linspace.default not in targets + assert exir_ops.edge.aten.arange.start_step in targets + + program = edge_manager.to_executorch().executorch_program + operators = {(op.name, op.overload) for op in program.execution_plan[0].operators} + assert not any("linspace" in name for name, _ in operators) + + @common.parametrize("test_data", LinspaceAdd.test_data) def test_linspace_tosa_FP(test_data: test_data_t): input_data, init_data = test_data diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index a3262af8bc0..96c8286f664 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -787,6 +787,22 @@ def ops_to_not_decompose( # noqa: C901 torch.ops.aten.linspace.default, torch.ops.aten.silu.default, } + ops_to_not_decompose = ( + ops_to_not_decompose_always + | ops_to_not_decompose_if_quant_op + | ops_to_not_decompose_if_fp + | ops_to_not_decompose_if_integer + ) + + if not self.tosa_spec.is_U55_subset: + # Tosa operator "RESIZE" is not supported on U55. Since + # upsample_bilinear2d and upsample_nearest2d decompose into that it + # will not be possible to delegate those operators on U55. If we + # have said here to not decompose them there will be an error saying + # the operator was not decomposed. It will not be possible for it + # to end up on either CPU or NPU. + ops_to_not_decompose.add(torch.ops.aten.upsample_nearest2d.vec) + ops_to_not_decompose.add(torch.ops.aten.upsample_bilinear2d.vec) def filter_fn(node: torch.fx.Node) -> bool: """Return True if an op should not be decomposed. @@ -803,6 +819,11 @@ def filter_fn(node: torch.fx.Node) -> bool: """ if _is_custom_partition_op(self._custom_partition_ops, node.target): return True + if ( + node.target in ops_to_not_decompose + and get_first_fake_tensor(node).dtype == torch.float64 + ): + return False if ( self.tosa_spec.support_float() and node.target in ops_to_not_decompose_if_fp @@ -873,21 +894,5 @@ def filter_fn(node: torch.fx.Node) -> bool: return True return False - ops_to_not_decompose = list( - ops_to_not_decompose_always - | ops_to_not_decompose_if_quant_op - | ops_to_not_decompose_if_fp - | ops_to_not_decompose_if_integer - ) - ops_to_not_decompose.extend(self._custom_partition_ops) - - if not self.tosa_spec.is_U55_subset: - # Tosa operator "RESIZE" is not supported on U55. Since upsample_bilinear2d - # and upsample_nearest2d decompose into that it will not be possible to - # delegate those operators on U55. If we have said here to not decompose - # them there will be an error saying the operator was not decomposed. It - # will not be possible for it to end up on either CPU or NPU. - ops_to_not_decompose.append(torch.ops.aten.upsample_nearest2d.vec) - ops_to_not_decompose.append(torch.ops.aten.upsample_bilinear2d.vec) - - return (ops_to_not_decompose, filter_fn) + ops_to_not_decompose.update(self._custom_partition_ops) + return (list(ops_to_not_decompose), filter_fn) From 9363f00175dce8239d6bce0c6c3f136e6eb6db99 Mon Sep 17 00:00:00 2001 From: Oscar Andersson <87121123+oscarandersson8218@users.noreply.github.com> Date: Fri, 4 Sep 2026 15:09:22 +0200 Subject: [PATCH 038/190] Arm backend: handle Python SymInt mod and div (#22557) Teach exact symbolic value analysis to evaluate the python_mod and python_floordiv expressions emitted by symbolic integer arithmetic. cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Oscar Andersson --- backends/arm/_passes/symbolic_value_range.py | 9 ++++++- .../test/passes/test_symbolic_value_range.py | 24 ++++++++++++++++++- 2 files changed, 31 insertions(+), 2 deletions(-) diff --git a/backends/arm/_passes/symbolic_value_range.py b/backends/arm/_passes/symbolic_value_range.py index 609a84edc54..36bff1ed47a 100644 --- a/backends/arm/_passes/symbolic_value_range.py +++ b/backends/arm/_passes/symbolic_value_range.py @@ -145,12 +145,20 @@ def mod(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: return None return _combine_values(lhs, rhs, lambda a, b: sympy.Mod(a, b)) + @staticmethod + def python_mod(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: + return _ExactValueAnalysis.mod(lhs, rhs) + @staticmethod def floordiv(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: if rhs is None or any(value == 0 for value in rhs): return None return _combine_values(lhs, rhs, lambda a, b: sympy.floor(a / b)) + @staticmethod + def python_floordiv(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: + return _ExactValueAnalysis.floordiv(lhs, rhs) + @staticmethod def pow(lhs: _ExactValues, rhs: _ExactValues) -> _ExactValues: return _combine_values(lhs, rhs, lambda a, b: a**b) @@ -181,7 +189,6 @@ def evaluate_symbolic_expr_values( """ root_expr = expr.node.expr if isinstance(expr, torch.SymInt) else expr - constant_values = _constant_expr_values(root_expr) if constant_values is not None: return constant_values diff --git a/backends/arm/test/passes/test_symbolic_value_range.py b/backends/arm/test/passes/test_symbolic_value_range.py index 99dfafc93a6..698be1a8bd2 100644 --- a/backends/arm/test/passes/test_symbolic_value_range.py +++ b/backends/arm/test/passes/test_symbolic_value_range.py @@ -9,6 +9,7 @@ evaluate_symbolic_expr_values, ) from torch.fx.experimental.symbolic_shapes import ShapeEnv +from torch.utils._sympy.functions import PythonMod def _make_shape_env( @@ -70,7 +71,7 @@ def test_evaluate_symbolic_expr_values_bails_out_for_large_symbol_ranges() -> No assert evaluate_symbolic_expr_values(symint, shape_env) is None -def test_evaluate_symbolic_expr_values_does_not_require_shape_env_bounds( +def test_evaluate_symbolic_expr_values_bails_out_on_recursive_bounds( monkeypatch, ) -> None: shape_env, symint = _make_shape_env(hint=3, compiler_min=2, compiler_max=6) @@ -81,3 +82,24 @@ def raise_recursion(_expr): monkeypatch.setattr(shape_env, "bound_sympy", raise_recursion) assert evaluate_symbolic_expr_values(symint, shape_env) == {2, 3, 4, 5, 6} + + +def test_evaluate_symbolic_expr_values_handles_python_mod() -> None: + shape_env, symint = _make_shape_env(hint=3, compiler_min=2, compiler_max=6) + + assert evaluate_symbolic_expr_values( + PythonMod(16 * symint.node.expr - 7, 4), shape_env + ) == {1} + + +def test_evaluate_symbolic_expr_values_handles_python_floordiv() -> None: + class PythonFloorDiv(sympy.Function): + _torch_handler_name = "python_floordiv" + is_integer = True + nargs = (2,) + + shape_env, symint = _make_shape_env(hint=3, compiler_min=2, compiler_max=6) + + assert evaluate_symbolic_expr_values( + PythonFloorDiv(symint.node.expr, 2), shape_env + ) == {1, 2, 3} From e6d1319cf1ff4fdc191e2ba24d14c5f0497ea148 Mon Sep 17 00:00:00 2001 From: Per Held Date: Wed, 19 Aug 2026 12:59:06 +0200 Subject: [PATCH 039/190] Arm backend: Pass memory mode to Corstone Pass the platform helper's explicit memory mode through to the Corstone configuration instead of relying on an ambient CMake variable. Update the MobileSAM caller to use the same explicit interface. This commit was authored with Codex. Signed-off-by: Per Held Change-Id: I5df5f88a1addf41240f3ebebc5a8acfc6face32b --- backends/arm/cmake/ArmRunnerUtilsInternal.cmake | 6 ++++-- backends/arm/scripts/corstone_utils.cmake | 2 +- .../runtime/CMakeLists.txt | 6 ++++-- 3 files changed, 9 insertions(+), 5 deletions(-) diff --git a/backends/arm/cmake/ArmRunnerUtilsInternal.cmake b/backends/arm/cmake/ArmRunnerUtilsInternal.cmake index 04c979b04ce..d7d66c45450 100644 --- a/backends/arm/cmake/ArmRunnerUtilsInternal.cmake +++ b/backends/arm/cmake/ArmRunnerUtilsInternal.cmake @@ -328,8 +328,10 @@ function(arm_runner_configure_ethos_u_platform) arm_ensure_ethos_u_content( "${ARG_SDK_PATH}" "${EXECUTORCH_ROOT}" ${FETCH_ETHOS_U_CONTENT} ) - add_corstone_subdirectory(${ARG_SYSTEM_CONFIG} ${ARG_SDK_PATH}) - configure_timing_adapters(${ARG_SYSTEM_CONFIG} ${ARG_MEMORY_MODE}) + add_corstone_subdirectory( + "${ARG_SYSTEM_CONFIG}" "${ARG_SDK_PATH}" "${ARG_MEMORY_MODE}" + ) + configure_timing_adapters("${ARG_SYSTEM_CONFIG}" "${ARG_MEMORY_MODE}") foreach(_platform_variable TARGET_BOARD ETHOSU_MODEL ETHOSU_ARENA) if(DEFINED ${_platform_variable}) set(${_platform_variable} diff --git a/backends/arm/scripts/corstone_utils.cmake b/backends/arm/scripts/corstone_utils.cmake index 72a0c27a8c4..95ef6f8b866 100644 --- a/backends/arm/scripts/corstone_utils.cmake +++ b/backends/arm/scripts/corstone_utils.cmake @@ -126,7 +126,7 @@ function(get_corstone_linker_script OUT_VAR SYSTEM_CONFIG) ) endfunction() -function(add_corstone_subdirectory SYSTEM_CONFIG ETHOS_SDK_PATH) +function(add_corstone_subdirectory SYSTEM_CONFIG ETHOS_SDK_PATH MEMORY_MODE) if(MEMORY_MODE MATCHES "^Dedicated_Sram($|_)") # Both model and scratch in DRAM. set(MEMORY_MODEL dram) diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt index 6f30110b95c..ce6786f6a58 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt @@ -100,8 +100,10 @@ find_package( executorch REQUIRED HINTS "${ET_BUILD_DIR_PATH}/lib/cmake/ExecuTorch" ) -add_corstone_subdirectory(${SYSTEM_CONFIG} ${ETHOS_SDK_PATH}) -configure_timing_adapters(${SYSTEM_CONFIG} ${MEMORY_MODE}) +add_corstone_subdirectory( + "${SYSTEM_CONFIG}" "${ETHOS_SDK_PATH}" "${MEMORY_MODE}" +) +configure_timing_adapters("${SYSTEM_CONFIG}" "${MEMORY_MODE}") add_executable(mobilesam_prompt_segmentation_example main.cpp) target_sources( From 090f5de54466c603636c0c898ab280e41fcfbb23 Mon Sep 17 00:00:00 2001 From: Jiri Ocenasek Date: Wed, 22 Jul 2026 14:23:22 +0200 Subject: [PATCH 040/190] NXP backend: Building MCUXpresso example --- .github/workflows/pull.yml | 19 +++ .../backends/nxp/nxp-mcuxpresso-example.md | 139 +++++++++++++++++ docs/source/backends/nxp/nxp-overview.md | 5 + docs/source/backends/nxp/terminal.png | Bin 0 -> 45222 bytes .../executorch_cifarnet/CMakeLists.txt | 144 ++++++++++++++++++ .../executorch_cifarnet/build_example.sh | 48 ++++++ .../executorch_cifarnet/prepare_model.sh | 32 ++++ .../test_build_from_scratch.sh | 58 +++++++ 8 files changed, 445 insertions(+) create mode 100644 docs/source/backends/nxp/nxp-mcuxpresso-example.md create mode 100644 docs/source/backends/nxp/terminal.png create mode 100644 examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/CMakeLists.txt create mode 100755 examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/build_example.sh create mode 100755 examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh create mode 100755 examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index a190dc132d1..30d3849a25e 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -1727,3 +1727,22 @@ jobs: echo "Neutron backend library not found!" exit 1 fi + + nxp-mcuxpresso-test: + name: nxp-mcuxpresso-test + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + permissions: + id-token: write + contents: read + with: + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + submodules: 'recursive' + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + timeout: 90 + script: | + # The generic Linux job chooses to use base env, not the one setup by the image + CONDA_ENV=$(conda env list --json | jq -r ".envs | .[-1]") + conda activate "${CONDA_ENV}" + + ./examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh diff --git a/docs/source/backends/nxp/nxp-mcuxpresso-example.md b/docs/source/backends/nxp/nxp-mcuxpresso-example.md new file mode 100644 index 00000000000..8f8bb93fda3 --- /dev/null +++ b/docs/source/backends/nxp/nxp-mcuxpresso-example.md @@ -0,0 +1,139 @@ +# Using the MCUXpresso Example + +This example demonstrates how to build and run the ExecuTorch CifarNet application for the NXP RT700 platform using the MCUXpresso SDK and the GNU Arm Embedded Toolchain. Before building the project, make sure that all required dependencies are installed and that the necessary environment variables are configured correctly. + +> **Tip:** The `test_build_from_scratch.sh` script automates all the steps described in this guide, including downloading the ARM GNU toolchain, preparing the model, and downloading the MCUXpresso SDK using the `west` tool. If you prefer a fully automated setup, you can run it directly instead of following the manual steps below. + +All scripts described in this guide are located in the following directory of the ExecuTorch repository: + +```text +examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/ +``` + +## 1. Install the Arm GNU Toolchain + +First, download the Arm GCC cross-compilation toolchain that is supported by the RT700 platform: + +```text +https://developer.arm.com/-/media/Files/downloads/gnu/15.2.rel1/binrel/arm-gnu-toolchain-15.2.rel1-x86_64-arm-none-eabi.tar.xz +``` + +After extracting the archive, create an environment variable called `ARMGCC_DIR` that points to the root directory of the toolchain installation. The build scripts use this variable to locate the compiler, linker, and other required tools. + +Example on Linux: + +```bash +export ARMGCC_DIR=/path/to/arm-gnu-toolchain-15.2.rel1-x86_64-arm-none-eabi +``` + +To verify the installation, you can run: + +```bash +$ARMGCC_DIR/bin/arm-none-eabi-gcc --version +``` + +The command should print the installed compiler version. + +## 2. Download the MCUXpresso SDK + +Next, download MCUXpresso SDK for the RT700 device family using the west tool: + +```bash +pip install west +west init -m https://github.com/nxp-mcuxpresso/mcuxsdk-manifests.git mcuxpresso-sdk +pushd mcuxpresso-sdk +west update_board --set board mimxrt700evk +popd +``` + +Afterwards, configure the `SdkRootDirPath` environment variable to point to the mcuxsdk directory in the downloaded dir. + +Example on Linux: + +```bash +export SdkRootDirPath=/path/to/mcuxpresso-sdk/mcuxsdk +``` + +The build system relies on this variable to locate board support packages, middleware components, startup code, linker scripts, and device-specific libraries. + +## 3. Prepare the Model Header File + +Before building the application, a compiled model must be provided as a C header file named `model_pte.h` and placed in the current directory. Run the provided helper script to generate it: + +```bash +./prepare_model.sh +``` + +The script performs the following steps: + +1. Installs ExecuTorch and its Python dependencies. +2. Installs the `eiq-neutron-sdk` Python package in the version that has been tested with the current ExecuTorch release. +3. Compiles the CifarNet model using the NXP ExecuTorch ahead-of-time (AoT) pipeline and produces a `.pte` model file. +4. Converts the `.pte` file into the `model_pte.h` C header, with the correct memory-section attributes for the RT700 target. + +> **Important:** The MCUXpresso SDK package includes a pre-built CifarNet model and a set of Neutron libraries, but this build flow deliberately does **not** use either of them. Instead, `prepare_model.sh` installs the `eiq-neutron-sdk` version that was tested with the current ExecuTorch release, compiles the model from scratch, and the linker later picks up the matching Neutron libraries from that same installation. This keeps the ExecuTorch AoT compiler, the model bytecode, the Neutron driver, the Neutron firmware, and the ExecuTorch runtime all in sync. + +Once the script finishes, verify that `model_pte.h` was created in the project directory before proceeding to the build step. + +## 4. Build the Application + +Once the environment variables have been configured and `model_pte.h` is present in the project directory, set the `NEUTRON_LIB_DIR` variable to the directory that contains the Neutron static libraries shipped with the eiq-neutron-sdk: + +```bash +export NEUTRON_LIB_DIR=/path/to/eiq_neutron_sdk/libs +``` + +The build script expects the following libraries to exist in that directory: + +- `libNeutronDriver.a` +- `libNeutronFirmware.a` + +Then build the project by executing the provided script: + +```bash +./build_example.sh +``` + +The script validates all required inputs, configures CMake, compiles the source code, links the application, and generates the executable image: + +```text +flash_release/executorch_cifarnet.elf +``` + +If the build completes successfully, the ELF file will be available and ready for programming onto the target board. + +## 5. Flash the Application + +The generated application can be programmed onto the RT700 device using SEGGER J-Link tools. + +### Linux + +```bash +echo "loadfile flash_release/executorch_cifarnet.elf" | \ +/opt/SEGGER/JLink_V796k/JLinkExe \ + -IF SWD \ + -speed auto \ + -Device MIMXRT798S_M33_0 +``` + +Before flashing, ensure that: + +- The board is powered on. +- The JLink debugger probe is flashed on device, if not see [documentation](https://mcuxpresso.nxp.com/mcuxsdk/latest/html/boards/RT/mimxrt700evk/gettingStartedXplorer/topics/program_lpc-link2_with_segger_j-link.html) how to flash it. +- The J-Link debugger is connected to the target. +- The SWD interface is available and correctly wired. +- No other debugging application is currently using the J-Link connection. + +The programming process typically takes only a few seconds. Once the image has been loaded successfully, the application can be started directly from flash memory. + +## 6. Running the Example + +After the firmware is programmed, reset the board and open a serial terminal connected to the device's debug UART interface. The application will initialize the hardware, load the embedded CifarNet model, and begin performing image inference. + +During execution, inference results and diagnostic messages are printed to the terminal. The included demonstration image contains a cat, and the model is expected to classify the image accordingly. + +A successful run produces output similar to the following: + +![example](terminal.png "Example") + +This example serves as a basic validation that the ExecuTorch runtime, model integration, SDK configuration, and hardware platform are all functioning correctly. It can also be used as a starting point for evaluating custom neural network models and experimenting with on-device machine learning workloads on the RT700 platform. diff --git a/docs/source/backends/nxp/nxp-overview.md b/docs/source/backends/nxp/nxp-overview.md index 581c375d038..604b480e614 100644 --- a/docs/source/backends/nxp/nxp-overview.md +++ b/docs/source/backends/nxp/nxp-overview.md @@ -52,6 +52,11 @@ For more finegrained tutorial, visit [this manual page](https://mcuxpresso.nxp.c For guideline how to update the eIQ Neutron Runtime on MCUXpresso SDK, follow the instructions from the eIQ Neutron SDK package `docs/NeutronSDKUserGuide.md` available here https://www.nxp.com/design/design-center/software/eiq-ai-development-environment/eiq-toolkit-for-end-to-end-model-development-and-deployment:EIQ-TOOLKIT. +## Using the MCUXpresso Example + +[This page](nxp-mcuxpresso-example.md) demonstrates how to build and run the ExecuTorch CIFARNet example from MCUXpresso SDK. + + ## Reference **→{doc}`nxp-partitioner` — Partitioner options.** diff --git a/docs/source/backends/nxp/terminal.png b/docs/source/backends/nxp/terminal.png new file mode 100644 index 0000000000000000000000000000000000000000..9a4f1d70ee84311ae5838c767c03cc908a0a5d05 GIT binary patch literal 45222 zcmbTe2T)U8_dbgHiWRV-(nUZNgebj7>Ae>z5)qKzrG*ldCS5>?)JX53C{0QPlolb< zrAP}Q^pX%DBqTS%_xt|8nfseN_uez(Fga(ReRf^zSo~-H1+p6Pjz!2 zIy%Nq+P@Rs9;J44bRT6klph-hSg#*3RpMqKTzet8XIzHvJ3VK5Ja%r^H}~e0Of>=X z{B?0rtA^PLulK(bvBY~yi)mTVV(#4AhJ$n+cW#ndPKn;_`B?D`bXP^zqq4*Xz2(67 zsY)uWzN?{MOy^N53&%SKCDQNl_!RrP9XS_gNhOC3oU~k_`YKZyWZ(K}D+rXm#e;nt z9Uc8{Q(A6uz8SK;53K+Z>p%sg&MK502uW*~i82o?qhZDyI*c~i+_@%eR#a@3W zB0+7xJOo3|xJRF&Jqt@+cUe@Qc5(VAW*>ei2-k$KNK%ZzSB~GC+X$iv!m)zRFB@%- zprOMxvn8uO(DFn(E5D0$b>a^bT#%E$%6o?buP4EalmB4*mFI%0PvmRI@_YJje`XYGx^q z7wj^8)EQrUu7;h3@B3%$a*@tZ+Hxm(y=IS^DOupv0RH{6M?FdEO!~jig73sxYu$7)WfYG07@8%)NzuD2TSxHv zC;|mfDNdJ;6a&|b2j>&I#bU_|W{?ohi4Eu)xD^ulMABc$4U617&aa)8G@rKA9=C zDj|&PsA2ew=@`*=mKeR<89keaiOBC7LsCe{73pqb14R;e9fM~B?gD$=vXv3!IT*7BEkoyI9e>^6Vs9B9mEK#BD}ZbDlxCrK=&FqyyAasjdrKk{D& zxg4#niqabStqvg(HSp)`l}I*4n@)!!R>>i{J_QTjaK4CE#UHS%lVeX})RizG64ZHj z3K;y$XN(+41ip2C7QB%}*-4n;+E)vG(hqU!@(+%7{>bYwh|uLyVl@Zun5lY&ds<2# zZ=USumNm#C@ax)?gghU~*WcTbXG2hIz{=dy=V52gYw$R-IbE|43aZ^~!^z@HbiCQ7 zC59RU4kJOojJTcjH99q^{c~p~mQYqwGDFV|{Z6(-wGm-Nj}>VQzI7I4wxXr?8@548 z!dub0)EKe&)PUGwdBtQ0Rxyv8PUI99W@YD|xWG%fqh!saZpS93(VRPRE_COHB&)u+ z%lqAvG4NkR*waGEyMZZN;5{{WOzOcFIfrubsIC|8|0$PM!;BRIK~bdIKU2*x9o>yP z_{rPzRXq}+*7*kSS~!FQv!m#b25V_l7!}HxaXBll?4(pe1=_qSG{p)1hckzzR9s@y)m3+W^Aa&JQX zK}Qpt(i!oH{K2`1;KbK&E$0)P0EVisK0~(kirr=7++9)!2MeS5SorOB8{F(1omiYL z$RTd|syXYOZBr+XXL?{BvRgx{Feb}WT+4-XF2Zpzf^gkWm)l7jU&!JP1(&bIuqr{8 zP_Qwd#Qn`03Lf4jKxqSh;BQEMJ-K(^;6P-wE+2XNxUO$&+ZN#0w}&A1=xWSH=6Yu^z7-649mumPJb4>5Crd}S zEyKqqizS${0X5Vv!WfS4)#!|T>y_~K69?A|1!>Jd$Ru;fKeQ)dE}tuQeAXCpjz#zYGH_vg5#YWP$zj(`1i`(M*v{iOcHm#T-? zggy&xn&xpISMC|f==6KeJYQ!uvK+F_M1CxP?r&WjrWMQ~1gwf42faKM4NxynpI?`b zy;)dRc+_#H>89DmB95QUE!oj)J?Qn89TA3(jCd>D2A1mMdy6;UOSyM=2{^lvABo4< zB}(Rj8=W&PqIz<^+_Mgm&d-$pb@ivWQ;LRF{p=c}V>_(C4WK4ed**@&N7CV7m(-Ez zuP1u^D;tIBdgeo=_E-4>kFMC?OdEFL?XJDtGDC9vj3f2tb^{iRkXw#KTy*APRY0KW zL6A?+PE*Pj6SZZPIU-L^YAC&-7_@u0rrBohi0wGIPghzAa^~oruHLPbQb+vkH~<-Kx`MH@J!ZERSol{|iHiYQgvPy(TL-t; zx%P5E01oR;f>2-&*|f#=qv_`5?lZGo|In_8wQ~zN3{ou;U*%P-#DI1BsBU!{jXXtc z6l}4|xK*v~$$F~x>bdv(%;!lb5=u^(Y)XfgoKA@5xs3|XW$b^U%uwyN;c(e3tix@s zqCM_aEB^e}H-$Oprh6;8jw0v6F*4V-PC#KLn2Pn@QG4-?=i3=O~rs5{mIY$*Svqrg*Gq3D0xQ&Jt%-{cD+kaibCB&a~tN85ehez+xg|a zVnynjQa&9?aHMVuo?s|?D#mH3-!%L5^24T~;d#|zvElPU>)t+9pFW?^FW+xUNo_p6 z6`%bgEo%Bf^Gy<Pk6+(!6x7Ia)YOo{|0DBl1n{Nu{FuYAy{p8nI`V_8dMB6 z^d9IFB4eZRdH#{i>X*~<17plTT4NU}%5krd2~&P83A*9SdH%m~qCW77Jxlc(8*Qwf zm4ZRxuIb?!PIL-~+F>qMKa9?R~ZTDOwD zu-NVrA|A-mxKMKA@=~-wS|K3J9mDW#s9?8O5`FZHtCsDK;rS(snBk_(AJe$~ zq|$zdvGMtAG@MA93?#+n0uN`$SL-kyB@WWZ3p%p}ntc75yH(FwuQH6#h&ub zEvSCzC%Ej6qn&1wKV4kjG5RefCKR2S)-vsoRyy=@vHG!4%jkok{W(7Z0npKSK~k8o zFH%5~UbPy`NOh0kcmGZA9Z0!&K*~1Y+i6MA&_#u@-9j3Yvr;hQGborP`3k5+cY)}z z>R^!qdKH3(L*K!Yn~L5IL#=bWck!_tvFNSmSS&EtP~XY$<)Ig^dB5hu{=6IpNL_v{ z99yA>B@i%PeikPWPN`r@)XD)%S3bIr5yM9BY^Y-&;8p*)V74b>Y=U&Ezu*g0wZve8 z20Ndx%XwVW&uL)ORaKhtwjW*>_r}(*GP9T^yc!%0OdNQ3m2v|TrI)D2k%vg{0S#{9+a||faD;s z^Y37qomw97$DIke-B10X5b}mYAja$y36loDNdGGjd^>b;f2`ru(mM#5kfRIfofU<_ zWo<}%s8mkIAb)ZS`g0J)U&A#BNtor>LIAbWnR;>N(>L94FBkUh{0SbjlztJagAE&B zd^i7Dl2vgVtdVp9Kd}3=hwYgcy1Cs1ozK8cbDaT}x=v@GKKuAjW5maXDrG!bnZYu+ zqQmBU=dTAZMwEDtU3+2}=W$`WM$P5Re8O;@%e02GRmZ=IPO#bbC#=dsD2@KajY?mfc31e)DKbOTxG+2lidn~UCo({Y+080gu zdEYYHt7+xv91cmecasM0r2!!zdx8us%Z&qIgSxpc2yemh?;o^{VQCb}@7tq1yseK9 zo$Gh$y;S&pBLfS%e_X6N>ti2Zf))j^bw@jSDsfxSc!|p2t8HrzR3RL2=EY3rvY;QM zJ#r`sf(ujC(-oXBd9$Yd?>8oaGk45xo0wMwV@13|FGM{>f2JI$KQNkFh%(QY)0Gb; ztif3A;7!SwJe2THiu-J4$tvhBG=~o2lhhqZpu9rh&eZE{vta9<9Z?QxwT`3sw&+^z zWSn4)nj?sib9kg+Nw|IM?{)1WX6^D76!+;g1@U+m-)?a31CH{46d109E$(-TF$law zigVSL+x?6ZOGMqs9*Iu}YgU1DcMH{*8_KLCM9hlxJ{1~xJc_%^lntN%`C*kH@p0+| zv4I`mv(#ZiSd{A`$g9lkTwanFmt7-?Uz7)BgGtyiU@AdLa;7VQHG=ClaW6@TUI8|> z5x^zGkbg!2WSxN4=79gW3dz+Q4GoR;+tHqd<)ESdttuH#h5F09}&A4uoz#`vgvas{+wWet3ooJq3XVcQclEM^T zGiYl9Mq;2$*|$0%Lx!mstzUBGmzaT(*k*@%3~tgMPms&CrR?>{DlBoZAi_9-k&oW43J@^%yIqQ6TLQHL$qm1BNcDO@?7 zO*(4ZqphhhOPp@*<-^=90w*FW-*zRnbopx=w6HRMZF~R>Q+A)M|@$g>*2gVa%sUgr|viDmGdEId&`0Co(EZ9 zYxZ@cQ8IP{60F@_&^Gf5osDXI-mk*U>2a?hHIWPP6zG+m?{@v&&t?J91^P7i9@TDu z2ZbHUN3+cspik)HSzI;$Z8l=DLN(ulzZSctG_2U9&Z8D2zgik{9+{erb+{ubCfnj3 zaCVGHg;$FQx!5Cxy5m`Ath1O|7fb;}lU8o`XIn8%SM2R2jB&c9or z8Uz?5<*CY*f!!wU^Uf4C4auP->-}j5Z|xV=q|12GFc;mNn>1Qpu{kJ0 z+Xo$BU4BMs!APDpJ!|!0%0bNToA>iy(078~V^OJyjZC;lNNr<6>4}le#hi3~To3BDKv&9Gg;F5uR!W?z;nin1yv$DL*%1s*h& z^3@z`p4x+31l^f!u)gH;2Vaf3?a9+ne;h zXF&N%r*w6)2;oP1vGuTz?091U_l-{$cS!z^L<`Du@1f{R!G7I~l=7~4&H}N18szogS zqFWNq=vu$sRo9PPXkXG!&F+%M-PlVE(6t`^b?d3Jos4kcRWp3bqM^J%f-m{Bq`(4a z8r904m()fDq*lJLM`aWipxxkfEAt|mLq51t1v7!^?TpnI`GIO(A|JZ-Y@U0Q#RShR zl81=2z@JXiUjB*V^AE2N=KZQ&WvpSVg(ay>6$qvEnD(>a5rc4+v*M5Fiq0tX#5J7Z z>Ns&>B%60ew@M4P`|OC%G&&k5W^ZS z{7c->oq6;hBQ_AUqWc#=Hu7yjsF>d3FPlWy$gl<3bAXuspPz11V_cCf`onS2ustiN z9CFYL+3g)Tl~`8d{-$L_LV#!nr(jPS54g1BloD6$76zsL5SrNqDa9@;XN3Ud6m7LO z(Rn}6>dm(4{cc{ZJrJ}NBO}M}1|jrZh)ab7Cr&|c&}ParY$xD8)48%|GBbU?Kk;AC zm5toj^lqF>q(Vq=t+=kDwVYxAdt*wY_@?h^7c{45s8nvl!q=^_#*X%PK0{Bvm=7Bd zwCE8!T#~s_?cj{Zf1QuBkoUe~)?&|gI>l7?DdERD^=7*u*VdOwZn{a&jl0S`9}yPKFS^#Av6ZhY(4v)0hyiB;E2YK`Z_h`C7&NJf228xS79@s z=;EcGfG5QQGxu$@WSUNi@OG7oRBI2ngIZYmCQiO0cPHNnyvAr^|dHXC95OJW+Dl{8yt(>eC8?1SJn=eRx65%`8Tfz8Jxc zD~E0Z(1Z@NpGcccf2T^o`J|G^2anasj;i5zy+3_;1!Mxk`cqzXzT&MCTK0OOzY|C1 zulATR`utg#5|xMHU%$`g=0VtyZ8J@apLk$xmLu57vNgz^F1E_I2YFaL{}}E6`qaQI zx>U{2pM5D6*Qj6(aCx|Wuk;HlWLfMDXQNzP3Mbo)kjZsx_~XyHI^mA91A9kHTRHu3 zuBz)0HnvExt?d#=G7=7KX>;>;x_R>@%t{`>N!=53bX)SsVJ=CbE~8B7OBKf*<#z8v zR`r!{%r13eR_QA4aK%@%v4}QF*-Dn$aZ6wuvKGQ8qOhElJf{JXihk%@l*~<7icjUREE@zw?UN~FdYHx6Z5=^x4o5hUCE< zco#^OilH39yjzA3^m{pM>pNfPkD;f~e^1pipJ4Axrq9?sbDT2=x z@QV`po{VPCp00YHHWlJt+YdPUo`(6eQ+p)gW*QQhVF6W;=ESef`(c>GSFf1e^;HR! zc22Hx_A~s160>`#5u`fuAumzz5$E--!f}srgx8_rAkQp+LGqJ3pCQC`5QIWTJbZA~ z{$qAh4<^>7SyyZ!b^~=p#;zEYJR+dkXYB~TQigY$!U8P7S5m^AUAbwfT#MB!&s1vI z=1K+@Ap~J@<*^czfPBdOIxEQj_Bn_|KAN&7msgH~!mNfGw|sJce6<>1@b1~`xi%Ld z% zS^>@fO$Xc=yaW!q?=B_QcV!}FehAvIs zR3J`j)I8e#=->I?V9Pb(){5I*gI1&ye>A%=?Hz^(TV;lssY%CKly7lC!~goMZNLAI zSp2GFM2Ag@y!RNu9V+=CO346b)94V$Qd;Qy{Q}#pjT!Yd}3Yw58j;*R$`YLaH zvI5#SGPZL?0vxv4fi@8}T$Sv8mR|;|KFj8mdXKI;p~*ef+7{4|yx=D%&cnadSN% zbzvAoF{+6!B~Ge5A^c?1h%78VCxy&65M`^pzbbtp4b3rx5|hKK;)E) zoG-!Q8mMeb1LW*~uZBy11B|S2zWC^C#)l(RW{P7iu;VYYy#J!~mxoc4t7|K>JCvXV z?TXvR(co=FA=iP}FR#Am+QT0d_lMk`{|BdjmMQrBXED7Y#D->UXK2PYu7!k7n=RTu zhHiAvEmadN^~bA|b3%mw+P>UR8fV;pFRRQ*#7&1ro`f@R1$p>A5=Ac|j^CKW|E8`r zzIBeBTI3Jc2g|zqAKoWw&3xAZp81S;PDQ}UoOJ_VN}Yc$8-O*sjLr+_)Gp9aWuzu} z)h=*Ka;xVF{qDiN_1cFt&h`744-RS+8o7_WsFsTImM`RIJf7yUnWGWCqmfmbV5E}% z!CApSTC8=htZ$X9;nkS_mB)B1PxSPLesg672EiwE3I_i3P{sdcM!HsxpR|>pYW7-W zLr3rS#L-vXjXc$--cB8e7<8#h{a1b+dckG=rzLQ!#GZyTYn)eEv_H@S_1={OWw+e) z-1O<>jetMr02L);`5{+{4Y28GT2yT_Ty6BkT7VcxK-_~X5^rp%iOe3iv%#UxVh{=d z?T*+i{F|uzCN5pX{qkFt*Kq+lk5m8mVD^8>>3#T{Bczi&^Ovink|AYLZ&5mqA{z8b zR^20?#yh?F7h9jTfxyJ$ia&yuBFnEgr3&b{JY}tLN@Z_Q68qxxQD1Z}%|{NJyS{vI zyrf&6mbB=LmvDdA;qES2O4WIq%Pr5bAu{MaavPR8QVK2ml0B?sydx2Kq?^^pS0^BZ z;aT-P!|Qf;#jPBgF2;J=I2qq@+Hqv9n$0|2$-RJTx`u)mjr>1W4`JMAy`64fcwpY8 z`(?%sc!kyRaW94t(wHD=7z4m?BY664L4j?=@+T9!T`m)eR?yDOc7FeGMz1BGR z5z8(r4#)&1Z>*q&BwMp5&9wBgQX~hq(qCdbcBW=Q3bmo-a*g+Er!=4rG7*FYY;1AS z{Z{xvr5qyq8m5V{ptQhg$dXT9{50WrWEWJLz*HC&ohz z@->IKMLG5R3yGSKjGXRe_}kq-1twjSXwgRGz2GB$N1PU-7b`gbBfQn?PLAq93lcYB zq>$-#LsgRya2Qx97QTD?IFy7Mos-U0t1c=joLY%-WZ9+|YI~a`C1)-7YRQPxAQs&m zx`gnM-b;~NHDSCH7X1z{aY?Bl~r+N|1Y6MXBBP4s_4Le|9J4}=4AetJ^o+rE0A&lzJ1}4 z=lFBq+yeH4AX{5ck2c^(4WkigU3EChhnm^1!?PL9kNyG(|InYJmGmsoWmO>ZL$Lg2 zki{g@WqB5AI$s8PaC?xq+h?VB^Gc?*il=Q)bAj+*Bjg#LTg1(aqi$6pj8DFI#z->| z>I43jkiH%HL5nrvWckSWPJuO)H^cFWpp;M;es49zBGF>rPdmCLRNe}}ZK0qyv!~*n zr669$jo{mRrS1OY<)raghI*TRo%*4j^rFNn_=vAxO~Lof+D}+*2p=y< z)~8yLf4yx8VGhVjJ%^d9>p*KC)l_bM6%=i{-`$jLEtUgo%XpKD%_!VA>lmEEAY*Ew zZW;u=7GLteISlF__j_-UFd!E*Bd54KEkbF&+ckWPbjfnaRqBJvbmXTG-5HNf-0#UH zJ)Z(4f~hp#y!VBpnic&bt75X9V(Q)9DPv)_FV^z>A-LSk9Q+}J5)vupOgs)(go7sb z0;Evv$J6{R??JGi-H!6Bi*ouE_KJg;w14UE7WZQ>y0~$y8+>Ho;C_~tZR`=l^JcsR zopa1A?Pc?~DRh!vbh~K4e{m@8?MmaKj+JgvDQRO%;y;_Vj@p34c>705unu9i3HS>Y z;JT6ez+U&UpS?0wpP9C8UUV(hYP4YS4!HJsG2OZ+{-mRA{&PV)V~4r;YKqp?R!_0H zSExYcHa%4(;4e7U0eFN6ax^fGF%h+<<4d0cy4RxP8b{_|j;7bHVSd9yX6i3W3UOT> z3V-N{>e%pPi6CnH^}60{YQ`cY zXP}bDu`H&=#(_(G^GSK$uT)KV?#@r))+6Z5UxM~Y5HmG>7T_ema}q#k+1;ZuWq(`0 zF}nNjc9Mfqx6@7WzX0KX=ZH5ejz^cSk?-H!U07DC)7gy0MoyhOSDEF;`H(&@r`Nav zp?G#9y;er7h{-)sDs^lvAGjmHg;7eZL02Oa%uLD174cli@dkTutVGo)*!?44KI$y6 zTl^}~oa$XGenE|RkGvXG_b#-KzBD6WKb2)sP?Y54URzA=CQPT=-DZPlQYQr#!1X&5bA( zsQnTy`fB59x$bEh$nCYY-H|rlbRR0-{+Cek(vl%|w2{g-^!8k5ea(u=l_$5F5%%nQ zTNjEuyCZQ(r>w+xHjNw;s>WT;lq>dSL;D9e%bf@<1C}o<`PsRIKw^;QJ|JMRxMtsA zLl65o5orm$AE(^LYYolN)@n1t@WA3Gj-=3}DG>RP^xG%&KY6=b?AQNI3qUq8p^3w9 z?mJG`^(J~Gh?dSu&#_#g*^F(LE9YodV>1NjGr`$8jM=33Mp1%K1Klv zjHGb)i-LCKj}An%2;O(%FXCUEiDkC*CQEPpO~v^ssARp}Jq28sj`5t(4L|?KM<=^d zfaSAIwxn-mv}mk{x2zg@?_LV|`6E9Be}Rx!@E41{uabBCL*l5kG*s^z-4>*wI4D(E zsD@ORgE^L8%T??H_iTa1KzEDy5-|XC)mh6$Z?yQf<(fnM3y*#ClyI*#z9)$bwdT=g z-Yz3WII{;KIHferX*o-5nQgH8ms_#bH56I~mhLy^h^o_NT1%RKn-Q@1B&NJDP%Ua# zZn^jEV}J1FVXxk-MjNzI`cqmY&@%02!6+^`!7m*9JDFT|{Zys_m9}o9($-YW(QT91 z5Be=TC+)BTn&2fmh>a2Xe}0568Ze%dbPm@{9g804tuWV+yXrGfVpgAhhjC<8W)jMB zwaf*r@~b)^EI|uvBO2Gj?jf%Qt|XH#hHNjAc6V_KPHLR9RPD^qi)uL3hWqTW5sJ1r zAMw&_y^Xpwt~mO6`YX5vZR0-dUB<(-f{6IOFJ*1EYphbwSuMy^OSJNG8`l7?3$HZX zv&3u5#}D|Gz(YcMx0R@@_LV7JpewTXJ5#;ZkN3lGfADUwjMq5@{8FMcW1-Xgf2NQM zKjDS;b4GgZS1)4L-~QewoEj}QNvkiMD$sp?DU63+%Mtfw_e<%Ct9xh<2%&@x(Y&q; z8!!xAIf5P}&sxqhW=umuNT0*)*@LKxTo~}l>))-A+mi#;-txzIYi>iz{$s>QsaHeS zLZ1`r{UDMyCvYez_u z-I=QCi&yH2<6yFqoa8(MNW1*|r7~gWECRoPQy*{|AG#cf9ROy-XwC6Z{U3CIMZqy0 zFy3$mJiy&GQYHV(N{?A&)3D04V9c@1p+!Cs<9&6BJ@M|WRK+VxwIWl3Gt1t*(a>FRZpKKjXtp{fv%x>U zM5G-U%hhbjWG}F6XuyOiABOgpSZ z;l`K{yvcwVj99)%axnjJT%Snp(%bt;SE>sqXB)Kza^E)-oQdisXjQ|Yl$Il898z%h zHeu>U&1?2?4tvAJ1NE~FpQVB4Wl@qtqWhd4|4n-izdj%la0c^74O9`V59_&qB?{(; z>9>dTH0Oq&<(2&S%RN|$@>OVdxvHb2LfbX^@3zILj{*3rN3@O=3IjeBPm*_UoT{}{ zU~!A7_*I$K=q5ljN#)vCPJbiLR!yvkNSn@3qs4}>AKx&xQ;7|OPM73&V`ui>H&CQ^ zPABSFlXVdk(YdXG%AxrE=*l}Li|ifoEq#Qk7=IX#Pxm=^yF!`l^z)`ssjoUuOa|Us*9Ep>%39ncZJ1 zcg*l0V``u<^U%X$lmQ5Q*+N2cx3FtcPmU+0q3!#4S?XOnBOq?!j_JeAj(}nB9vA2+ zh89p|b3t2TD%hB$r;tLk!Bpy%dF7w$!M}3fq@SPwZ2ZZ%J-9lFj7o*VZxfKIkYL@y zph$KHGmp0{iF@7MA^yih@vW%!f1r&st(}_q6xC_{o0rV4q1%1LJ#?8?`HXY#=2 zfF|~252mp+`LYAQ$!n?rC-S$11PMJYy|7iAl;{=Nk+C(chMywaW#b&fuPY)q|5Ys9 z_#_)rtzrbP-d{SVSMmir5C(DU+vF;F!qvZKs^IZU*(crgea54-WR3a(ZEZ17uCcqU;~UZu_`WmCVxQ=Ue2E z9%%}r6wv-wqJ(5)LvDkCjO7>&=Y@q=nhG9b(Q&ylKbzr)W|Y-cV%rvRmzOR(o=tJ) z&jWW7%_w}c&mOuY5bDBCx6aJaDiZ#EJ+`ZrUBXuv7n!VcIAI zci@TD!^o23!gQEq7IERD4g@+_IaEJyHC1O&_8u&$C7@Qb)55IP)hOyt`aVn*u>Vfg zaQF4da#2PWV8m)Dc%(oWz*SYVNtM@ZVt4)HPblU2N~1Q13J3GEv0px2A-b%uyr>^C}W+k~|QeHf*3yD|B>rkG;wnsFv2 zso77Aa>Cy1{HY@{6F40#bTSuReMf5VKXfMGUtF|iiPDZ{T^xRv+}JSyJL89SZqS(i zcoxLk`-+YGd)oIWyNVi3IsM*N3$c7Ua{jIH_sn+jTKx9XPHu4dRYpEP{u{(Ed2msQK}gd$;}8`Mq&Yjv6Wf zS7usK_G~}5R3Pmo2XGThif#f|gm%6+tobqAdM{+?f>ggi4QF-57?*-Us^FS=@$fBM zG41IqY1_Zit4;-TmPDPeHH({Lulh!o&btR&KJ6$Y*7P;pNb^W49;Hviy9EuYJ0_%q zd-3%tdqrGTpoPwLtu`q#_k+^pyXsEL#wD!*N32{RU%>mf^jsbu(yKoj5{2mHS%omW z`g8Ri3xyAFO5A3Rs&NmWj%5?b>5Gqw@kPFZZGg4E2mg_q^ zjylQ}z0bcVAEmUV^Q&=Xz;hf6Z)-+sfRsK`utozg`0c=X;a|YwN>00et$mX zqCkjygogQY;5jcPHN|3Xfdrj8MhtRI(y=RH5|!GLcTGTK8>T713`0uCh=iYrg zlTdnvCHiHg^eV%oWsLR3QppN^9~ZyQD9!hN3#R36KV?OByOkmvyg2j`LX8wPS))T@ z7Q?jqp!}U$Uas<1hOt!7HLq_B?d#G`eL~Uw{H&8%#3FES{*ilvqclWn{pP*tx2UN1 zA+7+Z^FLmx64J}Z`hgaiwWv#y!gdjj{SLZ|=FH@uJEA;3TCCk#fDYTd^Pt~Bi_$A( zn3!mT2VF<1rxS?%3#<9W zsCJ{+cUx5UGe;)ollgVnDxoqR&VPwoWBe3tO+ zuQ4BcnDOLcqoFq^gh`ct)Ta;7^yhtO*9F=-7})wND+chp{?qg9L?C`_id+7COw=!z zb_VV(_xH2kXr<^HKS&N;dR=I|)gicIqP`#Cy5@CkAp#Ws%kfe30>;4yuE)j0Bv1aA z=>Knl|3v_oTgA-_(8|%ivuFYC>RtSib2T#?*MEe~*Vjl5$Giu>P_!mx6(eNNH=%R= zgM~dw=ilC-56WE+_y7Lz-M-Aq5eEN9Xbn>Y2U=UdVGGq8k*-H0ye$q(xZI?N_*0=C}dHjEvd2ikI{xIuj!i>UVq&Gk2KEdF}61?zRHWLr~HYNS3W|Dme~-t(*40dY^v>@V)k90p9418<{z+?g0`!^$(8vi zH=E>8eW2y>W-ag<==UEeToK(2Az~Gc#bJ`A)xijFw(e&@T$|!6;D;Tt4}Bp^u`g$K zdSd|TK%;$aj&u(j6NXIrJ2P~|d$$V8jRw(zQ)YzNAj(*`!?P)Bb`F2z0ysO8izjH< zA#kV9=#ZhV>(AGCJsrw@Li2~#DM%D56*Ghma+aeGFZ&($v%Z!*7S!`HVDW{d5%;Tv zDpvV_=v5#tIn*;ZehJQX4}N#@JA2!zn_ui%3M@cpbG~ee*p!M-e76>*X|}+-p*nqj zZnu?jhkk5@>S_AY2%hJvmHo@b(4@#6`AkzbV^gow30zJ$YdnTW96e0O^gaS9Ez|?c z<-^wOAXI?Biuxx2=q@i8I$EE>iqN^Ou( zsSuJba`tk3^ET)nh`_iFId3L{#*x~KALnnfpkf0{L@pqn!|onEy|NScN^fL)q)Lxd zk}j(&9&9O@>-p%=GKUr7G`I&^Jy>jdNyufTr4eO|V3#chjEkPP<9#yqG=!_YLwLS3 z$7wcchlH4zMuukf*`q!P;&9Q`nr7-XirPa=q?bi69@1#g(Vb`i$!@MKd@i*xjt|u9 zACRF?$rWI93(4o6Nz-#S(5LI{vuJv~mVtW@mc_XtmxrTzu;bf^L4x{Y`sn~`wf zCw(W$f^e9BFr!o`5!Wy$k{WdlrUo9EDDA7K9@-h$INxtE)cFCG(pdS2QF3|Xm2jNs z!zRPEDZ8oZH+9Wo7vmFb_rYFDX0V18&j92fT|fve3oxg1PvObQhQ!xsbx^-v)CFU8 zz-vKgt3a9rdVgm9!;2B~_3H7qMAC+>3Bp*Q0j25QY?Y5S`IBixJdd9H92{vxy# zYp7-6J;=r-k(3U8(n?TPaxZnw-49&ZN@~aqCIMH@-MPU#)WDUSPF~p|suw5_nlD8v zg>qpVAQS&{xh;elpgLww1NNI(!}GgYjf?CSHp@u-_vyt%S=TW35iEQ8rBe%k`QwEE zBNO&i|MU&}e-Fk4B>euu?reI#CIHz zhRzv*j7|i$8IWOKLFUbX5~*@Ls{B z$r$PKldDvsYN+sOF*U6#zp4+nUYNkjoj0@olel_LKV8d$UIG8|htm6WEe>llp-)tb zyQacD+A!^=u3G7r9+Rm{wH%gzhWXMm7)N~%5oq;4Ggl~O&?u!69h)z5T;J02XfY5o z0Lq_1&HoNBMLj*1i;O7hXdv>0>qGl zTlC(wTS!2fhmzKlknBLRux63M%9mZeB>0!7Q;VGKy82Axc#ofLgL8J{^_#4$`!PQM zJ%_3bR$bO3#Q^hw$H0SbI%=Ow!EW}~Um@gneJgbH_-!es!2OBCNF?31MJrJdg087- zqY&XBJ!{=v#qa$$>^QaHPegGYovBZztg)F?$(+N>8B|snHMo8V*vB#mLMY z%?!;VzLG)N2MN0ulKoph`hpIjiNFYIW}@}`r>UD_e6s52VMd%VBj;Cq{ei?kRa4`b z+!iYt6NSJvOqxe1Ve6fg|8vC83u*QgSwr6NgVW=fl;DZV?B>C^dVDS<6ES7f2(ESW(8{ajoIT|B1&Z!G z>F3)60S3YWqHSS|>GQ!sL+4hv_jham|2OybT5?HqkVdB)-lzY)H_2aJOXO&x52F!JOXqw&?d+h^_OOyOaVX7ec3| zdpJ56HnI=+)Rw7BC1=P*D5i5Qo8Rk(XlG^ZF?_GtUB-u|Bz5F-o(gAgJ;`!HSHaBM zn$WPT5<{sXJ;A^f%wE$e$x9-$X$TUzy7Scvwc9B^Z{Vn<`o@;GK4gw-^*OTw>k2^J z9%z6IX;T@A;rFFaV0RgBTBpJw!n_fwpH{(BB*Jy6SrC+9kHXF{O?VS5;`W2L z8$Y%Zq${IP7vdqhA(x(NJu5BtN=}+!W)?BNLsENNoCzJ4lbZM5Q5P77;82`Tm+B@zj! zaQ?Oa@Dmg4q)0iKO$OqXySNG)P+PquUxfJ9XsNcT1A8h%zNA~Aja#vy(jZU#3lPI` z)YFz7Ng6}8wgV)ol`qFY2}GW{lb;NtqrtTEbL5Bpzi9i;u%@=HUEFRP3aE%kR}qmK0qG?oQlwWS(o}j4 zMOuJFML|)DRFNjVDJTM=N2PbAN)6IN4-g=Pw7U{@?|r`S+;i`93%|@G4_R4RbIdW{ z@fNj78aGSf-G%-|LxJ_eTqHW>VBy4A5G8eHH{A(JV3S9Ig5mD=<4>sClv%|tst`$2 zTcy)o`5WJ8^Cv-YamVxNTavXmy$f(U$16og0K|%;*@Dwo< z-fzcIJ#hmg5jbYDhDBb2QC2I+B$Wly3%6iYR#?kp>cQPwsQWao_9VBpc{4$1%sEDa zlQ)=%HUqwpdrDIZn$LN8HBOyj+E%>8G|4W<#$Y^2{T@wzMQKZv+SYuCD1JIm;&00! zVAZwYXF8vW*wY!WOq6GrJSguQ90R(**Cndr0|KPLpXZ6C(+Z4oa!& zh@l7xY~hGG26cCVfQd;$i02AJf~Q3~9?jH}LfJ}vv?<<5RNmU^oXyuKqMe-uh~BHh zFx6b141OB_!vJz&vnS(AtN&;Fi{PzRbP#HA-yE`EpttnAseLa?Y0b9C6<7Oo_sfLn zL}*+$V#U2yr}0z3OUCy00C0|zta!w7$2HR%1>1kPTtSi)Ha_?%9%NVFGk?OO@l4xD z{U+?-u1`&?$=C+~luOSf&2e-FP-k!XmnomLWlR*1^kY6)^X6gEia`2=TpWf#$vBH8 zetXU$#yr9@xY;-D#p}|yU?$nm7SqKum;S|5AZ}D z7gHw202{Kk4i83(&*4F6JvEbW?;Keo+_L?Kwb1j%WRVeC{Mknz~%KEMzyTlXRL{YIE;K0I+5}E{17WJhf2EL7c1)?td-$@ zGY*$jVAXygT6?L~&6>8SNXcU)A)<44ZAuB(8@t@EUfVu=9h19ttnJRhWXA4>ZCr;e zip0{yCQ$O+J#o?J8{1Ar5`~hnL%uZqM!GT;k>3Uev+Tmfsy?3+G?=sQ)sh@^_tGp6 z?ekbCy?}-G`0KN8D^MsatTLYNO9u`PF1x$8KX@jvR!#zn! z2s60v(uLEPx0PF$THkn5{NkZ&JH7m{Jk?E%y~Yu5rE2JH$E!8TyRzOP(E&XKligij z{0qa9FLTIj5-0>NQH@fWlrv#lzV5E>k+#@r)|0rU(B^)NJiunNIH7lw(xn)BGNi0} z(MF-+a$)&8cm&n%O%dn?Jr@y}+a(>vqM&`3SC|syEv=iW$1$bA<`cMaXa@JOTH<8F z_{|Dk=_=zaN86}^hgy{7XS~dd(w78;ul3@{=+yyiD_kChcB2yf9{>+adL7w-htI5G zdtl>vGX8#-yBC+o8tF>(GSf-u6xKlEZgg@@eW|qK9*5BbpNR0@#E#sep^li;{vu8X z3Hf98ax2mZ0&63x!@hzxLxUW9+zz&*H+hyU8HJR?i+t({r`@?ajD)#(iZmNL#Wk09 z;&zKKA_5fL8MEgo9u(W5dFSfbkG_XH%k6j*ee_chyQ98+jS%lqN$$J|z1Se{xgtaB zy0&)2Me7deEZ*r2aSV?!=$KZWbj4gN?ZBQ5iGUc7GmBm?4jK zp(kzaOD~9mJ_*dvl~QBlj)iGs==YHqF5V# zi5c&%FQ$TF1`)RKvStk{fk6M2wn>ih(FT=WM6k@*#;d?b&{((N%3$UXWD>KxGPt51 z4T{4J>t-O@Q4=SGis$>!A()*46OkY#mTW7x8KK<2teFBTaDU%K19F>&o`z)R1L^^xA4Yi59r zSFD9qCCYPCW6sh&e0g)e2MUb*JkRBC9sXs7!T>qKl*xm~mh$-Vt#qprXv5b}Y}fvh z&H##={b9XypphYf>%V)Uq#_58$Fe83TNKK>h1doF5XeQ4=~2AZf!=R8|8JdGOK}iI za`uTwx*+i^)C;_I#yfqm=EvSz;x?`iwX>_E=uS*hQQ7=*-5|dgwa;SG8sYNmVA^|$ zc+ITU2V$wL)m>W&2xTHfI4ou^%r@Id)&&q=_jKS=*&v6H){(kzSJ$oV$Q=pXaAZzQ zQSnf!C(Z?@s9FbLZo@*Pg7;awkn6Tel_`BcnpSQ0cH8jjT#&^@sXT=DZ-vZfRkt8> zmzVF?PQ;#z7i2A>(Uu{)0_%KB(o%uVH%kf>PFeTAfE}W_i`Z8SD<=!yklPb3;D=n_ z-aG04Ro$qw-BMz*E-Q#LE_Ig8w$~`oM#sh|aq{@^Oi0AeT^mY7JhW#=4AP5Xx@{!C zv?;GH2v8rYa4IkNWrSCqrL8UMdTIfi)=&Ot~jV^s8J{8Pu4i9JP6 zmp#XJKHytaO7T^5k~mv?Gy9wIB!@}%YvzfjPFtFt$c{n}FMxV;sa>~Ss?;bKXQJHmYP0zyg|<-yWZR!{E)x2 zo_BPaysT`Lz{T&S*t%*+_y6F|fDw8cNNhsr)Ya%(CF))MTv#|~%N5a98m6)AEl(Fl z#JfiUl2UbGx}S!0ta*nB(%nA9cU-&&W!)~Z-J?^p;f-H$y{bE)_5Or1poLI< z@+FG(d6A`8MCq-gQCRO#L2FDrY0d-TF2D4$TkMiqR9#n%lFu!@$(DE>D^b(Y?t+zf z00AUZ!6EYiGQq*9LK>`c7eP5~>ChKb^XceqGy5@s2ZC{U;vlwo_p8?p%OQ`}v*>vr zquxIzifd0WW2?s&bEBi%Bbm_`$4a~w`YvCVX|MpkQBTDP*~eABtfHDVvS&HO4KpVq ze!&@atCMaP(xfuK;S_6tZStYW9st?B1dvwZMuZ}dZlRYOs)!eV_AOaUflUJo1;nZs zB3V*2b(G&u24|gvJeKVpqpc7)pYwnqYS>6_mY9YjCn>4ij6-E)Z2cK`RAzBZ?XNp4 z@9un<X59xdT-Ua42CX@WK z0L6?3Ke;QGW9N@2$n8!FFnhiAX`r{OU<}QOfoc-wJ*cL!Lr-rPYI{7_(2%a1ru%Zl zvO&?NY^?z-o6|=XDIA#_p4Yv8NSJ-lOH~2}x+vAHud&4xZ!S}?gd}8_WEK~87n;!Z zV=fMPb3PZbx{;KKBm)JePch9cINAT8fZ|&K4W!df1FW%j165cm1v?yPDgWMbsl4;< z0e|`YaBEpF#rR{VzZxkY>F~bM`AN|F@#{Fh-MThF^4%}2Cl~oQ&>_vp$@>x;_if5v z;92Ar-h@p^U*{)FcBq6kvg)4nt+sc{3M@jKx>1%^J}YIecMR=5@BbuY+*-K}h8P^R4s^Tx1Kc~Y?Bj#vr`*xO2g>gVT9 z_2tg!!aqk8=X^`J|GUhFo4>ocHP^|m?E5M*?5lu0uk2&kGe;LXG{F#cUt4XV{w;~y zZg&M%NcZVb{At{^ZahhzLy?joh^=S}aHr_M@5!TIBfVyJpVyjeK1`sV&_9d7gomI&Y$X`zo$X$%KaH(wm5nRGo11tcvZ5+RWf z5%+JV+j>OC0rIl;_0t{5fL{PcfEKS5B^&Ae&Qk?)@n%$YHL1ii3h-AXeHh<@De|C7TiK?SXTZlz z{@H3!6L>*k+EUls66g+y7Ze&46*efD?8-|^$~dc}DVhx3i{2wU$(9(lj5{Wdq( z{D{}_N9%>f4-&yDf+~v~N7RSVkOpmQ(!<2hagESh##B+ZGWE?q z!Bl@m0?}(4FN-I&jZycnC{XcfsK8Ae*0BGTvH#`71*Rtre1~efabsBy%5SWvv4a-7 zMCtearI)Kq|D1)6{TXEfge zYFi&oq5X8<)U{v~kIy?dBa; zh~syFx^;H%TV+89TgGlg^3<+4b*gDOH5N6#T3Q%{A+Q@J@Pj_CE=@}0WU-r$Zb+;F zl3Umw1yeTzrHx{S@u9d6JS>hz6#{o@2*rRSTy^bV0E*W%fTEr4Hjkig$*Db=5mM`D zC2fGhUTB#50FV zVaFiKSFnGk6&-QMoMz-V{@##1VE~WOMHAa_!Bm47B3Z5|O(iKK>q_oA??AaE;k zsBXp3{1t9dXni|cE_`z3EyCgG2svd- zVbh<}<9i|a8ux9n9Y2T{Zi0Q)Z^0mO3E=84DJ!1cv<>*ljoeUb`%v^O zNhjUrG(@;#gTOusMgyz`E@7Xlx(MpIZX1BRiJN9=!8rQZUJ$BwcS&~f%)Q>*!fg~9 zLmFmPcJnqv#ydGF_2Lo8xCTB!1OnF+?L#8>`=wIz+rcRIP~OOb29=gVbgK!<#$tS& z;LbbDb)jSm3+bu@lParjeYef?z`Axv!COvErOJf$$`HZ}8rkf&*={?8q2fV^d%L%) zMUlsXr_~FD)NXVracuhBHfy|cG5kZkIPT4l6M*`O$@*sOcP#{XJKT`^UL24=5yFeD z1eHk_WVN%Da@v$Dyv+nb)@^*D%!At$v7|btpqSk+`7`!ZO?dP{*r6Q&tLPC!PXJH_ zLI^^sH9JOqUz2uyq?7FfEFg7OIP)KP#m)>L<%e#rfj3nJa&uHk3^E!3zV;1%p|k!h zfj`-$&c44*5gyw-F^3o>n7trvEcqlqpZ%KM?p|!Gy&Pf%N*Y0lHOgYO!QP9zS+g=n z8hg`gN61_o)w$4KiRvHeoVOLrS4l|^uV(dM=&{ZG3AFi-%K58t*B^x^2d}#y zDmlEyVg22*nw1pR@XRwCLmg8_ZNJCtdYnI7rzK~S!B7dpIziqL&O0FSsf3v z5z4i8t^}V7?#aXyMW6@S|LI9;AODuy@2Q(l4F!<75aWiKQPJCZZ%DMRbn4)59<@gg zfM!f%e|gCT@c$b$qw3s+MhkQb!?}K z;qZt0f2=CZ4^@qyC4L@@uk-oI?&w5&5x@$!JYc9m7`_uD6I%DRyJBn4VfPVoVDCxa zFevO}6}J?KbTnFLn$*UdU$+I~K`eSs?*PN`&FQ&auNe{L5G}&oT&W@Dx5r!$dsAgn z?&%|o7BlnwuQ|+&ho9-u*n+z1*J*r?cx6FLnR^Z;0Z_{r{T3+JNG>7JCi4^D*8on# zIhD9QRl?W{q-7e^<8HGue|n1e{z-mcX?Sc}MsE@4l_1?Zmo<~C8fN+D)$EGz3*K#* zp8V(wn*3B*s5rC?I=%HuC4ksJEQ(YMgZ59wrN4ZC=DQt>AIq0C+7J%pj@_~Oym2CM zV=G=OgTwzBR;ve(tKg#{9S_v%lFng<(wzV_ASVkwd$9S6)uQig#aFcymx|Yne!w7R z{0yUrj^(o@Q|Hm6GcFapdc{n_#Z@u=Z#S2WluE(DSr;LdJ*_nTrCR$=fR^fT;oa7H+E7Dm4|G+3x|Hdei7Xykm#wKc+ zWy5^!tMBfPUmmU8czk`&0+@aEWv!wf&BU>z3~f9NCTixr_BS|p6(qHd8=-rL;dJ?; zA0Vexl|H*?*`AgwVruDC{3pD|7)kR!idhQyKpSYBvRa?l(swy+cC>N3(|&2WxLcU6 zj5~x+U9ppvk2aYsyH!vpo!x}nOk$^5BLT*9=51w`2ba*_AJr6{4}8e}s))|6=qLd^ zdA2SuR9z5vMcT$dS?-&k!hODER5(7=npdKtgN2MfI%sU}QQ_1wxYoTp9Klu7?gL^D z$B(aWeQ5>E{*uQpZMc^ac1Qe@!XBQ0``GgZz3EufBj5nhTMpU!)YS=WKfJ90F`Z z+@K+@a!s}0oXU%amzPwr@B2_WCm^#2XJZ7B!#_Rk!g3{te^SH?$AYGdqLkbwoK4P% z^=f+BzXv?+a?NR;cFDs@>74e}bUlmNvAH1uYnc}}o>ta}CwCS1P|qh@mmA0v#!{&H z?+7X20ZAlSeM~F{E0)vx#%8#hw_;jkuPD}v1#jpcXZgCE^&jYgVEdcn3X5Bk7x}0T z)4RQVTOTYsX`6gfr7N*?de^CTMhb1SD50HJM{LN*{P_e=Rkj5j0%k|wOv0aWT?0NL z@B_Glh8b@BMTt24Hzk5PFN-Sto9Tceb)w$>{SrZnrh1W2(ErAiOdYinN|gEgC2X@=J@8lZ+|&~sy&8c zqx(IrHY?q1tmOhN_mTubIa^yFr>4pxDp>6rs^PwNa5sODY*u5)@9MM~63y8LxhOlLxfpa@ALpm7m4Dejlv1$Z{s6eQSCyP&=ph zMQu}uJP|325C#1j^WQRH^GRIAitu^4{PT4DRfNV_^(?}J331>;aD*TypieF_x9Da( zkQ`Akf}UWUpafrpvy02-MM`8xxwrJYsBt#y%M*K&U;_^ z(N2bf<}VF~&a_%pOwiOyZm+=$qzYcAYRkR?q=%N30@Z@U{|D-*$3TTf`7HHR$!sTv z5ULk?eUQ6WN^1Kkyyq#O*f$1PTbax}zl0oS z`UX$o@@v#R2e6SZnJw^&(G&|0;_9HtvAt{=z}0)vBviy}!j1>LZAC%xXOBookHb8CwIe;vu#rM|qV4 zUMfr!jEUTK@J>RsoOHv3s&nG1z%gxXMmF!~DUEjXA7s}JQ!jb~C3^NGGCnQTvfVmr zD4|Zy^3CTrbQ>n-w{685hTq$?WcSkfsTRUj$rjna4on_*Ir*52?Tf7&t}j3Gea@QEyn(v&;TC72w&Q@5|s7H?{IV%S&a=li-}XzVSOI~xtYdI=K^ z<>FDzO+GUxTvuO`lRv9nxy2jUle@<*>OB{jJbT%q@8ji+)6fF9E01QAlM6HpohZ>X zIyCi7laxTwy|gN@Zu;FN*{8g6h)OZ!-EIFD@XCq&vl9W5*kjY?w<)v>0ySKAj^6-> z4^a&!D3nckWy&2A*=1-L{U&ByrV5ys(k^bwlZFPv>PvW=Sc)xmIB%#u$T2b;#?5cN zfgQAa|M)LU(f1xsZH+r=GEcs3k{BUFV@hJmOhw^0!r0dgElNQXoSmQYt^IY64#xbj zZ+!LGX)Mq##6zI#S%=qNT?x;-l#GUB#S=Z^$1zrbC;f#{{cpkf?d!c{*ekj#%37jv z5?%2mHTO6xdzECHou9c(0-TKfYBzgU)pq6HyDZOV(m zSvW}d_gjZX;}j z+4ymhLkvMa;>HzK6y`LQts}rI)ySO*#orVH4Z_Kc-Q= z0V7+^>ty-1Omqeb1+7bOT`t77#m^LN9TnOW7%u zG)Wh38tiMn!B+zqFDywxMj<8YOHzC*3ASMw&;JA_oph3CIoM9GR{U4*5zaDV(le3!mVpZ9lkJ37ZmP))?u!P@ z_+av;0mYA($1edb^*98D4o`%@tGcjpvpzP2IbiadI)j}tf0SBlfkM?HP`C!v-Rg^Q zwB#I#8`TTI9zoKr_{R^h1`j>E@PdqrJs|VmimhNNUBMvn6lMfwSJAHnC{cnYv~^llXR7 z-HUf)-{1+1?bky+_D&MKIYFU0gT=UJ!L0P^7Qt)We1|4m&mb3u379LEf|Emmb{$Dz z9gbGOA2b59IdQx0Vzuz36<@7XtjJneFQyBaN8G*RFBBNE(Qq=xW-CH{R;a$t3n~8a zV@}40W9eJ@jdGEa-TBEn7@Fy}%dIt1h;xjA6x)dr9QkBmD}EY-866^b{^X@cwdOx~ zvGmaKNBT@@DktL23PFMwV6It~j3Z}e#9POGwl^?X8R}3w_Cwcmod7h!XcliV{-m>T z*KV}q83vmEoJwwM2e(>03m}Y5!6x@|dx@KXIUpXIA;R$^c=FAhn|#9|041Bi6Vvcy zrsdzq z{=c|W$yUkkaz-@LAsLg2TH=NHf!A8k# z1ZDT#izL&75Q>Y$#Wj9wr4H?W=J|dNm z2uC%NZ8Ab*{UkH;(@}&UN(iuDGlm>;v{Uj!dQ*OsFC6eAgT6$SdV9153ZevDEb6){ zQC`H**&r#tBFijvml!IexdG)Zpg@AJkjWFVaO;3UG?-_H7af8|ok!acCSuWV1JF5d z{h?5)DIzWwogsHh9<}8!lHxB<8NYR`mpXuTzUXEF7EHwu6W9^eo{7AfmgAi7{)F$r z+vua`wd(StbiNvL9`>5^HizN>sMEuINlsg@`l+ z8e{ddrzj9%91}gLavb=xKh;v3$V9_2^R3<_()P#c&$FJ6G9b zVHGkDs*BQ-!it>Y{+@l6X4cp`{BsWNmmZ{<+dIo~Sl-{%iC}mCB;;+xbZ|#u{dU&Z zfVS4EUJON){8oJuFM{1wX(VurArVRO{v?A)&{-g&C4|jm>RUU zz7*NtEb6#KGH+9_BbcU%WIVdOWQl1d-3{s84d4{$(N{)1ubM?i)8zz%{xTmJTK!$y zQ6~gvLHbZ#t=wiso2l$D6 z1S*O|sekXtE!}e-*$aj@s66F{eGH`><8_lr@2k6iM?J>id^AqJllk~(l>S4d?_6mY zyyP-fGf{@Fq{g#|>ldsjYJGqUNOst`KqbOq#^b!rJEPspu%hGqc{O6wUs0hq<9aii zN~1}i3+xbc@iTqsPvBTfh$1fOHCI;Ad3fHTCVq$~{J1;eq0-b)W8DnGh4i^65H5Of z<3Q?Yqowi-Zk|TdP{J*U7kmXrLVznrLMrQ`NNMM(c;W*rXjuHv<{S_p!f~T?9Z%>+ zP}j-T<VP56~EcQi2>zzerPG>ZJd2DYB;$82KrBnUY6%C`XmRZU_24&e2Q=sE;X1kU$upf9!Eiq=}f*D zqc^k32M28p^5G92Mt@oTC8Nzceb%@vV1Gn2Qc>x?^RT87h?A? zoF@9HXu=UYM>C)<0448Ttb9~(#+O?k&TKhd-`@5Ytw0c#`Bn(i^(ncw`vP|Of zWGb&QL^S6;Yom$JvS|%E7OfW7zo0c-rfV05j@!@5#8ID$t06Uho!1LjW7JX0(@Wk* zx+TL>+qTDKp4y>+a3Oki%7H^fB$|A6OQxG3SVF6^1WBBoLDn>QX^|CV0}Bxf-CED0 zSzh)aTz2PES36mB!tqAxfTDbgwIWcVg2hPwl-O_uY4eCc(wIzFUh1K4K!ttp+RIG? zmgt6%-5IMzY=zfacofe?Q(h10;1Fpe>G##fZ|-d!AB|6C=Xst=yXO=Ba3Jj{Tx&Js zCeT=ct;w(|!%8;;FeSeQYqUj*u?<-|p0)XAi-)Qz5JfZ+C?bf`I8!9Z z{C<_lHG(JBV~HZt_w$n}aChQZz&A98oj`dv$Aex)gR>kSVxeExp9W{V01{GexaFYu z8o(eTd4z};Gt;~^)@HLQd!?GHK8&8TM3(|X*{E`b*7XrUDeZV{({TPF%YGZPfMue; z-4xi*g72!v)c~^}k#Gch6A!h$%j4St3?Q%$vp~B zpu0pt<>TAN&z8u~IQ=-Szl%}GokY0R1|B;RP2G(yrF`cj z+WwBIR7mx9<)#^gt$WWlY>MHLr2-jh3=d8PBN%>z)nh)|YKMFOgHgXRpo{M>6#xPb zF3#Syw4^oi3<253OqqZ$#Efc^8=?fG1a9?d@5=!F=v$;KBR?(s&jY$}Z_DV5anBgj zF-?)o{YNmhu&}u$`>b?9jz41amjOk=_Fntk+>?88&ora9ZmDmK7nqmAZ|`O5QV1Vc z)L!2xiO>wCC%sL%0EUkw1BF{_5nRVpa6XS_0!2Lh-(r$!cs84?`oI<0n>WYTZ?!^4 zq5$rL$DN${cQY47dgr#L-Cq^X8R*#z;kjMYY^)Jq2nHa>lB#p%(_vqG?*DBfc3hl~ z25|U?p|+AJv5aHWJ;VHVwEQYg(XH}ZFPu3|7p(m|Qaykyrw6AuWKM6%Nfp5IN|KvD z<-h66Ygjw6M4OxL=ixFpMBdEiv^Jy@9M_A;eIPRM=k5$t#Y_0QJI{JLRTKc;iT-8C z78*#w(1RXqUg!YVqKd|s=ou8O~l6E(_R zjN$0Lyx0I_VtkYK(cX)~ymVa@TJ(dRzkQH|I`vbu%iy^>S$fku$m$a3ojYHco*LKl z!6(m%-a$T`;CH!rkZWQ5+W@cO0P}67j(3iCYsg`$6sM9(Ru-pEyyOSMWKS~uNgLF6 zW_$3RSHGZ(-+h((i^|(W$bz(V_*K)@$Iiu9(NXVZNTLz4_wqK;Y;8@tYC(ABph3z{fv5#UDIapH>Z~kOCb& ztx8zDzn{ia^5IMQa*@le38U+bsO4C8pj26E`|T9)<^py)xsuy6_S`(3y3fb2mp;TO z6!L}_Xllbwh%N>iE(=zEa3r56%;W}eJ)0EnNdYvz{XiNu0DT~{+W7{Yk9_zCim5NNkbZa z>k}SM#ZQ+(%I_w4zY1dA2!>7K{M?Zufu|uEzqN*HL&ix!wJF3 zy==II+@NDsMjxtEmri*1;qH4@RkRI`ZY%XqpSE)_Urc!ASyRHcqq{C?#G(#RqX==% z!`icuA&k41m4zOAhmHQFKK}BG;N0>1*1DPD(5O!sQHvN zz+@ht)UtP;O7jkm05XQ+S(Vj)dbmfXKmX~nzN0I;g*q^dP6N~GO^s2`RC%p*ifMlH zK0@uQIo&z?jiwBYNd;d)deCbDje9B@B(BN>8iSU9UlL)fUx@W z))e8Hesu-xT z;A%sBogO?fyKh@}BF+i0>0EmyvJH5Z3n&-U%!jkst)p(+v7C0+Z%?PIGRrZ3oU&6XLdctWoK=5P!}71MxF; zr8j0uYI}K0X$hiCI{OmtuTJ3@A;{7P%7Atpq2(AOZd97GzF20 zQr)SdNlOn6LZ76cRw{MxDi~ypE>*s;7c+iKzs~HP;rdxGF-GCR?%+-N*+afA=@I>-ZZZ=DrxGXNHcK=(q$U`cx$N2 zsI*1C&!9PI(bv}DyU7*#Jk>Hm>X+EXrk%KYVf7s>16_r9%isxc@|&kTIBhiZl=@C) zNcmpi#?0c27nQTc1}$zKdhL3^;P05VV2CB{GN$H(aVhBdZzoImZ>_mZ-&VxPm()%( zt||1C-z-^vn=cQ@P^W8xT&%g1KB~x`zgc3RHzPP1!F6(@Zl)j8Tjr6PSXkK_$);SN z-piE9Ci~dW!x7Tfi>j7nCElxMwT3H;0m)@-EM*g7i)|lf$xOIUMC>|6Fsbu7=FV#+ zw0D(deFEmmT)ePP#p;%X+V6@|)kBkr)=5=aQWM{ct*stQyrW<$lvJ3p@a`3TG~5i+ z_Z~(9e&7X)`s1AyF=N9kTdP*Jv$=Fcz&VZd=x3JB=p&>M65)&730S##mP+?ZrbU!d z$Ga3`bHGnbTz?f7dt%5W<;8saA8pV4>yEVg78priIvl(iBxa_+*eLbSDMQNrZKiY8 zRxTW#39yYOw(*xRzY9E!GG`b|5L;qq`SVy^KMeXj5}*M@N<6tkOzuh zQ`ty8Yo}j?A)!PS6_H2gubLu4{Sf-P%_hmilPpaW_6O(dIXb z6P%ZM==NQQ0+)4y(-b!LOV>5P^8~F%R3ynjz>c<*1Y`P*n;)+UC|CHcNL`aBBw4?v zd%03r`@O5&WRVR)ROYjh*C>t_`MT%*?W&m=Ne|PR=bCjwOUzjoR5^9tac~T}37Ol) zCth%O1g?dl$_hKnky#4#2Cz-GxbSV__Wi~?zv$%2rcpph@N(`?Ieg8`?#?B^&#+%S ze384Z;RbW-FzW_CN;9-+)j=>9#1yBW@?}Awn=d$vrhSj)G20DHKGMstPn?>}AroYQ znBeWp9`bxIdzp}aX-)&T(JDv;H6m{8`!SS8-Lk7BNg8U{XYD--msza~aIfbx4)cR? zdGq={g;}O++Huh?*81hYtt-R*#~&e|aGAy}#k&;V?3goOdtB77yPv+mOZouz3{vXW zAC|OXfrjnOlOlF|zyVO>%5{EPTG-S)ASnVl++pYT!dlMe)>R46{bA4ipPH`!$Nnq? zSb~-zKpTdM2AtPr5*i6dd|(B>&MyL*_<1*AYk(~fum?SGE&mmb=F$(b3CX39cQ>;4 zWQD%{s-9K6`Bxr$0(NAypMS(U|5Yphwt)XnSM#5GS3AyY7tKpFoammS4NR#oWaBZO!(XB;wwmQ@K!chigX=rl+)vL%iG)Qt=-VUdyrXp6ZAOvS**N~Qi- zc-r!^=i(NKw~R3Gc|)nEG#&aH2IBz{=39(luujUk-`QPfl{AlMfDw}%mPyMq+?CDm zfV43{y7QFgsty|pxYMs@n12zeRMl4u+n1%!;Bnv(rf3 zrR#KTa~*{)$}R<`xOohW(kL&g_yzj&^g!g4?Bmnke5P~O2Jg%VB@ZNzWtKqx;Chs6#8JD>E)Tc#Bi z+|i>gmwbavw{)NWb;$f?%U@)y?8VC8OSi6m;rNQ?Tn($-!qakNO1n&(RjM~w(*(&$ zFpV}q`(JRQaaUd<&_jf9Y$*SM4z&?R1@^pv)3ne3Y4FAbr|eo~TmHfAZFe&!Ta9;3 zt47Z?0oQw*=54NnglzeQCGvn_S2X-;0=x>2x;Bffc{Wy@cqK)}=wmi?Z1eTjfCXMD zH670EvBFecR9e}cd16Z+(Q;`~bkVUkzDJPr!M!g`xlg8e>tKT!ZcIK(iE`DdPEL62K2^sC&bJ5j@~$Kk#v}&)0FWeZeq#g(k~6YFYsA5M;#hxRQS{fk z>mWXy=@0lSFqW|_Z{FOGdR#}6w+?MfB2RBRZ(l07Z#-wKz*BUv{&Kn+$jPpQ{_&?2 zfK+K9Ki2lTpTDlr1YGI;;<*BoIFVW-ke2C2#kg+o!GA7n5XjU_88`@?X!gl=sNZv5 zWdAwwx-Cxgc{p@oT~v3%_5R4gIhp-qd26?ldhe-5e_XfiFe$UYVK>gxr%MM|0(hpw zVghZ>c?c{NR&3la5w1PzaY`UavSosEkoTnW(Z81USLK5ztE_U5xjhr1sSt&;z2yXDl9Mqv;dUtckGxS|`6r6~aG)`t{4nq(q5wvDYk=-=(|kAGd!FyAyukV~dV zM)6?H{aaj6GRa7R``gE+$g?!Y8}38j$5U}!kqz%W#L8F2nf8aSefxC< z6|Z5u{p*1qQoi_ii%>u9oekIv8$!e_CW}Fy3_M z=udVx$HfWAAFtEJ$Yh#*kBF4n2Llh1SwNxVV;W5NTJqInxV~=sn6^sjw$M}`Ku!kJ zC}niZXHQ>aWJ_eN*k++>NQi|Q5J4p36*MsY4S>ENUUc{f&}4$8PvNrp40HKP(4}o> z?IxO5&2(PtWfY8jL%Ij*!Is_y@#-6Z56SL?Sq?(R=D!|><%4bGk29@b{&6sM)*@tw zo5#YmvU4eTXTG;CU%UGHoRdm8y#7>+^K;ecM=pX<*fT1=8!wL2z7ESwe>7=c7V?iJ z-FeJ8WM5z0w@ls!%e9d>#0v{b`Ov@%6`hwd>6ar ztCD6t??$)IA2q`tJ;dF%^VrR2af`lFPI*s~7osT(ONS8ahqmc1ISikCfz!4YmF^

UQ4}l8| z(TZ;+W&^|_`JjBKJ;1mXe-5{K6G+TpYOl1Xx|L( zH`w8&f!Pj^fukqQg2iO}KYm^gQnCUi6ut{O7Tei*)Xa<}VoU4#Z+ut%7CoFVv6D92 z-gjj741MD!|I|MUuPpo2A$rND zFF*m3FmXUOB37AW#v{^&@8`>Cq~+V)mvQP_zVnqs z(cb!y$Mq_n(BKrjS+B6iUq!&>R}m-|pw5y1WRb!eFN5bBuM_DG;rq_;2!C%_9IX+^b00BE4|EDPLz+h z2?d-xgZo(BKC44ib!go&1Muh2>s)ZPU1y22;DX|ppR$VH%>@j|XcUfRXE|Qe?L1d% zbu**?SOhPZQNq%*lrwIb*Aw#89F!VqtmOWWxhQmiTvmX6p72h4WJVKhPZtW6SiroU zti9s}u_HhyTA6gHym<+F0Bu8ON22E={^W zd#=ptCWe*is>F|jdGW=2EAn^k9+E)+eRq$9Z)qY-TYBM%Qvs6^V^=I@dN_1jExcx& zT`gb6(02E~K(eb>+gND+gNmWcHNtMsb+!*%#U9H59@PTzdU+#I(}5?-yeS~l2@tpE z8!dyL-+e)6prGTKNEj#@ENJ=6B0PX+c0DGa-Z{A1DfF7+ zJun^e-ueAksmF81;bC@TWeAxO13hn(XOaJ?5kM74LdUOMC-@pOJ)2(Sv;9Ynh#Ee- zDS%Ez(dvibM=Fpw6PV;#7)*au>dNsp-CsUbRG&8ZZ*4&*tu64X9M_HBR|---w{4s} zS|fX|I{5q8)XTU9iDb`z+ZDvvex9>uX@M=>%3G6GK!keJO&ZBIq6=Q9;V+J@h|3NR z(cs*|M0%P21l;aJ@LrYz7C`k?mGDMIFf!jqa!uE)-C*bVi(2U-Lorpq>unqg=<3NsBthq3TyMoL*HkiFX59 zuS^~q`;q}V2t>0qKOnu)Pj_t&<-ZlB>24ce6Qb=A$C@`6Gb~=*>wIug{fjofN$*&t zAw4AuxTGH{e{}`Lfs@uOt+=*|tSvo#lzwg~*X`d8zfh?90_M|f_1ky7x03NyV)VoK zWH}yzp_2W#Gy}!+)m5YmN)MfPIwR#G;E{1#X%7gw6D~ntZsY3)s2TYlRM^)3qOHbl zA{$4}+mtG@%Z-nVtkQ~iNlmVOCeR5fFW8=Z)d$(xya%H);Tm5R$sE~_RZBHhE3*5O zPVDj)_^uxWNxeqG(FeL{12Dna^Zgd$S_&i}VyW-?6ZMdv4XdNkYmP(5)!Vu$4| zrNw5K3$|I7)A|^eeC{3SDojCPKoR!lQQbBu?dREW*5a33y<)Mr^vj7dFq~9C(SuPz zA;YK~uyx{#oSVPPLo-kw&eO|7@Q=lk^w^x=O-2&VV;99&@|s$omrnm47^;{Bw-1xE z0jBxigBm!A`bmxE!wLYU7L$lpD1rHJ);Ws0USqxcOZ5(D){E#Xh++Jxu38(tU-0-} z_6rSLY6?8Vef=-*E!spl-?m#MqZ>kcGK5n+(c_W!Wf)CAD*eq%)c!Az2YNOV^3hC; zyc*@n+8h7cGQ>hR?NL?1x0M-%(@Y7k^uvo(K`$_0k8&n&(g(v;%*$ah|L=N-Q+KL) zB!fN9hV)4>J1kjecgkEs(s6*A;L~*DR+IETG4RYQm=$b^#r&VG5JAlMTSmSDaViYp`dfNLh2$wxg zskN=}{u{RXumSy^e%DV-z54bVjZF;zEy}3zDD>Toh8w{5hW?N0t~?&fg>83IIx3VD zvV~CAia{7cD9em(vWzStN!ILp3CWf{yRz?T?7Iqs>@;B*OJj`5mMtTE&kUXS^gDg$ zJ>Tzpzdt^Ijo-|e8S^~XbKlo}-PeUsC08vIL!L9ubaA50G?@swXrm+k-RO%eclJE1 zP3>Tf#WL(W0=@EMGP(r-HLwm?v;Tq`_*E^Tfj$ACghp)F%d(>T36DR_YoH3Qw;tKP zOec+}&C#yc)3k$GU+>(c81DIKMknJ-ZzoF=1j-#Wp(B?UytS!Ag&NPRuy#Kc6t0Wo z0cmF_zWUdkr4q;=G1$WoFrNp{`!bfVBNl5{yV7S%Pr~w2tAd-)>^XFZkPTSCaN<^2 z5Z5!^yZLw0(BBG{w0g+lOdyL(RNKif&~x!#lmzXjtXZq#W`e5CuYCr6bPH zPR_!o|58C@0xFh~CBN9P*=1MUu(MVI1j0E=N90!8 zAn_|vsi8D$;SGX<|E{$FkD5?j@p7rxc~!GIbG`lq13f@{drabmCjBUeQ=i)i?cQ*0z{ri&CtI(t)|CUfV(zBsnNG;N$)cf# zMJ5JMOJcS{s*&t*p_~^w0^6#W)k6g8cSkFGHnlCdnz>xf625Fkj+63B4)_&tyt7;Z zf&~`LTVh8~>-xU4YJ;mw^QH>XvGf4+pqn-Ns(j>X^~*O?&vx!6W3zE9qA2V@r~Rku z1{$epb>&!0>vuJz>h2T=oJJA|6TrU2C1s}?wW&5#JH(`pYx(?$4)}fF&cT%ug(b=a z**Fl`I0gWtV$NEr?0{4d+o#nNgL&lw+t34q6mb4HUx*qOL3jBWFgt+8fqFEu1|4Lb*K2|D}w%1(5J4>g!zFiQ0oE@8Y9FcTU-H+E^Sf~M?| z)_eP>mZE?Vb2ZMkP%*5`KsY_nDTmb@*vs&Vr?H+cN|RLy;_`^AI{De<_*Wa5TOM(D z95ebV+h>%D_Oh-cX!Z|V2Zeha8*;FPwZX){wiO3~UAg4iHmjUc+iMo@OsV-mFRlZN zaOmFlV#~e8BGfek%2f7*bRbL+3q7CwLlB6>{S62Nvp|(z!8aQ9ApRXk#LJB4m^h+Y zP#HYb)42DatBXF1=%pPL+NLFZ6(Q|bkY4(LF0ZpSx@wd&?d%|`5(WxoDyOSd7G(22 z%IAF=WyknAVn!!_LtEjpEb%uOagw-wG!iq%V_Ra-|A2VCE4;y`lUzxx*zZ>o`3?W7 zBp!`rB=qU!2xDEzWyH^&#XrFXFN`thEH&6&zu!zAoFsJHHI9S{NEt^3cY?8`H~IY| zc_|vbH@_-Ef1(?Hg$WNDG@Kjln;t!YBE{k%?Br();q>>jnhtbh`Sy1u4j3K#W3lZI zF5wq*gm%#bC_(fbs zWplG?oN4K*aEqwmo7EjVj0{e=Lwqr*wp8t36azokUopW!QRW|E!q!yv%W_9nA%%w` zLPr7X18MV4xg%{WRnn_ofO!nTtzBF^yng`RLHpZ}pUJD_Frw~>8bfj1YeCyFK!XL$ zgMiG&vR*_hMH1U%Nl-PmIo4?C*a_!@_8AwyoPO#-?dBapR+BmZ17Ntu)@?>|MrpEm z@#%vpKrw#qXL@)GNnROd!m+vmUhhE3AaxgwpC;EU_KLqC7|WGZ!SOg(YaP!Qgo_o7 z&u-2>A5b}F+eDhpyB2|94+Gyb+_AUV!P=lu5QZXTY!t7?F~dxGo(VkbNunCe_K zf5>)_L_qL#xL#fgJ_{_0u`a?%g@e6Q$R3~CBE4clMQSC2$La9Smu8az#<9IlBv+?6 z`m%Kkb$#uYLK?zALl^Y*?-)pRG4cCk%g}UBC-UW{>GU&Z`D4|z`Hm$y9}E9dIsm>y zk*ph00<-P8(%NRGh0Gem_Imy*)I5EFv6r~h8-hNtO$S(b!R`BCy>uiz^W!gDI0<;h z(9YdI|2p@+^;Nd=FwMF3+}VT^peqKog)%uz^kA`9d_$+i`?$3BdJfZmZ|-#ba6iUQ znnknjiuUQP&#llhlFrrHF=K0ztb>b+;Kq9XpeIE`(@sU>Dx!K!ypL9`$`*D)J(_+! zY$$$T`JaM}n8VZ_wAaJz*?9oJUC)HkPDhK{!oQt8+*A*Hr>JZ??ZNxb zQk)DM;kc7)4{Dso^gXUf4yvMGxxDyqx{wB~U7zpYGy$)X2Y?04(lANv*`IW3P2N!B z({$-}Ku&D{1&z6Ui`ZK)t-@k;f1}BZr}dRC4;|%i$5fDkZca~D6GfwF%d9o z@+wMseBji=yP2hr+0#+D0Ec#l^1!!0iV0ofmFsRgE}A!Tw1284LN@kG2~+$g8OTuF zLnq##WkATBTSZDB8X!NZ!Bx);xolVrF^2T5#%G*aRG7gQ0lJs7Q8t}B!{)M}i?`$1 zwc|HxK?Za9v5IUcz$94knu?@M;f>$kxtzV6u390Bdf3lWBKNt|$6*Kl=MH1M(Wn<0 zL&nf1XFc&hri`y@w)dxux-J<;&vU1&s}oYGjWE-Ak32P-F|+4}GG7Fjk)%2zl^H}B zf~xKh82cPI>7Q$Y&pl$HsO?lwDF+BMpPDQkKVa4Kc3}9z(w>=K27DeJycDiEhA!6vtFR@t=c$SzxdC#`F4ftoasVZAYWF*hk zVr6wEHF;THo*o)4si!SsfPqI}a2xv^E?`^&T18u)@LA7F7p6!2Y0 zJNs<2k2bfT7PY8re&UVo)OtR5S2lF;%j{@H9ZBd4rp?Jo^Jn%{D4r{@d3f(o{Xp8~ zk*RN^Q`bq;iC4c};6XRVIKtAvyrD6z@P9X_Xjm*c8y1*aC(yR?)MJ<ww}AY7yOxQR6$|e2C+a;Qk;iJr42rb#2G+iQ=S<0 zN=R<2>It-pvN^S{<2)dV)1aqb!@}LZ1 z#K4D~?>kzE&m9itYlut24^OOtQ;SI!tYAu#-aJ45h$2Z+U^`f-D=WCoS(q<}XuIBt z1vH^+eutj9sZ`WG4Xss3jC?y-uM}^iVJq_5W-i=$XD=)1QAL0@O_m~6l zw2~2NQ+$ZqS+QzzIPb%~7Z)(%<-Vlz+k)esM%sl-d*q@J>&I5OtZxjB@~VIX7jmAlj^C?rRRAlnH^Eu@3TWTu7yR?lggj7z zLH4U*$0Pim}m}t zbdB?AuI*|VH z=Ks$|k*^t7F8~+sRjEMen19;Kb{G8=sOW)~C(JLjisulxCy_H?SO+(s^2*g!Le?dS zRL|bwS6Aoy(>cCJGyXq4$G!16ojn~f#y@oh|9OGIyhELADO@VDpR2OU26Rky^N={8E|?{hzcC! z5U0oBReRZPEI*K)J?xk!mYGLet9T~UZK+6^KY@Bo_T#2f7N4t_na23ImjpTS4NxG=da1n9@Z-#hF8UhheoGTvba3jsbAmNb}O_-D6unZ6b~Fb zXhr*_-I|L{dGm8reE*tuH&(`!#7!=NK2UIMw4Ms1`+|NuVU-O20ujByXL*HkK5=kC zH6h?8^s#|DaiA1E8v5wsSzQZ{NBX)adv0c#iSU4;uUKJezXNszw%u7N66nv|vH6rT zwr7~Dv}nEaZY4u$?}f`#^u&or1@*6`tkRv7?KUnNI9GuW%|yh|k=h){J6_tjOR;?l z7tJ_%2ouPY(N}y?>4KG9(3;!PefVc3kKVd1R!;B?6CRcuz}G{T&umA;bxK`PJ$pz| zP&w+_V(jf02hcsW$6totp1aqnFo&e3BGD4b>F(!5FE7fui@Vsq^tm1DgmumxfSts3 zK+g0tNI9_ZAYkSNj{&=;F_CollDb_@hcAl5!s#6V(lj9bDJ2WGhLTbfD87av{Uf{5q*ru3=<$9YV*f)jk; zSQ)*ZJfrObtz*Kh)GNTN`VAmf@1xVb`GddpmzTg2bZc!@OO(~wPTj1xttci)MXk1J zSvkcA-tke>T7IMZ&WX9_xco<Y{Lp;>_d{v9i}>QI35LVP?HN z&^_rL+9B#@JzQgElO%npf3$XIC}wDxoLJ>=j*UR3GOtXc6h?e!cYw!z+Yme~Uq4~gN?dHX6>V@!GD@-t z+!*Wc%3OvZ_>88F9Ns!MePhmrs6#zcE-ouptGXpNPil`D`7tN2uz#uPW}Qaj&(C%dI_4r4&%K9v31p$c^~C(5>Jy&%0q)?TU0OSyf>ug3JbI!>tN0pV2B08bR=p zt!8cTZmC<8%fa@v0?d8P`n2&~_Lvlv#1?D1?@~8HfzR$G804jVL+xCq?}NOaSEix`J-9J_#x_K+ z;7s{o#;`qO&s46=!!AeE&X99S#gMZmOShp^T2V=@OG=)nM2LcvaSZiZ(u%%Xry50$ zzE0(etwe~k>_abcQRjP7PBm;TAvP5;UZee+FIFm@6F0-=M<>7OdvcBW7zI^RVhJe> z-uHaHZ83b?v2XZk^XXr}mU=2=L+c`~#_58gOAeggqZ9@X$_ z6B1S{_Zf5f`f6aB_i~U2WCb=lgdJhJSt}u{5HVIf{Gg>`oe3Tx3xJ(nA6o))8W~Bj zS@d9<`8JgAgxgS3-GJo7lzY661T8dn39Ugr=+Hb<=mck?byb&0SejLG8Di_s@l+LY z%LoiF@MSENLAzWTR0z1<+ua8FwhK6D%tKO?B9wt~Iq6bso}sSxf6 zFWF|T*%g!exM*}+RgvOt-nb-guG7c{_(R{P!^OSUGkLk%jM?;>yAT2qJaJbDE$~9u z+f>_ivTYEwszY|3v<&W10-QDrU$oA6^-24JXeR8wpEa@Wy+W{>iWz*%MeAe|!1>;b|c#sGCD4~al|Hipd`RI@w=aehv za?L1wUaGgd5vu!2VA`4JuLZoVWJFdo-G`;&F!u+VbMI68(wmO|Zh||inliM;;ectT z^KyU9+E%G)q{WR5*Gu)lJvyc&^I**qmh}ub3yJp7DBfL`W^ksab zD`t~)DRBCwg8tku0+_?oh^6=#?r`H8%9JTs(b_!fDAp$xrqBAply5*ZGTc4vlJToj z{qurgmx>mx?bh|}mEZVpEp8W&7aRK;u6V|eg>0Ho)#!^cHV@IvV+I?#P>b$_Wp(Z( zpR|A}iB24#Ve{*8MVe6YXS8}3*T^Ij2};Ua$^oU6@+&828#t*~r;6RUlx~W2&u;(& zSo`}lR8PV;?xVvy*lpHQBAvQs*oL$A=mEB^FC7{i|8OoKu!AQ4b9#|ErYI+L_AHy9Au@hz02=p z>OGnDAeqkp%iTL$>+?w|B*gVwxu&%D)J/dev/null || sysctl -n hw.ncpu) + 1 )) executorch_cifarnet.elf diff --git a/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh new file mode 100755 index 00000000000..9a5c29022f5 --- /dev/null +++ b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/prepare_model.sh @@ -0,0 +1,32 @@ +#!/bin/bash +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -ue + +pushd "$(dirname "$0")/../../../../.." + +./install_executorch.sh +./devtools/install_requirements.sh + +pip install -r backends/nxp/requirements-eiq.txt + +python3 -m examples.nxp.aot_neutron_compile -m cifar10 -d -q --use_channels_last_dim_order --remove-quant-io-ops +mv cifar10_nxp_delegate.pte model.pte + +popd + +cat > model_pte.h <<'EOF' +#ifdef __MCUXPRESSO +#define __PLACEMENT __attribute__((section(".data.$modeldata"))) +#else +#define __PLACEMENT __attribute__((section(".modeldata"))) +#endif + +static const uint8_t model_pte[] __ALIGNED(16) __PLACEMENT = { +EOF + + +xxd -i "$(dirname "$0")/../../../../../model.pte" | grep -v unsigned >> model_pte.h diff --git a/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh new file mode 100755 index 00000000000..f734d1b79a1 --- /dev/null +++ b/examples/nxp/mcuxpresso/imxrt700/executorch_cifarnet/test_build_from_scratch.sh @@ -0,0 +1,58 @@ +#!/bin/bash +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -ue +ARM_TOOLCHAIN_URL="${ARM_TOOLCHAIN_URL:-https://developer.arm.com/-/media/Files/downloads/gnu/15.2.rel1/binrel/arm-gnu-toolchain-15.2.rel1-x86_64-arm-none-eabi.tar.xz}" + +# Get arm gcc +echo Downloading ARM GCC toolchain +if [ ! -d arm-toolchain ]; then + mkdir -p arm-toolchain + pushd arm-toolchain + wget $ARM_TOOLCHAIN_URL + tar -xvf *.tar.xz + rm *.tar.xz + popd +fi +export ARMGCC_DIR=$(pwd)/$(find arm-toolchain -maxdepth 1 -type d | tail -1) + +# Prepare model +# Side effect: the neutron SDK is installed +echo Preparing model and installing neutron SDK +$(dirname $0)/prepare_model.sh +# Check the model exists +if [ ! -f model_pte.h ]; then + echo "Cannot create the model_pte.h!" + exit 1; +fi + +# Locate Neutron SDK +NEUTRON_LIB_DIR=$(python3 -c "import eiq_neutron_sdk; print(eiq_neutron_sdk.__path__[0])") +export NEUTRON_LIB_DIR=${NEUTRON_LIB_DIR}/target/imxrt700/rt700/cm33 + +# Get MCUX SDK +echo Downloading MCUXpresso SDK +if [ ! -d mcuxpresso-sdk ]; then + pip install west + west init -m https://github.com/nxp-mcuxpresso/mcuxsdk-manifests.git mcuxpresso-sdk + pushd mcuxpresso-sdk + west update_board --set board mimxrt700evk + popd +fi +export SdkRootDirPath=$(pwd)/mcuxpresso-sdk/mcuxsdk + +# Build now +echo Building the example +$(dirname $0)/build_example.sh + +# Test the result +if [ ! -f cmake-out/flash_release/executorch_cifarnet.elf ]; then + echo "Build not successful!" + exit 1; +else + echo "Build successful." + exit 0; +fi From df1402a8d15d97cae2638cac0baadfdf96d6f2f8 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 3 Sep 2026 21:31:08 -0700 Subject: [PATCH 041/190] [ET-VK][q8ta] Avoid dynamic im2col vector stores Pull Request resolved: https://github.com/pytorch/executorch/pull/22538 PowerVR can lose padding lanes when q8ta im2col writes a local vector through a runtime component index. Construct the vector from four statically addressed loads to avoid the compiler/driver failure. Authored with Codex. ghstack-source-id: 424669153 @exported-using-ghexport Differential Revision: [D118585511](https://our.internmc.facebook.com/intern/diff/D118585511/) --- .../runtime/graph/ops/glsl/q8ta_im2col.glsl | 100 ++++++++++++------ 1 file changed, 70 insertions(+), 30 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl index b0cc4866a03..b58035f59a5 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl @@ -39,6 +39,35 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; #include "dispatch.glslh" +int load_packed_input( + const int x, + const int y, + const int z4, + const int n, + const int input_W, + const int input_H, + const int input_Z4, + const int zp_packed) { + if (x < 0 || x >= input_W || y < 0 || y >= input_H || z4 < 0 || + z4 >= input_Z4) { + return zp_packed; + } + + const int x4 = div_4(x); + const int x_mod = mod_4(x); + if (get_outer_packed_dim_block_size(inp_layout) == 1) { + const int scalar_idx = n * int(inp.strides[0][3]) + + y * int(inp.strides[0][1]) + + x * int(inp.strides[0][0]) + z4 * int(inp.strides[0][2]); + return t_packed_int8_input[scalar_idx]; + } + + const int scalar_idx = mul_4( + n * int(inp.strides[0][3]) + y * int(inp.strides[0][1]) + + x4 * int(inp.strides[0][0]) + z4) + x_mod; + return t_packed_int8_input[scalar_idx]; +} + void main() { const int out_buf_idx = int(linear_idx_from_gid()); @@ -96,36 +125,47 @@ void main() { const int zp_packed = pack_into_int32(ivec4(zp)); const int z4 = div_4(input_z); - // Check if y and z are in bounds (constant for all 4 width elements) - const bool y_z_in_bounds = - (input_y >= 0 && input_y < input_H && z4 >= 0 && z4 < input_Z4); - - // Load 4 elements from input, one for each output width position. - // Each loaded int contains 4 packed int8 channel values. - ivec4 im2col_block; - for (int i = 0; i < 4; i++) { - const int x = input_x_base + i * conv2d_params.stride.x; - if (!y_z_in_bounds || x < 0 || x >= input_W) { - im2col_block[i] = zp_packed; - } else { - const int x4 = div_4(x); - const int x_mod = mod_4(x); - int scalar_idx; - if (get_outer_packed_dim_block_size(inp_layout) == 1) { - scalar_idx = n_idx * int(inp.strides[0][3]) - + input_y * int(inp.strides[0][1]) - + x * int(inp.strides[0][0]) - + z4 * int(inp.strides[0][2]); - } else { - scalar_idx = mul_4( - n_idx * int(inp.strides[0][3]) - + input_y * int(inp.strides[0][1]) - + x4 * int(inp.strides[0][0]) - + z4) + x_mod; - } - im2col_block[i] = t_packed_int8_input[scalar_idx]; - } - } + const int stride_x = conv2d_params.stride.x; + + // Keep lane loads and their complete bounds checks static; some mobile + // drivers lose lanes otherwise. + const ivec4 im2col_block = ivec4( + load_packed_input( + input_x_base, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed), + load_packed_input( + input_x_base + stride_x, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed), + load_packed_input( + input_x_base + 2 * stride_x, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed), + load_packed_input( + input_x_base + 3 * stride_x, + input_y, + z4, + n_idx, + input_W, + input_H, + input_Z4, + zp_packed)); // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1) const int buffer_idx = n_idx * int(im2col_outp.strides[0][3]) From 0c3f3a964991d9cfe4b98b2bf62c5922788b5e56 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 3 Sep 2026 21:31:08 -0700 Subject: [PATCH 042/190] [ET-VK][q8ta] Route pointwise convolution through unsigned dot Pull Request resolved: https://github.com/pytorch/executorch/pull/22539 Use unsigned accumulating-saturating packed dot when Vulkan accelerates unsigned packed dot but not signed packed dot. Keep weight prepack and execution routing consistent while preserving signed routing elsewhere. Authored with Codex. ghstack-source-id: 424669155 @exported-using-ghexport Differential Revision: [D118585646](https://our.internmc.facebook.com/intern/diff/D118585646/) --- .../vulkan/runtime/graph/ops/PrepackNode.h | 4 + .../ops/glsl/pack_q8_conv2d_weights.glsl | 7 +- .../ops/glsl/pack_q8_conv2d_weights.yaml | 3 + .../graph/ops/glsl/q8ta_conv2d_pw.glsl | 53 +++++++++ .../graph/ops/glsl/q8ta_conv2d_pw.yaml | 3 + .../runtime/graph/ops/impl/Q8taConv2d.h | 14 +++ .../graph/ops/impl/Q8taConv2dIm2Col.cpp | 1 + .../runtime/graph/ops/impl/Q8taConv2dPW.cpp | 52 +++++++-- backends/vulkan/runtime/vk_api/Adapter.h | 30 ++++- .../test/custom_ops/impl/TestQ8taConv2d.cpp | 88 +++++++++++++- .../test/custom_ops/test_q8ta_conv2d_pw.cpp | 107 +++++++++++++++--- 11 files changed, 338 insertions(+), 24 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/PrepackNode.h b/backends/vulkan/runtime/graph/ops/PrepackNode.h index 8a301ef1e0a..fee7cf6fa60 100644 --- a/backends/vulkan/runtime/graph/ops/PrepackNode.h +++ b/backends/vulkan/runtime/graph/ops/PrepackNode.h @@ -50,6 +50,10 @@ class PrepackNode final { node_id_ = node_id; } + inline const std::string& name() const { + return shader_.kernel_name; + } + protected: uint32_t node_id_; const vkapi::ShaderInfo shader_; diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl index 5682f044b1d..8e1a07024ed 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.glsl @@ -9,6 +9,7 @@ #version 450 core #define PRECISION ${PRECISION} +#define ADD_UNSIGNED_OFFSET ${ADD_UNSIGNED_OFFSET} ${define_active_storage_type(STORAGE)} @@ -69,7 +70,11 @@ void main() { weight_vals[col] = (t_int8_weight[word_idx >> 2][word_idx & 3] >> (byte_pos * 8)) & 0xFF; } } - packed_block[row] = pack_into_int32(weight_vals); + int packed = pack_into_int32(weight_vals); +#if ADD_UNSIGNED_OFFSET == 1 + packed = int(uint(packed) ^ 0x80808080u); +#endif + packed_block[row] = packed; } buf_idx += oc_stride; } diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml index 9331de6e758..190eff31f4f 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_conv2d_weights.yaml @@ -7,9 +7,12 @@ pack_q8_conv2d_weights: parameter_names_with_default_values: STORAGE: buffer + ADD_UNSIGNED_OFFSET: 0 generate_variant_forall: STORAGE: - VALUE: buffer - VALUE: texture2d shader_variants: - NAME: pack_q8_conv2d_weights + - NAME: pack_q8_conv2d_weights_unsigned + ADD_UNSIGNED_OFFSET: 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl index aeb98f7a41b..ca7cd978176 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl @@ -11,6 +11,7 @@ ${define_required_extensions("buffer", DTYPE)} #define USE_INT8_DOT_PRODUCT_EXT ${USE_INT8_DOT_PRODUCT_EXT} +#define USE_UNSIGNED_DOT_PRODUCT ${USE_UNSIGNED_DOT_PRODUCT} #extension GL_EXT_control_flow_attributes : require $if USE_INT8_DOT_PRODUCT_EXT == 1: @@ -126,10 +127,19 @@ void main() { const int inp_n_stride = int(inp.strides[0][3]); // Initialize int32 accumulator +#if USE_UNSIGNED_DOT_PRODUCT == 1 + uvec4 out_accum[TILE_M][TILE_N4]; + uvec4 input_sums = uvec4(0u); +#else ivec4 out_accum[TILE_M][TILE_N4]; +#endif [[unroll]] for (int m = 0; m < TILE_M; ++m) { [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { +#if USE_UNSIGNED_DOT_PRODUCT == 1 + out_accum[m][n4] = uvec4(0u); +#else out_accum[m][n4] = ivec4(0); +#endif } } @@ -149,6 +159,10 @@ void main() { // Load the packed int8 input tile for the current width and K sub-block. // Each int contains 4 packed int8s (one per width position in the tile) ivec4 int8_input_tile = t_packed_int8_input[input_idx]; +#if USE_UNSIGNED_DOT_PRODUCT == 1 + const uvec4 uint8_input_tile = + uvec4(int8_input_tile) ^ uvec4(0x80808080u); +#endif // Load the int8 weight tile for the current K and output-channel sub-block. ivec4 int8_weight_tile[TILE_N4]; @@ -158,17 +172,34 @@ void main() { ivec2(oc_block_idx + n4, k4), 0); } +#if USE_UNSIGNED_DOT_PRODUCT == 1 + uvec4 uint8_weight_tile[TILE_N4]; + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + uint8_weight_tile[n4] = uvec4(int8_weight_tile[n4]); + } +#endif // Accumulate using int8 dot product // Input tile indexed as input[m] where m is the width index within tile // Weight tile indexed as weight[n4][n4i] where n4i is the channel index within block [[unroll]] for (int m = 0; m < TILE_M; ++m) { +#if USE_UNSIGNED_DOT_PRODUCT == 1 + input_sums[m] = dotPacked4x8AccSatEXT( + uint8_input_tile[m], 0x01010101u, input_sums[m]); +#endif [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { [[unroll]] for (int n4i = 0; n4i < 4; ++n4i) { +#if USE_UNSIGNED_DOT_PRODUCT == 1 + out_accum[m][n4][n4i] = dotPacked4x8AccSatEXT( + uint8_input_tile[m], + uint8_weight_tile[n4][n4i], + out_accum[m][n4][n4i]); +#else out_accum[m][n4][n4i] = dotPacked4x8AccSat( int8_input_tile[m], int8_weight_tile[n4][n4i], out_accum[m][n4][n4i]); +#endif } } } @@ -188,6 +219,14 @@ void main() { weight_sums[n4] = ivec4(t_weight_sums[oc_block_idx + n4]); } +#if USE_UNSIGNED_DOT_PRODUCT == 1 + ivec4 unsigned_weight_correction[TILE_N4]; + [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { + unsigned_weight_correction[n4] = + (128 + input_zp) * weight_sums[n4]; + } +#endif + // Initialize int8 output tile ivec4 int8_out_tile[TILE_M4][TILE_N4]; [[unroll]] for (int m4 = 0; m4 < TILE_M4; ++m4) { @@ -197,7 +236,9 @@ void main() { } // Compute int8 output tile from int32 accumulator +#if USE_UNSIGNED_DOT_PRODUCT == 0 ivec4 input_zp_vec = ivec4(-input_zp); +#endif if (apply_bias > 0) { // Load bias tile @@ -211,8 +252,14 @@ void main() { [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { const int m = mul_4(m4) + m4i; // Compute floating point output values +#if USE_UNSIGNED_DOT_PRODUCT == 1 + ivec4 accum_adjusted = ivec4(out_accum[m][n4]) + - ivec4(int(input_sums[m]) * 128) + - unsigned_weight_correction[n4]; +#else ivec4 accum_adjusted = input_zp_vec * weight_sums[n4] + out_accum[m][n4]; +#endif vec4 float_out_texel = fma(vec4(accum_adjusted), vec4(weight_scales[n4]) * input_scale, @@ -236,8 +283,14 @@ void main() { [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { const int m = mul_4(m4) + m4i; // Compute floating point output values +#if USE_UNSIGNED_DOT_PRODUCT == 1 + ivec4 accum_adjusted = ivec4(out_accum[m][n4]) + - ivec4(int(input_sums[m]) * 128) + - unsigned_weight_correction[n4]; +#else ivec4 accum_adjusted = input_zp_vec * weight_sums[n4] + out_accum[m][n4]; +#endif vec4 float_out_texel = vec4(accum_adjusted) * vec4(weight_scales[n4] * input_scale); // Apply ReLU if enabled diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml index 46670b8d2aa..6e34bdaf5c0 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml @@ -8,10 +8,13 @@ q8ta_conv2d_pw: parameter_names_with_default_values: DTYPE: float USE_INT8_DOT_PRODUCT_EXT: 1 + USE_UNSIGNED_DOT_PRODUCT: 0 generate_variant_forall: DTYPE: - VALUE: float shader_variants: - NAME: q8ta_conv2d_pw + - NAME: q8ta_conv2d_pw_unsigned + USE_UNSIGNED_DOT_PRODUCT: 1 - NAME: q8ta_conv2d_pw_fallback USE_INT8_DOT_PRODUCT_EXT: 0 diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h index 5d16cb3b78c..34556348019 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h @@ -111,6 +111,7 @@ void add_q8ta_conv2d_node( void add_q8ta_conv2d_pw_node( ComputeGraph& graph, + const bool use_unsigned_dot, const ValueRef packed_int8_input, const ValueRef input_scale, const ValueRef input_zp, @@ -130,6 +131,19 @@ void add_q8ta_conv2d_pw_node( const ValueRef padding = kDummyValueRef, const ValueRef dilation = kDummyValueRef); +constexpr int64_t kMaxUnsignedDotAccumulatorBytes = 33025; + +bool can_use_unsigned_pw_dot( + const vkapi::Adapter& adapter, + int64_t k_per_group); + +void q8ta_conv2d_pw_impl( + ComputeGraph& graph, + bool use_unsigned_dot, + const std::vector& args); + +void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args); + std::vector calculate_q8ta_im2col_sizes( ComputeGraph* graph, const ValueRef& input, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp index e93723c5125..334782d7af5 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp @@ -318,6 +318,7 @@ void q8ta_conv2d_im2col( add_q8ta_conv2d_pw_node( graph, + /*use_unsigned_dot=*/false, packed_int8_im2col, input_scale, input_zp, diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp index ee234319e8c..8b8bc63fc37 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp @@ -118,7 +118,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( const QuantizationConfig& weight_quant_config, const ValueRef weight_data, const ValueRef input, - const ValueRef output) { + const ValueRef output, + const bool use_unsigned_dot) { VK_CHECK_COND(weight_quant_config.nbits == 8); VK_CHECK_COND(weight_quant_config.is_symmetric); @@ -149,7 +150,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( std::vector packed_weight_sizes{output_height, output_width}; utils::StorageType storage_type = utils::kTexture2D; - uint32_t max_extent = graph.context()->adapter_ptr()->max_texture2d_dim(); + const uint32_t max_extent = + graph.context()->adapter_ptr()->max_texture2d_dim(); if (output_width > max_extent * 4 || output_height > max_extent) { storage_type = utils::kBuffer; } @@ -166,7 +168,8 @@ ValueRef prepack_quantized_conv2d_pw_weight( 1u}, kTiledWorkGrid); - std::string kernel_name = "pack_q8_conv2d_weights"; + std::string kernel_name = use_unsigned_dot ? "pack_q8_conv2d_weights_unsigned" + : "pack_q8_conv2d_weights"; add_storage_type_suffix(kernel_name, storage_type); graph.prepack_nodes().emplace_back(new PrepackNode( @@ -257,6 +260,7 @@ void resize_q8ta_conv2d_pw_im2col_node( void add_q8ta_conv2d_pw_node( ComputeGraph& graph, + const bool use_unsigned_dot, const ValueRef packed_int8_input, const ValueRef input_scale, const ValueRef input_zp, @@ -311,8 +315,15 @@ void add_q8ta_conv2d_pw_node( const bool use_hw_dot = graph.context()->adapter_ptr()->supports_int8_dot_product(); - std::string kernel_name = - use_hw_dot ? "q8ta_conv2d_pw" : "q8ta_conv2d_pw_fallback"; + std::string kernel_name; + if (use_unsigned_dot) { + VK_CHECK_COND( + use_hw_dot, + "Unsigned q8ta pointwise convolution requires integer dot product"); + kernel_name = "q8ta_conv2d_pw_unsigned"; + } else { + kernel_name = use_hw_dot ? "q8ta_conv2d_pw" : "q8ta_conv2d_pw_fallback"; + } add_dtype_suffix(kernel_name, graph.dtype_of(packed_weight_scales)); vkapi::ParamsBindList param_buffers = { @@ -364,7 +375,18 @@ void add_q8ta_conv2d_pw_node( // High level operator impl // -void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { +bool can_use_unsigned_pw_dot( + const vkapi::Adapter& adapter, + const int64_t k_per_group) { + return adapter.accelerates_unsigned_packed4x8_dot() && + !adapter.accelerates_signed_packed4x8_dot() && + k_per_group <= kMaxUnsignedDotAccumulatorBytes; +} + +void q8ta_conv2d_pw_impl( + ComputeGraph& graph, + const bool use_unsigned_dot, + const std::vector& args) { int32_t idx = 0; const ValueRef packed_int8_input = args.at(idx++); const ValueRef input_scale = args.at(idx++); @@ -385,6 +407,12 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { const ValueRef activation_ref = args.at(idx++); const ValueRef packed_int8_output = args.at(idx++); + VK_CHECK_COND( + !use_unsigned_dot || + graph.size_at(-1, weight_data) <= + kMaxUnsignedDotAccumulatorBytes, + "Unsigned q8ta pointwise convolution exceeds the accumulator bound"); + uint32_t activation_type_val = static_cast( activation_type_from_string(graph.extract_string(activation_ref))); @@ -396,7 +424,8 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { weight_quant_config, weight_data, packed_int8_input, - packed_int8_output); + packed_int8_output, + use_unsigned_dot); ValueRef packed_weight_sums = prepack_standard( graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); @@ -422,6 +451,7 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { add_q8ta_conv2d_pw_node( graph, + use_unsigned_dot, packed_int8_input, input_scale, input_zp, @@ -436,6 +466,14 @@ void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { packed_int8_output); } +void q8ta_conv2d_pw(ComputeGraph& graph, const std::vector& args) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + const ValueRef weight_data = args.at(3); + const int64_t k_per_group = graph.size_at(-1, weight_data); + const bool use_unsigned_dot = can_use_unsigned_pw_dot(*adapter, k_per_group); + q8ta_conv2d_pw_impl(graph, use_unsigned_dot, args); +} + REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.q8ta_conv2d_pw.default, q8ta_conv2d_pw); } diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index 6eae09e8eb6..968ca798369 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -231,7 +231,7 @@ class Adapter final { #endif /* VK_KHR_shader_float16_int8 */ } - inline bool supports_int8_dot_product() { + inline bool supports_int8_dot_product() const { #ifdef ETVK_FORCE_NO_EXTENSIONS return false; #endif @@ -243,6 +243,34 @@ class Adapter final { #endif /* VK_KHR_shader_integer_dot_product */ } + inline bool accelerates_signed_packed4x8_dot() const { +#ifdef ETVK_FORCE_NO_EXTENSIONS + return false; +#endif +#ifdef VK_KHR_shader_integer_dot_product + return supports_int8_dot_product() && + physical_device_.shader_int_dot_product_properties + .integerDotProductAccumulatingSaturating4x8BitPackedSignedAccelerated == + VK_TRUE; +#else + return false; +#endif /* VK_KHR_shader_integer_dot_product */ + } + + inline bool accelerates_unsigned_packed4x8_dot() const { +#ifdef ETVK_FORCE_NO_EXTENSIONS + return false; +#endif +#ifdef VK_KHR_shader_integer_dot_product + return supports_int8_dot_product() && + physical_device_.shader_int_dot_product_properties + .integerDotProductAccumulatingSaturating4x8BitPackedUnsignedAccelerated == + VK_TRUE; +#else + return false; +#endif /* VK_KHR_shader_integer_dot_product */ + } + inline bool supports_nv_cooperative_matrix2() { #ifdef ETVK_FORCE_NO_EXTENSIONS return false; diff --git a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp index 679ac33d11b..e1957c4c590 100644 --- a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp @@ -10,10 +10,74 @@ #include #include +#include #include namespace vkcompute { +namespace { + +void assert_pw_kernel_selection( + ComputeGraph& graph, + const bool expect_unsigned, + const bool expect_buffer_weights) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + std::string expected_execute; + std::string expected_prepack; + if (expect_unsigned) { + expected_execute = "q8ta_conv2d_pw_unsigned_float"; + expected_prepack = expect_buffer_weights + ? "pack_q8_conv2d_weights_unsigned_buffer" + : "pack_q8_conv2d_weights_unsigned_texture2d"; + } else { + VK_CHECK_COND(!expect_buffer_weights); + expected_execute = adapter->supports_int8_dot_product() + ? "q8ta_conv2d_pw_float" + : "q8ta_conv2d_pw_fallback_float"; + expected_prepack = "pack_q8_conv2d_weights_texture2d"; + } + + int32_t execute_matches = 0; + std::string execute_names; + for (const auto& node : graph.execute_nodes()) { + const ExecuteNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + execute_names += node_name + " "; + if (node_name.find("q8ta_conv2d_pw") == 0) { + VK_CHECK_COND( + node_name == expected_execute, + "Expected ", + expected_execute, + " but selected execute kernel ", + node_name); + ++execute_matches; + } + } + VK_CHECK_COND(execute_matches > 0, "Execute kernels: ", execute_names); + + int32_t prepack_matches = 0; + std::string prepack_names; + for (const auto& node : graph.prepack_nodes()) { + const PrepackNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + prepack_names += node_name + " "; + if (node_name.find("pack_q8_conv2d_weights") == 0) { + VK_CHECK_COND( + node_name == expected_prepack, + "Expected ", + expected_prepack, + " but selected prepack kernel ", + node_name); + ++prepack_matches; + } + } + VK_CHECK_COND(prepack_matches > 0, "Prepack kernels: ", prepack_names); +} + +} // namespace + void test_q8ta_conv2d_dw( ComputeGraph& graph, const std::vector& args) { @@ -287,7 +351,29 @@ void test_q8ta_conv2d_pw( groups, activation, packed_int8_output}; - VK_GET_OP_FN("et_vk.q8ta_conv2d_pw.default")(graph, conv_args); + if (impl_selector == "pw_signed" || impl_selector == "pw_unsigned" || + impl_selector == "pw_auto") { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + bool expect_unsigned = impl_selector == "pw_unsigned"; + if (impl_selector == "pw_auto") { + VK_GET_OP_FN("et_vk.q8ta_conv2d_pw.default")(graph, conv_args); + expect_unsigned = can_use_unsigned_pw_dot( + *adapter, graph.size_at(-1, weight_data)); + } else { + q8ta_conv2d_pw_impl(graph, expect_unsigned, conv_args); + } + const int64_t packed_height = + utils::div_up_4(graph.size_at(-1, weight_data)); + const int64_t packed_width = + utils::div_up_4(graph.size_at(-2, weight_data)) * 4; + const int64_t max_texture_extent = adapter->max_texture2d_dim(); + const bool expect_buffer_weights = + packed_width > max_texture_extent * 4 || + packed_height > max_texture_extent; + assert_pw_kernel_selection(graph, expect_unsigned, expect_buffer_weights); + } else { + VK_GET_OP_FN("et_vk.q8ta_conv2d_pw.default")(graph, conv_args); + } } // Dequantize packed int8 output to floating point diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp index ee7d8c9e5bf..1003c0540ce 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_pw.cpp @@ -4,10 +4,12 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. +#include #include #include #include +#include #include #include @@ -23,13 +25,19 @@ using namespace vkcompute; static constexpr int64_t kRefDimSizeLimit = 100; +struct PointwiseTestOptions { + int32_t input_zero_point = 2; + bool has_bias = true; +}; + // Utility function to create a test case from a Conv2dConfig static TestCase create_test_case_from_config( const Conv2dConfig& config, vkapi::ScalarType input_dtype, utils::StorageType fp_storage_type, utils::GPUMemoryLayout int8_memory_layout, - const std::string& impl_selector = "") { + const std::string& impl_selector = "", + const PointwiseTestOptions& options = {}) { TestCase test_case; // Calculate output dimensions @@ -92,7 +100,7 @@ static TestCase create_test_case_from_config( float input_scale_val = 0.008123; ValueSpec input_scale(input_scale_val); - int32_t input_zero_point_val = 2; + const int32_t input_zero_point_val = options.input_zero_point; ValueSpec input_zero_point(input_zero_point_val); // Quantized weight tensor (int8) - [C_out, C_in_per_group * K_h * K_w] @@ -109,6 +117,15 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::RANDINT8); quantized_weight.set_constant(true); + std::vector& weight_data = quantized_weight.get_int8_data(); + for (int64_t out_channel = 0; out_channel < config.channels.out; + ++out_channel) { + const auto padding_begin = + weight_data.begin() + out_channel * in_features + in_channels_per_group; + const auto padding_end = + weight_data.begin() + (out_channel + 1) * in_features; + std::fill(padding_begin, padding_end, 0); + } if (debugging()) { print_valuespec_data(quantized_weight, "weight_tensor"); @@ -145,6 +162,7 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::ZEROS); bias.set_constant(true); + bias.set_none(!options.has_bias); // Output quantization parameters float output_scale_val = 0.05314; @@ -215,8 +233,12 @@ static TestCase create_test_case_from_config( return test_case; } -// Generate test cases for quantized pointwise conv2d operation -static std::vector generate_quantized_conv2d_pw_test_cases() { +// Generate test cases for quantized pointwise conv2d operation. When +// pw_selector is non-empty ("pw_signed", "pw_unsigned", "pw_auto"), every +// non-legacy case carries that selector so the test graph builder can force or +// verify the unsigned-dot routing instead of using the default automatic path. +static std::vector generate_quantized_conv2d_pw_test_cases( + const std::string& pw_selector = "") { std::vector test_cases; if (!vkcompute::api::context()->adapter_ptr()->supports_int8_dot_product()) { return test_cases; @@ -342,7 +364,11 @@ static std::vector generate_quantized_conv2d_pw_test_cases() { config.test_case_name = make_test_case_name( config, is_performance, fp_storage_type, utils::kBuffer); test_cases.push_back(create_test_case_from_config( - config, vkapi::kFloat, fp_storage_type, int8_memory_layout)); + config, + vkapi::kFloat, + fp_storage_type, + int8_memory_layout, + pw_selector)); // For 4W4C layout, also test the legacy implementation if (int8_memory_layout == utils::kPackedInt8_4W4C) { @@ -390,13 +416,40 @@ static std::vector generate_quantized_conv2d_pw_test_cases() { config.test_case_name = make_test_case_name( config, is_performance, utils::kTexture3D, utils::kBuffer); test_cases.push_back(create_test_case_from_config( - config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C1W)); + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C1W, + pw_selector)); if (config.batch == 2) { test_cases.push_back(create_test_case_from_config( - config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4W4C)); + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + pw_selector)); } } + Conv2dConfig edge_config{ + OutInChannels(13, 7), + InputSize2D(7, 5), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}; + edge_config.op_name = "conv2d_q8ta_q8csw_q8to"; + edge_config.test_case_name = make_test_case_name( + edge_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + edge_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + pw_selector, + {.input_zero_point = -128, .has_bias = false})); + return test_cases; } @@ -484,6 +537,7 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { auto& weight_data = weight_spec.get_int8_data(); auto& weight_scales_data = weight_scales_spec.get_float_data(); auto& bias_data = bias_spec.get_float_data(); + const bool has_bias = !bias_spec.is_none(); const float output_scale = output_scale_spec.get_float_value(); const int32_t output_zero_point = output_zeros_spec.get_int_value(); @@ -498,7 +552,7 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { auto& ref_data = output_spec.get_ref_float_data(); ref_data.resize(num_output_elements); - const int in_features = utils::align_up_4(C_in_per_group * K_h * K_w); + const int64_t in_features = utils::align_up_4(C_in_per_group * K_h * K_w); // Perform activation, weight, and output quantized conv2d operation for (int64_t n = 0; n < N; ++n) { @@ -578,7 +632,9 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { accum_adjusted * input_scale * weight_scales_data[out_c]; // Add bias and store result - float_result += bias_data[out_c]; + if (has_bias) { + float_result += bias_data[out_c]; + } // Quantize the output to int8 float quant_output_f = @@ -635,6 +691,30 @@ static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { } int main(int argc, char* argv[]) { + const vkapi::Adapter& adapter = *vkcompute::api::context()->adapter_ptr(); + const bool prefers_unsigned_dot = + adapter.accelerates_unsigned_packed4x8_dot() && + !adapter.accelerates_signed_packed4x8_dot(); + VK_CHECK_COND( + can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes) == + prefers_unsigned_dot); + VK_CHECK_COND( + !can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes + 1)); + + std::string pw_impl_selector; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--pw-path=signed") { + pw_impl_selector = "pw_signed"; + } else if (arg == "--pw-path=unsigned") { + pw_impl_selector = "pw_unsigned"; + } else if (arg == "--pw-path=auto") { + pw_impl_selector = "pw_auto"; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } + } set_debugging(false); set_print_output(false); #ifdef DEBUG_MODE @@ -653,12 +733,11 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; // Execute test cases using the new framework with custom FLOP calculator + const auto test_case_generator = [pw_impl_selector]() { + return generate_quantized_conv2d_pw_test_cases(pw_impl_selector); + }; auto results = execute_test_cases( -#ifdef DEBUG_MODE - generate_quantized_conv2d_pw_test_cases, -#else - generate_quantized_conv2d_pw_test_cases, -#endif + test_case_generator, quantized_conv2d_flop_calculator, "QuantizedConv2dPW", /*warmup_runs = */ 1, From de3f49dee04a966b3432c931458f077706b8f307 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 3 Sep 2026 21:31:09 -0700 Subject: [PATCH 043/190] [ET-VK][q8ta] Route im2col convolution through unsigned dot Pull Request resolved: https://github.com/pytorch/executorch/pull/22540 Use corrected unsigned accumulating-saturating packed dot for im2col convolution when the adapter accelerates unsigned but not signed packed dot and the accumulator bound is safe. Support both texture- and buffer-backed packed weights. Preserve signed routing elsewhere. Authored with Codex. ghstack-source-id: 424669159 @exported-using-ghexport Differential Revision: [D118585647](https://our.internmc.facebook.com/intern/diff/D118585647/) --- .../graph/ops/glsl/pack_q8_linear_weight.glsl | 5 + .../graph/ops/glsl/pack_q8_linear_weight.yaml | 7 + .../graph/ops/glsl/q8ta_conv2d_pw.glsl | 10 +- .../graph/ops/glsl/q8ta_conv2d_pw.yaml | 4 + .../runtime/graph/ops/impl/Q8taConv2d.h | 5 + .../graph/ops/impl/Q8taConv2dIm2Col.cpp | 24 +- .../runtime/graph/ops/impl/Q8taConv2dPW.cpp | 3 + .../graph/ops/impl/QuantizedLinear.cpp | 15 +- .../runtime/graph/ops/impl/QuantizedLinear.h | 3 +- .../test/custom_ops/impl/TestQ8taConv2d.cpp | 84 ++++- .../test/custom_ops/test_q8ta_conv2d.cpp | 305 +++++++++++++++++- 11 files changed, 437 insertions(+), 28 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl index f2c74b67283..7b17dbe4990 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.glsl @@ -9,6 +9,7 @@ #version 450 core #define PRECISION ${PRECISION} +#define ADD_UNSIGNED_OFFSET ${ADD_UNSIGNED_OFFSET} ${define_active_storage_type(STORAGE)} @@ -53,6 +54,10 @@ void main() { load_block_data_with_checks(block , k4, n, K4, N); } +#if ADD_UNSIGNED_OFFSET == 1 + block.data = ivec4(uvec4(block.data) ^ uvec4(0x80808080u)); +#endif + // The weight blocks are stored in a tranposed manner, such that weight blocks // are indexed like packed_weight[k4][n4]. This is to optimize memory // coalescing when computing tiled GEMM. diff --git a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml index 13e6d43b2c5..d05ffb467bc 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/pack_q8_linear_weight.yaml @@ -7,8 +7,15 @@ pack_q8_linear_weight: parameter_names_with_default_values: STORAGE: buffer + ADD_UNSIGNED_OFFSET: 0 shader_variants: - NAME: pack_q8_linear_weight_buffer STORAGE: buffer - NAME: pack_q8_linear_weight_texture2d STORAGE: texture2d + - NAME: pack_q8_linear_weight_unsigned_buffer + STORAGE: buffer + ADD_UNSIGNED_OFFSET: 1 + - NAME: pack_q8_linear_weight_unsigned_texture2d + STORAGE: texture2d + ADD_UNSIGNED_OFFSET: 1 diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl index ca7cd978176..b68497a96be 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl @@ -13,6 +13,9 @@ ${define_required_extensions("buffer", DTYPE)} #define USE_INT8_DOT_PRODUCT_EXT ${USE_INT8_DOT_PRODUCT_EXT} #define USE_UNSIGNED_DOT_PRODUCT ${USE_UNSIGNED_DOT_PRODUCT} +$if WEIGHT_STORAGE == "buffer": + #define WEIGHT_BUFFER + #extension GL_EXT_control_flow_attributes : require $if USE_INT8_DOT_PRODUCT_EXT == 1: #extension GL_EXT_integer_dot_product : require @@ -44,7 +47,7 @@ layout(std430) buffer; ${layout_declare_tensor(B, "w", "t_packed_int8_output", "int", "buffer", is_scalar_array=True)} ${layout_declare_tensor(B, "r", "t_packed_int8_input", "int", "buffer", is_scalar_array=False)} -${layout_declare_tensor(B, "r", "t_packed_int8_weight", "int", "texture2d", is_scalar_array=False)} +${layout_declare_tensor(B, "r", "t_packed_int8_weight", "int", WEIGHT_STORAGE, is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_weight_sums", "int", "buffer", is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_weight_scales", DTYPE, "buffer", is_scalar_array=False)} ${layout_declare_tensor(B, "r", "t_bias", DTYPE, "buffer", is_scalar_array=False)} @@ -167,10 +170,15 @@ void main() { // Load the int8 weight tile for the current K and output-channel sub-block. ivec4 int8_weight_tile[TILE_N4]; [[unroll]] for (int n4 = 0; n4 < TILE_N4; ++n4) { +#ifdef WEIGHT_BUFFER + int8_weight_tile[n4] = t_packed_int8_weight[ + k4 * OC4 + oc_block_idx + n4]; +#else int8_weight_tile[n4] = texelFetch( t_packed_int8_weight, ivec2(oc_block_idx + n4, k4), 0); +#endif } #if USE_UNSIGNED_DOT_PRODUCT == 1 uvec4 uint8_weight_tile[TILE_N4]; diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml index 6e34bdaf5c0..f2979add927 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.yaml @@ -9,6 +9,7 @@ q8ta_conv2d_pw: DTYPE: float USE_INT8_DOT_PRODUCT_EXT: 1 USE_UNSIGNED_DOT_PRODUCT: 0 + WEIGHT_STORAGE: texture2d generate_variant_forall: DTYPE: - VALUE: float @@ -16,5 +17,8 @@ q8ta_conv2d_pw: - NAME: q8ta_conv2d_pw - NAME: q8ta_conv2d_pw_unsigned USE_UNSIGNED_DOT_PRODUCT: 1 + - NAME: q8ta_conv2d_pw_unsigned_buffer + USE_UNSIGNED_DOT_PRODUCT: 1 + WEIGHT_STORAGE: buffer - NAME: q8ta_conv2d_pw_fallback USE_INT8_DOT_PRODUCT_EXT: 0 diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h index 34556348019..690868e9f8c 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h @@ -165,6 +165,11 @@ void add_q8ta_im2col_node( void q8ta_conv2d_im2col(ComputeGraph& graph, const std::vector& args); +void q8ta_conv2d_im2col_impl( + ComputeGraph& graph, + bool use_unsigned_dot, + const std::vector& args); + // Transposed convolution void q8ta_conv2d_transposed( diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp index 334782d7af5..24538b8736a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp @@ -236,8 +236,9 @@ void add_q8ta_im2col_node( // High level operator impl // -void q8ta_conv2d_im2col( +void q8ta_conv2d_im2col_impl( ComputeGraph& graph, + const bool use_unsigned_dot, const std::vector& args) { int32_t idx = 0; const ValueRef packed_int8_input = args.at(idx++); @@ -260,8 +261,8 @@ void q8ta_conv2d_im2col( QuantizationConfig weight_quant_config(8, kPerChannel, {}); // Prepack weight using linear weight packing (for im2col approach) - ValueRef packed_weight = - prepack_quantized_linear_weight(graph, weight_quant_config, weight_data); + ValueRef packed_weight = prepack_quantized_linear_weight( + graph, weight_quant_config, weight_data, use_unsigned_dot); ValueRef packed_weight_sums = prepack_standard( graph, weight_sums_data, utils::kBuffer, utils::kWidthPacked); @@ -315,10 +316,15 @@ void q8ta_conv2d_im2col( // Step 2: Perform pointwise convolution on the im2col result const int32_t groups_val = graph.extract_scalar(groups); + VK_CHECK_COND( + !use_unsigned_dot || + graph.size_at(-1, weight_data) <= + kMaxUnsignedDotAccumulatorBytes, + "Unsigned q8ta im2col convolution exceeds the accumulator bound"); add_q8ta_conv2d_pw_node( graph, - /*use_unsigned_dot=*/false, + use_unsigned_dot, packed_int8_im2col, input_scale, input_zp, @@ -341,6 +347,16 @@ void q8ta_conv2d_im2col( dilation); } +void q8ta_conv2d_im2col( + ComputeGraph& graph, + const std::vector& args) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + const ValueRef weight_data = args.at(3); + const int64_t k_per_group = graph.size_at(-1, weight_data); + const bool use_unsigned_dot = can_use_unsigned_pw_dot(*adapter, k_per_group); + q8ta_conv2d_im2col_impl(graph, use_unsigned_dot, args); +} + REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.q8ta_conv2d_im2col.default, q8ta_conv2d_im2col); } diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp index 8b8bc63fc37..217afddb78f 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp @@ -321,6 +321,9 @@ void add_q8ta_conv2d_pw_node( use_hw_dot, "Unsigned q8ta pointwise convolution requires integer dot product"); kernel_name = "q8ta_conv2d_pw_unsigned"; + if (graph.storage_type_of(packed_weight) == utils::kBuffer) { + kernel_name += "_buffer"; + } } else { kernel_name = use_hw_dot ? "q8ta_conv2d_pw" : "q8ta_conv2d_pw_fallback"; } diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index 25a8d0b89ef..dc0a0a0837b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -331,9 +331,11 @@ vkapi::ShaderInfo pick_linear_dqa_qw_shader( ValueRef prepack_quantized_linear_weight( ComputeGraph& graph, const QuantizationConfig& weight_quant_config, - const ValueRef qmat2_data) { + const ValueRef qmat2_data, + const bool use_unsigned_dot) { VK_CHECK_COND( weight_quant_config.nbits == 8 || weight_quant_config.nbits == 4); + VK_CHECK_COND(!use_unsigned_dot || weight_quant_config.nbits == 8); std::vector qmat2_orig_sizes = graph.sizes_of(qmat2_data); const int64_t ndim = graph.dim_of(qmat2_data); @@ -410,10 +412,13 @@ ValueRef prepack_quantized_linear_weight( if (output_width > max_extent * 4 || output_height > max_extent) { storage_type = utils::kBuffer; } - - std::string kernel_name = weight_quant_config.nbits == 4 - ? "pack_q4_linear_weight" - : "pack_q8_linear_weight"; + std::string kernel_name; + if (weight_quant_config.nbits == 4) { + kernel_name = "pack_q4_linear_weight"; + } else { + kernel_name = use_unsigned_dot ? "pack_q8_linear_weight_unsigned" + : "pack_q8_linear_weight"; + } add_storage_type_suffix(kernel_name, storage_type); // Check prepack cache before creating a new prepack node. This avoids diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h index 7cb0e172c4a..4ab87f8ac33 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.h @@ -24,6 +24,7 @@ LocalWorkGroup quantized_linear_lwg( ValueRef prepack_quantized_linear_weight( ComputeGraph& graph, const QuantizationConfig& weight_quant_config, - const ValueRef qmat2_data); + const ValueRef qmat2_data, + const bool use_unsigned_dot = false); } // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp index e1957c4c590..e97ef8d61e0 100644 --- a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp @@ -17,6 +17,67 @@ namespace vkcompute { namespace { +void assert_im2col_kernel_selection( + ComputeGraph& graph, + const bool expect_unsigned, + const bool expect_buffer_weights) { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + std::string expected_execute; + std::string expected_prepack; + if (expect_unsigned) { + expected_execute = expect_buffer_weights + ? "q8ta_conv2d_pw_unsigned_buffer_float" + : "q8ta_conv2d_pw_unsigned_float"; + expected_prepack = expect_buffer_weights + ? "pack_q8_linear_weight_unsigned_buffer" + : "pack_q8_linear_weight_unsigned_texture2d"; + } else { + VK_CHECK_COND(!expect_buffer_weights); + expected_execute = adapter->supports_int8_dot_product() + ? "q8ta_conv2d_pw_float" + : "q8ta_conv2d_pw_fallback_float"; + expected_prepack = "pack_q8_linear_weight_texture2d"; + } + + int32_t execute_matches = 0; + std::string execute_names; + for (const auto& node : graph.execute_nodes()) { + const ExecuteNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + execute_names += node_name + " "; + if (node_name.find("q8ta_conv2d_pw") == 0) { + VK_CHECK_COND( + node_name == expected_execute, + "Expected ", + expected_execute, + " but selected execute kernel ", + node_name); + ++execute_matches; + } + } + VK_CHECK_COND(execute_matches > 0, "Execute kernels: ", execute_names); + + int32_t prepack_matches = 0; + std::string prepack_names; + for (const auto& node : graph.prepack_nodes()) { + const PrepackNode* const node_ptr = node.get(); + VK_CHECK_COND(node_ptr != nullptr); + const std::string& node_name = node_ptr->name(); + prepack_names += node_name + " "; + if (node_name.find("pack_q8_linear_weight") == 0) { + VK_CHECK_COND( + node_name == expected_prepack, + "Expected ", + expected_prepack, + " but selected prepack kernel ", + node_name); + ++prepack_matches; + } + } + VK_CHECK_COND(prepack_matches > 0, "Prepack kernels: ", prepack_names); +} + void assert_pw_kernel_selection( ComputeGraph& graph, const bool expect_unsigned, @@ -250,8 +311,27 @@ void test_q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { groups, activation, packed_int8_output}; - if (impl_selector == "im2col") { - VK_GET_OP_FN("et_vk.q8ta_conv2d_im2col.default")(graph, conv_args); + if (impl_selector == "im2col" || impl_selector == "im2col_unsigned" || + impl_selector == "im2col_auto") { + const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); + bool expect_unsigned = impl_selector == "im2col_unsigned"; + if (impl_selector == "im2col_auto") { + VK_GET_OP_FN("et_vk.q8ta_conv2d_im2col.default")(graph, conv_args); + expect_unsigned = can_use_unsigned_pw_dot( + *adapter, graph.size_at(-1, weight_data)); + } else { + q8ta_conv2d_im2col_impl(graph, expect_unsigned, conv_args); + } + const int64_t packed_height = + utils::div_up_4(graph.size_at(-1, weight_data)); + const int64_t packed_width = + utils::div_up_4(graph.size_at(-2, weight_data)) * 4; + const int64_t max_texture_extent = adapter->max_texture2d_dim(); + const bool expect_buffer_weights = + packed_width > max_texture_extent * 4 || + packed_height > max_texture_extent; + assert_im2col_kernel_selection( + graph, expect_unsigned, expect_buffer_weights); } else if (impl_selector == "general") { VK_GET_OP_FN("et_vk.q8ta_conv2d_general.default")(graph, conv_args); } else { diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index b30212feb72..308f2eac86a 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -4,10 +4,14 @@ // This source code is licensed under the BSD-style license found in the // LICENSE file in the root directory of this source tree. +#include +#include #include +#include #include #include +#include #include #include @@ -22,6 +26,17 @@ using namespace executorch::vulkan::prototyping; using namespace vkcompute; static constexpr int64_t kRefDimSizeLimit = 100; +static constexpr int64_t kRefOperationLimit = 2 * 1024 * 1024; + +struct Im2colUnsignedTestOptions { + int32_t input_zero_point = 2; + int32_t output_zero_point = -1; + bool use_extreme_values = false; + bool use_accumulator_limit_values = false; + bool has_bias = true; + const char* activation = "none"; + float weight_scale = 1.0f / 256.0f; +}; // Utility function to create a test case from a Conv2dConfig static TestCase create_test_case_from_config( @@ -29,7 +44,8 @@ static TestCase create_test_case_from_config( vkapi::ScalarType input_dtype, utils::StorageType fp_storage_type, utils::GPUMemoryLayout int8_memory_layout, - const std::string& impl_selector = "") { + const std::string& impl_selector = "", + const Im2colUnsignedTestOptions* im2col_options = nullptr) { TestCase test_case; // Calculate output dimensions @@ -91,9 +107,28 @@ static TestCase create_test_case_from_config( float input_scale_val = 0.008123; ValueSpec input_scale(input_scale_val); - int32_t input_zero_point_val = 2; + const int32_t input_zero_point_val = + im2col_options == nullptr ? 2 : im2col_options->input_zero_point; ValueSpec input_zero_point(input_zero_point_val); + if (im2col_options != nullptr && + im2col_options->use_accumulator_limit_values) { + input_tensor.ensure_data_generated(2401); + std::fill( + input_tensor.get_float_data().begin(), + input_tensor.get_float_data().end(), + (127.0f - input_zero_point_val) * input_scale_val); + } else if (im2col_options != nullptr && im2col_options->use_extreme_values) { + input_tensor.ensure_data_generated(2401); + constexpr std::array values = {-128, -1, 0, 1, 127}; + std::vector& input_data = input_tensor.get_float_data(); + for (size_t i = 0; i < input_data.size(); ++i) { + input_data.at(i) = (static_cast(values.at(i % values.size())) - + input_zero_point_val) * + input_scale_val; + } + } + // Quantized weight tensor (int8) - [C_out, C_in_per_group * K_h * K_w] // Memory layout: height, width, then channels - in_c is innermost (stride 1) // in the second dimension @@ -109,6 +144,22 @@ static TestCase create_test_case_from_config( DataGenType::RANDINT8); quantized_weight.set_constant(true); + if (im2col_options != nullptr && + im2col_options->use_accumulator_limit_values) { + quantized_weight.ensure_data_generated(2402); + std::fill( + quantized_weight.get_int8_data().begin(), + quantized_weight.get_int8_data().end(), + 127); + } else if (im2col_options != nullptr && im2col_options->use_extreme_values) { + quantized_weight.ensure_data_generated(2402); + constexpr std::array values = {-128, -1, 0, 1, 127}; + std::vector& weight_data = quantized_weight.get_int8_data(); + for (size_t i = 0; i < weight_data.size(); ++i) { + weight_data.at(i) = values.at((i * 3 + 1) % values.size()); + } + } + if (debugging()) { print_valuespec_data(quantized_weight, "weight_tensor"); } @@ -123,6 +174,13 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::RANDOM_SCALES); weight_scales.set_constant(true); + if (im2col_options != nullptr) { + weight_scales.ensure_data_generated(2403); + std::fill( + weight_scales.get_float_data().begin(), + weight_scales.get_float_data().end(), + im2col_options->weight_scale); + } ValueSpec weight_sums( {aligned_out_channels}, // Per output channel @@ -144,12 +202,16 @@ static TestCase create_test_case_from_config( utils::kWidthPacked, DataGenType::ZEROS); bias.set_constant(true); + if (im2col_options != nullptr && !im2col_options->has_bias) { + bias.set_none(true); + } // Output quantization parameters float output_scale_val = 0.05314; ValueSpec output_scale(output_scale_val); - int32_t output_zero_point_val = -1; + const int32_t output_zero_point_val = + im2col_options == nullptr ? -1 : im2col_options->output_zero_point; ValueSpec output_zero_point(output_zero_point_val); // Stride and padding parameters @@ -188,7 +250,8 @@ static TestCase create_test_case_from_config( test_case.add_input_spec(groups); // Activation (none = no activation) - ValueSpec activation = ValueSpec::make_string("none"); + ValueSpec activation = ValueSpec::make_string( + im2col_options == nullptr ? "none" : im2col_options->activation); test_case.add_input_spec(activation); // Add memory layout parameter for the quantized tensors @@ -201,7 +264,9 @@ static TestCase create_test_case_from_config( test_case.add_output_spec(output); - test_case.set_abs_tolerance(output_scale_val + 1e-4f); + test_case.set_abs_tolerance( + im2col_options == nullptr ? output_scale_val + 1e-4f + : output_scale_val * 0.25f); // Filter out quantize/dequantize shaders from timing measurements test_case.set_shader_filter({ @@ -271,10 +336,21 @@ std::vector generate_quantized_conv2d_easy_cases() { return test_cases; } +static std::vector generate_im2col_unsigned_test_cases( + const std::string& impl_selector); + // Generate test cases for quantized conv2d operation static std::vector generate_quantized_conv2d_test_cases() { std::vector test_cases; - if (!vkcompute::api::context()->adapter_ptr()->supports_int8_dot_product()) { + api::Context* const context = vkcompute::api::context(); + if (!context->adapter_ptr()->supports_int8_dot_product()) { + for (const std::string& impl_selector : {"im2col", "im2col_auto"}) { + std::vector im2col_cases = + generate_im2col_unsigned_test_cases(impl_selector); + for (TestCase& test_case : im2col_cases) { + test_cases.push_back(std::move(test_case)); + } + } return test_cases; } @@ -542,6 +618,171 @@ static std::vector generate_quantized_conv2d_test_cases() { } } + for (const std::string& impl_selector : + {"im2col", "im2col_unsigned", "im2col_auto"}) { + std::vector im2col_cases = + generate_im2col_unsigned_test_cases(impl_selector); + for (TestCase& test_case : im2col_cases) { + test_cases.push_back(std::move(test_case)); + } + } + + return test_cases; +} + +static std::vector generate_im2col_unsigned_test_cases( + const std::string& impl_selector) { + api::Context* const context = vkcompute::api::context(); + if (impl_selector == "im2col_unsigned" && + !context->adapter_ptr()->supports_int8_dot_product()) { + return {}; + } + + std::vector> configs = { + {{OutInChannels(5, 4), + InputSize2D(5, 5), + KernelSize(3, 3), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}, + {.input_zero_point = 2, + .output_zero_point = -1, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = true, + .activation = "none"}}, + {{OutInChannels(8, 8), + InputSize2D(5, 5), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + {.input_zero_point = -7, + .output_zero_point = 3, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = false, + .activation = "none"}}, + {{OutInChannels(12, 8), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1}, + {.input_zero_point = 127, + .output_zero_point = -5, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = true, + .activation = "relu"}}, + {{OutInChannels(12, 8), + InputSize2D(9, 9), + KernelSize(3, 3), + Stride(1, 1), + Padding(2, 2), + Dilation(2, 2), + 1}, + {.input_zero_point = -128, + .output_zero_point = 5, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = false, + .activation = "none"}}, + {{OutInChannels(8, 8), + InputSize2D(6, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2}, + {.input_zero_point = 11, + .output_zero_point = -3, + .use_extreme_values = true, + .use_accumulator_limit_values = false, + .has_bias = true, + .activation = "relu"}}, + {{OutInChannels(1, 4), + InputSize2D(90, 91), + KernelSize(90, 91), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}, + {.input_zero_point = 0, + .output_zero_point = -1, + .use_extreme_values = false, + .use_accumulator_limit_values = true, + .has_bias = true, + .activation = "none", + .weight_scale = 1.0f / 1000000.0f}}, + }; + + std::vector test_cases; + test_cases.reserve(configs.size()); + for (auto& [config, options] : configs) { + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + impl_selector, + &options)); + } + + if (impl_selector == "im2col_auto") { + Conv2dConfig config{ + OutInChannels(1, 4), + InputSize2D(91, 91), + KernelSize(91, 91), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}; + Im2colUnsignedTestOptions options; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + impl_selector, + &options)); + } + + const vkapi::Adapter& adapter = *context->adapter_ptr(); + if (impl_selector == "im2col_unsigned" || + (impl_selector == "im2col_auto" && can_use_unsigned_pw_dot(adapter, 4))) { + const int32_t buffer_output_channels = utils::safe_downcast( + static_cast(adapter.max_texture2d_dim()) * 4 + 1); + Conv2dConfig config{ + OutInChannels(buffer_output_channels, 4), + InputSize2D(1, 1), + KernelSize(1, 1), + Stride(1, 1), + Padding(0, 0), + Dilation(1, 1), + 1}; + Im2colUnsignedTestOptions options; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kBuffer, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kBuffer, + utils::kPackedInt8_4W4C, + impl_selector, + &options)); + } + return test_cases; } @@ -565,7 +806,6 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { const ValueSpec& dilation_spec = test_case.inputs()[idx++]; const ValueSpec& groups_spec = test_case.inputs()[idx++]; const ValueSpec& activation_spec = test_case.inputs()[idx++]; - (void)activation_spec; // Not used in reference implementation const ValueSpec& layout_spec = test_case.inputs()[idx++]; (void)layout_spec; // Not used in reference implementation const ValueSpec& impl_selector_spec = test_case.inputs()[idx++]; @@ -606,10 +846,12 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { int64_t dilation_w = dilation_data[1]; int64_t groups = groups_spec.get_int_value(); - // Skip for large tensors since computation time will be extremely slow - if (N > kRefDimSizeLimit || C_in > kRefDimSizeLimit || - H_in > kRefDimSizeLimit || W_in > kRefDimSizeLimit || - C_out > kRefDimSizeLimit) { + const int64_t reference_operations = + N * C_out * H_out * W_out * (C_in / groups) * K_h * K_w; + const bool has_large_dimension = N > kRefDimSizeLimit || + C_in > kRefDimSizeLimit || H_in > kRefDimSizeLimit || + W_in > kRefDimSizeLimit || C_out > kRefDimSizeLimit; + if (has_large_dimension && reference_operations > kRefOperationLimit) { throw std::invalid_argument( "One or more dimensions exceed the allowed limit for reference implementation."); std::cout @@ -643,7 +885,7 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { auto& ref_data = output_spec.get_ref_float_data(); ref_data.resize(num_output_elements); - const int in_features = utils::align_up_4(C_in_per_group * K_h * K_w); + const int64_t in_features = utils::align_up_4(C_in_per_group * K_h * K_w); // Perform activation, weight, and output quantized conv2d operation for (int64_t n = 0; n < N; ++n) { @@ -722,8 +964,12 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { float float_result = accum_adjusted * input_scale * weight_scales_data[out_c]; - // Add bias and store result - float_result += bias_data[out_c]; + if (!bias_spec.is_none()) { + float_result += bias_data[out_c]; + } + if (activation_spec.get_string_value() == "relu") { + float_result = std::max(float_result, 0.0f); + } // Quantize the output to int8 float quant_output_f = @@ -780,6 +1026,30 @@ static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { } int main(int argc, char* argv[]) { + const vkapi::Adapter& adapter = *vkcompute::api::context()->adapter_ptr(); + const bool prefers_unsigned_dot = + adapter.accelerates_unsigned_packed4x8_dot() && + !adapter.accelerates_signed_packed4x8_dot(); + VK_CHECK_COND( + can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes) == + prefers_unsigned_dot); + VK_CHECK_COND( + !can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes + 1)); + + std::string im2col_impl_selector; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--im2col-path=signed") { + im2col_impl_selector = "im2col"; + } else if (arg == "--im2col-path=unsigned") { + im2col_impl_selector = "im2col_unsigned"; + } else if (arg == "--im2col-path=auto") { + im2col_impl_selector = "im2col_auto"; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } + } set_debugging(false); set_print_output(false); #ifdef DEBUG_MODE @@ -798,11 +1068,16 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; // Execute test cases using the new framework with custom FLOP calculator + const auto test_case_generator = [im2col_impl_selector]() { + return im2col_impl_selector.empty() + ? generate_quantized_conv2d_test_cases() + : generate_im2col_unsigned_test_cases(im2col_impl_selector); + }; auto results = execute_test_cases( #ifdef DEBUG_MODE generate_quantized_conv2d_easy_cases, #else - generate_quantized_conv2d_test_cases, + test_case_generator, #endif quantized_conv2d_flop_calculator, "QuantizedConv2dQ8ToQ8To", From 0dd8f878e8421d2d83dd3fd9ebef836983731309 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 07:52:20 -0700 Subject: [PATCH 044/190] Cortex-M: add explicit layout convolution kernels (#22378) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Add and register NHWC implementations for regular, depthwise, and transpose convolution while preserving the existing channels-last entry points. ### Why separate operators The legacy kernels interpret logical NCHW tensors through channels-last dim order. Explicit layout instead passes ordinary contiguous tensors whose logical shape is NHWC. Separate experimental symbols make that contract unambiguous, prevent either mode from guessing layout, and allow validation to fail closed. The implementations share their CMSIS-NN setup paths where the two contracts are genuinely identical. The registration entries and FVP selected-op list live in this PR because this is the first layer containing the corresponding C++ kernels. A serialized program using any of these symbols can therefore load as soon as this PR lands. AI-assisted: Codex. --- Stack created with GitHub Stacks CLI • Give Feedback 💬 --- backends/cortex_m/ops/BUCK | 1 + backends/cortex_m/ops/cortex_m_ops_common.h | 6 + backends/cortex_m/ops/op_quantized_conv2d.cpp | 132 +++++++-- .../ops/op_quantized_depthwise_conv2d.cpp | 137 +++++++-- .../ops/op_quantized_transpose_conv2d.cpp | 134 +++++++-- backends/cortex_m/ops/operators.py | 277 ++++++++++++++++++ backends/cortex_m/ops/operators.yaml | 18 ++ .../cortex_m/passes/scratch_buffer_sizes.py | 34 ++- backends/cortex_m/test/build_test_runner.sh | 3 + backends/cortex_m/test/ops/nhwc_test_utils.py | 42 +++ backends/cortex_m/test/ops/test_nhwc_conv.py | 145 +++++++++ 11 files changed, 861 insertions(+), 68 deletions(-) create mode 100644 backends/cortex_m/test/ops/nhwc_test_utils.py create mode 100644 backends/cortex_m/test/ops/test_nhwc_conv.py diff --git a/backends/cortex_m/ops/BUCK b/backends/cortex_m/ops/BUCK index e538e9e1dd6..3c232ba048d 100644 --- a/backends/cortex_m/ops/BUCK +++ b/backends/cortex_m/ops/BUCK @@ -32,6 +32,7 @@ fbcode_target(_kind = runtime.python_library, "fbcode//caffe2:torch", "//executorch/backends/cortex_m/passes:passes_utils", "//executorch/backends/cortex_m/quantizer:quantization_configs", + "//executorch/exir:_warnings", ], ) diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index 2e3f49dd861..1e2eeb99d1d 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -14,6 +14,7 @@ #include #include +#include #include #include @@ -36,6 +37,11 @@ using KernelRuntimeContext = torch::executor::KernelRuntimeContext; // 16-byte alignment for MVE vector operations. constexpr size_t kCortexMMveAlignment = 16; +enum class ActivationLayout { + NCHWLogical, + NHWCLogical, +}; + // Basic tensor type / layout validation and dimension order checking inline void validate_cmsis_nn_tensor_requirements( const Tensor& input1, diff --git a/backends/cortex_m/ops/op_quantized_conv2d.cpp b/backends/cortex_m/ops/op_quantized_conv2d.cpp index 204a2b8369b..91cc893fba7 100644 --- a/backends/cortex_m/ops/op_quantized_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_conv2d.cpp @@ -1,10 +1,14 @@ /* * Copyright 2025-2026 Arm Limited and/or its affiliates. + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ +#include + #include "cortex_m_ops_common.h" namespace cortex_m { @@ -25,7 +29,8 @@ bool validate_conv2d_arguments( const Int64ArrayRef& padding, const Int64ArrayRef& dilation, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvDim || weight.dim() != kConvDim || output.dim() != kConvDim) { ET_LOG(Error, "quantized_conv2d_out: tensors must be 4-D"); @@ -33,20 +38,22 @@ bool validate_conv2d_arguments( return false; } - // Check for channels_last dim_order (NHWC: 0, 2, 3, 1) - // Skip check if channels == 1, as dim_order is ambiguous in that case - if (input.size(1) > 1 && !is_channels_last_tensor(input)) { - ET_LOG( - Error, - "quantized_conv2d_out: input must have channels_last dim_order (NHWC)"); - context.fail(Error::InvalidArgument); - return false; - } - - if (output.size(1) > 1 && !is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( Error, - "quantized_conv2d_out: output must have channels_last dim_order (NHWC)"); + "quantized_conv2d_out: input and output must have channels_last dim_order"); context.fail(Error::InvalidArgument); return false; } @@ -78,7 +85,8 @@ bool validate_conv2d_arguments( return false; } - const int64_t out_channels = output.size(1); + const int64_t out_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (requantize_multipliers.size(0) != out_channels || requantize_shifts.size(0) != out_channels) { ET_LOG( @@ -94,7 +102,7 @@ bool validate_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_conv2d_out( +static Tensor& quantized_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -109,6 +117,7 @@ Tensor& quantized_conv2d_out( const int64_t activation_min, const int64_t activation_max, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { if (!validate_conv2d_arguments( context, @@ -120,23 +129,30 @@ Tensor& quantized_conv2d_out( padding, dilation, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t kernel_output_channels = static_cast(weight.size(0)); const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); const int32_t kernel_input_channels = static_cast(weight.size(3)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t input_offset_val = static_cast(input_offset); const int32_t output_offset_val = static_cast(output_offset); @@ -228,5 +244,77 @@ Tensor& quantized_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_conv2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp index 0793606de44..296fda24b56 100644 --- a/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_depthwise_conv2d.cpp @@ -1,10 +1,14 @@ /* * Copyright 2025-2026 Arm Limited and/or its affiliates. + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. */ +#include + #include "cortex_m_ops_common.h" namespace cortex_m { @@ -26,7 +30,8 @@ bool validate_depthwise_conv2d_arguments( const Int64ArrayRef& dilation, const int64_t depth_multiplier, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvDim || weight.dim() != kConvDim || output.dim() != kConvDim) { ET_LOG(Error, "quantized_depthwise_conv2d_out: tensors must be 4-D"); @@ -55,7 +60,8 @@ bool validate_depthwise_conv2d_arguments( } const int64_t weight_output_channels = weight.size(3); - const int64_t output_channels = output.size(1); + const int64_t output_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (weight_output_channels != output_channels) { ET_LOG( Error, @@ -66,16 +72,22 @@ bool validate_depthwise_conv2d_arguments( return false; } - if (!is_channels_last_tensor(input)) { - ET_LOG( - Error, "quantized_depthwise_conv2d_out: input must be channels_last"); - context.fail(Error::InvalidArgument); - return false; - } - - if (!is_channels_last_tensor(output)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_depthwise_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( - Error, "quantized_depthwise_conv2d_out: output must be channels_last"); + Error, + "quantized_depthwise_conv2d_out: input and output must be channels_last"); context.fail(Error::InvalidArgument); return false; } @@ -108,7 +120,8 @@ bool validate_depthwise_conv2d_arguments( return false; } - const int64_t input_channels = input.size(1); + const int64_t input_channels = + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); // output_channels already extracted above for weight validation if (output_channels != input_channels * depth_multiplier) { ET_LOG( @@ -136,7 +149,7 @@ bool validate_depthwise_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_depthwise_conv2d_out( +static Tensor& quantized_depthwise_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -152,6 +165,7 @@ Tensor& quantized_depthwise_conv2d_out( const int64_t activation_min, const int64_t activation_max, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { if (!validate_depthwise_conv2d_arguments( context, @@ -164,23 +178,30 @@ Tensor& quantized_depthwise_conv2d_out( dilation, depth_multiplier, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); // Weight is in IHWO layout after permutation in the pass: [1, H, W, C_OUT] // For depthwise conv, this matches CMSIS-NN's expected format const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t depth_multiplier_val = static_cast(depth_multiplier); @@ -272,5 +293,81 @@ Tensor& quantized_depthwise_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_depthwise_conv2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t depth_multiplier, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_depthwise_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_depthwise_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const int64_t depth_multiplier, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + Tensor& out) { + return quantized_depthwise_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp index 04d57d4c693..4ac9b2338e6 100644 --- a/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp +++ b/backends/cortex_m/ops/op_quantized_transpose_conv2d.cpp @@ -24,7 +24,8 @@ bool validate_transpose_conv2d_arguments( const std::optional& bias, const Tensor& output, const Tensor& requantize_multipliers, - const Tensor& requantize_shifts) { + const Tensor& requantize_shifts, + ActivationLayout layout) { if (input.dim() != kConvTransposeDim || weight.dim() != kConvTransposeDim || output.dim() != kConvTransposeDim) { ET_LOG(Error, "quantized_transpose_conv2d_out: tensors must be 4-D"); @@ -32,16 +33,22 @@ bool validate_transpose_conv2d_arguments( return false; } - if (!is_channels_last_tensor(input)) { + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG( + Error, + "quantized_transpose_conv2d_nhwc_out: input and output must have contiguous dim_order"); + context.fail(Error::InvalidArgument); + return false; + } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { ET_LOG( - Error, "quantized_transpose_conv2d_out: input must be channels_last"); - context.fail(Error::InvalidArgument); - return false; - } - - if (!is_channels_last_tensor(output)) { - ET_LOG( - Error, "quantized_transpose_conv2d_out: output must be channels_last"); + Error, + "quantized_transpose_conv2d_out: input and output must be channels_last"); context.fail(Error::InvalidArgument); return false; } @@ -68,7 +75,8 @@ bool validate_transpose_conv2d_arguments( return false; } - const int64_t out_channels = output.size(1); + const int64_t out_channels = + output.size(layout == ActivationLayout::NHWCLogical ? 3 : 1); if (requantize_multipliers.size(0) != out_channels || requantize_shifts.size(0) != out_channels) { ET_LOG( @@ -84,7 +92,7 @@ bool validate_transpose_conv2d_arguments( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_transpose_conv2d_out( +static Tensor& quantized_transpose_conv2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Tensor& weight, @@ -101,6 +109,7 @@ Tensor& quantized_transpose_conv2d_out( const int64_t activation_max, const Tensor& scratch, const Tensor& output_scratch, + ActivationLayout layout, Tensor& out) { if (!validate_transpose_conv2d_arguments( context, @@ -109,23 +118,30 @@ Tensor& quantized_transpose_conv2d_out( bias, out, requantize_multipliers, - requantize_shifts)) { + requantize_shifts, + layout)) { return out; } const int32_t batch = static_cast(input.size(0)); - const int32_t input_channels = static_cast(input.size(1)); - const int32_t input_height = static_cast(input.size(2)); - const int32_t input_width = static_cast(input.size(3)); + const int32_t input_channels = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t input_height = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t input_width = static_cast( + input.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); const int32_t kernel_output_channels = static_cast(weight.size(0)); const int32_t kernel_height = static_cast(weight.size(1)); const int32_t kernel_width = static_cast(weight.size(2)); const int32_t kernel_input_channels = static_cast(weight.size(3)); - const int32_t output_channels = static_cast(out.size(1)); - const int32_t output_height = static_cast(out.size(2)); - const int32_t output_width = static_cast(out.size(3)); + const int32_t output_channels = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 3 : 1)); + const int32_t output_height = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 1 : 2)); + const int32_t output_width = static_cast( + out.size(layout == ActivationLayout::NHWCLogical ? 2 : 3)); if (kernel_output_channels != output_channels) { ET_LOG( @@ -246,5 +262,85 @@ Tensor& quantized_transpose_conv2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_transpose_conv2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef output_padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + const Tensor& output_scratch, + Tensor& out) { + return quantized_transpose_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_transpose_conv2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Tensor& weight, + const std::optional& bias, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef output_padding, + const Int64ArrayRef dilation, + const int64_t input_offset, + const int64_t output_offset, + const Tensor& requantize_multipliers, + const Tensor& requantize_shifts, + const int64_t activation_min, + const int64_t activation_max, + const Tensor& scratch, + const Tensor& output_scratch, + Tensor& out) { + return quantized_transpose_conv2d_out_impl( + context, + input, + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index 44e47087c11..ea67600cff6 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -23,6 +23,7 @@ CMSIS_SOFTMAX_SCALE, CMSIS_SOFTMAX_ZERO_POINT, ) +from executorch.exir._warnings import experimental from executorch.exir.dialects._ops import ops as exir_ops # To provide the implementation of the operators @@ -31,6 +32,11 @@ # New operator library with a custom namespace to allow fusion etc. lib = Library("cortex_m", "DEF") +_EXPLICIT_LAYOUT_EXPERIMENTAL = ( + "This explicit-layout Cortex-M operator may change while the legacy " + "dim-order operators remain supported." +) + SOFTMAX_INPUT_INTEGER_BITS = 5 @@ -915,6 +921,92 @@ def quantized_conv2d_impl( return result.to(torch.int8, memory_format=torch.channels_last) +lib.define( + "quantized_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch) -> Tensor" +) +lib.define( + "quantized_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_conv2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED DEPTHWISE CONV2D OPERATION DEFINITION # =================================================================== @@ -1062,6 +1154,96 @@ def quantized_depthwise_conv2d_impl( return result.to(torch.int8, memory_format=torch.channels_last) +lib.define( + "quantized_depthwise_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int depth_multiplier, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch) -> Tensor" +) +lib.define( + "quantized_depthwise_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] dilation, int depth_multiplier, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_depthwise_conv2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_depthwise_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + depth_multiplier: int, + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_depthwise_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_depthwise_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_depthwise_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + depth_multiplier: int, + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_depthwise_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + dilation, + depth_multiplier, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED TRANSPOSE_CONV2D OPERATION DEFINITION # =================================================================== @@ -1270,6 +1452,101 @@ def quantized_transpose_conv2d_impl( return result.to(torch.int8).to(memory_format=torch.channels_last) +lib.define( + "quantized_transpose_conv2d_nhwc(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] output_padding, int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "Tensor output_scratch) -> Tensor" +) +lib.define( + "quantized_transpose_conv2d_nhwc.out(" + "Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, " + "int[] output_padding, int[] dilation, int input_offset, int output_offset, " + "Tensor requantize_multipliers, Tensor requantize_shifts, " + "int activation_min, int activation_max, Tensor scratch, " + "Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_transpose_conv2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_transpose_conv2d_nhwc_meta( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + output_padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, + output_scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_transpose_conv2d_meta( + input.permute(0, 3, 1, 2), + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_transpose_conv2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_transpose_conv2d_nhwc_impl( + input: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + stride: Sequence[int], + padding: Sequence[int], + output_padding: Sequence[int], + dilation: Sequence[int], + input_offset: int, + output_offset: int, + requantize_multipliers: torch.Tensor, + requantize_shifts: torch.Tensor, + activation_min: int, + activation_max: int, + scratch: torch.Tensor, + output_scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_transpose_conv2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + weight, + bias, + stride, + padding, + output_padding, + dilation, + input_offset, + output_offset, + requantize_multipliers, + requantize_shifts, + activation_min, + activation_max, + scratch, + output_scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED AVG_POOL2D OPERATION DEFINITION # =================================================================== diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index 2c85325f854..a2912bfd660 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -83,6 +83,12 @@ - arg_meta: null kernel_name: cortex_m::quantized_conv2d_out +- func: cortex_m::quantized_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_conv2d_nhwc_out + - func: cortex_m::quantized_depthwise_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function @@ -90,12 +96,24 @@ - arg_meta: null kernel_name: cortex_m::quantized_depthwise_conv2d_out +- func: cortex_m::quantized_depthwise_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] dilation, int depth_multiplier, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_depthwise_conv2d_nhwc_out + - func: cortex_m::quantized_transpose_conv2d.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_transpose_conv2d_out +- func: cortex_m::quantized_transpose_conv2d_nhwc.out(Tensor input, Tensor weight, Tensor? bias, int[] stride, int[] padding, int[] output_padding, int[] dilation, int input_offset, int output_offset, Tensor requantize_multipliers, Tensor requantize_shifts, int activation_min, int activation_max, Tensor scratch, Tensor output_scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_transpose_conv2d_nhwc_out + - func: cortex_m::quantized_avg_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/passes/scratch_buffer_sizes.py b/backends/cortex_m/passes/scratch_buffer_sizes.py index b247e2be944..964b5530e13 100644 --- a/backends/cortex_m/passes/scratch_buffer_sizes.py +++ b/backends/cortex_m/passes/scratch_buffer_sizes.py @@ -4,6 +4,7 @@ # LICENSE file in the root directory of this source tree. from collections.abc import Callable +from functools import partial from typing import Any, cast import executorch.backends.cortex_m.ops.operators # noqa @@ -37,6 +38,7 @@ def _shape_from_node(node: torch.fx.Node) -> torch.Size: def _get_common_conv_buffer_size_inputs( conv_node: torch.fx.Node, *, + nhwc_logical: bool = False, stride_arg_idx: int = 3, padding_arg_idx: int = 4, dilation_arg_idx: int = 5, @@ -54,13 +56,14 @@ def _get_common_conv_buffer_size_inputs( padding = cast(list[int], conv_node.args[padding_arg_idx]) dilation = cast(list[int], conv_node.args[dilation_arg_idx]) - # Input is NCHW (PyTorch); CMSIS-NN wants NHWC dims. - n, c_in, height, width = _shape_from_node(x) - weight_shape = _shape_from_node(weight) - # Output is NCHW; convert to NHWC dims. - out_n, out_c, out_h, out_w = _shape_from_node(conv_node) + if nhwc_logical: + n, height, width, c_in = _shape_from_node(x) + out_n, out_h, out_w, out_c = _shape_from_node(conv_node) + else: + n, c_in, height, width = _shape_from_node(x) + out_n, out_c, out_h, out_w = _shape_from_node(conv_node) input_nhwc = [n, height, width, c_in] output_nhwc = [out_n, out_h, out_w, out_c] @@ -81,6 +84,7 @@ def _get_common_conv_buffer_size_inputs( def cmsis_nn_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -89,7 +93,9 @@ def cmsis_nn_conv_buffer_size( stride_hw, padding_hw, dilation_hw, - ) = _get_common_conv_buffer_size_inputs(conv_node=conv_node) + ) = _get_common_conv_buffer_size_inputs( + conv_node=conv_node, nhwc_logical=nhwc_logical + ) input_offset = cast(int, conv_node.args[6]) output_offset = cast(int, conv_node.args[7]) output_qmin = cast(int, conv_node.args[10]) @@ -122,6 +128,7 @@ def cmsis_nn_conv_buffer_size( def cmsis_nn_depthwise_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -130,7 +137,9 @@ def cmsis_nn_depthwise_conv_buffer_size( stride_hw, padding_hw, dilation_hw, - ) = _get_common_conv_buffer_size_inputs(conv_node=conv_node) + ) = _get_common_conv_buffer_size_inputs( + conv_node=conv_node, nhwc_logical=nhwc_logical + ) depth_multiplier = cast(int, conv_node.args[6]) input_offset = cast(int, conv_node.args[7]) output_offset = cast(int, conv_node.args[8]) @@ -185,6 +194,7 @@ def cmsis_nn_batch_matmul_buffer_size( def cmsis_nn_transpose_conv_buffer_size( backend: cmsis_nn.Backend, conv_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: ( input_nhwc, @@ -195,6 +205,7 @@ def cmsis_nn_transpose_conv_buffer_size( dilation_hw, ) = _get_common_conv_buffer_size_inputs( conv_node=conv_node, + nhwc_logical=nhwc_logical, stride_arg_idx=3, padding_arg_idx=4, dilation_arg_idx=6, @@ -270,9 +281,18 @@ def cmsis_nn_avgpool_buffer_size( _target_to_buffer_sizes_registry: dict[Any, BufferSizeFunction] = { exir_ops.edge.cortex_m.quantized_conv2d.default: cmsis_nn_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default: partial( + cmsis_nn_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default: cmsis_nn_depthwise_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default: partial( + cmsis_nn_depthwise_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_batch_matmul.default: cmsis_nn_batch_matmul_buffer_size, exir_ops.edge.cortex_m.quantized_transpose_conv2d.default: cmsis_nn_transpose_conv_buffer_size, + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default: partial( + cmsis_nn_transpose_conv_buffer_size, nhwc_logical=True + ), exir_ops.edge.cortex_m.quantized_avg_pool2d.default: cmsis_nn_avgpool_buffer_size, } diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index a38c6d53256..47b27da137e 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -67,8 +67,11 @@ ops_list=( cortex_m::transpose.out cortex_m::pad.out cortex_m::quantized_conv2d.out + cortex_m::quantized_conv2d_nhwc.out cortex_m::quantized_depthwise_conv2d.out + cortex_m::quantized_depthwise_conv2d_nhwc.out cortex_m::quantized_transpose_conv2d.out + cortex_m::quantized_transpose_conv2d_nhwc.out cortex_m::quantized_avg_pool2d.out cortex_m::quantized_max_pool2d.out cortex_m::quantized_batch_matmul.out diff --git a/backends/cortex_m/test/ops/nhwc_test_utils.py b/backends/cortex_m/test/ops/nhwc_test_utils.py new file mode 100644 index 00000000000..6cc02d85602 --- /dev/null +++ b/backends/cortex_m/test/ops/nhwc_test_utils.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math + +import torch + +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.passes.scratch_buffer_sizes import ( + required_cmsis_nn_buffer_sizes, +) +from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.test.harness.stages import RunPasses, StageType + + +def int8_values(shape): + values = torch.arange(math.prod(shape), dtype=torch.int32) + return (values.remainder(7) - 3).to(torch.int8).reshape(shape) + + +def run_on_fvp(module, x, target, target_config, scratch_count, atol=0): + sizing_inputs = (x,) + tuple( + torch.empty(0, dtype=torch.uint8) for _ in range(scratch_count) + ) + tester = CortexMTester(module, sizing_inputs, target_config=target_config) + tester.export().to_edge() + program = tester.get_artifact(StageType.TO_EDGE).exported_program() + [node] = [n for n in program.graph.nodes if n.target == target] + scratch_sizes = required_cmsis_nn_buffer_sizes(node, target_config.backend) or [] + assert len(scratch_sizes) == scratch_count + + inputs = (x,) + tuple( + torch.empty(size, dtype=torch.uint8) for size in scratch_sizes + ) + tester = CortexMTester(module, inputs, target_config=target_config) + tester.export().to_edge() + tester.run_passes(RunPasses(CortexMPassManager, pass_list=[])) + tester.to_executorch().serialize() + tester.run_method_and_compare_outputs(inputs=inputs, atol=atol) diff --git a/backends/cortex_m/test/ops/test_nhwc_conv.py b/backends/cortex_m/test/ops/test_nhwc_conv.py new file mode 100644 index 00000000000..da7602447c2 --- /dev/null +++ b/backends/cortex_m/test/ops/test_nhwc_conv.py @@ -0,0 +1,145 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.cortex_m.test.ops.nhwc_test_utils import ( + int8_values, + run_on_fvp, +) +from executorch.exir.dialects._ops import ops as exir_ops + +# Direct-op coverage is temporary until CortexMTester uses explicit layout by default. + + +class Conv2dNhwc(torch.nn.Module): + def __init__(self, grouped=False): + super().__init__() + in_channels = 4 if grouped else 3 + self.register_buffer( + "weight", int8_values((4, 2, 3, 2 if grouped else in_channels)) + ) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [2, 1], + [1, 0], + [1, 1], + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + ) + + +class DepthwiseConv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", int8_values((1, 3, 2, 4))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_depthwise_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [2, 1], + [1, 0], + [1, 1], + 1, + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + ) + + +class TransposeConv2dNhwc(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("weight", int8_values((4, 2, 4, 2))) + self.register_buffer("bias", torch.arange(4, dtype=torch.int32) - 2) + self.register_buffer( + "multipliers", torch.full((4,), 1 << 30, dtype=torch.int32) + ) + self.register_buffer("shifts", torch.full((4,), -1, dtype=torch.int32)) + + def forward(self, x, scratch, output_scratch): + return torch.ops.cortex_m.quantized_transpose_conv2d_nhwc.default( + x, + self.weight, + self.bias, + [1, 1], + [0, 0], + [0, 0], + [1, 1], + 0, + 0, + self.multipliers, + self.shifts, + -128, + 127, + scratch, + output_scratch, + ) + + +def test_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + Conv2dNhwc(), + int8_values((1, 7, 10, 3)), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + cortex_m_target, + 1, + ) + + +def test_grouped_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + Conv2dNhwc(grouped=True), + int8_values((1, 7, 10, 4)), + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + cortex_m_target, + 1, + ) + + +def test_depthwise_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + DepthwiseConv2dNhwc(), + int8_values((1, 7, 10, 4)), + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default, + cortex_m_target, + 1, + ) + + +def test_transpose_conv2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + TransposeConv2dNhwc(), + int8_values((1, 5, 6, 2)), + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default, + cortex_m_target, + 2, + ) From 738f197f32c37b019c928c2fcf808c5429333da8 Mon Sep 17 00:00:00 2001 From: Xingguo Li <100689130+xingguo01@users.noreply.github.com> Date: Fri, 4 Sep 2026 16:42:31 +0100 Subject: [PATCH 045/190] Arm backend: profile SmolLM2 Ethos-U KV cache on FVP (#22560) Capture per-inference Ethos-U85 PMU counters from the semihosting server. Separate token executions into prompt processing and decoding summaries. Disable FVP fast mode automatically during profiling. Support JSON and CSV reports, with optional NPU-only token-rate estimates. Document the CS-320 command and clarify that Cortex-M timing and end-to-end throughput cannot be inferred from the simulator. Generated with assistance from OpenAI Codex. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Xingguo Li --- .../arm/executor_runner/arm_perf_monitor.cpp | 6 + .../arm/smollm2_example_ethos_u/README.md | 39 +++ .../generate_sampled.py | 247 +++++++++++++++++- 3 files changed, 284 insertions(+), 8 deletions(-) diff --git a/examples/arm/executor_runner/arm_perf_monitor.cpp b/examples/arm/executor_runner/arm_perf_monitor.cpp index 9d92a878a2b..31fa5937fd3 100644 --- a/examples/arm/executor_runner/arm_perf_monitor.cpp +++ b/examples/arm/executor_runner/arm_perf_monitor.cpp @@ -187,6 +187,12 @@ void EthosUBackend_execute_end() { } void StartMeasurements() { +#if defined(__ARM_ARCH_8_1M_MAIN__) + // StopMeasurements() disables the cycle counter after each measurement. + // Server mode starts a new measurement for every inference. + ARM_PMU_Enable(); + ARM_PMU_CNTR_Enable(PMU_CNTENSET_CCNTR_ENABLE_Msk); +#endif ethosu_delegation_count = 0; ethosu_ArmBackendExecuteCycleCount = 0; ethosu_ArmWhenNPURunCycleCount = 0; diff --git a/examples/arm/smollm2_example_ethos_u/README.md b/examples/arm/smollm2_example_ethos_u/README.md index 25403fb0e7d..6571264af2b 100644 --- a/examples/arm/smollm2_example_ethos_u/README.md +++ b/examples/arm/smollm2_example_ethos_u/README.md @@ -234,6 +234,45 @@ How to interpret the main options: - `--repetition-penalty 1.1` still matters in greedy mode because it modifies the logits before `argmax`. +### 5.1 Profile prompt processing and decoding + +The runner exposes Ethos-U85 PMU counters for every server-mode inference. +Capture those counters with FVP fast mode disabled: + +```bash +python examples/arm/smollm2_example_ethos_u/generate_sampled.py \ + --fvp examples/arm/arm-scratch/FVP-corstone320/models/Linux64_GCC-9.3/FVP_Corstone_SSE-320 \ + --runner smollm2_ethosu_static_kvq_seq64_w8a16_wikitext/cmake-out/arm_executor_runner \ + --embedded-pte \ + --tokenizer data/tokenizers/smollm2/tokenizer.json \ + --prompt "Once upon a time in a small village," \ + --window 64 \ + --max-context-length 64 \ + --use-kv-cache \ + --max-new-tokens 2 \ + --temperature 0 \ + --no-topk-print \ + --profile-output outputs/ethosu_u85_profile.json \ + --timeout 24000 +``` + +`--profile-output` accepts `.json` or `.csv` and automatically removes the +Ethos-U `--fast` FVP option. The report contains one NPU PMU sample per model +execution, split into `prefill` and `decode` phases. Passing +`--npu-frequency-mhz` also reports an estimated NPU-only token rate. + +This KV-cache demo processes the prompt one token at a time; its prefill result +is therefore the aggregate and average of those token executions, not a +batched-prefill measurement. The final prompt execution supplies the logits for +the first generated token, so steady-state decode executions start with the +second generated token. + +CS-320 provides useful Ethos-U85 NPU cycle and event estimates when its timing +adapters match the target configuration. Cortex-M85 CPU timing and simulator +wall time are not cycle accurate, so the report must not be interpreted as +end-to-end latency or measured device tokens/s. Use FPGA or hardware for those +measurements. + ## 6. Optional: evaluate Wikitext perplexity The KV-cache generation artifact can also be used for step-wise perplexity scoring over the same 64-token context. diff --git a/examples/arm/smollm2_example_ethos_u/generate_sampled.py b/examples/arm/smollm2_example_ethos_u/generate_sampled.py index 67877d21fc3..fe4311a9dec 100644 --- a/examples/arm/smollm2_example_ethos_u/generate_sampled.py +++ b/examples/arm/smollm2_example_ethos_u/generate_sampled.py @@ -4,6 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import argparse +import csv +import json import re import secrets import select @@ -13,8 +15,9 @@ import time from collections import deque +from dataclasses import asdict, dataclass from pathlib import Path -from typing import Deque, List, Optional, Sequence +from typing import Deque, Dict, List, Optional, Sequence import numpy as np from pytorch_tokenizers import ( # type: ignore[import-not-found, import-untyped] @@ -26,6 +29,142 @@ re.MULTILINE, ) +ETHOSU_PMU_CYCLE_PATTERN = re.compile(r"ethosu_pmu_cycle_cntr\s*:\s*(\d+)") +ETHOSU_PMU_COUNTER_PATTERN = re.compile(r"ethosu_pmu_cntr(\d+)\s*:\s*(\d+)") +ETHOSU_DELEGATIONS_PATTERN = re.compile(r"NPU delegations:\s*(\d+)") +ETHOSU85_EVENT_NAMES = [ + "sram_read_beats", + "sram_write_beats", + "external_read_beats", + "external_write_beats", + "npu_idle", + "mac_active", + "weight_decoder_active", +] + + +@dataclass +class EthosUPmuMeasurement: + npu_cycles: int + delegations: int + events: Dict[str, int] + + +@dataclass +class ProfileSample: + prompt_no: int + phase: str + input_pos: int + npu_cycles: int + delegations: int + events: Dict[str, int] + + +def parse_ethosu_pmu(lines: Sequence[str]) -> EthosUPmuMeasurement: + text = "".join(lines) + cycle_match = ETHOSU_PMU_CYCLE_PATTERN.search(text) + if cycle_match is None: + raise RuntimeError("Ethos-U PMU cycle count was not found in FVP output") + + counter_values = { + int(index): int(value) + for index, value in ETHOSU_PMU_COUNTER_PATTERN.findall(text) + } + events = { + name: counter_values.get(index, 0) + for index, name in enumerate(ETHOSU85_EVENT_NAMES) + } + delegations_match = ETHOSU_DELEGATIONS_PATTERN.search(text) + return EthosUPmuMeasurement( + npu_cycles=int(cycle_match.group(1)), + delegations=( + int(delegations_match.group(1)) if delegations_match is not None else 0 + ), + events=events, + ) + + +def summarize_profile( + samples: Sequence[ProfileSample], npu_frequency_mhz: Optional[float] +) -> Dict[str, Dict[str, float]]: + summary: Dict[str, Dict[str, float]] = {} + for phase in ("prefill", "decode"): + phase_samples = [sample for sample in samples if sample.phase == phase] + if not phase_samples: + continue + total_cycles = sum(sample.npu_cycles for sample in phase_samples) + mean_cycles = total_cycles / len(phase_samples) + values = { + "executions": float(len(phase_samples)), + "total_npu_cycles": float(total_cycles), + "mean_npu_cycles": mean_cycles, + } + if npu_frequency_mhz is not None: + values["estimated_npu_tokens_per_second"] = ( + npu_frequency_mhz * 1_000_000 / mean_cycles + ) + summary[phase] = values + return summary + + +def write_profile( + path: Path, + samples: Sequence[ProfileSample], + npu_frequency_mhz: Optional[float], +) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + summary = summarize_profile(samples, npu_frequency_mhz) + if path.suffix.lower() == ".json": + path.write_text( + json.dumps( + { + "metadata": { + "ethosu_fast": False, + "npu_frequency_mhz": npu_frequency_mhz, + "timing_scope": "Ethos-U85 NPU only", + }, + "samples": [asdict(sample) for sample in samples], + "summary": summary, + }, + indent=2, + ) + + "\n", + encoding="utf-8", + ) + elif path.suffix.lower() == ".csv": + fieldnames = [ + "prompt_no", + "phase", + "input_pos", + "npu_cycles", + "delegations", + *ETHOSU85_EVENT_NAMES, + ] + with path.open("w", encoding="utf-8", newline="") as output: + writer = csv.DictWriter(output, fieldnames=fieldnames) + writer.writeheader() + for sample in samples: + row = asdict(sample) + row.update(row.pop("events")) + writer.writerow(row) + else: + raise ValueError("--profile-output must end in .json or .csv") + + print("\nEthos-U85 NPU profile (FVP estimate):") + for phase, values in summary.items(): + line = ( + f" {phase}: executions={int(values['executions'])} " + f"total_cycles={int(values['total_npu_cycles'])} " + f"mean_cycles={values['mean_npu_cycles']:.2f}" + ) + if "estimated_npu_tokens_per_second" in values: + line += ( + " estimated_npu_tokens/s=" + f"{values['estimated_npu_tokens_per_second']:.3f}" + ) + print(line) + print(f"Profile written to {path}") + def prepare_input( ids: List[int], @@ -220,12 +359,17 @@ def __init__( timeout: int, input_names: Optional[Sequence[str]] = None, server_mode: bool = False, + ethosu_fast: bool = True, + collect_profile: bool = False, ) -> None: self._fvp = fvp self._runner = runner self._pte = pte self._timeout = timeout self._server_mode = server_mode + self._ethosu_fast = ethosu_fast + self._collect_profile = collect_profile + self.last_pmu: Optional[EthosUPmuMeasurement] = None self._proc: Optional[subprocess.Popen[str]] = None self._recent_stdout: Deque[str] = deque(maxlen=400) self._tmpdir: Optional[tempfile.TemporaryDirectory[str]] = None @@ -246,7 +390,7 @@ def _init_paths(self, input_names: Sequence[str]) -> None: def _build_command(self, cmd_line: str) -> List[str]: assert self._tmpdir_path is not None - return [ + command = [ self._fvp, "-C", "mps4_board.subsystem.ethosu.num_macs=256", @@ -271,14 +415,19 @@ def _build_command(self, cmd_line: str) -> List[str]: "-C", f"mps4_board.subsystem.cpu0.semihosting-cwd={self._tmpdir_path}", "-C", - "mps4_board.subsystem.ethosu.extra_args='--fast'", - "-C", f"mps4_board.subsystem.cpu0.semihosting-cmd_line='{cmd_line}'", "-a", self._runner, "--timelimit", str(self._timeout), ] + if self._ethosu_fast: + insert_at = command.index("-a") + command[insert_at:insert_at] = [ + "-C", + "mps4_board.subsystem.ethosu.extra_args='--fast'", + ] + return command def close(self) -> None: if self._proc is not None: @@ -347,6 +496,9 @@ def _run_server_once(self, output_path: Path) -> np.ndarray: self._proc.stdin.write("go\n") self._proc.stdin.flush() + self.last_pmu = None + profile_lines: List[str] = [] + deadline = time.monotonic() + self._timeout while time.monotonic() < deadline: self._check_proc() @@ -362,7 +514,11 @@ def _run_server_once(self, output_path: Path) -> np.ndarray: f"\n\n[FVP stdout tail]\n{''.join(self._recent_stdout)}" ) self._recent_stdout.append(line) + if self._collect_profile: + profile_lines.append(line) if "SERVER_INFERENCE_DONE" in line: + if self._collect_profile: + self.last_pmu = parse_ethosu_pmu(profile_lines) if output_path.exists() and output_path.stat().st_size > 0: return np.fromfile(output_path, dtype=np.float32) raise RuntimeError( @@ -434,6 +590,8 @@ def __init__( runner: str, pte: Optional[str], timeout: int, + ethosu_fast: bool = True, + collect_profile: bool = False, ) -> None: self._runner = FvpRunnerSession( fvp, @@ -442,7 +600,10 @@ def __init__( timeout, input_names=["i0.bin", "i1.bin"], server_mode=True, + ethosu_fast=ethosu_fast, + collect_profile=collect_profile, ) + self.samples: List[ProfileSample] = [] def close(self) -> None: self._runner.close() @@ -453,13 +614,41 @@ def __enter__(self) -> "KvFvpRunnerSession": def __exit__(self, exc_type, exc, tb) -> None: # type: ignore[no-untyped-def] self.close() - def run(self, token_id: int, input_pos: int) -> np.ndarray: + def run( + self, + token_id: int, + input_pos: int, + *, + phase: str, + prompt_no: int, + ) -> np.ndarray: logits = self._runner.run_inputs( [ np.array([[token_id]], dtype=np.int32), np.array([input_pos], dtype=np.int32), ] ) + if self._runner.last_pmu is not None: + measurement = self._runner.last_pmu + if measurement.npu_cycles <= 0: + raise RuntimeError( + "Ethos-U PMU returned no cycles; ensure FVP fast mode is disabled" + ) + self.samples.append( + ProfileSample( + prompt_no, + phase, + input_pos, + measurement.npu_cycles, + measurement.delegations, + measurement.events, + ) + ) + print( + f"\n[Ethos-U profile phase={phase} input_pos={input_pos} " + f"npu_cycles={measurement.npu_cycles}]", + flush=True, + ) return logits.reshape(1, -1)[0] @@ -573,7 +762,7 @@ def run_one_prompt_kv( logits = None for pos, token_id in enumerate(ids): - logits = runner.run(token_id, pos) + logits = runner.run(token_id, pos, phase="prefill", prompt_no=prompt_no) if topk_print: token_text = tokenizer.decode_token(int(token_id)) print( @@ -601,7 +790,12 @@ def run_one_prompt_kv( print(tokenizer.decode_token(next_id), end="", flush=True) if next_id == eos_id or step == max_new_tokens - 1: break - logits = runner.run(next_id, len(ids) - 1) + logits = runner.run( + next_id, + len(ids) - 1, + phase="decode", + prompt_no=prompt_no, + ) print("\n=== Generation complete ===") decoded = tokenizer.decode(ids) @@ -728,6 +922,23 @@ def main() -> None: default=120, help="FVP time limit in seconds for each runner call.", ) + parser.add_argument( + "--no-ethosu-fast", + action="store_true", + help="Disable Ethos-U FVP fast mode. Required for NPU PMU profiling.", + ) + parser.add_argument( + "--profile-output", + type=Path, + default=None, + help="Write per-token Ethos-U85 PMU samples and summaries to .json or .csv.", + ) + parser.add_argument( + "--npu-frequency-mhz", + type=float, + default=None, + help="Optional assumed NPU frequency for NPU-only token/s estimates.", + ) parser.add_argument( "--full-logits", action="store_true", @@ -757,11 +968,25 @@ def main() -> None: pte_path = None if args.embedded_pte else args.pte if not args.embedded_pte and pte_path is None: raise ValueError("--pte is required unless --embedded-pte is set") + if args.profile_output is not None and not args.use_kv_cache: + raise ValueError("--profile-output requires --use-kv-cache") + if args.profile_output is not None and args.profile_output.suffix.lower() not in { + ".csv", + ".json", + }: + raise ValueError("--profile-output must end in .json or .csv") + if args.npu_frequency_mhz is not None and args.npu_frequency_mhz <= 0: + raise ValueError("--npu-frequency-mhz must be greater than zero") max_context_length = args.max_context_length or args.window if args.use_kv_cache: with KvFvpRunnerSession( - args.fvp, args.runner, pte_path, args.timeout + args.fvp, + args.runner, + pte_path, + args.timeout, + ethosu_fast=not (args.no_ethosu_fast or args.profile_output is not None), + collect_profile=args.profile_output is not None, ) as runner: for i, prompt in enumerate(prompts): run_one_prompt_kv( @@ -779,6 +1004,12 @@ def main() -> None: save_generations_path=args.save_generations, topk_print=not args.no_topk_print, ) + if args.profile_output is not None: + write_profile( + args.profile_output, + runner.samples, + args.npu_frequency_mhz, + ) else: with FvpRunnerSession(args.fvp, args.runner, pte_path, args.timeout) as runner: for i, prompt in enumerate(prompts): From 372aa3b99b162b8f6981db30d405f770ea108244 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Fri, 4 Sep 2026 09:52:06 -0700 Subject: [PATCH 046/190] Qualcomm: bounds-check the delegate (#22237) Qualcomm: bounds-check the delegate argument walk instead of running off the end execute() binds delegate arguments positionally. It walks the input and output tensor lists recovered from the context binary and, for every tensor the name prefixes mark as bindable, consumes one entry from args with a running counter. Nothing relates that counter to args.size(). So when the binary and the program disagree on the delegate signature -- a stale binary, or an AOT bug that publishes extra graph I/O -- the walk indexes past the end of the Span and dereferences whatever is there. In the case that prompted this, a context binary declaring 54 graph inputs and 56 graph outputs met a program passing 4 tensors, and the result was a null dereference at 0x8 with the two counts sitting in registers. Reading that back to a cause took days. Count the bindable tensors with the same prefix rules the loops use, then check once before either loop runs. A mismatch in either direction is fatal. A shortfall is the memory-safety case, since the walk reads past the end of args; a surplus does not read out of bounds, but it still means the binary and the program disagree on the signature, which is a defect either way. This started out warning on a surplus and was changed to fail at review request. Verified against 14 lowerings -- single I/O, multi-input, multi-output, partially-consumed multi-output, topk with both outputs used, a mutable buffer and a conv, each under an fp16 and a quantized spec -- with no false positives. Mutable buffers are excluded from the count by the same rule the binding loops use, so they do not create a surplus. Deliberately not included: a matching "input_" prefix filter on the input loop, for symmetry with the output loop. Inputs of a model built by from_context_binary carry names straight from the QNN converter with no such prefix, and the runtime only renames outputs (QnnManager.cpp SetName("output_" + tensor_name)). Filtering on it would skip every input of those models and leave the counter at zero when the output loop starts, writing outputs into input buffers. The count check gives the same protection without that risk. cc @cbilgin --------- Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> --- .../qualcomm/runtime/QnnExecuTorchBackend.cpp | 36 ++++++++++++++++++- 1 file changed, 35 insertions(+), 1 deletion(-) diff --git a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp index 8bbe047a967..76b67f911cd 100644 --- a/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp +++ b/backends/qualcomm/runtime/QnnExecuTorchBackend.cpp @@ -139,7 +139,41 @@ Error QnnExecuTorchBackend::execute( std::vector input_tensor_structs; std::vector output_tensor_structs; - int args_index = 0; + // The loops below walk the tensor lists recovered from the context binary and + // index args[] with a running counter, so the number of bindable tensors the + // binary declares has to agree with what the program passes. When it does not + // -- a stale binary, or an AOT bug that publishes extra graph I/O -- the walk + // runs off the end of the Span. Count first and fail with both numbers rather + // than reading out of bounds. + size_t bindable_inputs = 0; + for (const auto& input_tensor : input_tensors) { + const auto& name = input_tensor->GetName(); + if (name.find("mutbuf_") == std::string::npos) { + ++bindable_inputs; + } + } + size_t bindable_outputs = 0; + for (const auto& output_tensor : output_tensors) { + const auto& name = output_tensor->GetName(); + if (name.rfind("output_", 0) == 0 && + name.find("mutbuf_") == std::string::npos) { + ++bindable_outputs; + } + } + ET_CHECK_OR_RETURN_ERROR( + bindable_inputs + bindable_outputs == args.size(), + Internal, + "Method %s: the QNN context binary binds %zu tensors (%zu bindable inputs, " + "%zu bindable outputs) but ExecuTorch passed %zu arguments. The binary and " + "the program disagree on the delegate signature; the model has to be " + "re-exported.", + method_name.c_str(), + bindable_inputs + bindable_outputs, + bindable_inputs, + bindable_outputs, + args.size()); + + size_t args_index = 0; input_tensor_structs.reserve(input_tensors.size()); for (const auto& input_tensor : input_tensors) { if (input_tensor->GetName().find("mutbuf_") == std::string::npos) { From 6975be19f91756409511dc5f6799dcfc49b81aba Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 09:56:49 -0700 Subject: [PATCH 047/190] Run a recipe's edge-manager passes after partitioning (#22474) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary `LoweringRecipe.edge_manager_transform_passes` has been accepted since the field was added, but nothing ran it in a default pipeline: the stage that applies it was never scheduled, and it could only follow `TO_EDGE`, never the partitioning stage. A recipe declaring such passes had them silently dropped. Both halves are needed together. Scheduling the stage without allowing it after `TO_EDGE_TRANSFORM_AND_LOWER` fails pipeline validation, and allowing it there without scheduling leaves the passes unreachable. Running after partitioning is the point rather than an incidental ordering. The passes a backend wants here act on what the partitioner left outside the delegates -- rewriting the boundary quantize/dequantize pairs, for instance -- so running them earlier would either find nothing or feed the partitioner a graph it no longer recognises. The commented-out predecessor this enables was raised in #21516; NXP Neutron wants the same hook. Authored with Claude Code. --- Stack created with GitHub Stacks CLI • Give Feedback 💬 --- export/export.py | 18 ++++++++++------- export/stages.py | 2 +- export/tests/test_export_session.py | 30 +++++++++++++++++++++++++++++ export/tests/test_export_stages.py | 9 +++++++++ 4 files changed, 51 insertions(+), 8 deletions(-) diff --git a/export/export.py b/export/export.py index f569a4196ee..fa8e534a430 100644 --- a/export/export.py +++ b/export/export.py @@ -283,13 +283,17 @@ def _get_default_pipeline(self) -> List[StageType]: if self._input_model_type != "ExportedProgram": stages.append(StageType.TORCH_EXPORT) - # Always include edge and executorch stages - stages.extend( - [ - StageType.TO_EDGE_TRANSFORM_AND_LOWER, - StageType.TO_EXECUTORCH, - ] - ) + stages.append(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + + # This is the only stage that runs edge_manager_transform_passes, so a + # recipe declaring them would otherwise have them silently dropped. + if ( + self._lowering_recipe + and self._lowering_recipe.edge_manager_transform_passes + ): + stages.append(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM) + + stages.append(StageType.TO_EXECUTORCH) return stages diff --git a/export/stages.py b/export/stages.py index 3d2c8c86a4d..c08a72cdba7 100644 --- a/export/stages.py +++ b/export/stages.py @@ -702,7 +702,7 @@ def stage_type(self) -> str: def valid_predecessor_stages(self) -> List["StageType"]: return [ StageType.TO_EDGE, - # StageType.TO_EDGE_TRANSFORM_AND_LOWER, # TODO + StageType.TO_EDGE_TRANSFORM_AND_LOWER, ] @property diff --git a/export/tests/test_export_session.py b/export/tests/test_export_session.py index 2e9112f7995..93902fd86d7 100644 --- a/export/tests/test_export_session.py +++ b/export/tests/test_export_session.py @@ -642,6 +642,36 @@ def test_dict_exported_program_input_type_detection(self) -> None: pipeline = session._get_default_pipeline() self.assertNotIn(StageType.TORCH_EXPORT, pipeline) + def test_edge_manager_transform_passes_get_their_stage(self) -> None: + # Nothing else in the default pipeline runs them, so without this the + # recipe's passes are accepted and then silently never applied. + session = ExportSession( + model=self.model, + example_inputs=[self.example_inputs], + export_recipe=ExportRecipe( + name="t", + lowering_recipe=LoweringRecipe( + edge_manager_transform_passes=[lambda epm: []] + ), + ), + ) + pipeline = session._get_default_pipeline() + self.assertIn(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, pipeline) + self.assertGreater( + pipeline.index(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM), + pipeline.index(StageType.TO_EDGE_TRANSFORM_AND_LOWER), + ) + + def test_no_transform_passes_means_no_stage(self) -> None: + session = ExportSession( + model=self.model, + example_inputs=[self.example_inputs], + export_recipe=self.recipe, + ) + self.assertNotIn( + StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, session._get_default_pipeline() + ) + def test_example_inputs_required_for_nn_module(self) -> None: """Test that example_inputs are required for nn.Module.""" with self.assertRaises(ValueError) as cm: diff --git a/export/tests/test_export_stages.py b/export/tests/test_export_stages.py index 46eb000b743..351ae748b86 100644 --- a/export/tests/test_export_stages.py +++ b/export/tests/test_export_stages.py @@ -1231,6 +1231,15 @@ def _manager(self) -> Mock: manager.exported_program.return_value = Mock() return manager + def test_edge_program_manager_stage_may_follow_partitioning(self) -> None: + # The point of the stage for a delegate recipe: its passes act on what + # the partitioner left outside the delegates, so it has to be able to + # run after TO_EDGE_TRANSFORM_AND_LOWER and not only after TO_EDGE. + self.assertEqual( + set(EdgeProgramManagerTransformStage().valid_predecessor_stages), + {StageType.TO_EDGE, StageType.TO_EDGE_TRANSFORM_AND_LOWER}, + ) + def test_edge_program_manager_stage_skips_empty_transform(self) -> None: manager = self._manager() stage = EdgeProgramManagerTransformStage( From fcc3eb5a20449c8e2e549ecb7a314a8d3a3ac87a Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 10:09:06 -0700 Subject: [PATCH 048/190] Cortex-M: add explicit layout pooling kernels (#22379) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Add and register NHWC average- and max-pooling implementations using the shared layout-aware CMSIS-NN validation and configuration path. ### Why separate operators The legacy pooling symbols consume logical NCHW tensors represented with channels-last dim order. Explicit layout consumes ordinary contiguous NHWC tensors. Separate experimental symbols keep those contracts visible and prevent silent interpretation of one representation as the other, while shared helpers retain one implementation of the common CMSIS-NN mechanics. The registration entries and FVP selected-op list live beside these C++ kernels so this PR is independently loadable and testable through serialized programs. AI-assisted: Codex. --- Stack created with GitHub Stacks CLI • Give Feedback 💬 --- backends/cortex_m/ops/cortex_m_ops_common.h | 42 ++++-- .../cortex_m/ops/op_quantized_avg_pool2d.cpp | 63 +++++++- .../cortex_m/ops/op_quantized_max_pool2d.cpp | 70 ++++++++- backends/cortex_m/ops/operators.py | 139 ++++++++++++++++++ backends/cortex_m/ops/operators.yaml | 13 ++ .../cortex_m/passes/scratch_buffer_sizes.py | 14 +- backends/cortex_m/test/build_test_runner.sh | 2 + backends/cortex_m/test/ops/test_nhwc_pool.py | 67 +++++++++ 8 files changed, 390 insertions(+), 20 deletions(-) create mode 100644 backends/cortex_m/test/ops/test_nhwc_pool.py diff --git a/backends/cortex_m/ops/cortex_m_ops_common.h b/backends/cortex_m/ops/cortex_m_ops_common.h index 1e2eeb99d1d..bcaed0a1bc7 100644 --- a/backends/cortex_m/ops/cortex_m_ops_common.h +++ b/backends/cortex_m/ops/cortex_m_ops_common.h @@ -209,7 +209,7 @@ inline bool prepare_cmsis_pool2d_config( int64_t activation_min, int64_t activation_max, CmsisPool2DConfig& config, - bool require_channels_last = true, + ActivationLayout layout, bool allow_ceil_mode = false) { if (input.dim() != 4 || output.dim() != 4) { ET_LOG(Error, "%s: tensors must be 4-D", op_name); @@ -224,7 +224,9 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (input.size(0) != output.size(0) || input.size(1) != output.size(1)) { + const int64_t channel_dim = layout == ActivationLayout::NHWCLogical ? 3 : 1; + if (input.size(0) != output.size(0) || + input.size(channel_dim) != output.size(channel_dim)) { ET_LOG( Error, "%s: batch and channel dimensions must match between input and output", @@ -233,13 +235,21 @@ inline bool prepare_cmsis_pool2d_config( return false; } - if (require_channels_last) { - if (!is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { - ET_LOG( - Error, "%s: tensors must use channels_last dimension order", op_name); + if (layout == ActivationLayout::NHWCLogical) { + if (!executorch::runtime::is_contiguous_dim_order( + input.dim_order().data(), input.dim_order().size()) || + !executorch::runtime::is_contiguous_dim_order( + output.dim_order().data(), output.dim_order().size())) { + ET_LOG(Error, "%s: tensors must use contiguous dimension order", op_name); context.fail(Error::InvalidArgument); return false; } + } else if ( + !is_channels_last_tensor(input) || !is_channels_last_tensor(output)) { + ET_LOG( + Error, "%s: tensors must use channels_last dimension order", op_name); + context.fail(Error::InvalidArgument); + return false; } auto check_tuple_len = [&](const Int64ArrayRef& arr, @@ -318,19 +328,29 @@ inline bool prepare_cmsis_pool2d_config( return false; } + const int64_t height_dim = layout == ActivationLayout::NHWCLogical ? 1 : 2; + const int64_t width_dim = layout == ActivationLayout::NHWCLogical ? 2 : 3; int32_t batch, channels, input_h, input_w, output_h, output_w; if (!check_int32_within_range( context, op_name, input.size(0), "input batch", batch) || !check_int32_within_range( - context, op_name, input.size(1), "input channels", channels) || + context, + op_name, + input.size(channel_dim), + "input channels", + channels) || !check_int32_within_range( - context, op_name, input.size(2), "input height", input_h) || + context, op_name, input.size(height_dim), "input height", input_h) || !check_int32_within_range( - context, op_name, input.size(3), "input width", input_w) || + context, op_name, input.size(width_dim), "input width", input_w) || !check_int32_within_range( - context, op_name, output.size(2), "output height", output_h) || + context, + op_name, + output.size(height_dim), + "output height", + output_h) || !check_int32_within_range( - context, op_name, output.size(3), "output width", output_w)) { + context, op_name, output.size(width_dim), "output width", output_w)) { return false; } diff --git a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp index 39b6432c45a..ba8b51d18fd 100644 --- a/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_avg_pool2d.cpp @@ -1,5 +1,7 @@ /* * Copyright 2025-2026 Arm Limited and/or its affiliates. + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. @@ -66,7 +68,7 @@ bool validate_avg_pool2d_output_size( } // namespace // cppcheck-suppress unusedFunction -Tensor& quantized_avg_pool2d_out( +static Tensor& quantized_avg_pool2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef kernel_size, @@ -77,6 +79,7 @@ Tensor& quantized_avg_pool2d_out( const int64_t multiplier, const int64_t shift, const Tensor& scratch, + ActivationLayout layout, Tensor& out) { constexpr int32_t activation_min = std::numeric_limits::min(); constexpr int32_t activation_max = std::numeric_limits::max(); @@ -97,7 +100,7 @@ Tensor& quantized_avg_pool2d_out( activation_min, activation_max, pool_config, - true, + layout, true)) { return out; } @@ -153,5 +156,61 @@ Tensor& quantized_avg_pool2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_avg_pool2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const bool ceil_mode, + const int64_t zero_point, + const int64_t multiplier, + const int64_t shift, + const Tensor& scratch, + Tensor& out) { + return quantized_avg_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_avg_pool2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const bool ceil_mode, + const int64_t zero_point, + const int64_t multiplier, + const int64_t shift, + const Tensor& scratch, + Tensor& out) { + return quantized_avg_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp index ca1b00ff340..75986e99d9c 100644 --- a/backends/cortex_m/ops/op_quantized_max_pool2d.cpp +++ b/backends/cortex_m/ops/op_quantized_max_pool2d.cpp @@ -1,5 +1,7 @@ /* - * Copyright 2026 Arm Limited and/or its affiliates. + * Copyright 2026 Arm Limited + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. * * This source code is licensed under the BSD-style license found in the * LICENSE file in the root directory of this source tree. @@ -11,7 +13,7 @@ namespace cortex_m { namespace native { // cppcheck-suppress unusedFunction -Tensor& quantized_max_pool2d_out( +static Tensor& quantized_max_pool2d_out_impl( KernelRuntimeContext& context, const Tensor& input, const Int64ArrayRef kernel_size, @@ -23,6 +25,7 @@ Tensor& quantized_max_pool2d_out( const int64_t output_zero_point, const int64_t activation_min, const int64_t activation_max, + ActivationLayout layout, Tensor& out) { CmsisPool2DConfig pool_config; if (!prepare_cmsis_pool2d_config( @@ -37,7 +40,8 @@ Tensor& quantized_max_pool2d_out( ceil_mode, activation_min, activation_max, - pool_config)) { + pool_config, + layout)) { return out; } @@ -95,5 +99,65 @@ Tensor& quantized_max_pool2d_out( return out; } +// cppcheck-suppress unusedFunction +Tensor& quantized_max_pool2d_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const bool ceil_mode, + const int64_t input_zero_point, + const int64_t output_zero_point, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_max_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ActivationLayout::NCHWLogical, + out); +} + +// cppcheck-suppress unusedFunction +Tensor& quantized_max_pool2d_nhwc_out( + KernelRuntimeContext& context, + const Tensor& input, + const Int64ArrayRef kernel_size, + const Int64ArrayRef stride, + const Int64ArrayRef padding, + const Int64ArrayRef dilation, + const bool ceil_mode, + const int64_t input_zero_point, + const int64_t output_zero_point, + const int64_t activation_min, + const int64_t activation_max, + Tensor& out) { + return quantized_max_pool2d_out_impl( + context, + input, + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ActivationLayout::NHWCLogical, + out); +} + } // namespace native } // namespace cortex_m diff --git a/backends/cortex_m/ops/operators.py b/backends/cortex_m/ops/operators.py index ea67600cff6..56af4cbe9fb 100644 --- a/backends/cortex_m/ops/operators.py +++ b/backends/cortex_m/ops/operators.py @@ -1642,6 +1642,73 @@ def quantized_avg_pool2d_impl( return output.to(torch.int8) +lib.define( + "quantized_avg_pool2d_nhwc(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "bool ceil_mode, int zero_point, int multiplier, int shift, " + "Tensor scratch) -> Tensor" +) +lib.define( + "quantized_avg_pool2d_nhwc.out(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "bool ceil_mode, int zero_point, int multiplier, int shift, " + "Tensor scratch, *, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_avg_pool2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_avg_pool2d_nhwc_meta( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + ceil_mode: bool, + zero_point: int, + multiplier: int, + shift: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_avg_pool2d_meta( + input.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_avg_pool2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_avg_pool2d_nhwc_impl( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + ceil_mode: bool, + zero_point: int, + multiplier: int, + shift: int, + scratch: torch.Tensor, +) -> torch.Tensor: + nchw = quantized_avg_pool2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + kernel_size, + stride, + padding, + ceil_mode, + zero_point, + multiplier, + shift, + scratch, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + # =================================================================== # QUANTIZED MAX POOL2D OPERATION DEFINITION # =================================================================== @@ -1797,3 +1864,75 @@ def quantized_max_pool2d_impl( ) result = torch.clamp(result, activation_min, activation_max) return result.to(torch.int8).contiguous(memory_format=torch.channels_last) + + +lib.define( + "quantized_max_pool2d_nhwc(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "int[] dilation, bool ceil_mode, int input_zero_point, " + "int output_zero_point, int activation_min, int activation_max) -> Tensor" +) +lib.define( + "quantized_max_pool2d_nhwc.out(" + "Tensor input, int[] kernel_size, int[] stride, int[] padding, " + "int[] dilation, bool ceil_mode, int input_zero_point, " + "int output_zero_point, int activation_min, int activation_max, " + "*, Tensor(a!) out) -> Tensor(a!)" +) + + +@register_fake("cortex_m::quantized_max_pool2d_nhwc") # type: ignore[misc] +@experimental(_EXPLICIT_LAYOUT_EXPERIMENTAL) # type: ignore[misc] +def quantized_max_pool2d_nhwc_meta( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + input_zero_point: int, + output_zero_point: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + nchw = quantized_max_pool2d_meta( + input.permute(0, 3, 1, 2), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ) + return nchw.permute(0, 2, 3, 1).contiguous() + + +@impl(lib, "quantized_max_pool2d_nhwc", "CompositeExplicitAutograd") # type: ignore[misc] +def quantized_max_pool2d_nhwc_impl( + input: torch.Tensor, + kernel_size: Sequence[int], + stride: Sequence[int], + padding: Sequence[int], + dilation: Sequence[int], + ceil_mode: bool, + input_zero_point: int, + output_zero_point: int, + activation_min: int, + activation_max: int, +) -> torch.Tensor: + nchw = quantized_max_pool2d_impl( + input.permute(0, 3, 1, 2).contiguous(), + kernel_size, + stride, + padding, + dilation, + ceil_mode, + input_zero_point, + output_zero_point, + activation_min, + activation_max, + ) + return nchw.permute(0, 2, 3, 1).contiguous() diff --git a/backends/cortex_m/ops/operators.yaml b/backends/cortex_m/ops/operators.yaml index a2912bfd660..ebe7590c2f4 100644 --- a/backends/cortex_m/ops/operators.yaml +++ b/backends/cortex_m/ops/operators.yaml @@ -119,12 +119,25 @@ kernels: - arg_meta: null kernel_name: cortex_m::quantized_avg_pool2d_out + +- func: cortex_m::quantized_avg_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, bool ceil_mode, int zero_point, int multiplier, int shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_avg_pool2d_nhwc_out + - func: cortex_m::quantized_max_pool2d.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, int input_zero_point, int output_zero_point, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: - arg_meta: null kernel_name: cortex_m::quantized_max_pool2d_out +- func: cortex_m::quantized_max_pool2d_nhwc.out(Tensor input, int[] kernel_size, int[] stride, int[] padding, int[] dilation, bool ceil_mode, int input_zero_point, int output_zero_point, int activation_min, int activation_max, *, Tensor(a!) out) -> Tensor(a!) + variants: function + kernels: + - arg_meta: null + kernel_name: cortex_m::quantized_max_pool2d_nhwc_out + - func: cortex_m::quantized_batch_matmul.out(Tensor lhs, int lhs_zero_point, Tensor rhs_transposed, int rhs_zero_point, int output_zero_point, int output_multiplier, int output_shift, Tensor scratch, *, Tensor(a!) out) -> Tensor(a!) variants: function kernels: diff --git a/backends/cortex_m/passes/scratch_buffer_sizes.py b/backends/cortex_m/passes/scratch_buffer_sizes.py index 964b5530e13..65a3a178757 100644 --- a/backends/cortex_m/passes/scratch_buffer_sizes.py +++ b/backends/cortex_m/passes/scratch_buffer_sizes.py @@ -259,13 +259,16 @@ def cmsis_nn_transpose_conv_buffer_size( def cmsis_nn_avgpool_buffer_size( backend: cmsis_nn.Backend, pool_node: torch.fx.Node, + nhwc_logical: bool = False, ) -> list[int]: x = cast(torch.fx.Node, pool_node.args[0]) - # Input is NCHW (PyTorch); CMSIS-NN's avgpool buffer sizer only needs the - # input channel count and output width. - _, c_in, _, _ = _shape_from_node(x) - _, _, _, out_w = _shape_from_node(pool_node) + if nhwc_logical: + _, _, _, c_in = _shape_from_node(x) + _, _, out_w, _ = _shape_from_node(pool_node) + else: + _, c_in, _, _ = _shape_from_node(x) + _, _, _, out_w = _shape_from_node(pool_node) return [ int( @@ -294,6 +297,9 @@ def cmsis_nn_avgpool_buffer_size( cmsis_nn_transpose_conv_buffer_size, nhwc_logical=True ), exir_ops.edge.cortex_m.quantized_avg_pool2d.default: cmsis_nn_avgpool_buffer_size, + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default: partial( + cmsis_nn_avgpool_buffer_size, nhwc_logical=True + ), } diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index 47b27da137e..1178e27aa27 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -73,7 +73,9 @@ ops_list=( cortex_m::quantized_transpose_conv2d.out cortex_m::quantized_transpose_conv2d_nhwc.out cortex_m::quantized_avg_pool2d.out + cortex_m::quantized_avg_pool2d_nhwc.out cortex_m::quantized_max_pool2d.out + cortex_m::quantized_max_pool2d_nhwc.out cortex_m::quantized_batch_matmul.out ) diff --git a/backends/cortex_m/test/ops/test_nhwc_pool.py b/backends/cortex_m/test/ops/test_nhwc_pool.py new file mode 100644 index 00000000000..4377df62f9f --- /dev/null +++ b/backends/cortex_m/test/ops/test_nhwc_pool.py @@ -0,0 +1,67 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.cortex_m.test.ops.nhwc_test_utils import ( + int8_values, + run_on_fvp, +) +from executorch.exir.dialects._ops import ops as exir_ops + +# Direct-op coverage is temporary until CortexMTester uses explicit layout by default. + + +class AvgPool2dNhwc(torch.nn.Module): + def forward(self, x, scratch): + return torch.ops.cortex_m.quantized_avg_pool2d_nhwc.default( + x, + [2, 3], + [2, 1], + [0, 1], + False, + 0, + 1 << 30, + 1, + scratch, + ) + + +class MaxPool2dNhwc(torch.nn.Module): + def forward(self, x): + return torch.ops.cortex_m.quantized_max_pool2d_nhwc.default( + x, + [2, 3], + [2, 1], + [0, 1], + [1, 1], + False, + 0, + 0, + -128, + 127, + ) + + +def test_avg_pool2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + AvgPool2dNhwc(), + int8_values((1, 7, 9, 3)), + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default, + cortex_m_target, + 1, + atol=1, + ) + + +def test_max_pool2d_nhwc_runs_on_fvp(cortex_m_target): + run_on_fvp( + MaxPool2dNhwc(), + int8_values((1, 7, 9, 3)), + exir_ops.edge.cortex_m.quantized_max_pool2d_nhwc.default, + cortex_m_target, + 0, + ) From 9ae78daaca89f544f6223dec0c1af7b125440891 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 11:44:46 -0700 Subject: [PATCH 049/190] Cortex-M: add opt-in explicit-layout lowering (#22544) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Add the opt-in explicit-layout lowering pipeline for Cortex-M. `CortexMQuantizer(use_explicit_layout=True)` selects the explicit convolution eligibility checks and includes Conv1d annotation before it is widened to Conv2d. `CortexMPassManager(use_explicit_layout=True)` selects a separate, ordered pass list that forms NHWC spatial regions, removes redundant local layout copies, validates the resulting graph, and lowers to the experimental NHWC Cortex-M operators. The default quantizer checks and pass list remain unchanged. The mode is selected only at the quantization and lowering boundaries; it is not threaded through individual passes or the target configuration. Unsupported quantized spatial anchors fail closed instead of mixing legacy and explicit operator families. ### Test plan `pytest --config-file=backends/arm/test/pytest.ini backends/cortex_m/test/test_explicit_layout_pipeline.py backends/cortex_m/test/misc/test_quantization.py backends/cortex_m/test/misc/test_quantizer_reporter.py` `lintrunner -m origin/main` AI-assisted: Codex. --- Stack created with GitHub Stacks CLI • Give Feedback 💬 --- backends/cortex_m/passes/BUCK | 8 + .../cortex_m/passes/aten_to_cortex_m_pass.py | 84 ++++++++-- .../cortex_m/passes/cortex_m_pass_manager.py | 49 +++++- .../cortex_m/passes/explicit_layout_pass.py | 156 +++++++++++++++++ .../cortex_m/quantizer/pattern_checkers.py | 21 ++- backends/cortex_m/quantizer/quantizer.py | 32 +++- .../cortex_m/quantizer/quantizer_support.py | 49 ++++++ backends/cortex_m/test/targets.bzl | 35 ++++ .../test/test_explicit_layout_pipeline.py | 158 ++++++++++++++++++ 9 files changed, 569 insertions(+), 23 deletions(-) create mode 100644 backends/cortex_m/passes/explicit_layout_pass.py create mode 100644 backends/cortex_m/test/test_explicit_layout_pipeline.py diff --git a/backends/cortex_m/passes/BUCK b/backends/cortex_m/passes/BUCK index f33ddbf9cf3..4d758703cb4 100644 --- a/backends/cortex_m/passes/BUCK +++ b/backends/cortex_m/passes/BUCK @@ -35,6 +35,7 @@ fbcode_target(_kind = runtime.python_library, "cortex_m_pass_manager.py", "decompose_hardswish_pass.py", "decompose_mean_pass.py", + "explicit_layout_pass.py", "matmul_to_bmm_pass.py", "quantized_clamp_activation_pass.py", "scratch_buffer_sizes.py", @@ -48,8 +49,15 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/cortex_m/passes:passes_utils", "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", "//executorch/backends/transforms:aten_to_dialect_pass", + "//executorch/backends/transforms:canonicalize_view_copy_permute_pass", + "//executorch/backends/transforms:channels_last_layout", + "//executorch/backends/transforms:channels_last_ops", + "//executorch/backends/transforms:convert_conv1d_to_conv2d_pass", "//executorch/backends/transforms:remove_getitem_op", + "//executorch/backends/transforms:remove_permutes_around_elementwise_ops", "//executorch/backends/transforms:replace_scalar_with_tensor", + "//executorch/backends/transforms:replace_ops_with_channels_last_variants", + "//executorch/backends/transforms:replace_squeeze_unsqueeze_with_view", "//executorch/backends/transforms:utils", "//executorch/exir:lib", "//executorch/exir:pass_base", diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index af60df686f5..7f0cd2dd434 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -10,6 +10,7 @@ from typing import cast, Optional import executorch.backends.cortex_m.ops.operators # noqa +import executorch.backends.transforms.channels_last_ops # noqa: F401 import executorch.exir as exir import torch import torch.fx @@ -75,6 +76,10 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: ) for node in result.graph_module.graph.nodes: + if getattr(node.target, "namespace", None) == "channels_last": + raise RuntimeError( + f"Cortex-M lowering left {node.target} in the graph." + ) self._initialize_alloc_node_size(node) return PassResult(result.graph_module, result.modified or max_pool_modified) @@ -448,6 +453,9 @@ def _get_linear_replacement( return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_linear.default, args) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.convolution.default +) @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.convolution.default) def _get_convolution_replacement( node: Node, dialect_pass: AtenToDialectPass @@ -455,6 +463,8 @@ def _get_convolution_replacement( if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.convolution.default + exported_program = dialect_pass.exported_program conv_args = node.args ( @@ -612,7 +622,11 @@ def _get_convolution_replacement( scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default, + ( + exir_ops.edge.cortex_m.quantized_depthwise_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default + ), depthwise_args, ) @@ -634,7 +648,14 @@ def _get_convolution_replacement( output_qmax, scratch, ) - return DialectNodeSpec(exir_ops.edge.cortex_m.quantized_conv2d.default, conv2d_args) + return DialectNodeSpec( + ( + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_conv2d.default + ), + conv2d_args, + ) def _get_transpose_conv2d_replacement( @@ -646,6 +667,7 @@ def _get_transpose_conv2d_replacement( if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.convolution.default exported_program = dialect_pass.exported_program conv_t_args = node.args ( @@ -758,7 +780,12 @@ def _get_transpose_conv2d_replacement( output_scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_transpose_conv2d.default, new_args + ( + exir_ops.edge.cortex_m.quantized_transpose_conv2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_transpose_conv2d.default + ), + new_args, ) @@ -825,12 +852,16 @@ def _get_bmm_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.avg_pool2d.default) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.avg_pool2d.default +) def _get_avg_pool2d_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: if not _has_qparams(node): return None + explicit_nhwc = node.target == exir_ops.edge.channels_last.avg_pool2d.default exported_program = dialect_pass.exported_program pool_args = node.args kernel_size = cast(list[int], pool_args[1]) @@ -851,8 +882,11 @@ def _get_avg_pool2d_replacement( avg_padding = padding if count_include_pad: pad_h, pad_w = padding - input_tensor = get_first_fake_tensor(input_node) - pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) + if explicit_nhwc: + pre_pad = post_pad = [0, pad_h, pad_w, 0] + else: + input_tensor = get_first_fake_tensor(input_node) + pre_pad = post_pad = to_physical_order([0, 0, pad_h, pad_w], input_tensor) with node.graph.inserting_before(node): input_node = node.graph.create_node( "call_function", @@ -875,7 +909,12 @@ def _get_avg_pool2d_replacement( scratch, ) return DialectNodeSpec( - exir_ops.edge.cortex_m.quantized_avg_pool2d.default, new_args + ( + exir_ops.edge.cortex_m.quantized_avg_pool2d_nhwc.default + if explicit_nhwc + else exir_ops.edge.cortex_m.quantized_avg_pool2d.default + ), + new_args, ) @@ -1108,10 +1147,14 @@ def _get_softmax_replacement( @AtenToCortexMPass.register_dialect_substitution(exir_ops.edge.aten.max_pool2d.default) +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.max_pool2d.default +) def _get_max_pool2d_replacement( node: Node, dialect_pass: AtenToDialectPass ) -> DialectNodeSpec | None: del dialect_pass + explicit_nhwc = node.target == exir_ops.edge.channels_last.max_pool2d.default input_qparams = node.meta.get("input_qparams", {}).get(0) cortex_m_meta = node.meta.get("custom", {}).get("cortex_m", {}) if input_qparams is None or cortex_m_meta.get("skip_quantized_max_pool2d", False): @@ -1169,6 +1212,12 @@ def _get_max_pool2d_replacement( activation_min, activation_max, ) + if explicit_nhwc: + quantized_op = getattr( + exir_ops.edge.cortex_m, "quantized_max_pool2d_nhwc", None + ) + if quantized_op is None: + return None return DialectNodeSpec(quantized_op.default, args) @@ -1194,6 +1243,14 @@ def _get_maximum_replacement( return DialectNodeSpec(exir_ops.edge.cortex_m.maximum.default, node.args) +def _transpose_spec(node: Node, input_tensor) -> DialectNodeSpec: + rank = len(input_tensor.shape) + perms = [p % rank for p in cast(tuple[int, ...], node.args[1])] + return DialectNodeSpec( + exir_ops.edge.cortex_m.transpose.default, (node.args[0], perms) + ) + + @AtenToCortexMPass.register_dialect_substitution( exir_ops.edge.aten.permute_copy.default ) @@ -1204,12 +1261,17 @@ def _get_permute_replacement( input_tensor = _get_input_tensor_data(node) if input_tensor.dtype != torch.int8: return None + return _transpose_spec(node, input_tensor) - rank = len(input_tensor.shape) - perms = [p % rank for p in cast(tuple[int, ...], node.args[1])] - return DialectNodeSpec( - exir_ops.edge.cortex_m.transpose.default, (node.args[0], perms) - ) + +@AtenToCortexMPass.register_dialect_substitution( + exir_ops.edge.channels_last.permute_copy.default +) +def _get_layout_permute_replacement( + node: Node, dialect_pass: AtenToDialectPass +) -> DialectNodeSpec | None: + del dialect_pass + return _transpose_spec(node, _get_input_tensor_data(node)) @AtenToCortexMPass.register_dialect_substitution( diff --git a/backends/cortex_m/passes/cortex_m_pass_manager.py b/backends/cortex_m/passes/cortex_m_pass_manager.py index 892baf136ed..a15a7262660 100644 --- a/backends/cortex_m/passes/cortex_m_pass_manager.py +++ b/backends/cortex_m/passes/cortex_m_pass_manager.py @@ -13,10 +13,19 @@ ScalarsToAttributePass, ) from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.transforms.convert_conv1d_to_conv2d_pass import ( + ConvertConv1dToConv2dPass, +) from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass +from executorch.backends.transforms.remove_permutes_around_elementwise_ops import ( + RemovePermutesAroundElementwiseOps, +) from executorch.backends.transforms.replace_scalar_with_tensor import ( ReplaceScalarWithTensorArgPass, ) +from executorch.backends.transforms.replace_squeeze_unsqueeze_with_view import ( + ReplaceSqueezeAndUnsqueezeWithViewPass, +) from executorch.exir.pass_base import ExportPass from executorch.exir.pass_manager import PassManager from executorch.exir.program._program import _transform, lift_constant_tensor_pass @@ -27,6 +36,11 @@ from .clamp_hardswish_pass import ClampHardswishPass from .decompose_hardswish_pass import DecomposeHardswishPass from .decompose_mean_pass import DecomposeMeanPass +from .explicit_layout_pass import ( + CortexMCanonicalizeViewCopyPermutePass, + CortexMReplaceOpsWithChannelsLastVariants, + ValidateCortexMExplicitLayoutPass, +) from .matmul_to_bmm_pass import MatmulToBmmPass from .quantized_clamp_activation_pass import QuantizedClampActivationPass from .replace_quant_nodes_pass import ReplaceQuantNodesPass @@ -35,7 +49,7 @@ class CortexMPassManager(PassManager): - pass_list: list[PassClass] = [ + legacy_pass_list: list[PassClass] = [ # Run before folding so qparams attach to max_pool2d values, not tuple + getitem. RemoveGetItemPass, FoldAndAnnotateQParamsPass, @@ -47,6 +61,26 @@ class CortexMPassManager(PassManager): AtenToCortexMPass, ] + explicit_layout_pass_list: list[PassClass] = [ + RemoveGetItemPass, + FoldAndAnnotateQParamsPass, + ReplaceScalarWithTensorArgPass, + ActivationFusionPass, + QuantizedClampActivationPass, + DecomposeHardswishPass, + ConvertConv1dToConv2dPass, + CortexMReplaceOpsWithChannelsLastVariants, + ReplaceSqueezeAndUnsqueezeWithViewPass, + CortexMCanonicalizeViewCopyPermutePass, + RemovePermutesAroundElementwiseOps, + CortexMCanonicalizeViewCopyPermutePass, + ValidateCortexMExplicitLayoutPass, + ReplaceQuantNodesPass, + AtenToCortexMPass, + ] + + pass_list = legacy_pass_list + pass_list_transform_for_annotation: list[PassClass] = [ ScalarsToAttributePass, ReplaceScalarWithTensorArgPass, @@ -61,6 +95,7 @@ def __init__( exported_program: ExportedProgram | None, passes: Optional[list[PassClass]] = None, target_config: Optional[CortexMTargetConfig] = None, + use_explicit_layout: bool = False, ) -> None: """Initialize the Cortex-M pass manager. @@ -69,17 +104,25 @@ def __init__( before calling ``transform()``; may be ``None`` for callers that only use ``transform_for_annotation()``. passes: Optional override of the pass list. Defaults to - ``CortexMPassManager.pass_list``. + the legacy or explicit-layout pass list selected by + ``use_explicit_layout``. target_config: Compilation target for passes that need it. Defaults to ``CortexMTargetConfig(cpu=CortexM.M55)``, which resolves through cmsis_nn to the MVE backend — matching the pre-config historical behaviour. + use_explicit_layout: Select the experimental explicit-layout pass + sequence. Legacy lowering remains the default. """ super().__init__(passes=[]) self.exported_program = exported_program # PassManager.passes is typed as callables; this manager stores pass classes which are initialized at transform time with the exported_program. + default_passes = ( + self.explicit_layout_pass_list + if use_explicit_layout + else self.legacy_pass_list + ) self.passes: list[PassClass] = ( # type: ignore[assignment] - passes if passes is not None else self.pass_list # type: ignore[assignment] + passes if passes is not None else default_passes # type: ignore[assignment] ) self.target_config: CortexMTargetConfig = target_config or CortexMTargetConfig( cpu=CortexM.M55 diff --git a/backends/cortex_m/passes/explicit_layout_pass.py b/backends/cortex_m/passes/explicit_layout_pass.py new file mode 100644 index 00000000000..19481bd0314 --- /dev/null +++ b/backends/cortex_m/passes/explicit_layout_pass.py @@ -0,0 +1,156 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import executorch.backends.transforms.channels_last_ops # noqa: F401 + +import torch +from executorch.backends.cortex_m.passes.passes_utils import ( + coerce_int_pair, + skips_quantized_max_pool2d, +) +from executorch.backends.transforms.canonicalize_view_copy_permute_pass import ( + CanonicalizeViewCopyPermutePass, +) +from executorch.backends.transforms.channels_last_layout import ( + LAYOUT_PERMUTE_COPY, + PERMUTE_COPY_TARGETS, +) +from executorch.backends.transforms.replace_ops_with_channels_last_variants import ( + ChannelsLastOpSpec, + ReplaceOpsWithChannelsLastVariants, +) +from executorch.exir import ExportedProgram +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node +from torch.fx.node import Target + + +def _is_rank4(node: Node) -> bool: + return len(node.meta["val"].shape) == 4 + + +def _has_per_tensor_qparam(node: Node, key: str, index: int) -> bool: + qparam = node.meta.get(key, {}).get(index) + return qparam is not None and not getattr(qparam, "per_channel", False) + + +def _has_input_and_output_qparams(node: Node) -> bool: + return _has_per_tensor_qparam(node, "input_qparams", 0) and _has_per_tensor_qparam( + node, "output_qparams", 0 + ) + + +def _supports_avg_pool2d(node: Node) -> bool: + divisor_override = node.args[6] if len(node.args) > 6 else None + return ( + _is_rank4(node) + and _has_input_and_output_qparams(node) + and divisor_override is None + ) + + +def _supports_max_pool2d(node: Node) -> bool: + if not _is_rank4(node) or not _has_per_tensor_qparam(node, "input_qparams", 0): + return False + if skips_quantized_max_pool2d(node): + return False + + dilation = coerce_int_pair(node.args[4] if len(node.args) > 4 else None, (1, 1)) + ceil_mode = bool(node.args[5]) if len(node.args) > 5 else False + if dilation != (1, 1) or ceil_mode: + return False + + input_qparams = node.meta["input_qparams"][0] + output_qparams = node.meta.get("output_qparams", {}).get(0) + return output_qparams is None or ( + not getattr(output_qparams, "per_channel", False) + and abs(float(input_qparams.scale) - float(output_qparams.scale)) <= 1e-6 + and int(input_qparams.zp) == int(output_qparams.zp) + ) + + +_EXPLICIT_LAYOUT_OP_MAP: dict[Target, ChannelsLastOpSpec] = { + exir_ops.edge.aten.convolution.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.convolution.default, + input_indices=[0], + output_indices=[0], + filter_fn=lambda node: _is_rank4(node) and _has_input_and_output_qparams(node), + ), + exir_ops.edge.aten.avg_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.avg_pool2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_supports_avg_pool2d, + ), + exir_ops.edge.aten.max_pool2d.default: ChannelsLastOpSpec( + target=exir_ops.edge.channels_last.max_pool2d.default, + input_indices=[0], + output_indices=[0], + filter_fn=_supports_max_pool2d, + ), +} + + +class CortexMReplaceOpsWithChannelsLastVariants(ReplaceOpsWithChannelsLastVariants): + def __init__(self, exported_program: ExportedProgram) -> None: + super().__init__(exported_program, op_map=dict(_EXPLICIT_LAYOUT_OP_MAP)) + + +class CortexMCanonicalizeViewCopyPermutePass(CanonicalizeViewCopyPermutePass): + def __init__(self) -> None: + super().__init__(permute_targets=PERMUTE_COPY_TARGETS) + + def _set_node_op(self, node, target, input_node, arg) -> None: + super()._set_node_op(node, target, input_node, arg) + if node.target != LAYOUT_PERMUTE_COPY: + return + input_value = input_node.meta.get("val") + if isinstance(input_value, torch.Tensor): + dims = [int(dim) % input_value.dim() for dim in arg] + node.meta["val"] = input_value.new_empty( + tuple(input_value.shape[dim] for dim in dims) + ) + + +class ValidateCortexMExplicitLayoutPass(ExportPass): + def call(self, graph_module: GraphModule) -> PassResult: + for node in graph_module.graph.nodes: + if node.target in _EXPLICIT_LAYOUT_OP_MAP: + raise RuntimeError( + "Cortex-M explicit layout requires every quantized spatial " + f"operator to be NHWC-eligible, but {node.target} was not. " + "Use the legacy layout pipeline for this model." + ) + + if node.target != LAYOUT_PERMUTE_COPY: + continue + input_node = node.args[0] + dims = node.args[1] if len(node.args) > 1 else None + input_value = ( + input_node.meta.get("val") if isinstance(input_node, Node) else None + ) + if ( + not isinstance(input_value, torch.Tensor) + or input_value.dtype != torch.int8 + ): + raise RuntimeError( + f"Cortex-M layout copy {node.name} must move an int8 tensor." + ) + rank = input_value.dim() + if ( + not 1 <= rank <= 4 + or not isinstance(dims, (list, tuple)) + or len(dims) != rank + or not all(isinstance(dim, int) for dim in dims) + or sorted(dim % rank for dim in dims) != list(range(rank)) + ): + raise RuntimeError( + f"Cortex-M layout copy {node.name} has invalid permutation " + f"{dims!r} for rank {rank}." + ) + + return PassResult(graph_module, False) diff --git a/backends/cortex_m/quantizer/pattern_checkers.py b/backends/cortex_m/quantizer/pattern_checkers.py index ef2a91e6c2c..d1b7f8b8255 100644 --- a/backends/cortex_m/quantizer/pattern_checkers.py +++ b/backends/cortex_m/quantizer/pattern_checkers.py @@ -131,6 +131,18 @@ def check_quantization_config( return is_int8 and is_ch_axis_0 +class CortexMExplicitConv2DCheck(CortexMConv2DCheck): + @classmethod + def check_pattern(cls, pattern): + return all(get_first_fake_tensor(node).dim() == 4 for node in pattern) + + +class CortexMExplicitConv1DCheck(CortexMConv2DCheck): + @classmethod + def check_pattern(cls, pattern): + return all(get_first_fake_tensor(node).dim() == 3 for node in pattern) + + class CortexMLinearCheck(PatternCheck): @classmethod def check_quantization_config( @@ -219,6 +231,8 @@ def check_quantization_config( class CortexMConvTranspose2DCheck(PatternCheck): + require_channels_last = True + @classmethod def _check_node(cls, node: Node) -> bool: if node is None: @@ -228,8 +242,7 @@ def _check_node(cls, node: Node) -> bool: if tensor is None: return False # Reject if no tensor found - # REJECT if using NCHW format (we need channels_last/NHWC) - if not is_channels_last(tensor): + if cls.require_channels_last and not is_channels_last(tensor): return False # Reject NCHW # For aten.conv_transpose2d.input: @@ -288,6 +301,10 @@ def check_quantization_config( return is_int8 and is_ch_axis_1 +class CortexMExplicitConvTranspose2DCheck(CortexMConvTranspose2DCheck): + require_channels_last = False + + class CortexMAvgPool2DCheck(PatternCheck): @classmethod def check_pattern(cls, pattern): diff --git a/backends/cortex_m/quantizer/quantizer.py b/backends/cortex_m/quantizer/quantizer.py index d3f49114144..c1f0778b336 100644 --- a/backends/cortex_m/quantizer/quantizer.py +++ b/backends/cortex_m/quantizer/quantizer.py @@ -25,8 +25,10 @@ ) from executorch.backends.cortex_m.quantizer.quantizer_support import ( __name__ as cortex_m_quantizer_support_module, + CONV1D_OP_PATTERNS, CONV_OP_PATTERNS, CONV_TRANSPOSE_OP_PATTERNS, + CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT, CORTEX_M_QUANTIZER_SUPPORT_DICT, ) from executorch.backends.cortex_m.quantizer_reporter import QuantizerReporter @@ -45,8 +47,11 @@ def mark_node_as_annotated( class CortexMQuantizer(ComposableQuantizer): - - def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> None: + def __init__( + self, + per_tensor_config: Optional[QuantizationConfig] = None, + use_explicit_layout: bool = False, + ) -> None: """Cortex-M PT2E quantizer. Args: @@ -57,20 +62,33 @@ def __init__(self, per_tensor_config: Optional[QuantizationConfig] = None) -> No ``INT8_PER_TENSOR_CONFIG``; pass ``INT16_PER_TENSOR_CONFIG`` to quantize the ops that support it (e.g. ``quantized_div``) with int16 activations. + use_explicit_layout: Select the support checks for the experimental + explicit-layout pipeline. Legacy mode continues to require + channels-last convolution inputs. """ per_tensor_config = per_tensor_config or INT8_PER_TENSOR_CONFIG - conv_targets: set[OpOverload] = set() - for key in CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys(): - conv_targets.update(key) - + support_dict = CORTEX_M_QUANTIZER_SUPPORT_DICT support_dict_name = ( cortex_m_quantizer_support_module + ".CORTEX_M_QUANTIZER_SUPPORT_DICT" ) + conv_patterns = CONV_OP_PATTERNS.keys() | CONV_TRANSPOSE_OP_PATTERNS.keys() + if use_explicit_layout: + support_dict = CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT + support_dict_name = ( + cortex_m_quantizer_support_module + + ".CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT" + ) + conv_patterns |= CONV1D_OP_PATTERNS.keys() + + conv_targets: set[OpOverload] = set() + for key in conv_patterns: + conv_targets.update(key) + pattern_matcher = PatternMatcher( cast( dict[tuple[OpOverload, ...], Optional[type[PatternCheck]]], - CORTEX_M_QUANTIZER_SUPPORT_DICT, + support_dict, ), support_dict_name=support_dict_name, ) diff --git a/backends/cortex_m/quantizer/quantizer_support.py b/backends/cortex_m/quantizer/quantizer_support.py index d0cce702cff..b00909c3b5f 100644 --- a/backends/cortex_m/quantizer/quantizer_support.py +++ b/backends/cortex_m/quantizer/quantizer_support.py @@ -12,6 +12,9 @@ CortexMConv2DCheck, CortexMConvTranspose2DCheck, CortexMDivCheck, + CortexMExplicitConv1DCheck, + CortexMExplicitConv2DCheck, + CortexMExplicitConvTranspose2DCheck, CortexMLinearCheck, CortexMMaxPool2DCheck, CortexMSoftmaxCheck, @@ -102,6 +105,42 @@ (torch.ops.aten.conv2d.default, torch.ops.aten.clamp_.default): CortexMConv2DCheck, } +CONV1D_OP_PATTERNS = { + (torch.ops.aten.conv1d.default,): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.relu.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.relu_.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardtanh.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardtanh_.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardsigmoid.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.hardsigmoid_.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.clamp.default, + ): CortexMExplicitConv1DCheck, + ( + torch.ops.aten.conv1d.default, + torch.ops.aten.clamp_.default, + ): CortexMExplicitConv1DCheck, +} + CONV_TRANSPOSE_OP_PATTERNS = { (torch.ops.aten.conv_transpose2d.input,): CortexMConvTranspose2DCheck, ( @@ -209,3 +248,13 @@ | BMM_OP_PATTERNS | ACTIVATION_OP_PATTERNS ) + +CORTEX_M_EXPLICIT_LAYOUT_QUANTIZER_SUPPORT_DICT = ( + CORTEX_M_QUANTIZER_SUPPORT_DICT + | CONV1D_OP_PATTERNS + | {pattern: CortexMExplicitConv2DCheck for pattern in CONV_OP_PATTERNS} + | { + pattern: CortexMExplicitConvTranspose2DCheck + for pattern in CONV_TRANSPOSE_OP_PATTERNS + } +) diff --git a/backends/cortex_m/test/targets.bzl b/backends/cortex_m/test/targets.bzl index 2ea3a5b3b99..06b2daed033 100644 --- a/backends/cortex_m/test/targets.bzl +++ b/backends/cortex_m/test/targets.bzl @@ -5,6 +5,8 @@ # LICENSE file in the root directory of this source tree. load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") +load("@fbcode_macros//build_defs:python_library.bzl", "python_library") +load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load("@fbsource//tools/build_defs:platform_defs.bzl", "CXX") @@ -36,6 +38,21 @@ def define_common_targets(is_fbcode = False): define_operator_test_target(op) if is_fbcode: + python_library( + name = "tester", + srcs = ["tester.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/arm/test:arm_tester", + "//executorch/backends/arm/test:common", + "//executorch/backends/cortex_m:edge_compile_config", + "//executorch/backends/cortex_m:target_config", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m/quantizer:quantizer", + "//executorch/backends/test/harness:tester", + ], + ) + python_unittest( name = "test_activation_lut", srcs = [ @@ -50,6 +67,24 @@ def define_common_targets(is_fbcode = False): ], ) + python_pytest( + name = "test_explicit_layout_pipeline", + srcs = ["test_explicit_layout_pipeline.py"], + compile = "with-source", + typing = False, + deps = [ + "//caffe2:torch", + "//executorch/backends/cortex_m:target_config", + "//executorch/backends/cortex_m/ops:ops", + "//executorch/backends/cortex_m/passes:cortex_passes", + "//executorch/backends/cortex_m/quantizer:quantizer", + "//executorch/backends/test/harness:tester", + "//executorch/exir/dialects:lib", + ":tester", + "fbsource//third-party/pypi/pytest:pytest", + ], + ) + python_unittest( name = "test_replace_quant_nodes", srcs = [ diff --git a/backends/cortex_m/test/test_explicit_layout_pipeline.py b/backends/cortex_m/test/test_explicit_layout_pipeline.py new file mode 100644 index 00000000000..7c93233e768 --- /dev/null +++ b/backends/cortex_m/test/test_explicit_layout_pipeline.py @@ -0,0 +1,158 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from functools import partial + +import pytest +import torch +from executorch.backends.cortex_m.passes.cortex_m_pass_manager import CortexMPassManager +from executorch.backends.cortex_m.quantizer.quantizer import CortexMQuantizer +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMTester +from executorch.backends.test.harness.stages import Quantize, RunPasses, StageType +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Node + +# Temporary opt-in coverage. Move these invariants into the standard Cortex-M +# tests when this pass manager becomes the default. + + +class Conv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class Conv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(2, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class ConvPadConv(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv1 = torch.nn.Conv2d(3, 4, 3, padding=1) + self.conv2 = torch.nn.Conv2d(4, 5, 3, padding=1) + + def forward(self, x): + return self.conv2(torch.nn.functional.pad(self.conv1(x), (1, 1, 1, 1))) + + +class UnsupportedAvgPool(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.avg_pool2d( + x, kernel_size=2, stride=2, divisor_override=3 + ) + + +def _count(exported_program, target) -> int: + return sum(node.target == target for node in exported_program.graph.nodes) + + +def _run_explicit_layout_pass_manager(tester: CortexMTester) -> CortexMTester: + target_config = CortexMTargetConfig(cpu=CortexM.M55) + tester.run_passes( + RunPasses( + partial( + CortexMPassManager, + target_config=target_config, + use_explicit_layout=True, + ), # type: ignore[arg-type] + CortexMPassManager.explicit_layout_pass_list, # type: ignore[arg-type] + ) + ) + return tester + + +def _run_explicit_layout_passes(tester: CortexMTester) -> CortexMTester: + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=True))) + tester.export().to_edge() + return _run_explicit_layout_pass_manager(tester) + + +def test_layout_pipelines_select_distinct_spatial_operators(): + legacy_input = torch.randn(1, 3, 8, 8).to(memory_format=torch.channels_last) + legacy = CortexMTester( + Conv2d().eval().to(memory_format=torch.channels_last), + (legacy_input,), + ) + legacy.quantize().export().to_edge().run_passes() + legacy_program = legacy.get_artifact(StageType.RUN_PASSES).exported_program() + + explicit = _run_explicit_layout_passes( + CortexMTester(Conv2d().eval(), (torch.randn(1, 3, 8, 8),)) + ) + explicit_program = explicit.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(legacy_program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 1 + assert ( + _count( + legacy_program, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ) + == 0 + ) + assert ( + _count(explicit_program, exir_ops.edge.cortex_m.quantized_conv2d.default) == 0 + ) + assert ( + _count( + explicit_program, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ) + == 1 + ) + assert _count(explicit_program, exir_ops.edge.cortex_m.transpose.default) == 2 + + +def test_conv1d_is_quantized_before_layout_conversion(): + tester = CortexMTester(Conv1d().eval(), (torch.randn(1, 2, 8),)) + tester.quantize(Quantize(CortexMQuantizer(use_explicit_layout=True))) + quantized = tester.get_artifact(StageType.QUANTIZE) + [conv1d] = [ + node + for node in quantized.graph.nodes + if node.target == torch.ops.aten.conv1d.default + ] + + weight = conv1d.args[1] + assert isinstance(weight, Node) + assert ( + weight.target == torch.ops.quantized_decomposed.dequantize_per_channel.default + ) + + tester.export().to_edge() + _run_explicit_layout_pass_manager(tester) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + assert _count(program, exir_ops.edge.aten.convolution.default) == 0 + + +def test_explicit_layout_reuses_pad(): + tester = _run_explicit_layout_passes( + CortexMTester(ConvPadConv().eval(), (torch.randn(1, 3, 8, 8),)) + ) + program = tester.get_artifact(StageType.RUN_PASSES).exported_program() + + assert _count(program, exir_ops.edge.cortex_m.pad.default) == 1 + + +def test_explicit_layout_rejects_unsupported_spatial_operator(): + tester = CortexMTester(UnsupportedAvgPool(), (torch.randn(1, 3, 8, 8),)) + + with pytest.raises(Exception) as caught: + _run_explicit_layout_passes(tester) + + assert caught.value.__cause__ is not None + assert "NHWC-eligible" in str(caught.value.__cause__) From 77fb78d26b83c6696424c2fa727603bc7cfee1fa Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 11:48:11 -0700 Subject: [PATCH 050/190] Add ExportRecipe support for Arm targets (#22368) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary `ArmRecipeProvider` reaches the Ethos-U, TOSA and VGF targets through `ExportRecipe` instead of `aot_arm_compiler.py`, mirroring the XNNPACK and QNN providers. Eight recipes: Ethos-U55/U65/U85 INT8 with `macs`, `system_config`, `memory_mode`, `extra_flags` and `config_ini` kwargs; TOSA FP, INT8 and A16W8; VGF FP and INT8. Each reproduces the default CLI invocation for its target byte for byte. Two things are easy to get wrong there and worth review attention: the CLI pins `extract_delegate_segments=False`, which is the delegate layout `arm_executor_runner` has actually been run against; and the partitioner snapshots the compile spec, so the pass pipeline config has to be populated first or the U55 and U65 delegates silently lose it. Accepted MAC counts come from the installed Vela rather than a local copy, so a Vela bump cannot leave the recipe rejecting a target its compiler accepts. That import is guarded — the target-less CI job has no Vela and a recipe still has to build there. Review order: the two target tables, then `_build_recipe`, then the tests. The tests are named so the existing target-less and TOSA suites each collect a disjoint subset pre-merge; no new CI job. Authored with Claude Code. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell --- backends/arm/README.md | 43 ++ backends/arm/recipes/BUCK | 61 +++ backends/arm/recipes/__init__.py | 15 + backends/arm/recipes/arm_recipe_provider.py | 287 ++++++++++ backends/arm/recipes/arm_recipe_types.py | 52 ++ backends/arm/test/recipes/test_arm_recipes.py | 491 ++++++++++++++++++ backends/arm/test/targets.bzl | 6 + 7 files changed, 955 insertions(+) create mode 100644 backends/arm/recipes/BUCK create mode 100644 backends/arm/recipes/__init__.py create mode 100644 backends/arm/recipes/arm_recipe_provider.py create mode 100644 backends/arm/recipes/arm_recipe_types.py create mode 100644 backends/arm/test/recipes/test_arm_recipes.py diff --git a/backends/arm/README.md b/backends/arm/README.md index ce0c49919e4..d5986b5ccd1 100644 --- a/backends/arm/README.md +++ b/backends/arm/README.md @@ -156,6 +156,49 @@ compile specs, see: Additional examples are available in `examples/arm`. +#### Export recipes + +An `ExportRecipe` bundles those steps for a target, so a standard export needs +no compile spec, quantizer or partitioner of its own. Each recipe carries the +settings its target expects: + +```python +from executorch.backends.arm.recipes.arm_recipe_types import ArmRecipeType +from executorch.export import export, ExportRecipe + +session = export( + model=model, + example_inputs=[example_inputs], + export_recipe=ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8), +) +session.save_to_pte("model") +``` + +The recipe quantizes the model, so pass `example_inputs` that are representative +of real data; they are used to calibrate. + +Available recipes: + +| Recipe | Target | +| --- | --- | +| `ETHOS_U55_INT8`, `ETHOS_U65_INT8`, `ETHOS_U85_INT8` | Ethos-U NPUs, int8 | +| `TOSA_FP`, `TOSA_INT8`, `TOSA_A16W8` | TOSA, for testing without hardware | +| `VGF_FP`, `VGF_INT8` | VGF, for the ML SDK for Vulkan | + +The Ethos-U recipes accept `macs`, `system_config`, `memory_mode`, +`extra_flags` and `config_ini`, matching the corresponding Vela options: + +```python +ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U85_INT8, macs=512) +``` + +`macs` is validated against the accelerator configurations the installed Vela +accepts, so an unsupported count fails at recipe construction rather than during +compilation. + +Reach for the step-by-step flow above when a recipe does not fit -- a custom +quantization scheme, extra passes, or a compile spec the recipe does not expose. + ### Direct Drive (experimental, Ethos-U85 on Linux) workflow Direct Drive enables execution on Ethos-U85 via the Linux driver stack. diff --git a/backends/arm/recipes/BUCK b/backends/arm/recipes/BUCK new file mode 100644 index 00000000000..c7dfa0abd55 --- /dev/null +++ b/backends/arm/recipes/BUCK @@ -0,0 +1,61 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target( + _kind = runtime.python_library, + name = "recipes", + srcs = [ + "__init__.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":arm_recipe_provider", + ":arm_recipe_types", + "//executorch/export:recipe_registry", + ], +) + +fbcode_target( + _kind = runtime.python_library, + name = "arm_recipe_provider", + srcs = [ + "arm_recipe_provider.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":arm_recipe_types", + # Imported lazily for the accelerator-config list; declared so the dep + # does not rely on reaching vela through :ethosu. + "fbsource//third-party/pypi/ethos-u-vela:ethos-u-vela", + "//executorch/backends/arm:_factory", + "//executorch/backends/arm:arm_compile_spec", + "//executorch/backends/arm:ethosu", + "//executorch/backends/arm:vgf", + "//executorch/backends/arm/quantizer:lib", + "//executorch/backends/arm/tosa:compile_spec", + "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", + "//executorch/exir:lib", + "//executorch/exir/backend:op_backend", + "//executorch/export:lib", + ], +) + +fbcode_target( + _kind = runtime.python_library, + name = "arm_recipe_types", + srcs = [ + "arm_recipe_types.py", + ], + visibility = ["PUBLIC"], + deps = [ + "//executorch/export:recipe", + ], +) diff --git a/backends/arm/recipes/__init__.py b/backends/arm/recipes/__init__.py new file mode 100644 index 00000000000..2b751645d68 --- /dev/null +++ b/backends/arm/recipes/__init__.py @@ -0,0 +1,15 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.export import recipe_registry + +from .arm_recipe_provider import ArmRecipeProvider +from .arm_recipe_types import ArmRecipeType + +recipe_registry.register_backend_recipe_provider(ArmRecipeProvider()) + + +__all__ = ["ArmRecipeProvider", "ArmRecipeType"] diff --git a/backends/arm/recipes/arm_recipe_provider.py b/backends/arm/recipes/arm_recipe_provider.py new file mode 100644 index 00000000000..35f1370856a --- /dev/null +++ b/backends/arm/recipes/arm_recipe_provider.py @@ -0,0 +1,287 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +import logging +from dataclasses import dataclass +from typing import Any, Callable, Optional, Sequence + +from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec +from executorch.backends.arm.ethosu import EthosUCompileSpec +from executorch.backends.arm.quantizer import ( + get_symmetric_a16w8_quantization_config, + get_symmetric_quantization_config, +) +from executorch.backends.arm.recipes.arm_recipe_types import ARM_BACKEND, ArmRecipeType +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.util._factory import create_partitioner, create_quantizer +from executorch.backends.arm.vgf import VgfCompileSpec +from executorch.exir.capture import EdgeCompileConfig, ExecutorchBackendConfig +from executorch.exir.pass_manager import PassType +from executorch.exir.program import EdgeProgramManager +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) + + +logger: logging.Logger = logging.getLogger(__name__) + +# (target prefix, default MAC count). Which counts are *accepted* is Vela's to +# say, so it is asked at build time rather than restated here. +_ETHOS_U_FAMILIES: dict[ArmRecipeType, tuple[str, int]] = { + ArmRecipeType.ETHOS_U55_INT8: ("ethos-u55", 128), + ArmRecipeType.ETHOS_U65_INT8: ("ethos-u65", 256), + ArmRecipeType.ETHOS_U85_INT8: ("ethos-u85", 256), +} + +_ETHOS_U_KWARGS: frozenset[str] = frozenset( + {"macs", "system_config", "memory_mode", "extra_flags", "config_ini"} +) + +# Prepended to any caller-supplied Vela flags, matching `_get_compile_spec` in +# aot_arm_compiler.py. +_VELA_DEFAULT_FLAGS: tuple[str, ...] = ( + "--verbose-operators", + "--verbose-cycle-estimate", +) + + +@dataclass(frozen=True) +class _TosaVersionedTarget: + """A target with no caller-tunable options: class plus TOSA version.""" + + compile_spec: Callable[[str], ArmCompileSpec] + tosa_spec: str + quant_mode: Optional[str] + replace_quant_nodes: bool + + +# VGF keeps the quantized_decomposed QDQ ops it is given; see +# `_apply_replace_quant_nodes` in aot_arm_compiler.py. +_TOSA_VERSIONED_TARGETS: dict[ArmRecipeType, _TosaVersionedTarget] = { + ArmRecipeType.TOSA_FP: _TosaVersionedTarget( + TosaCompileSpec, "TOSA-1.0+FP", None, False + ), + ArmRecipeType.TOSA_INT8: _TosaVersionedTarget( + TosaCompileSpec, "TOSA-1.0+INT", "INT8", True + ), + ArmRecipeType.TOSA_A16W8: _TosaVersionedTarget( + TosaCompileSpec, "TOSA-1.0+INT+int16", "A16W8", True + ), + ArmRecipeType.VGF_FP: _TosaVersionedTarget( + VgfCompileSpec, "TOSA-1.0+FP", None, False + ), + ArmRecipeType.VGF_INT8: _TosaVersionedTarget( + VgfCompileSpec, "TOSA-1.0+INT", "INT8", False + ), +} + + +def _reject_unsupported_accelerator( + recipe_type: ArmRecipeType, family: str, target: str, macs: int +) -> None: + """Ask Vela which accelerator configurations it accepts. + + A local copy would go stale on the next Vela bump. Without Vela there is + nothing to check against and the compile spec still has to build, so the + import is guarded the way `arm_vela` guards its own. + + """ + try: + from ethosu.vela.architecture_features import Accelerator # type: ignore + except ImportError: + logger.debug("ethos-u-vela is not installed; macs=%s unvalidated", macs) + return + + supported = {accelerator.value for accelerator in Accelerator} + if target not in supported: + allowed = sorted( + int(name.rsplit("-", 1)[1]) + for name in supported + if name.startswith(f"{family}-") + ) + raise ValueError( + f"Recipe '{recipe_type.value}' does not support macs={macs}. " + f"Allowed: {allowed}" + ) + + +def _replace_quant_nodes( + edge_program_manager: EdgeProgramManager, +) -> list[PassType]: + """Rewrite the QDQ ops left outside the delegate into cortex_m kernels. + + Matches `_apply_replace_quant_nodes` in aot_arm_compiler.py, which applies + the same pass to the whole manager once the partitioner has run. Without it + the boundary quantize/dequantize keep their `quantized_decomposed` targets, + which have no out variants, and `to_executorch` refuses to emit them. + + """ + # Function-local: an FP recipe must not pull in the cortex_m operator + # library, which registers its whole op set on import. + from executorch.backends.cortex_m.passes.replace_quant_nodes_pass import ( + ReplaceQuantNodesPass, + ) + + return [ReplaceQuantNodesPass()] + + +class ArmRecipeProvider(BackendRecipeProvider): + """Builds ExportRecipes for the delegated Arm targets: Ethos-U, TOSA, VGF. + + Each recipe is built to reproduce the default + ``backends/arm/scripts/aot_arm_compiler.py`` invocation for its target: the + same compile spec, quantizer, pass pipeline and backend config. The CLI + options with no recipe equivalent, debug mode and direct drive, are the + exceptions. + """ + + @property + def backend_name(self) -> str: + return ARM_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(_ETHOS_U_FAMILIES) + list(_TOSA_VERSIONED_TARGETS) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + if not isinstance(recipe_type, ArmRecipeType): + return None + + if recipe_type in _ETHOS_U_FAMILIES: + self._warn_unknown_kwargs(recipe_type, kwargs, _ETHOS_U_KWARGS) + return self._build_recipe( + recipe_type, + self._ethos_u_compile_spec(recipe_type, kwargs), + quant_mode="INT8", + replace_quant_nodes=True, + ) + + target = _TOSA_VERSIONED_TARGETS.get(recipe_type) + if target is None: + return None + + self._warn_unknown_kwargs(recipe_type, kwargs, frozenset()) + return self._build_recipe( + recipe_type, + target.compile_spec(target.tosa_spec), + quant_mode=target.quant_mode, + replace_quant_nodes=target.replace_quant_nodes, + ) + + @staticmethod + def _ethos_u_compile_spec( + recipe_type: ArmRecipeType, kwargs: dict[str, Any] + ) -> EthosUCompileSpec: + family, default_macs = _ETHOS_U_FAMILIES[recipe_type] + macs = kwargs.get("macs", default_macs) + if not isinstance(macs, int): + raise ValueError(f"macs must be an int, got {macs!r}") + + extra_flags = kwargs.get("extra_flags") or [] + # The list check comes first: a bare string would be iterated into one + # flag per character, and anything not iterable would raise TypeError + # out of `all` rather than reaching this message. + if not isinstance(extra_flags, list) or not all( + isinstance(flag, str) for flag in extra_flags + ): + raise ValueError( + f"extra_flags must be a list of strings, got {extra_flags!r}" + ) + + target = f"{family}-{macs}" + _reject_unsupported_accelerator(recipe_type, family, target, macs) + + return EthosUCompileSpec( + target=target, + system_config=kwargs.get("system_config"), + memory_mode=kwargs.get("memory_mode"), + extra_flags=list(_VELA_DEFAULT_FLAGS) + list(extra_flags), + # EthosUCompileSpec owns the default. + config_ini=kwargs.get("config_ini"), + ) + + @classmethod + def _build_recipe( + cls, + recipe_type: ArmRecipeType, + compile_spec: ArmCompileSpec, + quant_mode: Optional[str], + replace_quant_nodes: bool, + ) -> ExportRecipe: + # The partitioner snapshots the compile spec and the pipeline config is + # materialised on first read, which the CLI gets for free by quantizing + # before it partitions. + compile_spec.set_pass_pipeline_config(compile_spec._get_pass_pipeline_config()) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=cls._build_quantization_recipe( + compile_spec, quant_mode + ), + lowering_recipe=LoweringRecipe( + partitioners=[create_partitioner(compile_spec)], + # The CLI disables edge verification on every Arm path. + edge_compile_config=EdgeCompileConfig(_check_ir_validity=False), + edge_manager_transform_passes=( + [_replace_quant_nodes] if replace_quant_nodes else None + ), + ), + # The Arm runtime expects the delegate payload inline rather than + # in its own segment, as every other Arm AOT path asks for. + executorch_backend_config=ExecutorchBackendConfig( + extract_delegate_segments=False + ), + ) + + @staticmethod + def _build_quantization_recipe( + compile_spec: ArmCompileSpec, quant_mode: Optional[str] + ) -> Optional[QuantizationRecipe]: + if quant_mode is None: + return None + + if quant_mode == "INT8": + operator_config = get_symmetric_quantization_config(is_per_channel=True) + elif quant_mode == "A16W8": + if not compile_spec.tosa_spec.support_extension("int16"): + raise ValueError( + f"TOSA spec {compile_spec.tosa_spec} does not support int16 " + "(required for A16W8)" + ) + operator_config = get_symmetric_a16w8_quantization_config( + is_per_channel=True + ) + else: + raise ValueError(f"Unsupported quant_mode: {quant_mode}") + + quantizer = create_quantizer(compile_spec) + quantizer.set_global(operator_config) + return QuantizationRecipe(quantizers=[quantizer]) + + @staticmethod + def _warn_unknown_kwargs( + recipe_type: ArmRecipeType, + kwargs: dict[str, Any], + expected: frozenset[str], + ) -> None: + # Warn, as XNNPACK and QNN do: `_create_target_recipe` hands every + # recipe in a combination the same kwargs. + unexpected = set(kwargs.keys()) - expected + if unexpected: + allowed = sorted(expected) if expected else "none" + logger.warning( + "Arm recipe '%s' ignoring unexpected parameters: %s. Allowed: %s", + recipe_type.value, + sorted(unexpected), + allowed, + ) diff --git a/backends/arm/recipes/arm_recipe_types.py b/backends/arm/recipes/arm_recipe_types.py new file mode 100644 index 00000000000..91904a4d981 --- /dev/null +++ b/backends/arm/recipes/arm_recipe_types.py @@ -0,0 +1,52 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +from executorch.export import RecipeType + + +ARM_BACKEND: str = "arm" + + +class ArmRecipeType(RecipeType): + """Arm-specific recipe types. + + Covers the delegated targets of ``backends/arm/scripts/aot_arm_compiler.py``. + Its non-delegated Cortex-M/CMSIS-NN path is a separate backend and is not + reachable from these recipes. + + Ethos-U recipes accept the following kwargs: + macs (int): MAC count for the family, validated against the accelerator + configurations the installed Vela accepts -- today 32/64/128/256 for + U55, 256/512 for U65 and 128/256/512/1024/2048 for U85. Defaults to + 128 for U55 and 256 for U65 and U85. + system_config (str): Vela system config name. Defaults from + ``EthosUCompileSpec`` apply when omitted. + memory_mode (str): Vela memory mode. Defaults from + ``EthosUCompileSpec`` apply when omitted. + extra_flags (list[str]): Vela compiler flags, appended to the + ``--verbose-operators --verbose-cycle-estimate`` the CLI always + passes rather than replacing them. + config_ini (str): Path to a Vela .ini configuration file. Defaults to + ``"Arm/vela.ini"``. + + """ + + ETHOS_U55_INT8 = "arm_ethos_u55_int8" + ETHOS_U65_INT8 = "arm_ethos_u65_int8" + ETHOS_U85_INT8 = "arm_ethos_u85_int8" + + TOSA_FP = "arm_tosa_fp" + TOSA_INT8 = "arm_tosa_int8" + TOSA_A16W8 = "arm_tosa_a16w8" + + VGF_FP = "arm_vgf_fp" + VGF_INT8 = "arm_vgf_int8" + + @classmethod + def get_backend_name(cls) -> str: + return ARM_BACKEND diff --git a/backends/arm/test/recipes/test_arm_recipes.py b/backends/arm/test/recipes/test_arm_recipes.py new file mode 100644 index 00000000000..1544863d86d --- /dev/null +++ b/backends/arm/test/recipes/test_arm_recipes.py @@ -0,0 +1,491 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Tests for the Arm ExportRecipe provider. + +Building a recipe only assembles a compile spec, a quantizer and a partitioner, +so none of these tests need Vela, the model converter or an FVP. Class and +method names route each test to exactly one of the existing target-less, TOSA +and VGF suites; see the ``-k`` filters in +``backends/arm/test/test_arm_backend.sh``. + +""" + +# pyre-strict + +import unittest +from typing import Any, Optional + +import torch + +from executorch.backends.arm.recipes.arm_recipe_provider import ArmRecipeProvider +from executorch.backends.arm.recipes.arm_recipe_types import ARM_BACKEND, ArmRecipeType +from executorch.export import ExportRecipe, recipe_registry, StageType +from executorch.export.export import ExportSession + + +_PROVIDER_LOGGER = "executorch.backends.arm.recipes.arm_recipe_provider" + +try: + import ethosu.vela.architecture_features # type: ignore # noqa: F401 + + _VELA_INSTALLED = True +except ImportError: + # The target-less CI job installs the Arm deps without Vela, and a recipe + # has to build there; only the accelerator check needs it. + _VELA_INSTALLED = False + + +class _AddModule(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + +class _ConvReluModule(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = torch.nn.ReLU() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.conv(x)) + + +def _compile_spec_value(partitioner: Any, key: str) -> Optional[str]: + for spec in partitioner.delegation_spec.compile_specs: + if spec.key == key: + value = spec.value + return value.decode() if isinstance(value, (bytes, bytearray)) else value + return None + + +def _backend_id(recipe: ExportRecipe) -> str: + return _first_partitioner(recipe).delegation_spec.backend_id + + +def _first_partitioner(recipe: ExportRecipe) -> Any: + assert recipe.lowering_recipe is not None + # Arm recipes partition every method the same way, so the list form rather + # than the per-method dict `LoweringRecipe` also accepts. + partitioners = recipe.lowering_recipe.partitioners + assert isinstance(partitioners, list) and partitioners + return partitioners[0] + + +def _global_config(recipe: ExportRecipe) -> Any: + assert recipe.quantization_recipe is not None + assert recipe.quantization_recipe.quantizers is not None + return recipe.quantization_recipe.quantizers[0].global_config # type: ignore[attr-defined] + + +def _input_activation_dtype(recipe: ExportRecipe) -> Optional[torch.dtype]: + config = _global_config(recipe) + if config is None or config.input_activation is None: + return None + return config.input_activation.dtype + + +class _ArmRecipeTestCase(unittest.TestCase): + """Re-registers the Arm provider before each test. + + Registration happens when ``backends.arm.recipes`` is imported, so it cannot + repeat: the module is already loaded. Other suites in the same process clear + the singleton registry in teardown, and under ``pytest -n`` their tests can + interleave with these. + + """ + + def setUp(self) -> None: + super().setUp() + recipe_registry.register_backend_recipe_provider(ArmRecipeProvider()) + + +class TestArmRecipeRegistration(_ArmRecipeTestCase): + def test_backend_registered(self) -> None: + self.assertIn(ARM_BACKEND, recipe_registry.list_backends()) + + def test_supported_recipes_match_enum(self) -> None: + # Catches an enum member added but never wired into a target table. + supported = recipe_registry.get_supported_recipes(ARM_BACKEND) + self.assertEqual(set(supported), set(ArmRecipeType)) + + def test_unknown_recipe_returns_none(self) -> None: + from executorch.export import RecipeType + + class _StubRecipeType(RecipeType): + FOO = "stub_foo" + + @classmethod + def get_backend_name(cls) -> str: + return "stub" + + self.assertIsNone(ArmRecipeProvider().create_recipe(_StubRecipeType.FOO)) + + +class TestTosaRecipes(_ArmRecipeTestCase): + def test_tosa_construction(self) -> None: + cases = [ + (ArmRecipeType.TOSA_FP, "arm_tosa_fp", "TOSA-1.0+FP", None), + (ArmRecipeType.TOSA_INT8, "arm_tosa_int8", "TOSA-1.0+INT", torch.int8), + ( + ArmRecipeType.TOSA_A16W8, + "arm_tosa_a16w8", + "TOSA-1.0+INT+int16", + torch.int16, + ), + ] + for recipe_type, expected_name, expected_spec, expected_act_dtype in cases: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + self.assertEqual(recipe.name, expected_name) + # A VGF spec here would emit a container TOSA cannot consume. + self.assertEqual(_backend_id(recipe), "TOSABackend") + self.assertEqual( + _compile_spec_value(_first_partitioner(recipe), "tosa_spec"), + expected_spec, + ) + if expected_act_dtype is None: + self.assertIsNone(recipe.quantization_recipe) + else: + self.assertEqual( + _input_activation_dtype(recipe), expected_act_dtype + ) + + def test_weights_are_per_channel(self) -> None: + for recipe_type in (ArmRecipeType.TOSA_INT8, ArmRecipeType.TOSA_A16W8): + with self.subTest(recipe_type=recipe_type): + weight = _global_config(ExportRecipe.get_recipe(recipe_type)).weight + self.assertEqual(weight.qscheme, torch.per_channel_symmetric) + + def test_unexpected_kwarg_warns(self) -> None: + with self.assertLogs(_PROVIDER_LOGGER, level="WARNING"): + ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8, foo=1) + + +class TestVgfRecipes(_ArmRecipeTestCase): + """Named without the ``_vgf_`` token that routes tests to the VKML suite: + + constructing a compile spec needs no model converter, so these belong in the + target-less suite that runs on every PR. + + """ + + def test_construction(self) -> None: + cases = [ + (ArmRecipeType.VGF_FP, "arm_vgf_fp", "TOSA-1.0+FP", None), + (ArmRecipeType.VGF_INT8, "arm_vgf_int8", "TOSA-1.0+INT", torch.int8), + ] + for recipe_type, expected_name, expected_spec, expected_act_dtype in cases: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + self.assertEqual(recipe.name, expected_name) + self.assertEqual( + _compile_spec_value(_first_partitioner(recipe), "tosa_spec"), + expected_spec, + ) + # A TOSA spec here would emit a flatbuffer VKML cannot load. + self.assertEqual(_backend_id(recipe), "VgfBackend") + self.assertEqual( + _compile_spec_value(_first_partitioner(recipe), "output_format"), + "vgf", + ) + if expected_act_dtype is None: + self.assertIsNone(recipe.quantization_recipe) + else: + self.assertEqual( + _input_activation_dtype(recipe), expected_act_dtype + ) + + def test_keeps_quantized_decomposed_ops(self) -> None: + # VGF consumes the quantized_decomposed QDQ ops, so ReplaceQuantNodesPass + # must not run; see _apply_replace_quant_nodes in aot_arm_compiler.py. + for recipe_type in (ArmRecipeType.VGF_INT8, ArmRecipeType.VGF_FP): + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + assert recipe.lowering_recipe is not None + self.assertIsNone(recipe.lowering_recipe.edge_manager_transform_passes) + + +class TestEthosURecipes(_ArmRecipeTestCase): + def test_ethos_recipes_carry_their_pass_pipeline_config(self) -> None: + # Building the partitioner before the config is materialised silently + # drops this entry. Only the U55 subset has a non-default config. + for recipe_type in ( + ArmRecipeType.ETHOS_U55_INT8, + ArmRecipeType.ETHOS_U65_INT8, + ): + with self.subTest(recipe_type=recipe_type): + partitioner = _first_partitioner(ExportRecipe.get_recipe(recipe_type)) + self.assertIsNotNone( + _compile_spec_value(partitioner, "transform_pipeline_config") + ) + self.assertIsNone( + _compile_spec_value( + _first_partitioner( + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U85_INT8) + ), + "transform_pipeline_config", + ) + ) + + def test_default_macs(self) -> None: + cases = [ + (ArmRecipeType.ETHOS_U55_INT8, "ethos-u55-128"), + (ArmRecipeType.ETHOS_U65_INT8, "ethos-u65-256"), + (ArmRecipeType.ETHOS_U85_INT8, "ethos-u85-256"), + ] + for recipe_type, expected_target in cases: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + self.assertEqual(recipe.name, recipe_type.value) + self.assertEqual(_input_activation_dtype(recipe), torch.int8) + partitioner = _first_partitioner(recipe) + self.assertEqual( + _compile_spec_value(partitioner, "target"), expected_target + ) + + def test_custom_macs(self) -> None: + cases = [ + (ArmRecipeType.ETHOS_U55_INT8, 32, "ethos-u55-32"), + (ArmRecipeType.ETHOS_U55_INT8, 256, "ethos-u55-256"), + (ArmRecipeType.ETHOS_U65_INT8, 512, "ethos-u65-512"), + (ArmRecipeType.ETHOS_U85_INT8, 128, "ethos-u85-128"), + (ArmRecipeType.ETHOS_U85_INT8, 2048, "ethos-u85-2048"), + ] + for recipe_type, macs, expected_target in cases: + with self.subTest(recipe_type=recipe_type, macs=macs): + recipe = ExportRecipe.get_recipe(recipe_type, macs=macs) + partitioner = _first_partitioner(recipe) + self.assertEqual( + _compile_spec_value(partitioner, "target"), expected_target + ) + + @unittest.skipUnless(_VELA_INSTALLED, "accelerator configs come from Vela") + def test_invalid_macs_raises_u55(self) -> None: + cases = [ + (ArmRecipeType.ETHOS_U55_INT8, 512), + (ArmRecipeType.ETHOS_U65_INT8, 128), + (ArmRecipeType.ETHOS_U85_INT8, 64), + (ArmRecipeType.ETHOS_U55_INT8, 999), + ] + for recipe_type, macs in cases: + with self.subTest(recipe_type=recipe_type, macs=macs): + with self.assertRaises(ValueError): + ExportRecipe.get_recipe(recipe_type, macs=macs) + + def test_pass_through_kwargs(self) -> None: + recipe = ExportRecipe.get_recipe( + ArmRecipeType.ETHOS_U55_INT8, + macs=128, + system_config="Custom_System", + memory_mode="Custom_Memory", + extra_flags=["--user-flag"], + config_ini="custom/vela.ini", + ) + partitioner = _first_partitioner(recipe) + flags = _compile_spec_value(partitioner, "compile_flags") or "" + # Vela takes the last occurrence of a repeated flag, so the defaults + # have to come first for a caller override to win. + self.assertTrue( + flags.startswith("--verbose-operators --verbose-cycle-estimate"), + f"default flags must be prepended, got {flags}", + ) + self.assertLess(flags.index("--verbose-operators"), flags.index("--user-flag")) + self.assertIn("--system-config=Custom_System", flags) + self.assertIn("--memory-mode=Custom_Memory", flags) + self.assertIn("--verbose-operators", flags) + self.assertIn("--verbose-cycle-estimate", flags) + self.assertIn("--user-flag", flags) + self.assertIn("--config=custom/vela.ini", flags) + + def test_default_vela_flags(self) -> None: + # test_pass_through_kwargs supplies every kwarg, so the defaults would + # otherwise never be built. A wrong default config path only surfaces + # when Vela runs. + partitioner = _first_partitioner( + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8) + ) + flags = _compile_spec_value(partitioner, "compile_flags") or "" + self.assertIn("--config=Arm/vela.ini", flags) + self.assertIn("--verbose-operators", flags) + self.assertIn("--verbose-cycle-estimate", flags) + + def test_documented_kwargs_do_not_warn(self) -> None: + # Every one of these is honoured, so reporting it as ignored would be + # a lie about what the recipe did. + with self.assertNoLogs(_PROVIDER_LOGGER, level="WARNING"): + ExportRecipe.get_recipe( + ArmRecipeType.ETHOS_U55_INT8, + macs=128, + system_config="Custom_System", + memory_mode="Custom_Memory", + extra_flags=["--user-flag"], + config_ini="custom/vela.ini", + ) + + def test_unexpected_kwarg_warns(self) -> None: + # Flags typos like `mac=128` (instead of `macs=128`), which would + # otherwise silently produce a default-target binary. + with self.assertLogs(_PROVIDER_LOGGER, level="WARNING"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, mac=128) + + def test_extra_flags_must_be_a_list(self) -> None: + # A bare string is iterable, so it would reach Vela as one flag per + # character instead of failing. + with self.assertRaisesRegex(ValueError, "extra_flags must be a list"): + ExportRecipe.get_recipe( + ArmRecipeType.ETHOS_U55_INT8, extra_flags="--enable-debug-db" + ) + with self.assertRaisesRegex(ValueError, "extra_flags must be a list"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, extra_flags=[1]) + # Not iterable at all: checking the elements first raises TypeError out + # of `all` and the caller never sees the real complaint. + with self.assertRaisesRegex(ValueError, "extra_flags must be a list"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, extra_flags=7) + + def test_macs_must_be_an_int(self) -> None: + with self.assertRaisesRegex(ValueError, "macs must be an int"): + ExportRecipe.get_recipe(ArmRecipeType.ETHOS_U55_INT8, macs="128") + + def test_program_config_matches_the_cli(self) -> None: + # The Arm runtime has only ever been run against an inline delegate + # payload, and quantized Arm graphs do not survive the edge verifier. + recipe = ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8) + assert recipe.executorch_backend_config is not None + self.assertFalse(recipe.executorch_backend_config.extract_delegate_segments) + assert recipe.lowering_recipe is not None + assert recipe.lowering_recipe.edge_compile_config is not None + self.assertFalse(recipe.lowering_recipe.edge_compile_config._check_ir_validity) + + def test_fp_recipes_run_no_post_partition_transform(self) -> None: + # No QDQ ops to rewrite, so the extra stage must not be scheduled. + recipe = ExportRecipe.get_recipe(ArmRecipeType.TOSA_FP) + assert recipe.lowering_recipe is not None + self.assertIsNone(recipe.lowering_recipe.edge_manager_transform_passes) + + def test_recipes_do_not_share_config_objects(self) -> None: + # _combine_recipes compares configs by value on the assumption that + # each provider hands out a fresh one. + first = ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8) + second = ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8) + assert first.lowering_recipe is not None + assert second.lowering_recipe is not None + self.assertIsNot( + first.lowering_recipe.edge_compile_config, + second.lowering_recipe.edge_compile_config, + ) + self.assertIsNot( + first.executorch_backend_config, second.executorch_backend_config + ) + + +class TestQuantizedRecipeLowering(_ArmRecipeTestCase): + """Every quantized recipe has to rewrite the QDQ ops left outside the + delegate, and can only do so from a stage that runs after partitioning. + """ + + QUANTIZED_RECIPES = ( + ArmRecipeType.TOSA_INT8, + ArmRecipeType.TOSA_A16W8, + ArmRecipeType.ETHOS_U55_INT8, + ArmRecipeType.ETHOS_U65_INT8, + ArmRecipeType.ETHOS_U85_INT8, + ) + + def test_replace_quant_nodes_is_wired(self) -> None: + for recipe_type in self.QUANTIZED_RECIPES: + with self.subTest(recipe_type=recipe_type): + recipe = ExportRecipe.get_recipe(recipe_type) + assert recipe.lowering_recipe is not None + self.assertTrue( + recipe.lowering_recipe.edge_manager_transform_passes, + "quantized recipes must run ReplaceQuantNodesPass", + ) + + def test_session_schedules_the_stage(self) -> None: + session = ExportSession( + model=_ConvReluModule(), + example_inputs=[(torch.randn(1, 3, 8, 8),)], + export_recipe=ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8), + ) + stages = session._pipeline_stages + self.assertIn(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM, stages) + self.assertGreater( + stages.index(StageType.EDGE_PROGRAM_MANAGER_TRANSFORM), + stages.index(StageType.TO_EDGE_TRANSFORM_AND_LOWER), + ) + + +class TestTosaAOTRoundTrip(_ArmRecipeTestCase): + """End-to-end exports through the recipe pipeline. + + Ethos-U and VGF round-trips need a real compiler and are deferred to an FVP- + bearing follow-up. + + """ + + def _export( + self, + recipe: ExportRecipe, + model: torch.nn.Module, + example_inputs: tuple, + ): + from executorch.export import export + + session = export( + model=model, + example_inputs=[example_inputs], + export_recipe=recipe, + ) + return session.get_executorch_program() + + def _instruction_kinds(self, program) -> tuple[list, list]: + from executorch.exir.schema import DelegateCall, KernelCall + + instructions = program.execution_plan[0].chains[0].instructions + assert instructions is not None + operators = program.execution_plan[0].operators + delegate_calls = [ + i for i in instructions if isinstance(i.instr_args, DelegateCall) + ] + kernel_op_names = [ + operators[i.instr_args.op_index].name + for i in instructions + if isinstance(i.instr_args, KernelCall) + ] + return delegate_calls, kernel_op_names + + def test_tosa_fp_export(self) -> None: + # FP path: no quant ops, expect full delegation (Add is supported by TOSA). + program = self._export( + ExportRecipe.get_recipe(ArmRecipeType.TOSA_FP), + _AddModule(), + (torch.randn(2, 3), torch.randn(2, 3)), + ) + delegates, kernels = self._instruction_kinds(program) + self.assertEqual(len(delegates), 1, "Add should produce one TOSA delegate") + self.assertEqual( + kernels, [], f"Expected full delegation, got kernels {kernels}" + ) + + def test_tosa_int8_export(self) -> None: + # INT8 path: boundary quantize/dequantize remain outside the delegate + # and ReplaceQuantNodesPass rewrites them to cortex_m::*. + program = self._export( + ExportRecipe.get_recipe(ArmRecipeType.TOSA_INT8), + _ConvReluModule(), + (torch.randn(1, 3, 8, 8),), + ) + delegates, kernels = self._instruction_kinds(program) + self.assertGreaterEqual(len(delegates), 1, "Conv+ReLU should delegate") + for op_name in kernels: + self.assertTrue( + op_name.startswith("cortex_m::"), + f"Non-delegate kernels must be cortex_m boundary ops; got {op_name}", + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index aeb783045a9..53d6ddadd58 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -48,6 +48,11 @@ def define_arm_tests(): "ops/test_split.py", ] + # Export recipes + test_files += [ + "recipes/test_arm_recipes.py", + ] + # Quantization test_files += [ "quantizer/test_generic_annotater.py", @@ -131,6 +136,7 @@ def define_arm_tests(): "//executorch/backends/arm/test/misc:dw_convs_shared_weights_module", "//executorch/backends/arm:ao_ext", "//executorch/backends/arm:ethosu", + "//executorch/backends/arm/recipes:recipes", "//executorch/backends/arm/tosa:compile_spec", "//executorch/backends/arm/tosa:partitioner", "//executorch/backends/arm:vgf", From 758a8966eb015a2311e3a4a4eed025a907b2898e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 4 Sep 2026 12:20:11 -0700 Subject: [PATCH 051/190] Read a single tensor element as the type the caller asked for (#22567) ## Summary Running a model on the CUDA backend can fail the moment it loads: ``` symbol lookup error: executorch_cuda_0.so: undefined symbol: aoti_torch_item_int64 ``` The compiler that produces the shared library emits a call to read one element out of a tensor, and there are fourteen of those calls, one per type. This runtime defined a single one, for booleans, so a model reading an integer or a float referred to something that was never there. Nothing catches it when the library is built, because that name is resolved when the library loads. ## The type in the name is the return type A branch selector is stored as a boolean and read through the integer call, because the generated branch table compares integers. So these calls convert. The one that was already here insisted on an exact match, which is wrong for the same reason; it had not caused trouble only because nothing had reached it with anything else yet. Two things follow. The element is read as whatever the tensor actually holds, because reading it as the wrong width would copy more bytes than the tensor owns. And the value is refused only when it does not fit, which is what the reference implementation does and what this project already does elsewhere when narrowing a number. ## Test plan Compiled the file and read the exported names out of the object file. Before, one name. After, eight: ``` aoti_torch_item_bool aoti_torch_item_int8 aoti_torch_item_int16 aoti_torch_item_int32 aoti_torch_item_int64 aoti_torch_item_uint8 aoti_torch_item_float32 aoti_torch_item_bfloat16 ``` That check fails on the previous version. Two tests asserted that a mismatched type is refused. That is the normal case, not an error, so they are replaced with tests of what the generated code actually does: a boolean read as an integer returns one, an integer read as a boolean returns true, and an integer too large for a byte is refused. The first two fail on the previous version. The third fails on a version that converts without checking, where three hundred silently becomes forty four. Those conversions were measured, not assumed, including the two that must keep working: a boolean widened to an integer, and a fraction read as a boolean, which the reference implementation allows deliberately. The remaining calls upstream declares need types this runtime does not have, so a model cannot ask for them. Co-authored-by: PyTorch Bot --- backends/cuda/runtime/shims/memory.cpp | 109 +++++++++++++----- backends/cuda/runtime/shims/memory.h | 43 ++++++- .../shims/tests/test_aoti_torch_item_bool.cpp | 64 +++++----- 3 files changed, 148 insertions(+), 68 deletions(-) diff --git a/backends/cuda/runtime/shims/memory.cpp b/backends/cuda/runtime/shims/memory.cpp index 976b282a29c..bfd43c25090 100644 --- a/backends/cuda/runtime/shims/memory.cpp +++ b/backends/cuda/runtime/shims/memory.cpp @@ -6,7 +6,9 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include @@ -31,10 +33,63 @@ const PalInitializer kPalInitializer{}; } // namespace namespace c10 = executorch::backends::aoti::slim::c10; + using c10::Device; using c10::DeviceIndex; using c10::DeviceType; using c10::ScalarType; + +namespace { + +// Reads the one element as T. Dispatch on the dtype the tensor actually holds: +// item() copies sizeof(T) bytes and does not check, so asking for the wrong +// width reads past a one-element allocation. +// +// The dtype in each entry point's name is the type the caller wants back, not +// the type the tensor holds. Generated code reads a boolean branch selector +// through the int64 entry point, for instance. So convert, and refuse only when +// the value does not fit, which is what the reference implementation does. +template +AOTITorchError narrow_to(From value, To* ret_value) { + // Anything at all converts to a boolean, so there is nothing to check there. + if (!std::is_same_v && ::c10::overflows(value)) { + ET_CHECK_OR_RETURN_ERROR( + false, InvalidArgument, "reading a single element: value does not fit"); + } + *ret_value = static_cast(value); + return Error::Ok; +} + +template +AOTITorchError read_one_element(const SlimTensor* tensor, T* ret_value) { + switch (tensor->dtype()) { + case ScalarType::Byte: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Char: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Short: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Int: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Long: + return narrow_to(tensor->item(), ret_value); + case ScalarType::Float: + return narrow_to(tensor->item(), ret_value); + case ScalarType::BFloat16: + return narrow_to( + static_cast(tensor->item()), ret_value); + case ScalarType::Bool: + return narrow_to(tensor->item(), ret_value); + default: + ET_CHECK_OR_RETURN_ERROR( + false, + InvalidArgument, + "reading a single element: dtype %d is not supported", + static_cast(tensor->dtype())); + } +} + +} // namespace using executorch::backends::aoti::slim::empty_strided; using executorch::backends::aoti::slim::from_blob; using executorch::backends::aoti::slim::IntArrayRef; @@ -347,34 +402,36 @@ aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking) { return Error::Ok; } -AOTITorchError aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value) { - ET_CHECK_OR_RETURN_ERROR( - tensor != nullptr, - InvalidArgument, - "aoti_torch_item_bool: tensor is null"); - - ET_CHECK_OR_RETURN_ERROR( - ret_value != nullptr, - InvalidArgument, - "aoti_torch_item_bool: ret_value is null"); - - ET_CHECK_OR_RETURN_ERROR( - tensor->numel() == 1, - InvalidArgument, - "aoti_torch_item_bool: tensor must have exactly 1 element, got %zu", - tensor->numel()); - - ET_CHECK_OR_RETURN_ERROR( - tensor->dtype() == ScalarType::Bool, - InvalidArgument, - "aoti_torch_item_bool: tensor dtype must be Bool"); +#define ET_CUDA_DEFINE_ITEM_SHIM(SUFFIX, CTYPE) \ + AOTITorchError aoti_torch_item_##SUFFIX( \ + SlimTensor* tensor, CTYPE* ret_value) { \ + ET_CHECK_OR_RETURN_ERROR( \ + tensor != nullptr, \ + InvalidArgument, \ + "aoti_torch_item_" #SUFFIX ": tensor is null"); \ + ET_CHECK_OR_RETURN_ERROR( \ + ret_value != nullptr, \ + InvalidArgument, \ + "aoti_torch_item_" #SUFFIX ": ret_value is null"); \ + ET_CHECK_OR_RETURN_ERROR( \ + tensor->numel() == 1, \ + InvalidArgument, \ + "aoti_torch_item_" #SUFFIX \ + ": tensor must have exactly 1 element, got %zu", \ + tensor->numel()); \ + return read_one_element(tensor, ret_value); \ + } - // SlimTensor::item() handles both CPU and CUDA tensors. - // For CUDA tensors, it copies the value to CPU automatically. - *ret_value = tensor->item(); +ET_CUDA_DEFINE_ITEM_SHIM(uint8, uint8_t) +ET_CUDA_DEFINE_ITEM_SHIM(int8, int8_t) +ET_CUDA_DEFINE_ITEM_SHIM(int16, int16_t) +ET_CUDA_DEFINE_ITEM_SHIM(int32, int32_t) +ET_CUDA_DEFINE_ITEM_SHIM(int64, int64_t) +ET_CUDA_DEFINE_ITEM_SHIM(float32, float) +ET_CUDA_DEFINE_ITEM_SHIM(bfloat16, c10::BFloat16) +ET_CUDA_DEFINE_ITEM_SHIM(bool, bool) - return Error::Ok; -} +#undef ET_CUDA_DEFINE_ITEM_SHIM AOTITorchError aoti_torch_assign_tensors_out( SlimTensor* src, diff --git a/backends/cuda/runtime/shims/memory.h b/backends/cuda/runtime/shims/memory.h index 03e2b3ed18d..96c47c03aed 100644 --- a/backends/cuda/runtime/shims/memory.h +++ b/backends/cuda/runtime/shims/memory.h @@ -195,18 +195,49 @@ AOTI_SHIM_EXPORT AOTITorchError aoti_torch__reinterpret_tensor( AOTI_SHIM_EXPORT AOTITorchError aoti_torch_copy_(SlimTensor* self, SlimTensor* src, int32_t non_blocking); +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value); + /** - * Extracts a boolean scalar value from a single-element tensor. + * Extracts a scalar value from a single-element tensor. * - * The tensor must contain exactly one element and have Bool dtype. - * For CUDA tensors, this will synchronize to copy the value to CPU. + * The type in the name is the type returned, not the tensor's dtype: generated + * code reads a boolean branch selector through the int64 entry point. The value + * is converted, and rejected only when it does not fit. The tensor must contain + * exactly one element. For CUDA tensors, this will synchronize to copy the + * value to CPU. * - * @param tensor Single-element boolean tensor (must not be null) - * @param ret_value Output parameter for the extracted boolean value + * @param tensor Single-element tensor (must not be null) + * @param ret_value Output parameter for the extracted value * @return AOTITorchError error code (Error::Ok on success) */ AOTI_SHIM_EXPORT AOTITorchError -aoti_torch_item_bool(SlimTensor* tensor, bool* ret_value); +aoti_torch_item_uint8(SlimTensor* tensor, uint8_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int8(SlimTensor* tensor, int8_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int16(SlimTensor* tensor, int16_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int32(SlimTensor* tensor, int32_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_int64(SlimTensor* tensor, int64_t* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_float32(SlimTensor* tensor, float* ret_value); + +/// See aoti_torch_item_uint8. +AOTI_SHIM_EXPORT AOTITorchError +aoti_torch_item_bfloat16(SlimTensor* tensor, c10::BFloat16* ret_value); /** * Moves a tensor into a new handle and assigns it to the output parameter. diff --git a/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp b/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp index fa8b5bb9245..cc4b68d372c 100644 --- a/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp +++ b/backends/cuda/runtime/shims/tests/test_aoti_torch_item_bool.cpp @@ -155,49 +155,62 @@ TEST_F(AOTITorchItemBoolSlimTest, NullReturnValue) { EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } -TEST_F(AOTITorchItemBoolSlimTest, MultiElementTensor) { - std::vector sizes = {2, 3}; +TEST_F(AOTITorchItemBoolSlimTest, ConvertsFromLong) { + // Generated code reads a value through whichever entry point matches the type + // it wants back, not the type the tensor holds. + std::vector sizes = {1}; Tensor* tensor = createTestTensor( sizes, - static_cast(slim_c10::ScalarType::Bool), + static_cast(slim_c10::ScalarType::Long), static_cast(slim_c10::DeviceType::CPU), 0); ASSERT_NE(tensor, nullptr); - EXPECT_GT(tensor->numel(), 1); + *static_cast(tensor->data_ptr()) = 1; bool result = false; - AOTITorchError error = aoti_torch_item_bool(tensor, &result); + EXPECT_EQ(aoti_torch_item_bool(tensor, &result), Error::Ok); + EXPECT_TRUE(result); - EXPECT_EQ(error, Error::InvalidArgument); + EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); +} + +TEST_F(AOTITorchItemBoolSlimTest, BoolReadAsInt64) { + Tensor* tensor = createScalarBoolTensor( + true, static_cast(slim_c10::DeviceType::CPU), 0); + ASSERT_NE(tensor, nullptr); + + int64_t result = -1; + EXPECT_EQ(aoti_torch_item_int64(tensor, &result), Error::Ok); + EXPECT_EQ(result, 1); EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } -TEST_F(AOTITorchItemBoolSlimTest, WrongDtype_Float) { +TEST_F(AOTITorchItemBoolSlimTest, RejectsValueThatDoesNotFit) { std::vector sizes = {1}; Tensor* tensor = createTestTensor( sizes, - static_cast(slim_c10::ScalarType::Float), + static_cast(slim_c10::ScalarType::Long), static_cast(slim_c10::DeviceType::CPU), 0); ASSERT_NE(tensor, nullptr); + *static_cast(tensor->data_ptr()) = 300; - bool result = false; - AOTITorchError error = aoti_torch_item_bool(tensor, &result); - - EXPECT_EQ(error, Error::InvalidArgument); + uint8_t result = 0; + EXPECT_EQ(aoti_torch_item_uint8(tensor, &result), Error::InvalidArgument); EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } -TEST_F(AOTITorchItemBoolSlimTest, WrongDtype_Long) { - std::vector sizes = {1}; +TEST_F(AOTITorchItemBoolSlimTest, MultiElementTensor) { + std::vector sizes = {2, 3}; Tensor* tensor = createTestTensor( sizes, - static_cast(slim_c10::ScalarType::Long), + static_cast(slim_c10::ScalarType::Bool), static_cast(slim_c10::DeviceType::CPU), 0); ASSERT_NE(tensor, nullptr); + EXPECT_GT(tensor->numel(), 1); bool result = false; AOTITorchError error = aoti_torch_item_bool(tensor, &result); @@ -270,24 +283,3 @@ TEST_F(AOTITorchItemBoolSlimTest, MultiElementTensor_CUDA) { EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); } - -TEST_F(AOTITorchItemBoolSlimTest, WrongDtype_Float_CUDA) { - if (!isCudaAvailable()) { - GTEST_SKIP() << "CUDA not available"; - } - - std::vector sizes = {1}; - Tensor* tensor = createTestTensor( - sizes, - static_cast(slim_c10::ScalarType::Float), - static_cast(slim_c10::DeviceType::CUDA), - 0); - ASSERT_NE(tensor, nullptr); - - bool result = false; - AOTITorchError error = aoti_torch_item_bool(tensor, &result); - - EXPECT_EQ(error, Error::InvalidArgument); - - EXPECT_EQ(aoti_torch_delete_tensor_object(tensor), Error::Ok); -} From 32722a32f7d423646d8c5278d03f63ed6f72e51d Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 4 Sep 2026 13:00:41 -0700 Subject: [PATCH 052/190] Fine-tune MobileBert on smaller batches in its test (#22565) ## Summary The MobileBert test in the Samsung model suite is killed for using too much memory. That stops the whole suite, and because the suite is a required check, it stops the branch nightly builds are cut from moving forward. The test asks for a fine-tuned model, and fine-tuning means real training: five passes over a sentiment dataset, sixty four sentences at a time. Training keeps every intermediate value until the backward pass finishes, and a sentence of a hundred and twenty eight tokens through twenty four layers produces a lot of them. Measured on a single training step, same model and sequence length: ``` batch 64 peak memory 7.58 GiB batch 8 peak memory 1.43 GiB ``` Smaller batches take longer but hold far less at once, and they do not change what the model learns, only how often it steps. So the batch size becomes an argument, today's value stays the default, and the test asks for a smaller one. ## What this deliberately does not change The number of passes is left alone. Cutting it would shorten the test further, but the output range the model ends up with is not a falling function of it. Measured on a fixed probe, the largest logit is about two million untuned, about half a million after one pass, and about two and a half million after two. Fewer passes is therefore not obviously safer for the backend, so the memory problem is solved without touching it. ## Test plan Measured the peak memory of one training step at both batch sizes, on Linux, with the same model and sequence length: 7.58 GiB against 1.43 GiB, a bit over five times less. Ran the suite's own order in one process to find where the peak comes from. Two vision models ahead of this one reach half a gigabyte between them, so the training step is the peak rather than an accumulation across models. Confirmed the example script is unchanged. It passes only the artifacts directory, so the batch size falls back to its default and the accuracy it prints is produced exactly as before. Co-authored-by: PyTorch Bot --- backends/samsung/test/models/test_mobilebert_finetuning.py | 6 +++++- examples/samsung/scripts/mobilebert_finetune.py | 6 +++--- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/backends/samsung/test/models/test_mobilebert_finetuning.py b/backends/samsung/test/models/test_mobilebert_finetuning.py index 9ed7440d5ba..0c695f14045 100644 --- a/backends/samsung/test/models/test_mobilebert_finetuning.py +++ b/backends/samsung/test/models/test_mobilebert_finetuning.py @@ -18,7 +18,11 @@ class Test_Milestone_MobileBertFinetune(unittest.TestCase): def test_mobilebert_finetuning_fp16(self): mobilebert_finetune = MobileBertFinetune() - model, _ = mobilebert_finetune.get_finetune_mobilebert(None) + # Smaller batches, because training keeps every intermediate value until + # the backward pass is done and a test runner cannot hold them all. The + # number of passes is left alone: it is what the example script uses, and + # the range this model ends up with does not fall off with fewer of them. + model, _ = mobilebert_finetune.get_finetune_mobilebert(None, batch_size=8) example_input = mobilebert_finetune.get_example_inputs() tester = SamsungTester( model, example_input, [gen_samsung_backend_compile_spec(TestConfig.chipset)] diff --git a/examples/samsung/scripts/mobilebert_finetune.py b/examples/samsung/scripts/mobilebert_finetune.py index 76c1e9b03d3..bee2909ffb2 100644 --- a/examples/samsung/scripts/mobilebert_finetune.py +++ b/examples/samsung/scripts/mobilebert_finetune.py @@ -117,7 +117,7 @@ def build_loader_from_dataset(self, dataset, batch_size, usage="train"): return data_loader - def get_finetune_mobilebert(self, artifacts_dir): + def get_finetune_mobilebert(self, artifacts_dir, batch_size=64): # Pretrained bert's output ranges in a large scale. It is challenge for enn backend to support directly. # Please finetune mobilebert on specific tasks, make sure that bert's output and hidden states are friendly # to resource-constraint device. @@ -138,7 +138,7 @@ def get_finetune_mobilebert(self, artifacts_dir): labels_set = train_data.label.unique() train_data_loader = self.build_loader_from_dataset( - train_data, batch_size=64, usage="train" + train_data, batch_size=batch_size, usage="train" ) val_url = "https://raw.githubusercontent.com/clairett/pytorch-sentiment-classification/refs/heads/master/data/SST2/test.tsv" @@ -147,7 +147,7 @@ def get_finetune_mobilebert(self, artifacts_dir): BytesIO(content), delimiter="\t", header=None, names=["text", "label"] ) val_data_loader = self.build_loader_from_dataset( - val_data, batch_size=64, usage="val" + val_data, batch_size=batch_size, usage="val" ) artifacts_dir = artifacts_dir if artifacts_dir is not None else "./mobilebert" From b903c2ab1517036f6217b6773863851fdef1af87 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 4 Sep 2026 13:05:27 -0700 Subject: [PATCH 053/190] Detect ATen when a target names it by its resolved label (#22553) ## Summary A target can name a third-party dependency two ways. It can give the short name, which goes in `external_deps` and is resolved later, or it can call `external_dep_location`, which hands back the resolved label so the target puts it in an ordinary `deps` list. The check for whether a target compiles against ATen only looked at the first. A target using the second was treated as having no ATen dependency and was compiled at C++17. PyTorch's headers need C++20 now, so those targets are left compiling against headers they cannot. The Vulkan operator tests are what this reaches. They name `libtorch` through `external_dep_location`, and the open source Buck build cannot query that backend at all, so nothing here builds them and the mismatch does not surface until someone does. Those same tests already ask for C++20 in their CMake build, so the requirement is not in question, only the Buck side of it. This resolves the same names the check already knows and compares against the labels they resolve to. ## Test plan Added a unit test for the decision, next to the other tests of build script logic and for the same reason: the targets it matters for cannot be built here, while the decision itself is a plain function of a target's arguments. It covers a resolved label in `deps` and in `exported_deps`, every short name in `external_deps`, a plain target with no ATen dependency, and a target whose only dependency is a project label, since every one of those contains the word torch and an earlier version of this check matched on that substring. The test fails without the change: two of its five cases fail against the previous version of the file, and all five pass with it. Co-authored-by: PyTorch Bot --- .ci/scripts/tests/test_is_aten_target.py | 151 ++++++++++++++++++ .../executorch/build/runtime_wrapper.bzl | 27 ++++ 2 files changed, 178 insertions(+) create mode 100644 .ci/scripts/tests/test_is_aten_target.py diff --git a/.ci/scripts/tests/test_is_aten_target.py b/.ci/scripts/tests/test_is_aten_target.py new file mode 100644 index 00000000000..b8debfddcc9 --- /dev/null +++ b/.ci/scripts/tests/test_is_aten_target.py @@ -0,0 +1,151 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the ATen detection in the Buck macro layer. + +Here rather than as a build test, because the open source Buck build cannot query the +Vulkan backend at all, so the targets this decision matters most for are never built by +CI. The decision itself is a pure function of a target's keyword arguments, so it can be +exercised directly. + +The detection had a real defect that this covers. A target can name a third-party +dependency either by its short name, which lands in ``external_deps``, or through +``external_dep_location``, which hands back the resolved label and lands in an ordinary +``deps`` list. Only the first was checked, so the Vulkan operator tests, which use the +second, compiled at the older standard against headers that need the newer one. +""" + +import types +import unittest +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[3] +MACROS = ( + REPO_ROOT / "shim_et" / "xplat" / "executorch" / "build" / "runtime_wrapper.bzl" +) + +# What the open source dependency map resolves these names to. Kept here rather than +# imported so the test states the mapping it depends on. +RESOLVED = { + "c10": ["//third-party:libtorch"], + "libtorch": ["//third-party:libtorch"], + "libtorch_python": ["//third-party:libtorch_python"], + "torch-core-cpp": ["//third-party:libtorch"], + "gtest_aten": ["//third-party:gtest_aten"], + "gmock_aten": ["//third-party:gmock_aten"], +} +FALLTHROUGH = "@fallthrough@" + + +def _load_is_aten_target(): + """Execute the real macro text, with the little of Starlark it uses shimmed.""" + text = MACROS.read_text() + start = text.index("def _has_pytorch_dep") + end = text.index("def _cxx_library_common", start) + + env = types.SimpleNamespace( + EXTERNAL_DEP_FALLTHROUGH=FALLTHROUGH, + resolve_external_dep=lambda name: RESOLVED.get(name, FALLTHROUGH), + ) + + def _apply(obj, function): + """Stand-in for selects.apply: run over each list the object holds.""" + if isinstance(obj, dict): + return {key: function(value) for key, value in obj.items()} + return function(obj) + + namespace = { + # Starlark's type() returns a name, and the macros compare against "string". + "type": lambda value: "string" if isinstance(value, str) else "other", + "env": env, + "selects": types.SimpleNamespace(apply=_apply), + } + exec(compile(text[start:end], str(MACROS), "exec"), namespace) + return namespace["_is_aten_target"] + + +class TestIsAtenTarget(unittest.TestCase): + def setUp(self) -> None: + self.is_aten_target = _load_is_aten_target() + + def test_resolved_label_in_deps(self) -> None: + """The Vulkan operator tests name libtorch this way.""" + self.assertTrue( + self.is_aten_target( + { + "name": "compute_graph_op_tests_bin", + "deps": [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + "//third-party:libtorch", + ], + } + ) + ) + + def test_resolved_label_in_exported_deps(self) -> None: + self.assertTrue( + self.is_aten_target( + {"name": "some_lib", "exported_deps": ["//third-party:libtorch"]} + ) + ) + + def test_short_name_in_external_deps(self) -> None: + for name in RESOLVED: + with self.subTest(name=name): + self.assertTrue( + self.is_aten_target({"name": "some_test", "external_deps": [name]}) + ) + + def test_plain_target_is_not_aten(self) -> None: + """The embedded builds rely on these staying at the older standard.""" + self.assertFalse( + self.is_aten_target( + { + "name": "op_add_test", + "deps": [ + "//executorch/runtime/core:core", + "//third-party/googletest:gtest_main", + ], + } + ) + ) + + def test_executorch_label_alone_is_not_aten(self) -> None: + """Every label under the project contains the word torch.""" + self.assertFalse( + self.is_aten_target( + {"name": "evalue_test", "deps": ["//executorch/test/utils:utils"]} + ) + ) + + def test_resolved_label_inside_a_select(self) -> None: + """A dep list can be a select, which cannot be walked like a list.""" + self.assertTrue( + self.is_aten_target( + { + "name": "some_test", + "deps": { + "DEFAULT": ["//third-party:libtorch"], + "ovr_config//os:windows": [], + }, + } + ) + ) + + def test_select_without_aten_is_not_aten(self) -> None: + self.assertFalse( + self.is_aten_target( + { + "name": "some_test", + "deps": {"DEFAULT": ["//executorch/runtime/core:core"]}, + } + ) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/shim_et/xplat/executorch/build/runtime_wrapper.bzl b/shim_et/xplat/executorch/build/runtime_wrapper.bzl index 445c1aa3b98..85c390fe6c2 100644 --- a/shim_et/xplat/executorch/build/runtime_wrapper.bzl +++ b/shim_et/xplat/executorch/build/runtime_wrapper.bzl @@ -162,6 +162,33 @@ def _is_aten_target(kwargs): for dep in kwargs.get(key) or []: if dep in aten_external_deps: return True + + # A target can also name one of those through external_dep_location, which + # hands back the resolved label and puts it in an ordinary dep list. + aten_targets = [] + for name in aten_external_deps: + resolved = env.resolve_external_dep(name) + if resolved != env.EXTERNAL_DEP_FALLTHROUGH: + for target in resolved: + if target not in aten_targets: + aten_targets.append(target) + + # A dep list can be a select(), so collect through selects.apply rather than + # walking it. The lists it holds are the same shape either way. + found = [] + + def _note_aten_deps(targets): + for dep in targets: + if dep in aten_targets: + found.append(dep) + return targets + + for key in ["deps", "exported_deps"]: + if kwargs.get(key): + selects.apply(obj = kwargs.get(key), function = _note_aten_deps) + if found: + return True + for key in ["xplat_deps", "fbcode_deps"]: if _has_pytorch_dep(kwargs.get(key)): return True From f624c34665128dc92eb870d7c6e3dce5b69c091d Mon Sep 17 00:00:00 2001 From: Huy Do Date: Fri, 4 Sep 2026 13:20:26 -0700 Subject: [PATCH 054/190] Move the shared reusable workflows to linux_job_v3 (#22246) Second of six splitting up #22107. Stacked on the landed #22106. Moves the shared reusable workflows to `linux_job_v3`: `_test_backend.yml`, `_test_cadence.yml`, `_test_cortex_m_e2e.yml`, `_test_cortex_m_ops.yml`, `_test_riscv.yml`, `_test_arduino_library.yml`, `_llm_server.yml`, plus the nine `test-backend-*.yml` callers. Every job also changes runner fleet, from EC2 to OSDC: | EC2 | OSDC | |---|---| | `linux.2xlarge` | `mt-l-x86iavx512-8-64` | | `linux.4xlarge.memory` (`_test_backend` default) | `mt-l-x86iavx512-16-128` | | `linux.8xlarge.memory` (`_test_cadence` default, `test-backend-qnn`) | `mt-l-x86iavx512-32-256` | Also caps pytest-xdist's worker count. `-n auto` and `-n logical` size themselves from `psutil.cpu_count()`, which reports the node rather than the pod's cpuset, so on `mt-l-x86iavx512-16-128` the backend suite started 48 export workers on 116Gi instead of 8 on 128Gi. `PYTEST_XDIST_AUTO_NUM_WORKERS` from `nproc` pins it to what the pod actually has. **Test plan.** `test-backend-*` and arduino run on this pull request and are green. `_test_riscv`, `_llm_server` and the two cortex-m workflows are only reachable from `riscv64.yml`, trunk and nightly, none of which this change's paths match, so they were dispatched by hand. Fork pull requests are a known gap, see below. **Known gap, not introduced here.** A fork pull request gets no OIDC token, so v3's role assume fails and sccache dies against a bogus `the C compiler is not able to compile a simple test program`. This is already live on main from the unit test migration, and pytorch/test-infra#8735 fixes it for every repo on v3. Authored with Claude Code. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- .ci/scripts/pytest-parallelism.sh | 26 +++++++++++++++++++++ .ci/scripts/test_backend.sh | 3 +++ .ci/scripts/unittest-linux-cmake.sh | 3 +++ .github/workflows/_docker-image.yml | 16 +++++++++++-- .github/workflows/_llm_server.yml | 13 +++++++---- .github/workflows/_test_arduino_library.yml | 17 +++++++++----- .github/workflows/_test_backend.yml | 17 ++++++++++---- .github/workflows/_test_cadence.yml | 24 ++++++++++++++----- .github/workflows/_test_cortex_m_e2e.yml | 14 ++++++++--- .github/workflows/_test_cortex_m_ops.yml | 14 ++++++++--- .github/workflows/_test_riscv.yml | 14 ++++++++--- .github/workflows/riscv64.yml | 1 + .github/workflows/test-backend-arm.yml | 6 +++++ .github/workflows/test-backend-coreml.yml | 3 +++ .github/workflows/test-backend-cortex-m.yml | 3 +++ .github/workflows/test-backend-nxp.yml | 3 +++ .github/workflows/test-backend-openvino.yml | 3 +++ .github/workflows/test-backend-qnn.yml | 5 +++- .github/workflows/test-backend-vulkan.yml | 3 +++ .github/workflows/test-backend-webgpu.yml | 3 +++ .github/workflows/test-backend-xnnpack.yml | 3 +++ backends/arm/test/test_arm_backend.sh | 3 +++ backends/nxp/run_unittests.sh | 4 ++++ examples/riscv/README.md | 8 +++---- 24 files changed, 173 insertions(+), 36 deletions(-) create mode 100755 .ci/scripts/pytest-parallelism.sh diff --git a/.ci/scripts/pytest-parallelism.sh b/.ci/scripts/pytest-parallelism.sh new file mode 100755 index 00000000000..97d50061cf3 --- /dev/null +++ b/.ci/scripts/pytest-parallelism.sh @@ -0,0 +1,26 @@ +#!/bin/bash +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# Pin pytest-xdist's worker count, but only when the container is allowed less +# CPU than the machine it landed on. `auto` asks psutil for the machine's +# physical cores, which inside an OSDC pod is the whole node, so the workers +# exhaust the pod's memory. nproc honours the pod's cpuset, which is what +# pytorch relies on for OMP_NUM_THREADS on the same fleet. +# +# Left alone when unconstrained. `auto` already discounts hyperthreads there, +# while nproc counts them, so overriding it would double the workers and halve +# the memory each one gets. + +if [[ -z "${PYTEST_XDIST_AUTO_NUM_WORKERS:-}" ]] && command -v nproc >/dev/null 2>&1; then + cpus_allowed="$(nproc)" + cpus_installed="$(nproc --all)" + echo "pytest-xdist: ${cpus_allowed} of ${cpus_installed} CPUs available" + if [[ "${cpus_allowed}" -lt "${cpus_installed}" ]]; then + export PYTEST_XDIST_AUTO_NUM_WORKERS="${cpus_allowed}" + echo "PYTEST_XDIST_AUTO_NUM_WORKERS=${PYTEST_XDIST_AUTO_NUM_WORKERS}" + fi +fi diff --git a/.ci/scripts/test_backend.sh b/.ci/scripts/test_backend.sh index 45e3322a0c8..32c682ad5f0 100755 --- a/.ci/scripts/test_backend.sh +++ b/.ci/scripts/test_backend.sh @@ -7,6 +7,9 @@ # LICENSE file in the root directory of this source tree. set -eux +# Cap pytest-xdist's `auto` workers to the container's CPU quota. +source .ci/scripts/pytest-parallelism.sh + SUITE=$1 FLOW=$2 ARTIFACT_DIR=$3 diff --git a/.ci/scripts/unittest-linux-cmake.sh b/.ci/scripts/unittest-linux-cmake.sh index 0f750e1fe13..83f2f1464ee 100755 --- a/.ci/scripts/unittest-linux-cmake.sh +++ b/.ci/scripts/unittest-linux-cmake.sh @@ -7,6 +7,9 @@ # LICENSE file in the root directory of this source tree. set -eux +# Cap pytest-xdist's `auto` workers to the container's CPU quota. +source .ci/scripts/pytest-parallelism.sh + # Some ARM/TOSA-adjacent tests import modules that require tosa_serializer. # Install from a local tosa-tools checkout when available. If absent in this # checkout layout, clone the pinned upstream tag and install from there. diff --git a/.github/workflows/_docker-image.yml b/.github/workflows/_docker-image.yml index a328a8fc3ca..9dc4813e16e 100644 --- a/.github/workflows/_docker-image.yml +++ b/.github/workflows/_docker-image.yml @@ -37,10 +37,22 @@ jobs: - name: Checkout ExecuTorch uses: actions/checkout@v4 with: - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + # build-cadence-runner.yml reaches this workflow on pull_request_target, + # where github.sha is the base branch tip. Resolving the tag from there + # while the test jobs check out the fork head would test a docker change + # against the old image. + ref: ${{ (github.event_name == 'pull_request' || github.event_name == 'pull_request_target') && github.event.pull_request.head.sha || github.sha }} + # Nothing here runs the checked-out code, only `git rev-parse` over it, + # but on pull_request_target that code is the fork's and the token is + # the base repository's, so do not leave one next to the other. + persist-credentials: false - name: Compute the docker tag id: hash run: | set -eu - echo "ci-docker-hash=$(git rev-parse HEAD:.ci/docker)" >> "$GITHUB_OUTPUT" + # Assigned on its own line: set -e does not catch a failure inside a + # command substitution, so `echo "x=$(git rev-parse ...)"` would write + # git's error text as the tag and exit 0. + CI_DOCKER_HASH="$(git rev-parse HEAD:.ci/docker)" + echo "ci-docker-hash=${CI_DOCKER_HASH}" >> "$GITHUB_OUTPUT" diff --git a/.github/workflows/_llm_server.yml b/.github/workflows/_llm_server.yml index e1ef5a30db9..6492a3106ee 100644 --- a/.github/workflows/_llm_server.yml +++ b/.github/workflows/_llm_server.yml @@ -4,20 +4,25 @@ on: workflow_call: inputs: docker-image: - description: Docker image to use for Linux tests. + description: Name of the docker image to use, without registry or tag suffix. required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + linux: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ${{ inputs.docker-image }} + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 60 diff --git a/.github/workflows/_test_arduino_library.yml b/.github/workflows/_test_arduino_library.yml index a3e73958463..5001202c069 100644 --- a/.github/workflows/_test_arduino_library.yml +++ b/.github/workflows/_test_arduino_library.yml @@ -14,15 +14,20 @@ on: default: 90 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: job-name: arduino-library - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} @@ -53,9 +58,9 @@ jobs: # moving branch: reproducible, and one less thing to fail transiently. ARDUINO_CLI_VERSION=1.5.1 ARDUINO_LINT_VERSION=1.3.0 - # RUNNER_TEMP and GITHUB_WORKSPACE both point at host paths this - # container cannot write to, and HOME may too, so keep the toolchain - # and every arduino-cli directory somewhere local to the container. + # HOME is /github/home, which arduino-cli's directories should not share + # with the runner's own state, so keep the toolchain and every + # arduino-cli directory somewhere local to the container. ARDUINO_CI_DIR=/tmp/.arduino-ci export ARDUINO_DIRECTORIES_USER="${ARDUINO_CI_DIR}/user" export ARDUINO_DIRECTORIES_DATA="${ARDUINO_CI_DIR}/data" diff --git a/.github/workflows/_test_backend.yml b/.github/workflows/_test_backend.yml index 18c70a31531..dda19fa033f 100644 --- a/.github/workflows/_test_backend.yml +++ b/.github/workflows/_test_backend.yml @@ -40,15 +40,21 @@ on: description: 'Runner type for Linux jobs' required: false type: string - default: linux.4xlarge.memory + default: mt-l-x86iavx512-16-128 docker-image: - description: 'Docker image for Linux jobs' + description: 'Name of the docker image to use, without registry or tag suffix' required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 jobs: + docker-image: + name: Resolve CI docker image + if: ${{ inputs.run-linux }} + uses: ./.github/workflows/_docker-image.yml + test-backend-linux: + needs: docker-image if: ${{ inputs.run-linux }} strategy: fail-fast: false @@ -57,11 +63,14 @@ jobs: suite: [models, operators] exclude: ${{ fromJSON(inputs.exclude) }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: ref: ${{ inputs.ref }} runner: ${{ inputs.runner-linux }} - docker-image: ${{ inputs.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive timeout: ${{ inputs.timeout }} upload-artifact: test-report-${{ inputs.backend }}-${{ matrix.flow }}-${{ matrix.suite }} diff --git a/.github/workflows/_test_cadence.yml b/.github/workflows/_test_cadence.yml index 2e98d21db1c..d899e2ae585 100644 --- a/.github/workflows/_test_cadence.yml +++ b/.github/workflows/_test_cadence.yml @@ -8,7 +8,7 @@ on: workflow_call: inputs: docker-image: - description: 'Docker image to use' + description: 'Name of the docker image to use, without registry or tag suffix' required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 @@ -16,7 +16,7 @@ on: description: 'Runner type' required: false type: string - default: linux.8xlarge.memory + default: mt-l-x86iavx512-32-256 ref: description: 'Git ref to checkout' required: false @@ -29,12 +29,20 @@ on: default: 90 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + test-aot: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: job-name: test-aot runner: ${{ inputs.runner }} - docker-image: ${{ inputs.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive ref: ${{ inputs.ref }} timeout: ${{ inputs.timeout }} @@ -50,11 +58,15 @@ jobs: python -m pytest backends/cadence/aot/tests/ -v -n auto --reruns 2 --reruns-delay 1 test-ops: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: job-name: test-ops runner: ${{ inputs.runner }} - docker-image: ${{ inputs.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/${{ inputs.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: recursive ref: ${{ inputs.ref }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/_test_cortex_m_e2e.yml b/.github/workflows/_test_cortex_m_e2e.yml index 0510b017723..1957a2b0c56 100644 --- a/.github/workflows/_test_cortex_m_e2e.yml +++ b/.github/workflows/_test_cortex_m_e2e.yml @@ -23,8 +23,16 @@ on: default: 120 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read strategy: matrix: model: ${{ fromJSON(inputs.models) }} @@ -32,8 +40,8 @@ jobs: fail-fast: false with: job-name: ${{ matrix.model }}-${{ matrix.target }} - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/_test_cortex_m_ops.yml b/.github/workflows/_test_cortex_m_ops.yml index a9e2e6180c3..885f5fb64d9 100644 --- a/.github/workflows/_test_cortex_m_ops.yml +++ b/.github/workflows/_test_cortex_m_ops.yml @@ -18,16 +18,24 @@ on: default: 120 jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read strategy: matrix: target: ${{ fromJSON(inputs.targets) }} fail-fast: false with: job-name: cortex-m-ops-${{ matrix.target }} - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/_test_riscv.yml b/.github/workflows/_test_riscv.yml index 223a146e3d8..954cb94d596 100644 --- a/.github/workflows/_test_riscv.yml +++ b/.github/workflows/_test_riscv.yml @@ -37,11 +37,19 @@ on: type: string jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + run: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-24.04-gcc14 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-gcc14-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ inputs.timeout }} diff --git a/.github/workflows/riscv64.yml b/.github/workflows/riscv64.yml index 244db631021..34f4ab04874 100644 --- a/.github/workflows/riscv64.yml +++ b/.github/workflows/riscv64.yml @@ -10,6 +10,7 @@ on: pull_request: paths: - .github/workflows/riscv64.yml + - .github/workflows/_test_riscv.yml - .ci/scripts/test_riscv_qemu.sh - tools/cmake/preset/riscv64_linux.cmake - examples/riscv/** diff --git a/.github/workflows/test-backend-arm.yml b/.github/workflows/test-backend-arm.yml index d71696ee096..7d844748dac 100644 --- a/.github/workflows/test-backend-arm.yml +++ b/.github/workflows/test-backend-arm.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-arm: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: arm flows: >- @@ -35,6 +38,9 @@ jobs: test-arm-vgf: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: arm-vgf flows: >- diff --git a/.github/workflows/test-backend-coreml.yml b/.github/workflows/test-backend-coreml.yml index 86844ffb559..fca10b619ce 100644 --- a/.github/workflows/test-backend-coreml.yml +++ b/.github/workflows/test-backend-coreml.yml @@ -36,6 +36,9 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.github/workflows/test-backend-coreml.yml') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/_test_backend.yml') uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: coreml # The heavier coreml_static_int8 macOS matrix saturates the shared diff --git a/.github/workflows/test-backend-cortex-m.yml b/.github/workflows/test-backend-cortex-m.yml index 9e9b4cdcaa1..b6cd09d3073 100644 --- a/.github/workflows/test-backend-cortex-m.yml +++ b/.github/workflows/test-backend-cortex-m.yml @@ -49,6 +49,9 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.github/workflows/test-backend-cortex-m.yml') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/_test_backend.yml') uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: cortex_m flows: '["cortex_m"]' diff --git a/.github/workflows/test-backend-nxp.yml b/.github/workflows/test-backend-nxp.yml index fed7ab7d19b..0db7d7efeac 100644 --- a/.github/workflows/test-backend-nxp.yml +++ b/.github/workflows/test-backend-nxp.yml @@ -39,6 +39,9 @@ jobs: contains(needs.changed-files.outputs.changed-files, '.github/workflows/test-backend-nxp.yml') || contains(needs.changed-files.outputs.changed-files, '.github/workflows/_test_backend.yml') uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: nxp flows: '["nxp_neutron_imxrt700_int8_ptq"]' diff --git a/.github/workflows/test-backend-openvino.yml b/.github/workflows/test-backend-openvino.yml index aeeb01e3eb1..cff7f8ba6e9 100644 --- a/.github/workflows/test-backend-openvino.yml +++ b/.github/workflows/test-backend-openvino.yml @@ -21,6 +21,9 @@ concurrency: jobs: test-openvino: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: openvino flows: '["openvino"]' diff --git a/.github/workflows/test-backend-qnn.yml b/.github/workflows/test-backend-qnn.yml index 939b7c36aee..51207d139b3 100644 --- a/.github/workflows/test-backend-qnn.yml +++ b/.github/workflows/test-backend-qnn.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-qnn: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: qnn flows: >- @@ -28,4 +31,4 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true - runner-linux: linux.8xlarge.memory + runner-linux: mt-l-x86iavx512-32-256 diff --git a/.github/workflows/test-backend-vulkan.yml b/.github/workflows/test-backend-vulkan.yml index 80ac3ee73a1..d3c9829dfe8 100644 --- a/.github/workflows/test-backend-vulkan.yml +++ b/.github/workflows/test-backend-vulkan.yml @@ -21,6 +21,9 @@ jobs: # runners. Runs on every PR and nightly. test-vulkan: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: vulkan flows: >- diff --git a/.github/workflows/test-backend-webgpu.yml b/.github/workflows/test-backend-webgpu.yml index 99dae3ee76d..42e958855ac 100644 --- a/.github/workflows/test-backend-webgpu.yml +++ b/.github/workflows/test-backend-webgpu.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-webgpu: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: webgpu flows: '["webgpu"]' diff --git a/.github/workflows/test-backend-xnnpack.yml b/.github/workflows/test-backend-xnnpack.yml index e9f43608b16..c5abb15e837 100644 --- a/.github/workflows/test-backend-xnnpack.yml +++ b/.github/workflows/test-backend-xnnpack.yml @@ -19,6 +19,9 @@ concurrency: jobs: test-xnnpack: uses: ./.github/workflows/_test_backend.yml + permissions: + id-token: write + contents: read with: backend: xnnpack flows: >- diff --git a/backends/arm/test/test_arm_backend.sh b/backends/arm/test/test_arm_backend.sh index 0a209c8e356..7e2e28d9d51 100755 --- a/backends/arm/test/test_arm_backend.sh +++ b/backends/arm/test/test_arm_backend.sh @@ -12,6 +12,9 @@ script_dir=$(cd -- "$( dirname -- "${BASH_SOURCE[0]}" )" &> /dev/null && pwd) et_root_dir=$(cd ${script_dir}/../../.. && pwd) cd "${et_root_dir}" pwd + +# Cap pytest-xdist's `auto` workers to the container's CPU quota. +source .ci/scripts/pytest-parallelism.sh scratch_dir=${et_root_dir}/examples/arm/arm-scratch setup_path_script=${scratch_dir}/setup_path.sh _setup_msg="please refer to ${et_root_dir}/examples/arm/setup.sh to properly install necessary tools." diff --git a/backends/nxp/run_unittests.sh b/backends/nxp/run_unittests.sh index 66e51c39a1d..d515d12ae1c 100755 --- a/backends/nxp/run_unittests.sh +++ b/backends/nxp/run_unittests.sh @@ -10,6 +10,10 @@ EXECUTORCH_DIR=$(dirname $(dirname $SCRIPT_DIR)) cd $EXECUTORCH_DIR +# Cap pytest-xdist's workers to the container's CPU quota. Applies to +# `-n logical` as well, despite the variable's name. +source .ci/scripts/pytest-parallelism.sh + # '-c /dev/null' is used to ignore root level pytest.ini. pytest -c /dev/null -n "logical" backends/nxp/tests/ diff --git a/examples/riscv/README.md b/examples/riscv/README.md index 563ff4913fd..02a55351a8a 100644 --- a/examples/riscv/README.md +++ b/examples/riscv/README.md @@ -35,7 +35,7 @@ The driver does three steps: ## CI -`.github/workflows/_test_riscv_qemu.yml` is a reusable `workflow_call` -job (mirroring `_test_cortex_m_e2e.yml`) invoked from `pull.yml` to run on -every PR. It runs on the standard `linux.2xlarge` x86_64 runner using the -`executorch-ubuntu-22.04-gcc11` docker image. +`.github/workflows/_test_riscv.yml` is a reusable `workflow_call` +job (mirroring `_test_cortex_m_e2e.yml`) invoked from `riscv64.yml`. It runs on +the `mt-l-x86iavx512-8-64` x86_64 runner using the +`executorch-ubuntu-24.04-gcc14` docker image. From bf88c643ce9abd6197d30b41d1f69a3d1e6009cd Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 13:35:19 -0700 Subject: [PATCH 055/190] Cortex-M: expose explicit-layout AOT (#22545) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Expose the experimental explicit-layout Cortex-M pipeline through `--cortex-m-explicit-layout`. Explicit mode keeps model inputs contiguous and selects the matching quantizer mode and blessed pass sequence. Legacy mode remains the default and continues converting inputs to channels-last before quantization. Explicit mode requires quantization and never falls back to legacy spatial operators. The tests cover mode-specific operator selection, removal of intermediate channels-last dialect nodes, int8 layout-copy boundaries, copy ceilings for MobileNetV2, ResNet8, and Silero VAD, and serialized Conv1d execution on the M55 FVP. The documentation also records how the explicit mode becomes the default and how legacy AOT generation is subsequently removed without removing legacy runtime operators. ### Test plan `source examples/arm/arm-scratch/setup_path.sh && python -m pytest --config-file=backends/arm/test/pytest.ini backends/cortex_m/test/test_explicit_layout.py` `lintrunner -m origin/main` AI-assisted: Codex. --- Stack created with GitHub Stacks CLI • Give Feedback 💬 --- backends/arm/scripts/aot_arm_compiler.py | 38 +++- backends/cortex_m/README.md | 12 +- backends/cortex_m/test/build_test_runner.sh | 1 + .../cortex_m/test/test_explicit_layout.py | 180 ++++++++++++++++++ .../arm-cortex-m/arm-cortex-m-overview.md | 2 + 5 files changed, 226 insertions(+), 7 deletions(-) create mode 100644 backends/cortex_m/test/test_explicit_layout.py diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py index 250606f9a4b..81ca626a031 100644 --- a/backends/arm/scripts/aot_arm_compiler.py +++ b/backends/arm/scripts/aot_arm_compiler.py @@ -626,6 +626,14 @@ def _get_args(): choices=TARGETS, help=f"Target backend. For delegated models: Ethos-U/VGF/TOSA variants. For non-delegated: cortex-m (CMSIS-NN portable kernels). Valid targets: {TARGETS}", ) + parser.add_argument( + "--cortex-m-explicit-layout", + action="store_true", + help=( + "Use explicit NCHW/NHWC permutes for Cortex-M instead of dim-order " + "operators. This is an experimental Cortex-M-only option." + ), + ) # TODO: Remove --evaluate and --evaluate_config completely after a suitable time. # They are deprecated and no longer functional in this script. parser.add_argument( @@ -923,9 +931,16 @@ def _to_edge_cortex_m( """Cortex-M/CMSIS-NN compilation path with no delegation.""" logging.info( f"Using Cortex-M/CMSIS-NN compilation path for cpu={target_config.cpu.name} " - f"backend={target_config.backend.name}" + f"backend={target_config.backend.name} " + f"layout={'explicit' if args.cortex_m_explicit_layout else 'dim-order'}" ) + if args.cortex_m_explicit_layout and not args.quantize: + raise RuntimeError( + "--cortex-m-explicit-layout requires --quantize; explicit layout " + "does not fall back to portable float spatial operators." + ) + def _to_channels_last(x): if isinstance(x, torch.Tensor): if x.dim() == 4: @@ -949,17 +964,26 @@ def _to_channels_last(x): ) model_quant = None else: - model = model.to(memory_format=torch.channels_last) # type: ignore[call-overload] - example_inputs = tuple(_to_channels_last(x) for x in example_inputs) + if not args.cortex_m_explicit_layout: + model = model.to(memory_format=torch.channels_last) # type: ignore[call-overload] + example_inputs = tuple(_to_channels_last(x) for x in example_inputs) + # Refresh fake-tensor strides after changing the captured module's + # memory format so the legacy quantizer sees channels-last inputs. + model = torch.export.export( + model, example_inputs, strict=args.strict_export + ).module() + + quantizer = CortexMQuantizer(use_explicit_layout=args.cortex_m_explicit_layout) - quantizer = CortexMQuantizer() prepared = prepare_pt2e(model, quantizer) if calibration_samples is None: calibration_samples = [example_inputs] for sample in calibration_samples: - prepared(*tuple(_to_channels_last(x) for x in sample)) + if not args.cortex_m_explicit_layout: + sample = tuple(_to_channels_last(x) for x in sample) + prepared(*sample) model_quant = convert_pt2e(prepared) @@ -973,7 +997,9 @@ def _to_channels_last(x): ) pass_manager = CortexMPassManager( - edge.exported_program(), target_config=target_config + edge.exported_program(), + target_config=target_config, + use_explicit_layout=args.cortex_m_explicit_layout, ) edge._edge_programs["forward"] = pass_manager.transform() diff --git a/backends/cortex_m/README.md b/backends/cortex_m/README.md index f077814d8a5..caead2cb7d3 100644 --- a/backends/cortex_m/README.md +++ b/backends/cortex_m/README.md @@ -5,7 +5,17 @@ ## Overview -The Cortex-M backend is implemented as an operator dialect/library based on [CMSIS-NN](https://github.com/ARM-software/CMSIS-NN), together with the `CortexMQuantizer` which targets supported ops, and the `CortexMPassManager` which modifies the exported program to use Cortex-M operators where possible. It is intended for use with **channels-last input** since this is what the accelerated kernels are using. +The Cortex-M backend is implemented as an operator dialect/library based on [CMSIS-NN](https://github.com/ARM-software/CMSIS-NN), together with the `CortexMQuantizer` which targets supported ops, and the `CortexMPassManager` which modifies the exported program to use Cortex-M operators where possible. + +The default AOT path retains the established channels-last input and dim-order contract. An experimental explicit-layout path accepts ordinary contiguous inputs, inserts graph-visible NHWC copies, and lowers spatial kernels to the `cortex_m::*_nhwc` operator family. Enable it with `--cortex-m-explicit-layout`; the two modes do not fall back to or mix with each other. Explicit-layout compilation fails when a spatial operator is not eligible for NHWC lowering, so models using unsupported configurations must use the legacy mode. + +### Explicit-layout migration + +The `use_explicit_layout=True` modes on `CortexMQuantizer` and `CortexMPassManager` are temporary staging APIs. They keep the experimental path isolated while the default modes and `CortexMTester` continue to exercise the legacy path. + +When explicit layout becomes the default, its support table and pass list will become the defaults in `CortexMQuantizer` and `CortexMPassManager`. The existing legacy support table and pass list will remain temporarily behind an opt-out AOT flag. This keeps the public Python entry points stable and switches `CortexMTester` to explicit layout without changing its callers. The direct NHWC kernel tests and explicit-only model tests will then be folded into the normal operator and model suites. + +After the legacy AOT compatibility period, the legacy support table, pass list, input conversion, opt-out flag, and remaining dual-mode tests will be removed. Legacy runtime operators will remain registered so programs serialized by the old AOT path continue to load. For a detailed example of the full lowering flow, see `examples/arm/cortex_m_mv2_example.ipynb`. diff --git a/backends/cortex_m/test/build_test_runner.sh b/backends/cortex_m/test/build_test_runner.sh index 1178e27aa27..39a760ee39f 100755 --- a/backends/cortex_m/test/build_test_runner.sh +++ b/backends/cortex_m/test/build_test_runner.sh @@ -52,6 +52,7 @@ ops_list=( aten::full.out aten::ge.Tensor_out aten::unsqueeze_copy.out + aten::squeeze_copy.dim_out aten::select_copy.int_out aten::amax.out cortex_m::quantize_per_tensor.out diff --git a/backends/cortex_m/test/test_explicit_layout.py b/backends/cortex_m/test/test_explicit_layout.py new file mode 100644 index 00000000000..1e9ae4488d8 --- /dev/null +++ b/backends/cortex_m/test/test_explicit_layout.py @@ -0,0 +1,180 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from types import SimpleNamespace + +import pytest +import torch +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMSerialize +from executorch.exir.dialects._ops import ops as exir_ops + + +_LEGACY_SPATIAL_OPS = { + exir_ops.edge.cortex_m.quantized_conv2d.default, + exir_ops.edge.cortex_m.quantized_depthwise_conv2d.default, + exir_ops.edge.cortex_m.quantized_transpose_conv2d.default, + exir_ops.edge.cortex_m.quantized_avg_pool2d.default, + exir_ops.edge.cortex_m.quantized_max_pool2d.default, +} + +# Temporary dual-mode acceptance coverage; see "Explicit-layout migration" in +# the Cortex-M README. + + +class Conv2d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +class Conv1d(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.conv = torch.nn.Conv1d(2, 4, 3, padding=1) + + def forward(self, x): + return self.conv(x) + + +def _compile(module, inputs, *, explicit_layout: bool, quantize: bool = True): + from executorch.backends.arm.scripts.aot_arm_compiler import _to_edge_cortex_m + + exported_program = torch.export.export(module, inputs, strict=True) + return _to_edge_cortex_m( + exported_program, + SimpleNamespace( + cortex_m_explicit_layout=explicit_layout, + quantize=quantize, + strict_export=True, + ), + exported_program.module(), + inputs, + None, + CortexMTargetConfig(cpu=CortexM.M55), + ) + + +def _count(exported_program, target) -> int: + return sum(node.target == target for node in exported_program.graph.nodes) + + +@pytest.mark.parametrize( + "explicit_layout,expected,unexpected", + [ + ( + False, + exir_ops.edge.cortex_m.quantized_conv2d.default, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + ), + ( + True, + exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default, + exir_ops.edge.cortex_m.quantized_conv2d.default, + ), + ], +) +def test_aot_layout_mode_selects_operator_family(explicit_layout, expected, unexpected): + _, edge, _ = _compile( + Conv2d().eval(), + (torch.randn(1, 3, 8, 8),), + explicit_layout=explicit_layout, + ) + program = edge.exported_program() + + assert _count(program, expected) == 1 + assert _count(program, unexpected) == 0 + + +def test_aot_explicit_layout_requires_quantization(): + with pytest.raises(RuntimeError, match="requires --quantize"): + _compile( + Conv2d().eval(), + (torch.randn(1, 3, 8, 8),), + explicit_layout=True, + quantize=False, + ) + + +def _assert_explicit_copy_ceiling(module, inputs, ceiling): + _, edge, _ = _compile(module, inputs, explicit_layout=True) + program = edge.exported_program() + copies = [ + node + for node in program.graph.nodes + if node.target + in { + exir_ops.edge.cortex_m.transpose.default, + exir_ops.edge.aten.view_copy.default, + } + ] + + assert len(copies) <= ceiling + assert all( + node.args[0].meta["val"].dtype == torch.int8 + for node in copies + if node.target == exir_ops.edge.cortex_m.transpose.default + ) + assert not any(node.target in _LEGACY_SPATIAL_OPS for node in program.graph.nodes) + assert not any( + getattr(node.target, "namespace", None) == "channels_last" + for node in program.graph.nodes + ) + + +def test_mobilenet_v2_explicit_copy_ceiling(): + torchvision = pytest.importorskip("torchvision") + _assert_explicit_copy_ceiling( + torchvision.models.mobilenet_v2(weights=None).eval(), + (torch.randn(1, 3, 224, 224),), + ceiling=2, + ) + + +def test_resnet8_explicit_copy_ceiling(): + from executorch.examples.models.mlperf_tiny.resnet8 import ResNet8 + + _assert_explicit_copy_ceiling( + ResNet8().eval(), + (torch.rand(1, 3, 32, 32) * 2 - 1,), + ceiling=2, + ) + + +def test_silero_explicit_copy_ceiling(): + from executorch.examples.models.silero_vad.export_silero_vad import ( + CONTEXT_SIZE, + HIDDEN_DIM, + SileroVAD16k, + WINDOW_SIZE, + ) + + _assert_explicit_copy_ceiling( + SileroVAD16k().eval(), + ( + torch.randn(1, CONTEXT_SIZE + WINDOW_SIZE), + torch.zeros(2, 1, HIDDEN_DIM), + ), + ceiling=12, + ) + + +def test_explicit_conv1d_runs_on_fvp(): + inputs = (torch.linspace(-5, 5, steps=16).reshape(1, 2, 8),) + model_quant, edge, runtime_inputs = _compile( + Conv1d().eval(), inputs, explicit_layout=True + ) + program = edge.exported_program() + assert _count(program, exir_ops.edge.cortex_m.quantized_conv2d_nhwc.default) == 1 + + serialized = CortexMSerialize(CortexMTargetConfig(cpu=CortexM.M55)) + serialized.run(edge.to_executorch()) + [actual] = serialized.run_artifact(runtime_inputs) + expected = model_quant(*runtime_inputs) + torch.testing.assert_close(actual, expected, atol=0.05, rtol=1e-3) diff --git a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md index 0a6a250f968..9094ab96d51 100644 --- a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md +++ b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md @@ -6,6 +6,8 @@ This backend is in **beta**. It has been validated with a set of small models (e The Arm® Cortex®-M backend accelerates quantized model execution on Arm Cortex-M CPUs using [CMSIS-NN](https://arm-software.github.io/CMSIS-NN/latest/) optimized kernels. Unlike delegate-based backends, it operates as an operator library: quantized subgraphs are replaced with CMSIS-NN accelerated kernels during the pass-lowering stage, while unsupported operators fall back to portable fp32 kernels. +The default AOT flow uses channels-last inputs and the existing dim-order representation. The experimental explicit-layout flow uses ordinary contiguous inputs, represents NCHW/NHWC conversions as graph operators, and selects the experimental `cortex_m::*_nhwc` kernels. Enable it with `--cortex-m-explicit-layout`. Layout modes are selected independently of the Cortex-M CPU target and never mix operator families. + ## Target Support The backend targets Arm Cortex-M CPUs via CMSIS-NN, which provides optimized kernel implementations for three instruction set variants: From 4fc2acc7f21ce71b46d381d925d753c57e352e5c Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 4 Sep 2026 16:08:27 -0700 Subject: [PATCH 056/190] Lower the CMake version floor so 3.26 through 3.28 can build again (#22570) ### Summary You cannot build ExecuTorch from source unless you have CMake 3.29 or newer. CMake is the tool that drives the build. Every build file can name the oldest CMake it accepts, and if yours is older, CMake stops with an error instead of building. Five files in this repo named 3.29, and none of them uses anything that needs it. This hurts most on small Arm computers such as developer boards. Those ship CMake 3.28, and 3.28 is also the newest their system offers, so you cannot simply upgrade. You have to go and install a second, private copy of CMake before you can build at all. Those five files now name 3.26 instead, so anyone on 3.26 or newer can build. 3.26 rather than something lower, because `kernels/optimized` uses CMake syntax that only exists from 3.26 on Arm builds. Going lower means rewriting that file, which belongs in its own change. The docs and requirement files name a CMake version too, and some of those numbers could never have worked, so they are corrected here: the build requirements said 3.24 and now say 3.26, the WebGPU docs said 3.19 and now say 3.24, and the Supertonic example said 3.24 and now says 3.26. Installing a prebuilt wheel is untouched and still works on CMake 3.19. One more file, `examples/models/phi-3-mini`, named 3.24 and then asked CMake to turn on a behaviour that only exists from 3.27. Asking for a behaviour CMake does not have is a hard error, so that file could never build on the version it advertised. It now asks only when the running CMake has it. --- backends/apple/metal/CMakeLists.txt | 2 +- backends/cuda/CMakeLists.txt | 2 +- backends/webgpu/README.md | 2 +- docs/source/backends/webgpu/webgpu-overview.md | 2 +- docs/source/llm/run-with-c-plus-plus.md | 2 +- docs/source/raspberry_pi_llama_tutorial.md | 4 ++-- examples/models/phi-3-mini/CMakeLists.txt | 6 +++++- examples/models/supertonic/README.md | 2 +- examples/models/whisper/CMakeLists.txt | 2 +- extension/wasm/CMakeLists.txt | 2 +- extension/wasm/tokenizers/CMakeLists.txt | 2 +- pyproject.toml | 2 +- requirements-dev.txt | 2 +- 13 files changed, 18 insertions(+), 14 deletions(-) diff --git a/backends/apple/metal/CMakeLists.txt b/backends/apple/metal/CMakeLists.txt index 4d242eae235..34477cde833 100644 --- a/backends/apple/metal/CMakeLists.txt +++ b/backends/apple/metal/CMakeLists.txt @@ -14,7 +14,7 @@ # ~~~ # It should also be cmake-lint clean. # -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/backends/cuda/CMakeLists.txt b/backends/cuda/CMakeLists.txt index c3e5d17809e..6ed8f83678f 100644 --- a/backends/cuda/CMakeLists.txt +++ b/backends/cuda/CMakeLists.txt @@ -14,7 +14,7 @@ # ~~~ # It should also be cmake-lint clean. # -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) set(CMAKE_CXX_STANDARD 17) set(CMAKE_CXX_STANDARD_REQUIRED ON) diff --git a/backends/webgpu/README.md b/backends/webgpu/README.md index da23a74e8bf..5b9877b9601 100644 --- a/backends/webgpu/README.md +++ b/backends/webgpu/README.md @@ -218,4 +218,4 @@ backends/webgpu/ - **Linux:** Vulkan-capable GPU and drivers - **Browser:** A WebGPU-enabled browser; the benchmark harness uses Chrome Canary -- **Build:** CMake 3.19+ and a Python environment with ExecuTorch installed +- **Build:** CMake 3.24+ and a Python environment with ExecuTorch installed diff --git a/docs/source/backends/webgpu/webgpu-overview.md b/docs/source/backends/webgpu/webgpu-overview.md index 91a388ac576..ebc46031a53 100644 --- a/docs/source/backends/webgpu/webgpu-overview.md +++ b/docs/source/backends/webgpu/webgpu-overview.md @@ -68,7 +68,7 @@ depend on the selected WebGPU adapter. ## Development Requirements -- CMake 3.19 or later. +- CMake 3.24 or later. - A Python environment with ExecuTorch installed for model export. - Dawn's CMake package for native builds. - Emscripten for browser builds. diff --git a/docs/source/llm/run-with-c-plus-plus.md b/docs/source/llm/run-with-c-plus-plus.md index b6c6082c3a6..69d98325e4f 100644 --- a/docs/source/llm/run-with-c-plus-plus.md +++ b/docs/source/llm/run-with-c-plus-plus.md @@ -12,7 +12,7 @@ Before you begin, make sure you have: - For HuggingFace tokenizers, this is a JSON file `tokenizer.json` - For SentencePiece tokenizers, this is a `tokenizer.model` file and normally lives alongside the weights file 3. CMake and a C++ compiler installed - - CMake version 3.29 or higher + - CMake version 3.26 or higher - g++ or clang compiler ## Model Metadata diff --git a/docs/source/raspberry_pi_llama_tutorial.md b/docs/source/raspberry_pi_llama_tutorial.md index 6075e455c9b..9eb99711ef5 100644 --- a/docs/source/raspberry_pi_llama_tutorial.md +++ b/docs/source/raspberry_pi_llama_tutorial.md @@ -21,7 +21,7 @@ This tutorial demonstrates how to deploy **Llama models on Raspberry Pi 4/5 devi - **Python 3.10-3.14** (ExecuTorch requirement) - **conda** or **venv** for environment management -- **CMake 3.29.6+** +- **CMake 3.26+** - **Git** for repository cloning ### Target Device Requirements @@ -47,7 +47,7 @@ python3 --version # Should be 3.10-3.14 # Check required tools hash cmake git md5sum 2>/dev/null || echo "Missing required tools" -cmake --version # Should be 3.29.6+ at minimum +cmake --version # Should be 3.26+ at minimum ## Development Environment Setup diff --git a/examples/models/phi-3-mini/CMakeLists.txt b/examples/models/phi-3-mini/CMakeLists.txt index 3c7ed6a4acb..c5c5eae30ac 100644 --- a/examples/models/phi-3-mini/CMakeLists.txt +++ b/examples/models/phi-3-mini/CMakeLists.txt @@ -14,7 +14,11 @@ # cmake_minimum_required(VERSION 3.24) -cmake_policy(SET CMP0144 NEW) +# CMP0144 arrives in 3.27, and cmake_policy(SET) on a policy the running CMake +# does not know is a hard error, so the request has to be guarded. +if(POLICY CMP0144) + cmake_policy(SET CMP0144 NEW) +endif() project(phi_3_mini_runner) set(CMAKE_CXX_STANDARD 17) diff --git a/examples/models/supertonic/README.md b/examples/models/supertonic/README.md index d7089d2802c..2508faf02f9 100644 --- a/examples/models/supertonic/README.md +++ b/examples/models/supertonic/README.md @@ -110,7 +110,7 @@ exit cleanly; closing stdin also exits with status zero. ## Platform and model limits - This workflow requires an Apple silicon Mac, macOS, Xcode command-line - tools, CMake 3.24 or newer, and an ExecuTorch Python environment with the MLX + tools, CMake 3.26 or newer, and an ExecuTorch Python environment with the MLX backend and custom operations available. The native runner supports only arm64 Darwin and uses MLX GPU delegation with FP16 activations. - Exported programs use dynamic sequence lengths, five flow-matching steps, diff --git a/examples/models/whisper/CMakeLists.txt b/examples/models/whisper/CMakeLists.txt index 6a1c2902977..08baf3b825f 100644 --- a/examples/models/whisper/CMakeLists.txt +++ b/examples/models/whisper/CMakeLists.txt @@ -4,7 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) project(whisper_runner) set(CMAKE_CXX_STANDARD 17) diff --git a/extension/wasm/CMakeLists.txt b/extension/wasm/CMakeLists.txt index 8ffd1801c63..5b52622140c 100644 --- a/extension/wasm/CMakeLists.txt +++ b/extension/wasm/CMakeLists.txt @@ -9,7 +9,7 @@ # cmake-format -i CMakeLists.txt # ~~~ -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) project(executorch_wasm) diff --git a/extension/wasm/tokenizers/CMakeLists.txt b/extension/wasm/tokenizers/CMakeLists.txt index 03b7ea1ff6b..a1bead4bb80 100644 --- a/extension/wasm/tokenizers/CMakeLists.txt +++ b/extension/wasm/tokenizers/CMakeLists.txt @@ -9,7 +9,7 @@ # cmake-format -i CMakeLists.txt # ~~~ -cmake_minimum_required(VERSION 3.29) +cmake_minimum_required(VERSION 3.26) if(NOT CMAKE_CXX_STANDARD) set(CMAKE_CXX_STANDARD 17) diff --git a/pyproject.toml b/pyproject.toml index 8f630a71e42..105565862df 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [build-system] requires = [ - "cmake>=3.24,<4.0.0", # For building binary targets in the wheel. 4.0.0 breaks third-party CMake build so temporarily pin the version. + "cmake>=3.26,<4.0.0", # For building binary targets in the wheel. 4.0.0 breaks third-party CMake build so temporarily pin the version. "packaging>=24.2", # Lower bound required by setuptools "patchelf; sys_platform == 'linux'", # Writes the runtime search paths that let the shipped libraries find each other. "pip>=23", # For building the pip package. diff --git a/requirements-dev.txt b/requirements-dev.txt index c916ccfa472..1220547aeb5 100644 --- a/requirements-dev.txt +++ b/requirements-dev.txt @@ -1,6 +1,6 @@ # Pip packages needed to build from source. Mainly for development of ExecuTorch. -cmake>=3.24, <4.0.0 # For building binary targets in the wheel. +cmake>=3.26, <4.0.0 # For building binary targets in the wheel. packaging>=24.2 # Lower bound required by setuptools patchelf; sys_platform == 'linux' # Writes the runtime search paths that let the shipped libraries find each other. pip>=23 # For building the pip package. From 898795db89dcd5bce1e32bfd881b463c57938832 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 16:59:30 -0700 Subject: [PATCH 057/190] Fix Arm recipe BUCK dependency (#22574) ## Summary The Arm recipe provider moved from `OpBackend` to `edge_manager_transform_passes` before merge, but its BUCK target retained a dependency on the unlanded `//executorch/exir/backend:op_backend` target. This causes fbsource target graph analysis to fail for D118843589. Remove the unused dependency. This does not change source behavior. ## Test plan `buck2 build fbcode//executorch/backends/arm/recipes:arm_recipe_provider fbcode//executorch/backends/arm/recipes:recipes fbcode//executorch/backends/arm/test:arm_recipes` `buck2 test fbcode//executorch/backends/arm/test:arm_recipes` (25 passed, 0 failed) `git diff --check` Authored with Codex. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell --- backends/arm/recipes/BUCK | 1 - 1 file changed, 1 deletion(-) diff --git a/backends/arm/recipes/BUCK b/backends/arm/recipes/BUCK index c7dfa0abd55..df3c7874f9e 100644 --- a/backends/arm/recipes/BUCK +++ b/backends/arm/recipes/BUCK @@ -43,7 +43,6 @@ fbcode_target( "//executorch/backends/arm/tosa:compile_spec", "//executorch/backends/cortex_m/passes:replace_quant_nodes_pass", "//executorch/exir:lib", - "//executorch/exir/backend:op_backend", "//executorch/export:lib", ], ) From c570b6c146beed40b4f68f136c38787cb5af0861 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Fri, 4 Sep 2026 17:44:05 -0700 Subject: [PATCH 058/190] Declare dynamic quantization transform dependency in test harness (#22576) The //executorch/backends/test/harness:tester target includes stages/quantize.py, which imports duplicate_dynamic_quant_chain, but did not declare its Buck dependency. Existing consumers supplied the dependency transitively, masking the omission. The new Cortex-M explicit-layout test exposed it during test collection. Declare the dependency on the harness target that owns the importing source. Test plan: git diff --check; GitHub and internal CI. AI-assisted: Codex. --- backends/test/harness/BUCK | 1 + 1 file changed, 1 insertion(+) diff --git a/backends/test/harness/BUCK b/backends/test/harness/BUCK index d432aa6dc3d..3331a5176dd 100644 --- a/backends/test/harness/BUCK +++ b/backends/test/harness/BUCK @@ -8,6 +8,7 @@ fbcode_target(_kind = runtime.python_library, srcs = native.glob(["*.py", "stages/*.py"]), visibility = ["PUBLIC"], deps = [ + "//executorch/backends/transforms:duplicate_dynamic_quant_chain", "//executorch/exir:graph_module", ], ) From 5cdcc02a8f66f8d22740254ffa9ecf5fbdea6ee6 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Fri, 4 Sep 2026 23:07:27 -0700 Subject: [PATCH 059/190] Stop the backend suites logging captured output (#22581) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Since #22246 moved these suites to OSDC, `arm_ethos_u55 / models` and `arm_ethos_u85 / models` have died partway through reporting their failures. The cause is the runner agent's heap, not the pod: the log ends with `##[error]Exception of type 'System.OutOfMemoryException' was thrown.` Each Ethos-U failure captures a few hundred thousand lines of Vela operator listings, and the agent cannot process a step that size. `--show-capture=no` drops those sections from the terminal report. Verified on the default label and worker count: https://github.com/pytorch/executorch/actions/runs/33935989761 ``` success arm_ethos_u85, models <- never passed on OSDC before success arm_ethos_u55, models 9 failed, 13 passed, 20 skipped log: 12,596 lines <- was 458k, truncated where the agent died ``` Nothing is lost. `--show-capture` only affects what the terminal prints, so the captured streams still ship in the JSON report artifact (19MB, against 17MB before), and each failure keeps its traceback and exception message — `RuntimeError: Corstone simulation failed` and the arena sizes are still in the log. This applies to every backend suite, Linux and macOS, since `_test_backend.yml` sources the script in both jobs. Only the Ethos-U cells produce output at this scale today, but any of them could. The failures themselves are pre-existing and unrelated; the job was green with them before because `test_backend.sh` exits with the summary generator's status rather than pytest's. Authored with Claude Code. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- .ci/scripts/test_backend.sh | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/.ci/scripts/test_backend.sh b/.ci/scripts/test_backend.sh index 32c682ad5f0..95c5c8e9db9 100755 --- a/.ci/scripts/test_backend.sh +++ b/.ci/scripts/test_backend.sh @@ -158,7 +158,12 @@ GOLDEN_DIR="${ARTIFACT_DIR}/golden-artifacts" export GOLDEN_ARTIFACTS_DIR="${GOLDEN_DIR}" EXIT_CODE=0 -PYTEST_ARGS=(-c /dev/null -n auto) +# An Ethos-U failure captures a few hundred thousand lines of Vela operator +# listings, and the runner agent throws System.OutOfMemoryException processing +# a step that size, taking the whole job down before pytest can report. The +# reason for each failure is in its exception message and traceback, which are +# unaffected. +PYTEST_ARGS=(-c /dev/null -n auto --show-capture=no) if [[ ${#PYTEST_RETRY_ARGS[@]} -gt 0 ]]; then PYTEST_ARGS+=("${PYTEST_RETRY_ARGS[@]}") fi From 4ce2ec2ee639d3a3aa8b2503f757319dcbb0f733 Mon Sep 17 00:00:00 2001 From: Shamsudeen Saleem <104304558+ShamSaleem@users.noreply.github.com> Date: Sun, 6 Sep 2026 16:47:53 +0300 Subject: [PATCH 060/190] Use dynamic CPU count for -j in vendor/example scripts and docs (#21455) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Follow-up to #20436, which replaced the hardcoded `cmake --build -j` parallelism in the general build docs and `test/` scripts. This PR finishes the same job for the vendor backend scripts, example scripts, and their documentation — 65 sites across 37 files, all mechanical: -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) `nproc` on Linux, `sysctl -n hw.ncpu` on macOS (the mps and coreml scripts are Apple-only), and the arithmetic degrades to `-j1` if neither tool exists. "Core count + 1" is the guidance already in `docs/source/using-executorch-building-from-source.md`. The pinned values being removed ranged from `-j4` to `-j100`, including `-j64` in the Vulkan test scripts and `-j100` in `tools/cmake/preset/README.md`. Two Python sites differ. `extension/llm/export/quantizer_lib.py` holds a shell command inside a user-facing error string, so it takes the same shell expression. `backends/mlx/test/test_utils.py` builds an argv list handed to `subprocess.run` with no shell, where a shell expression would reach `cmake` as a literal string, so it uses `f"-j{(os.cpu_count() or 1) + 1}"` instead. Deliberately out of scope: `.ci/**` and `.github/workflows/**`, where the runners are fixed-size and the parallelism is a resource-tuning decision rather than a portability problem (`cuda.yml` pins `-j4`, likely to bound peak memory); the `-j4` in `backends/mlx`'s READMEs and `run_all_tests.py`, which is a test-worker count and not a build flag; and `docs/source/archive/`. Happy to take the CI files in a separate PR if you'd like them changed. Review order: the three groups are independent — vendor backend scripts under `backends/`, example scripts and docs under `examples/`, then the two Python files, which are the only sites that are not a pure token swap. Partial fix for #10887. ### Test plan ExecuTorch does not build on my Windows host, so verification is static and per-site: - `bash -n` passes on all 20 modified shell scripts. - Every edited command line was re-run with `cmake --build`/`make` swapped for `echo`, confirming all 65 sites expand to a single valid integer flag (`-j17` on this 16-core machine) with zero expansion failures. This covers the markdown sites too, including the two lines that begin with `&& ` and the one with a `$ ` prompt prefix. - `git diff` normalised on the `-j` token shows every removed line has a matching added line, so nothing outside the flag changed. Every changed markdown line contains a `-j` token. - Both Python files parse; the argv site renders `-j17`; the instruction string was extracted via `ast` and shell-expanded to confirm a user pasting it gets `-j17`. - `black --check` reports both Python files unchanged. Line endings are unchanged (still LF) and `git diff --check` reports no whitespace errors. - `lintrunner` was **not** run locally — it is not installed on this Windows host and is unavailable in my WSL environment, so CI lint is the gate for that. E501 is in the repo's flake8 ignore list, so the one long instruction string in `quantizer_lib.py` (already 165 chars before this change) is not a new violation. This PR was authored with AI assistance (Claude Code); the diff and every verification step above were reviewed by me. cc @GregoryComer @digantdesai @cbilgin @JakeStevens @larryliu0820 --- backends/apple/coreml/scripts/build_tests.sh | 4 ++-- backends/cadence/build_cadence_fusionG3.sh | 4 ++-- backends/cadence/build_cadence_hifi4.sh | 4 ++-- backends/cadence/build_cadence_runner.sh | 4 ++-- backends/cadence/build_cadence_vision.sh | 4 ++-- backends/cadence/runtime/executor_main.sh | 4 ++-- backends/mediatek/scripts/mtk_build.sh | 2 +- backends/mlx/test/test_utils.py | 9 ++++++++- backends/samsung/README.md | 2 +- backends/vulkan/test/custom_ops/build_and_run.sh | 8 ++++---- backends/vulkan/test/scripts/test_model.sh | 6 +++--- backends/vulkan/test/scripts/test_op.sh | 8 ++++---- docs/source/backends-cadence.md | 2 +- .../backends/vulkan/tutorials/etvk-llama-tutorial.md | 4 ++-- .../backends/vulkan/tutorials/etvk-profiling-tutorial.md | 2 +- docs/source/tutorial-xnnpack-delegate-lowering.md | 2 +- examples/mediatek/mtk_build_examples.sh | 2 +- examples/models/llama/README.md | 8 ++++---- examples/models/phi-3-mini-lora/README.md | 2 +- examples/models/phi-3-mini/README.md | 2 +- examples/nxp/run.sh | 2 +- examples/portable/README.md | 2 +- examples/portable/custom_ops/test_custom_ops.sh | 4 ++-- .../portable/scripts/test_demo_backend_delegation.sh | 2 +- examples/qualcomm/test_qualcomm.sh | 4 ++-- examples/selective_build/README.md | 4 ++-- examples/selective_build/test_selective_build.sh | 6 +++--- examples/wasm/README.md | 6 +++--- examples/wasm/test_build_wasm.sh | 2 +- examples/xnnpack/quantization/test_quantize.sh | 2 +- extension/llm/export/quantizer_lib.py | 2 +- extension/wasm/README.md | 2 +- extension/wasm/tokenizers/README.md | 2 +- tools/cmake/preset/README.md | 2 +- 34 files changed, 66 insertions(+), 59 deletions(-) diff --git a/backends/apple/coreml/scripts/build_tests.sh b/backends/apple/coreml/scripts/build_tests.sh index 0203e5027a2..f14463f032e 100755 --- a/backends/apple/coreml/scripts/build_tests.sh +++ b/backends/apple/coreml/scripts/build_tests.sh @@ -36,7 +36,7 @@ cmake "$EXECUTORCH_ROOT_PATH" -B"$CMAKE_EXECUTORCH_BUILD_DIR_PATH" \ -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=OFF \ -DEXECUTORCH_BUILD_XNNPACK=OFF -cmake --build "$CMAKE_EXECUTORCH_BUILD_DIR_PATH" -j9 -t executorch +cmake --build "$CMAKE_EXECUTORCH_BUILD_DIR_PATH" -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) -t executorch # Build protobuf echo "ExecuTorch: Building libprotobuf-lite" @@ -53,7 +53,7 @@ cmake "$PROTOBUF_DIR_PATH/cmake" -B"$CMAKE_PROTOBUF_BUILD_DIR_PATH" \ -DCMAKE_MACOSX_BUNDLE=OFF \ -DCMAKE_CXX_STANDARD=17 -cmake --build "$CMAKE_PROTOBUF_BUILD_DIR_PATH" -j9 -t libprotobuf-lite +cmake --build "$CMAKE_PROTOBUF_BUILD_DIR_PATH" -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) -t libprotobuf-lite # Copy required libraries echo "ExecuTorch: Copying libraries" diff --git a/backends/cadence/build_cadence_fusionG3.sh b/backends/cadence/build_cadence_fusionG3.sh index 47a0f9ff9bb..a95721833f4 100644 --- a/backends/cadence/build_cadence_fusionG3.sh +++ b/backends/cadence/build_cadence_fusionG3.sh @@ -55,7 +55,7 @@ if $STEPWISE_BUILD; then -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out/backends/cadence \ backends/cadence - cmake --build cmake-out/backends/cadence -j8 + cmake --build cmake-out/backends/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" @@ -80,7 +80,7 @@ else -DHAVE_FNMATCH_H=OFF \ -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out - cmake --build cmake-out --target install --config Release -j8 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) fi echo "Run simple model to verify cmake build" diff --git a/backends/cadence/build_cadence_hifi4.sh b/backends/cadence/build_cadence_hifi4.sh index 22775af7082..cbac5e5b5d6 100644 --- a/backends/cadence/build_cadence_hifi4.sh +++ b/backends/cadence/build_cadence_hifi4.sh @@ -54,7 +54,7 @@ if $STEPWISE_BUILD; then -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out/backends/cadence \ backends/cadence - cmake --build cmake-out/backends/cadence -j8 + cmake --build cmake-out/backends/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" @@ -78,7 +78,7 @@ else -DHAVE_FNMATCH_H=OFF \ -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out - cmake --build cmake-out --target install --config Release -j8 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) fi echo "Run simple model to verify cmake build" diff --git a/backends/cadence/build_cadence_runner.sh b/backends/cadence/build_cadence_runner.sh index 82968b196b3..57ce5a339e5 100755 --- a/backends/cadence/build_cadence_runner.sh +++ b/backends/cadence/build_cadence_runner.sh @@ -27,7 +27,7 @@ main() { -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -DEXECUTORCH_ENABLE_LOGGING=ON \ -Bcmake-out . - cmake --build cmake-out --target install --config Release -j16 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) local example_dir=backends/cadence local build_dir="cmake-out/${example_dir}" @@ -46,7 +46,7 @@ main() { -DPYTHON_EXECUTABLE="$(which python3)" \ -B"${build_dir}" \ "${example_dir}" - cmake --build "${build_dir}" --config Release -j16 + cmake --build "${build_dir}" --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) local runner="${PWD}/${build_dir}/cadence_runner" if [[ ! -f "${runner}" ]]; then diff --git a/backends/cadence/build_cadence_vision.sh b/backends/cadence/build_cadence_vision.sh index b3972db4f31..fe2a07974d1 100755 --- a/backends/cadence/build_cadence_vision.sh +++ b/backends/cadence/build_cadence_vision.sh @@ -54,7 +54,7 @@ if $STEPWISE_BUILD; then -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out/backends/cadence \ backends/cadence - cmake --build cmake-out/backends/cadence -j8 + cmake --build cmake-out/backends/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) else echo "Building Cadence toolchain with ExecuTorch packages" cmake_prefix_path="${PWD}/cmake-out/lib/cmake/ExecuTorch;${PWD}/cmake-out/third-party/gflags" @@ -78,7 +78,7 @@ else -DHAVE_FNMATCH_H=OFF \ -DFLATCC_ALLOW_WERROR=OFF \ -Bcmake-out - cmake --build cmake-out --target install --config Release -j8 + cmake --build cmake-out --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) fi echo "Run simple model to verify cmake build" diff --git a/backends/cadence/runtime/executor_main.sh b/backends/cadence/runtime/executor_main.sh index 7d6cba09b87..5e630bdb524 100644 --- a/backends/cadence/runtime/executor_main.sh +++ b/backends/cadence/runtime/executor_main.sh @@ -24,7 +24,7 @@ cmake_install_executorch_devtools_lib() { -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -Bcmake-out . - cmake --build cmake-out -j9 --target install --config Release + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release } test_cmake_devtools_example_runner() { @@ -40,7 +40,7 @@ test_cmake_devtools_example_runner() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running devtools/example_runner' ${build_dir}/example_runner --bundled_program_path="./CadenceDemoModel.bpte" diff --git a/backends/mediatek/scripts/mtk_build.sh b/backends/mediatek/scripts/mtk_build.sh index d42e5f7e10a..6d5b013f9f3 100755 --- a/backends/mediatek/scripts/mtk_build.sh +++ b/backends/mediatek/scripts/mtk_build.sh @@ -35,7 +35,7 @@ cmake -DCMAKE_INSTALL_PREFIX="${build_dir}" \ -B"${build_dir}" # Build the project -cmake --build "${build_dir}" --target install --config Release -j5 +cmake --build "${build_dir}" --target install --config Release -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) # Switch back to the original directory cd - > /dev/null diff --git a/backends/mlx/test/test_utils.py b/backends/mlx/test/test_utils.py index eb20095db97..adf4e56e3a2 100644 --- a/backends/mlx/test/test_utils.py +++ b/backends/mlx/test/test_utils.py @@ -587,7 +587,14 @@ def rebuild_op_test_runner(verbose: bool = False) -> bool: print(f"Rebuilding op_test_runner in {build_dir}...") - cmd = ["cmake", "--build", str(build_dir), "--target", "op_test_runner", "-j8"] + cmd = [ + "cmake", + "--build", + str(build_dir), + "--target", + "op_test_runner", + f"-j{(os.cpu_count() or 1) + 1}", + ] if verbose: print(f"Running: {' '.join(cmd)}") diff --git a/backends/samsung/README.md b/backends/samsung/README.md index bc48bad830a..f498bd0a098 100644 --- a/backends/samsung/README.md +++ b/backends/samsung/README.md @@ -64,7 +64,7 @@ cmake extension/android \ -DCMAKE_INSTALL_PREFIX=cmake-android-out \ -Bcmake-android-out/extension/android -cmake --build cmake-android-out/extension/android -j8 +cmake --build cmake-android-out/extension/android -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ``` ## Examples diff --git a/backends/vulkan/test/custom_ops/build_and_run.sh b/backends/vulkan/test/custom_ops/build_and_run.sh index b1195568b1b..2641970faf8 100755 --- a/backends/vulkan/test/custom_ops/build_and_run.sh +++ b/backends/vulkan/test/custom_ops/build_and_run.sh @@ -30,13 +30,13 @@ configure_and_build_main() { -B$CMAKE_OUT_DIR fi - cmake --build $CMAKE_OUT_DIR -j16 --target install + cmake --build $CMAKE_OUT_DIR -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install # -DCMAKE_CXX_FLAGS="-DVULKAN_DEBUG" \ } # Function to build main project only build_main() { - cmake --build $CMAKE_OUT_DIR -j16 --target install + cmake --build $CMAKE_OUT_DIR -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install } # Function to configure and build tests @@ -65,12 +65,12 @@ configure_and_build_tests() { -B$CMAKE_OUT_DIR/backends/vulkan/test/custom_ops fi - cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j16 --target all + cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target all } build_tests() { - cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j16 --target all + cmake --build $CMAKE_OUT_DIR/backends/vulkan/test/custom_ops -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target all } # Function to rebuild both main and tests diff --git a/backends/vulkan/test/scripts/test_model.sh b/backends/vulkan/test/scripts/test_model.sh index 40ec88bae70..4adbdc407df 100755 --- a/backends/vulkan/test/scripts/test_model.sh +++ b/backends/vulkan/test/scripts/test_model.sh @@ -95,7 +95,7 @@ clean_build_directory() { } recompile() { - cmake --build cmake-out -j64 --target install + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install } build_core_libraries_and_devtools() { @@ -119,7 +119,7 @@ build_core_libraries_and_devtools() { -DEXECUTORCH_BUILD_VULKAN=ON \ -DEXECUTORCH_BUILD_XNNPACK=ON \ -Bcmake-out && \ - cmake --build cmake-out -j64 --target install + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install # Build devtools example runner cmake examples/devtools \ @@ -127,7 +127,7 @@ build_core_libraries_and_devtools() { -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -DEXECUTORCH_BUILD_VULKAN=ON \ -Bcmake-out/examples/devtools && \ - cmake --build cmake-out/examples/devtools -j16 --config Release + cmake --build cmake-out/examples/devtools -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release } run_example_runner() { diff --git a/backends/vulkan/test/scripts/test_op.sh b/backends/vulkan/test/scripts/test_op.sh index 797089e54dc..d3967f98abf 100755 --- a/backends/vulkan/test/scripts/test_op.sh +++ b/backends/vulkan/test/scripts/test_op.sh @@ -149,7 +149,7 @@ build_core_libraries() { -DEXECUTORCH_BUILD_XNNPACK=ON \ -DEXECUTORCH_BUILD_TESTS=ON \ -Bcmake-out && \ - cmake --build cmake-out -j64 --target install + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install } build_operator_tests() { @@ -195,13 +195,13 @@ build_operator_tests() { # Build operator tests cmake "${CMAKE_ARGS[@]}" \ -Bcmake-out/backends/vulkan/test/op_tests && \ - cmake --build cmake-out/backends/vulkan/test/op_tests -j16 + cmake --build cmake-out/backends/vulkan/test/op_tests -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) } recompile() { echo "Recompiling..." - cmake --build cmake-out -j64 --target install - cmake --build cmake-out/backends/vulkan/test/op_tests -j16 + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install + cmake --build cmake-out/backends/vulkan/test/op_tests -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) } run_operator_test() { diff --git a/docs/source/backends-cadence.md b/docs/source/backends-cadence.md index c5a5fc8497a..ddd258cdca5 100644 --- a/docs/source/backends-cadence.md +++ b/docs/source/backends-cadence.md @@ -311,7 +311,7 @@ cmake -DCMAKE_BUILD_TYPE=Debug \ -Bcmake-out/examples/cadence \ examples/cadence -cmake --build cmake-out/examples/cadence -j8 -t cadence_executorch_example +cmake --build cmake-out/examples/cadence -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) -t cadence_executorch_example ``` After having succesfully run the above step you should see two binary files in their CMake output directory. diff --git a/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md b/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md index cb14c72331e..42a7a846d75 100644 --- a/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md +++ b/docs/source/backends/vulkan/tutorials/etvk-llama-tutorial.md @@ -93,7 +93,7 @@ cmake . \ -DEXECUTORCH_BUILD_VULKAN=ON \ -DEXECUTORCH_BUILD_TESTS=OFF \ -Bcmake-out-android-so && \ -cmake --build cmake-out-android-so -j16 --target install --config Release +cmake --build cmake-out-android-so -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` ## Build and push the llama runner binary to Android @@ -111,7 +111,7 @@ cmake examples/models/llama \ -DCMAKE_BUILD_TYPE=Release \ -DPYTHON_EXECUTABLE=python \ -Bcmake-out-android-so/examples/models/llama && \ -cmake --build cmake-out-android-so/examples/models/llama -j16 --config Release +cmake --build cmake-out-android-so/examples/models/llama -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` Once the binary is built, it can be pushed to your Android device. diff --git a/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md b/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md index 07982d81c1c..9330ba370e8 100644 --- a/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md +++ b/docs/source/backends/vulkan/tutorials/etvk-profiling-tutorial.md @@ -71,7 +71,7 @@ cmake . \ -DEXECUTORCH_BUILD_EXECUTOR_RUNNER=ON \ -DEXECUTORCH_ENABLE_EVENT_TRACER=ON \ -Bcmake-out-android-so && \ -cmake --build cmake-out-android-so -j16 --target install --config Release +cmake --build cmake-out-android-so -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` Once the build completes, we can push the runner binary to device. diff --git a/docs/source/tutorial-xnnpack-delegate-lowering.md b/docs/source/tutorial-xnnpack-delegate-lowering.md index 5c88246b0ba..81b26db2f12 100644 --- a/docs/source/tutorial-xnnpack-delegate-lowering.md +++ b/docs/source/tutorial-xnnpack-delegate-lowering.md @@ -166,7 +166,7 @@ cmake \ Then you can build the runtime componenets with ```bash -cmake --build cmake-out -j9 --target install --config Release +cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` Now you should be able to find the executable built at `./cmake-out/executor_runner` you can run the executable with the model you generated as such diff --git a/examples/mediatek/mtk_build_examples.sh b/examples/mediatek/mtk_build_examples.sh index afdd9f16d51..c8ea0e927c1 100755 --- a/examples/mediatek/mtk_build_examples.sh +++ b/examples/mediatek/mtk_build_examples.sh @@ -43,7 +43,7 @@ main() { -B"${example_build_dir}" \ $EXECUTORCH_ROOT/$example_dir - cmake --build "${example_build_dir}" -j5 + cmake --build "${example_build_dir}" -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) # Switch back to the original directory cd - > /dev/null diff --git a/examples/models/llama/README.md b/examples/models/llama/README.md index 5a65eaa0cb4..3ee568120a3 100644 --- a/examples/models/llama/README.md +++ b/examples/models/llama/README.md @@ -288,7 +288,7 @@ cmake -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \ -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ -Bcmake-out-android . -cmake --build cmake-out-android -j16 --target install --config Release +cmake --build cmake-out-android -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release ``` **1.2 Build llama runner for android** @@ -307,7 +307,7 @@ cmake -DCMAKE_TOOLCHAIN_FILE=$ANDROID_NDK/build/cmake/android.toolchain.cmake \ -Bcmake-out-android/examples/models/llama \ examples/models/llama -cmake --build cmake-out-android/examples/models/llama -j16 --config Release +cmake --build cmake-out-android/examples/models/llama -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` **2. Run on Android via adb shell** @@ -400,7 +400,7 @@ cmake -DPYTHON_EXECUTABLE=python \ -DEXECUTORCH_BUILD_EXTENSION_LLM=ON \ -DEXECUTORCH_BUILD_KERNELS_LLM=ON \ -Bcmake-out . -cmake --build cmake-out -j16 --config Release --target install +cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release --target install ``` Next install the llama runner with torchao kernels enabled (similar to step 3.2 above): @@ -410,7 +410,7 @@ cmake -DPYTHON_EXECUTABLE=python \ -DCMAKE_BUILD_TYPE=Release \ -Bcmake-out/examples/models/llama \ examples/models/llama -cmake --build cmake-out/examples/models/llama -j16 --config Release +cmake --build cmake-out/examples/models/llama -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` Finally run your model (similar to step 3.3 above): diff --git a/examples/models/phi-3-mini-lora/README.md b/examples/models/phi-3-mini-lora/README.md index 62efda6c3dc..fa00e17fd3c 100644 --- a/examples/models/phi-3-mini-lora/README.md +++ b/examples/models/phi-3-mini-lora/README.md @@ -24,7 +24,7 @@ python export_model.py (mkdir cmake-out && cd cmake-out && cmake ..) # Build the executor_runner target -cmake --build cmake-out --target executor_runner -j9 +cmake --build cmake-out --target executor_runner -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) # Run the model for inference. ./cmake-out/executor_runner --model_path phi3_mini_lora.pte diff --git a/examples/models/phi-3-mini/README.md b/examples/models/phi-3-mini/README.md index dac378213d8..1df5844db9c 100644 --- a/examples/models/phi-3-mini/README.md +++ b/examples/models/phi-3-mini/README.md @@ -39,7 +39,7 @@ cmake -DCMAKE_PREFIX_PATH=cmake-out \ -Bcmake-out/examples/models/phi-3-mini \ examples/models/phi-3-mini -cmake --build cmake-out/examples/models/phi-3-mini -j16 --config Release +cmake --build cmake-out/examples/models/phi-3-mini -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release ``` - Run model. Options available [here](https://github.com/pytorch/executorch/blob/main/examples/models/phi-3-mini/main.cpp#L16-L33) ``` diff --git a/examples/nxp/run.sh b/examples/nxp/run.sh index b8cc87a5964..dc73bb474a2 100755 --- a/examples/nxp/run.sh +++ b/examples/nxp/run.sh @@ -19,7 +19,7 @@ rm -rf ${SCRIPT_DIR}/executor_runner/build/* pushd ${SCRIPT_DIR}/executor_runner/build cmake -DCMAKE_BUILD_TYPE=Release .. -make -j8 nxp_executor_runner +make -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) nxp_executor_runner popd echo "** Export cifar10 model to executorch" diff --git a/examples/portable/README.md b/examples/portable/README.md index ef9b44a48a3..f3b097881dd 100644 --- a/examples/portable/README.md +++ b/examples/portable/README.md @@ -49,7 +49,7 @@ Use `-h` (or `--help`) to see all the supported models. (mkdir cmake-out \ && cd cmake-out \ && cmake -DEXECUTORCH_PAL_DEFAULT=posix ..) \ - && cmake --build cmake-out -j32 --target executor_runner + && cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner # Run the tool on the generated model. ./cmake-out/executor_runner --model_path mv2.pte diff --git a/examples/portable/custom_ops/test_custom_ops.sh b/examples/portable/custom_ops/test_custom_ops.sh index 58a7de3a5f2..a54761aefc7 100644 --- a/examples/portable/custom_ops/test_custom_ops.sh +++ b/examples/portable/custom_ops/test_custom_ops.sh @@ -30,7 +30,7 @@ test_cmake_custom_op_1() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running custom_ops_executor_runner' ${build_dir}/custom_ops_executor_runner --model_path="./${model_name}.pte" @@ -66,7 +66,7 @@ test_cmake_custom_op_2() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release EXT=$(get_shared_lib_ext) echo "Exporting ${model_name}.pte" diff --git a/examples/portable/scripts/test_demo_backend_delegation.sh b/examples/portable/scripts/test_demo_backend_delegation.sh index d1ecf9150f9..2aeb1f90689 100644 --- a/examples/portable/scripts/test_demo_backend_delegation.sh +++ b/examples/portable/scripts/test_demo_backend_delegation.sh @@ -24,7 +24,7 @@ build_cmake_executor_runner() { && cd ${CMAKE_OUTPUT_DIR} \ && retry cmake -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" ..) - cmake --build ${CMAKE_OUTPUT_DIR} -j4 + cmake --build ${CMAKE_OUTPUT_DIR} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) } test_demo_backend_delegation() { diff --git a/examples/qualcomm/test_qualcomm.sh b/examples/qualcomm/test_qualcomm.sh index 51a563863f3..24a9c224061 100644 --- a/examples/qualcomm/test_qualcomm.sh +++ b/examples/qualcomm/test_qualcomm.sh @@ -21,7 +21,7 @@ cmake_install_executorch_qnn_lib() { -DEXECUTORCH_BUILD_QNN=ON \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" \ -Bcmake-out . - cmake --build cmake-out -j9 --target install --config Release + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target install --config Release } test_cmake_qualcomm() { @@ -45,7 +45,7 @@ test_cmake_qualcomm() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release # Need to run on device # ${build_dir}/qnn_executor_runner --model_path="./mv2_qnn.pte" } diff --git a/examples/selective_build/README.md b/examples/selective_build/README.md index c6c8dc1ba57..1e1e556f415 100644 --- a/examples/selective_build/README.md +++ b/examples/selective_build/README.md @@ -18,7 +18,7 @@ python -m examples.portable.scripts.export --model_name="mv2" # Create a PTE fil cd examples/selective_build/basic mkdir cmake-out && cd cmake-out cmake .. -DEXECUTORCH_SELECT_OPS_MODEL="../../mv2.pte" # Build with kernels needed for mv2.pte -cmake --build . -j8 +cmake --build . -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ./selective_build_test --model_path="../../mv2.pte" # Run the model with the selective kernel library ``` @@ -78,7 +78,7 @@ python -m examples.portable.custom_ops.custom_ops_1 # Create a model PTE file cd examples/selective_build/basic mkdir cmake-out && cd cmake-out cmake .. -DEXECUTORCH_SELECT_OPS_MODEL="../../custom_ops_1.pte" -DEXECUTORCH_EXAMPLE_USE_CUSTOM_OPS=ON # Build with kernels needed for the model -cmake --build . -j8 +cmake --build . -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ./selective_build_test --model_path="../../custom_ops_1.pte" # Run the model with the selective kernel library ``` diff --git a/examples/selective_build/test_selective_build.sh b/examples/selective_build/test_selective_build.sh index c1b5c627c42..bd286f8db19 100755 --- a/examples/selective_build/test_selective_build.sh +++ b/examples/selective_build/test_selective_build.sh @@ -105,7 +105,7 @@ aten,aten::clone.out" \ ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running selective build test' ${build_dir}/selective_build_test --model_path="./mv2.pte" @@ -129,7 +129,7 @@ test_cmake_select_ops_in_yaml() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config Release + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config Release echo 'Running selective build test' ${build_dir}/selective_build_test --model_path="./custom_ops_1.pte" @@ -158,7 +158,7 @@ test_cmake_select_ops_in_model() { ${example_dir} echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --config $CMAKE_BUILD_TYPE + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --config $CMAKE_BUILD_TYPE echo "Verifying auto-right-sized MAX_KERNEL_NUM header was generated" local generated_header diff --git a/examples/wasm/README.md b/examples/wasm/README.md index 15ce07493d1..716e1115846 100644 --- a/examples/wasm/README.md +++ b/examples/wasm/README.md @@ -51,13 +51,13 @@ Use -h (or --help) to see all the supported models. For the browser example, mak (mkdir cmake-out-wasm \ && cd cmake-out-wasm \ && emcmake cmake -DEXECUTORCH_PAL_DEFAULT=posix ..) \ - && cmake --build cmake-out-wasm -j32 --target executor_runner + && cmake --build cmake-out-wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner ``` If you need to rebuild `executor_runner` after modifying the contents of `./models/`, you can run the following command ```bash -cmake --build cmake-out-wasm -j32 --target executor_runner --clean-first +cmake --build cmake-out-wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner --clean-first ``` 4. Run the model with Node.js. Emscripten should come preinstalled with a compatible version of Node.js. If you have an incompatible version of Node.js installed, you can use the Emscripten-provided version by running `$EMSDK_NODE` instead of `node`. @@ -91,7 +91,7 @@ echo $EMSDK_NODE The file may not have been present while building the Wasm binary. You can rebuild with the following command ```bash -cmake --build cmake-out-wasm -j32 --target executor_runner --clean-first +cmake --build cmake-out-wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner --clean-first ``` The path may also be incorrect. The files in the `WASM_MODEL_DIR` are placed into the root directory of the virtual file system, so you would use `--model_path mv2.pte` instead of `--model_path models/mv2.pte`, for example. diff --git a/examples/wasm/test_build_wasm.sh b/examples/wasm/test_build_wasm.sh index f7144a209df..ef836b45901 100644 --- a/examples/wasm/test_build_wasm.sh +++ b/examples/wasm/test_build_wasm.sh @@ -23,7 +23,7 @@ test_build_wasm() { retry emcmake cmake -DWASM_MODEL_DIR="$(realpath "${model_dir_name}")" -B${build_dir} . echo "Building ${example_dir}" - cmake --build ${build_dir} -j9 --target executor_runner + cmake --build ${build_dir} -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner echo "Removing ${model_dir_name}" rm -rf "${model_dir_name}" diff --git a/examples/xnnpack/quantization/test_quantize.sh b/examples/xnnpack/quantization/test_quantize.sh index 1f50667c788..1211470b084 100644 --- a/examples/xnnpack/quantization/test_quantize.sh +++ b/examples/xnnpack/quantization/test_quantize.sh @@ -56,7 +56,7 @@ test_cmake_quantization() { -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \ -DPYTHON_EXECUTABLE="$PYTHON_EXECUTABLE" ..) - cmake --build cmake-out -j4 + cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) EXT=$(get_shared_lib_ext) SO_LIB="cmake-out/kernels/quantized/libquantized_ops_aot_lib$EXT" diff --git a/extension/llm/export/quantizer_lib.py b/extension/llm/export/quantizer_lib.py index e4564f32360..9c900c13e63 100644 --- a/extension/llm/export/quantizer_lib.py +++ b/extension/llm/export/quantizer_lib.py @@ -109,7 +109,7 @@ def check_embedding_byte_registered(): "Need to specify shared library path to register quantized ops (and their out variants) into EXIR.\n" "Follow the following steps to build the needed lib via cmake.\n" "Then from root executorch dir do the following:\n" - "rm -rf cmake-out && mkdir cmake-out && (cd cmake-out && cmake -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON ..) && cmake --build . -j16\n" + "rm -rf cmake-out && mkdir cmake-out && (cd cmake-out && cmake -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON ..) && cmake --build . -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 ))\n" 'To find the location of the lib: find cmake-out -name "libquantized_ops_aot_lib*"\n' "Then specify the said library via -s /dev/null || sysctl -n hw.ncpu) + 1 )) ``` To reduce the binary size, you may also use the selective build options found in the [Kernel Library Selective Build guide](../../docs/source/kernel-library-selective-build.md). You may also use optimized kernels with the `EXECUTORCH_BUILD_KERNELS_OPTIMIZED` option. Portable kernels are used by default. diff --git a/extension/wasm/tokenizers/README.md b/extension/wasm/tokenizers/README.md index e1c48992e94..5e446e18acc 100644 --- a/extension/wasm/tokenizers/README.md +++ b/extension/wasm/tokenizers/README.md @@ -14,7 +14,7 @@ emcmake cmake . -DEXECUTORCH_BUILD_TOKENIZERS_WASM=ON \ -Bcmake-out-wasm # Build the Wasm extension -cmake --build cmake-out-wasm --target tokenizers_wasm -j32 +cmake --build cmake-out-wasm --target tokenizers_wasm -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) ``` Emscripten modules are loaded into the global `Module` object by default. This means you cannot have multiple modules in the same page. If you are also using the ExecuTorch Wasm bindings, it is recommended to use the `MODULARIZE` option to avoid conflicts. diff --git a/tools/cmake/preset/README.md b/tools/cmake/preset/README.md index 2eaad11ef37..a1dbbbe0f3b 100644 --- a/tools/cmake/preset/README.md +++ b/tools/cmake/preset/README.md @@ -12,7 +12,7 @@ See: https://github.com/pytorch/executorch/discussions/10661. tl;dr instead of t ```bash $ cmake --preset macos -$ cmake --build cmake-out -j100 --target executor_runner +$ cmake --build cmake-out -j$(( $(nproc 2>/dev/null || sysctl -n hw.ncpu) + 1 )) --target executor_runner ``` ## Working with Presets From 55a969e674022cc23575008902e9e74659c03dfb Mon Sep 17 00:00:00 2001 From: Youngsik Yang Date: Mon, 7 Sep 2026 13:06:06 +0900 Subject: [PATCH 061/190] Arm backend: fix round() decomposition to round-half-to-even (#21065) ### Summary `DecomposeRoundPass` lowers `aten.round` with round-half-**away-from-zero**, but `torch.round` is round-half-**to-even** (banker's rounding). (see [`torch.round`](https://pytorch.org/docs/stable/generated/torch.round.html)). | input | `torch.round` (reference) | delegated output | fixed | | --- | --- | --- | --- | | `0.5` | `0` | `1` | `0` | | `2.5` | `2` | `3` | `2` | | `-1.5` | `-2` | `-2` | `-2` | | `-2.5` | `-2` | `-3` | `-2` | ### The fix round(x) now picks whichever of floor(x) and ceil(x) is nearer, and the even one of the two when x is exactly halfway. ``` dist_to_floor = x - floor(x) halved = floor(x) * 0.5 floor_is_odd = (halved - floor(halved)) == 0.5 take_ceil = (dist_to_floor > 0.5) | ((dist_to_floor == 0.5) & floor_is_odd) result = where(take_ceil, ceil(x), floor(x)) ``` | `x` | `floor(x)` | `ceil(x)` | `dist_to_floor` | `floor_is_odd` | `take_ceil` | `result` | | --- | --- | --- | --- | --- | --- | --- | | `-2.5` | `-3` | `-2` | `0.5` | `true` | `true` | `-2` | | `-2.25` | `-3` | `-2` | `0.75` | ~~`true`~~ | `true` | `-2` | | `-1.5` | `-2` | `-1` | `0.5` | `false` | `false` | `-2` | | `0.49999997` | `0` | `1` | `0.49999997` | ~~`false`~~ | `false` | `0` | | `0.50000006` | `0` | `1` | `0.50000006` | ~~`false`~~ | `true` | `1` | | `3.5` | `3` | `4` | `0.5` | `true` | `true` | `4` | ### Testing | case | data | | --- | --- | | `halfway_ties` | exact `.5`, both signs and both parities | | `halfway_neighbors` | one ulp either side of a tie | ```bash FILES="backends/arm/test/ops/test_round.py backends/arm/test/misc/test_mixed_type_lowering.py" PYTEST="pytest --config-file=backends/arm/test/pytest.ini --numprocesses=auto" $PYTEST $FILES -k tosa # 17 passed $PYTEST $FILES -k "u55 or u65" # 7 passed $PYTEST $FILES -k u85 # 7 passed lintrunner backends/arm/_passes/decompose_round_pass.py \ backends/arm/test/ops/test_round.py \ backends/arm/test/misc/test_mixed_type_lowering.py # No lint issues. ``` cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --------- Signed-off-by: Youngsik Yang --- backends/arm/_passes/decompose_round_pass.py | 98 ++++++++----------- .../arm/test/misc/test_mixed_type_lowering.py | 15 +-- backends/arm/test/ops/test_pixel_shuffling.py | 12 ++- backends/arm/test/ops/test_round.py | 17 +++- 4 files changed, 74 insertions(+), 68 deletions(-) diff --git a/backends/arm/_passes/decompose_round_pass.py b/backends/arm/_passes/decompose_round_pass.py index 48b26f1d027..dcb72d4202c 100644 --- a/backends/arm/_passes/decompose_round_pass.py +++ b/backends/arm/_passes/decompose_round_pass.py @@ -7,47 +7,25 @@ from executorch.backends.arm._passes import ArmOpTargetedPass from executorch.exir.dialects._ops import ops as exir_ops -from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass -from torch._ops import OpOverload -Op = OpOverload | EdgeOpOverload - - -def _get_round_decomposition_ops(op) -> tuple[Op, Op, Op, Op, Op, Op, Op]: - """Returns the (full_op, ge_op, add_op, sub_op, floor_op, ceil_op, where_op) - for the given round operation. - - The ops depend on whether the round op is an aten or edge op. - - """ - if op == exir_ops.edge.aten.round.default: - return ( - exir_ops.edge.aten.full.default, - exir_ops.edge.aten.ge.Tensor, - exir_ops.edge.aten.add.Scalar, - exir_ops.edge.aten.sub.Scalar, - exir_ops.edge.aten.floor.default, - exir_ops.edge.aten.ceil.default, - exir_ops.edge.aten.where.self, - ) - raise RuntimeError(f"Can't get round decomposition ops for op {op}") +class DecomposeRoundPass(ArmOpTargetedPass): + """Decomposes round(x) into round-half-to-even, matching the semantics of + aten.round / torch.round. + x lies between floor(x) and ceil(x), and its distance above floor(x) says + which one is nearer: less than 0.5 takes floor(x), more takes ceil(x), and + exactly 0.5 is a tie that takes whichever of the two is even. -class DecomposeRoundPass(ArmOpTargetedPass): - """ - For inputs >= 0, round(x) is equivalent to floor(x + 0.5), and for inputs < 0, - round(x) is equivalent to ceil(x - 0.5). This pass decomposes the round operation into - a sequence of more primitive operations. Example: - %zero = full((1,), 0.0, dtype=torch.float32) - %is_non_negative = ge(x, %zero) - %plus_half = add(x, 0.5) - %minus_half = sub(x, 0.5) - %floor = floor(%plus_half) - %ceil = ceil(%minus_half) - %result = where(%is_non_negative, %floor, %ceil) + %dist_to_floor = sub(x, floor(x)) + %halved = mul(floor(x), 0.5) + %floor_is_odd = eq(sub(%halved, floor(%halved)), 0.5) + %tie_to_even = logical_and(eq(%dist_to_floor, 0.5), %floor_is_odd) + %take_ceil = logical_or(gt(%dist_to_floor, 0.5), %tie_to_even) + %result = where(%take_ceil, ceil(x), floor(x)) + """ _passes_required_after: Set[Type[ExportPass]] = set() @@ -60,26 +38,30 @@ def call_operator(self, op, args, kwargs, meta, updated=False): if op not in self.target_ops or self._is_quantized_meta(meta): return super().call_operator(op, args, kwargs, meta, updated) x = args[0] - input_dtype = x.node.meta["val"].dtype - full, ge, add, sub, floor, ceil, where = _get_round_decomposition_ops(op) - zero = super().call_operator( - full, - args=((1,), 0.0), - kwargs={"dtype": input_dtype}, - meta=meta, - updated=True, - ) - is_non_negative = super().call_operator( - ge, (x, zero), kwargs, meta, updated=True - ) - plus_half = super().call_operator(add, (x, 0.5), kwargs, meta, updated=True) - minus_half = super().call_operator(sub, (x, 0.5), kwargs, meta, updated=True) - floor = super().call_operator(floor, (plus_half,), kwargs, meta, updated=True) - ceil = super().call_operator(ceil, (minus_half,), kwargs, meta, updated=True) - return super().call_operator( - where, - (is_non_negative, floor, ceil), - kwargs, - meta, - updated=True, - ) + + def call(op, *op_args): + return super(DecomposeRoundPass, self).call_operator( + op, op_args, kwargs, meta, updated=True + ) + + sub = exir_ops.edge.aten.sub.Tensor + mul = exir_ops.edge.aten.mul.Scalar + floor = exir_ops.edge.aten.floor.default + ceil = exir_ops.edge.aten.ceil.default + eq = exir_ops.edge.aten.eq.Scalar + gt = exir_ops.edge.aten.gt.Scalar + logical_and = exir_ops.edge.aten.logical_and.default + logical_or = exir_ops.edge.aten.logical_or.default + where = exir_ops.edge.aten.where.self + + floor_x = call(floor, x) + dist_to_floor = call(sub, x, floor_x) + + # floor_x is odd iff floor_x / 2 has a .5 fractional part + halved = call(mul, floor_x, 0.5) + halved_frac = call(sub, halved, call(floor, halved)) + floor_is_odd = call(eq, halved_frac, 0.5) + + tie_to_even = call(logical_and, call(eq, dist_to_floor, 0.5), floor_is_odd) + take_ceil = call(logical_or, call(gt, dist_to_floor, 0.5), tie_to_even) + return call(where, take_ceil, call(ceil, x), floor_x) diff --git a/backends/arm/test/misc/test_mixed_type_lowering.py b/backends/arm/test/misc/test_mixed_type_lowering.py index 6a2a1e4cbd5..14663b6fc51 100644 --- a/backends/arm/test/misc/test_mixed_type_lowering.py +++ b/backends/arm/test/misc/test_mixed_type_lowering.py @@ -33,14 +33,17 @@ def repeat_op_dict(op_dict, times): } q_tosa_ops = { "CAST": {"INT8": 1}, - "MUL": {"FP32": 1}, # scale multiplication - "ADD": {"FP32": 2}, # zero-point addition, rounding - "SUB": {"FP32": 1}, # for rounding + "MUL": {"FP32": 2}, # scale multiplication + round()'s internal multiply + "ADD": {"FP32": 1}, # zero-point addition + "SUB": {"FP32": 2}, # for rounding + "CEIL": {"FP32": 1}, # for rounding "CLAMP": {"FP32": 1}, # clamp - "GREATER_EQUAL": {"BOOL": 1}, # for rounding "SELECT": {"FP32": 1}, # for rounding - "CEIL": {"FP32": 1}, # for rounding - "FLOOR": {"FP32": 1}, # for rounding + "FLOOR": {"FP32": 2}, # for rounding + "EQUAL": {"BOOL": 2}, # for rounding + "GREATER": {"BOOL": 1}, # for rounding + "LOGICAL_AND": {"BOOL": 1}, # for rounding + "LOGICAL_OR": {"BOOL": 1}, # for rounding } diff --git a/backends/arm/test/ops/test_pixel_shuffling.py b/backends/arm/test/ops/test_pixel_shuffling.py index 4980e24bab3..03f8e9a17fe 100644 --- a/backends/arm/test/ops/test_pixel_shuffling.py +++ b/backends/arm/test/ops/test_pixel_shuffling.py @@ -34,6 +34,10 @@ "rand_4d_channels_last": "Known U55 partitioning limitation for large 4D pixel shuffle layouts.", } +mixed_precision_xfails = { + "rand_4d_channels_last": "Permute propagation stops at f(x, g(x)) shapes such as the round decomposition.", +} + class PixelUnShuffle(nn.Module): @@ -110,7 +114,9 @@ def test_pixel_unshuffle_tosa_FP(test_data: input_t1): pipeline.run() -@common.parametrize("test_data", PixelUnShuffle.test_data_generators) +@common.parametrize( + "test_data", PixelUnShuffle.test_data_generators, xfails=mixed_precision_xfails +) def test_pixel_unshuffle_no_target_tosa_mixed_precision(test_data: input_t1): inputs, expected_transposes = test_data() pipeline = TosaPipelineINT[input_t1]( @@ -140,7 +146,9 @@ def test_pixel_shuffle_tosa_FP(test_data: input_t1): pipeline.run() -@common.parametrize("test_data", PixelShuffle.test_data_generators) +@common.parametrize( + "test_data", PixelShuffle.test_data_generators, xfails=mixed_precision_xfails +) def test_pixel_shuffle_no_target_tosa_mixed_precision(test_data: input_t1): inputs, expected_transposes = test_data() pipeline = TosaPipelineINT[input_t1]( diff --git a/backends/arm/test/ops/test_round.py b/backends/arm/test/ops/test_round.py index 1f4470cebb3..f06f7e37737 100644 --- a/backends/arm/test/ops/test_round.py +++ b/backends/arm/test/ops/test_round.py @@ -21,6 +21,7 @@ aten_op = "torch.ops.aten.round.default" exir_op = "executorch_exir_dialects_edge__ops_aten_round_default" + test_data_suite = { # (test_name, test_data) "zeros": lambda: torch.zeros(1, 10, 10, 10), @@ -29,6 +30,14 @@ "randn_pos": lambda: torch.randn(10) + 10, "randn_neg": lambda: torch.randn(10) - 10, "ramp": lambda: torch.arange(-16, 16, 0.2), + "halfway_ties": lambda: torch.arange(-8, 8, 0.5), +} + +# One ulp either side of a tie. +test_data_suite_fp = { + "halfway_neighbors": lambda: torch.nextafter( + torch.tensor([0.5, 0.5]), torch.tensor([-float("inf"), float("inf")]) + ), } test_data_suite_bf16 = { @@ -41,7 +50,9 @@ def forward(self, x: torch.Tensor): return x.round() -@common.parametrize("test_data", test_data_suite | test_data_suite_bf16) +@common.parametrize( + "test_data", test_data_suite | test_data_suite_fp | test_data_suite_bf16 +) def test_round_tosa_FP(test_data: torch.Tensor): pipeline = TosaPipelineFP[input_t1]( Round(), @@ -88,7 +99,9 @@ def test_round_u85_INT(test_data: torch.Tensor): pipeline.run() -@common.parametrize("test_data", test_data_suite | test_data_suite_bf16) +@common.parametrize( + "test_data", test_data_suite | test_data_suite_fp | test_data_suite_bf16 +) @common.SkipIfNoModelConverter def test_round_vgf_no_quant(test_data: torch.Tensor): pipeline = VgfPipeline[input_t1]( From 067c74b7f92547ae4293288b2f1b08fb0605c0d3 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Mon, 7 Sep 2026 15:10:41 +0100 Subject: [PATCH 062/190] Arm backend: Fix bad init performance in VGF for multi segments (#22588) Improvements to init performance in VGF. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Elena Zhelezina --- backends/arm/runtime/VGFSetup.cpp | 53 +++++++++++++++++++++++++++++-- backends/arm/runtime/VGFSetup.h | 12 +++++++ 2 files changed, 63 insertions(+), 2 deletions(-) diff --git a/backends/arm/runtime/VGFSetup.cpp b/backends/arm/runtime/VGFSetup.cpp index 9fca73d3551..4adcb406666 100644 --- a/backends/arm/runtime/VGFSetup.cpp +++ b/backends/arm/runtime/VGFSetup.cpp @@ -2668,9 +2668,44 @@ bool VgfRepr::process_vgf( } } + // Keep the pipeline cache local to this VgfRepr. This preserves reuse across + // the many segment pipelines in one VGF without sharing mutable cache state + // or cache lifetime across independently initialized delegate handles. + if (vk_pipeline_cache == VK_NULL_HANDLE) { + VkPipelineCacheCreateInfo pipeline_cache_info{ + .sType = VK_STRUCTURE_TYPE_PIPELINE_CACHE_CREATE_INFO, + .pNext = nullptr, + // Default mode: Vulkan internally synchronizes concurrent + // pipeline-cache access. The cache is per-VgfRepr anyway, so + // independent delegates do not contend on it. + .flags = 0, + .initialDataSize = 0, + .pInitialData = nullptr, + }; + + { + VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_CREATE_PIPELINE_CACHE"); + result = vkCreatePipelineCache( + vk_device, &pipeline_cache_info, nullptr, &vk_pipeline_cache); + } + + if (result != VK_SUCCESS) { + ET_LOG( + Info, + "Failed to create optional per-VgfRepr Vulkan pipeline cache, " + "error 0x%08X; continuing without pipeline caching", + result); + vk_pipeline_cache = VK_NULL_HANDLE; + } else { + ET_LOG(Info, "VGF per-VgfRepr Vulkan pipeline cache enabled"); + } + } + // Build per-segment pipelines and descriptor sets. segments.clear(); segments.reserve(segment_count); + size_t graph_segment_count = 0; + size_t compute_segment_count = 0; { VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_BUILD_SEGMENTS"); @@ -2683,6 +2718,12 @@ bool VgfRepr::process_vgf( return false; } + if (segment_type == vgflib::ModuleType::GRAPH) { + ++graph_segment_count; + } else { + ++compute_segment_count; + } + SegmentState segment; segment.segment_id = segment_id; segment.use_data_graph_pipeline = @@ -3137,7 +3178,7 @@ bool VgfRepr::process_vgf( result = vkCreateDataGraphPipelinesARM( vk_device, VK_NULL_HANDLE, - VK_NULL_HANDLE, + vk_pipeline_cache, 1, &graph_pipeline_info, nullptr, @@ -3429,7 +3470,7 @@ bool VgfRepr::process_vgf( VGF_PROFILE_SCOPE(event_tracer, "VGF_INIT_CREATE_COMPUTE_PIPELINE"); result = vkCreateComputePipelines( vk_device, - VK_NULL_HANDLE, + vk_pipeline_cache, 1, &compute_info, nullptr, @@ -3445,6 +3486,14 @@ bool VgfRepr::process_vgf( } } + ET_LOG( + Info, + "VGF segment counts: total=%d graph=%zu compute=%zu pipeline_cache=%s", + segment_count, + graph_segment_count, + compute_segment_count, + vk_pipeline_cache != VK_NULL_HANDLE ? "enabled" : "disabled"); + // Map model sequence inputs/outputs to IO indices auto input_handle = sequence_decoder->getModelSequenceInputBindingSlotsHandle(); diff --git a/backends/arm/runtime/VGFSetup.h b/backends/arm/runtime/VGFSetup.h index e4b60bbbf94..b87f648962c 100644 --- a/backends/arm/runtime/VGFSetup.h +++ b/backends/arm/runtime/VGFSetup.h @@ -177,6 +177,12 @@ class VgfRepr { ~VgfRepr() { free_vgf(); + if (vk_pipeline_cache != VK_NULL_HANDLE) { + // The cache is private to this VgfRepr, so no other delegate instance can + // be accessing it while this object is being destroyed. + vkDestroyPipelineCache(vk_device, vk_pipeline_cache, nullptr); + vk_pipeline_cache = VK_NULL_HANDLE; + } } private: @@ -188,6 +194,12 @@ class VgfRepr { VkCommandPool vk_command_pool; uint32_t vk_queue_family_index = UINT32_MAX; + // Owned by this VgfRepr. One cache is reused across all graph and compute + // segments in this loaded VGF, but is not shared with independent VgfRepr + // instances. flags=0 uses Vulkan's default internally synchronized cache + // mode. + VkPipelineCache vk_pipeline_cache = VK_NULL_HANDLE; + bool neural_statistics_requested_ = false; bool neural_statistics_device_enabled_ = false; int neural_statistics_mode_index_ = 1; From cfc6ecd5166bec70c5fde2bad836fed26018fe52 Mon Sep 17 00:00:00 2001 From: zhaoxul-qti Date: Mon, 7 Sep 2026 23:47:26 +0800 Subject: [PATCH 063/190] Qualcomm AI Engine Direct - Enable QNN Windows ARM64 build in CI (#22438) Follow up on #20802. With Windows ARM64 build CI enabled in #21817, add Qualcomm Windows ARM64 build validation to the CI pipeline. ### Summary - Adds a `build-qnn-windows-arm64` job to `qnn-windows-msvc.yml`, running on the native windows-11-arm, to build the QNN backend for Windows ARM64. - Switches Python environment setup from `conda` to `venv` (`py -3.12 -m venv et`) for both architectures, since `conda` isn't available on the windows-11-arm image. - Verifies build artifacts (pybind module, backend DLL, executor runner) per architecture after each build. ### Test plan - Passing .github/workflows/qnn-windows-msvc.yml on both `build-qnn-windows-x64` and `build-qnn-windows-arm64` jobs. --- .ci/scripts/build-qnn-windows-msvc.ps1 | 63 +++++++++++++++++--------- .github/workflows/qnn-windows-msvc.yml | 36 +++++++++++++-- 2 files changed, 75 insertions(+), 24 deletions(-) diff --git a/.ci/scripts/build-qnn-windows-msvc.ps1 b/.ci/scripts/build-qnn-windows-msvc.ps1 index a69316be6e3..c94daf453e7 100644 --- a/.ci/scripts/build-qnn-windows-msvc.ps1 +++ b/.ci/scripts/build-qnn-windows-msvc.ps1 @@ -4,13 +4,35 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +param( + [switch]$SkipX86Windows, + [switch]$SkipArm64Windows +) + $ErrorActionPreference = "Stop" -conda create --yes --quiet -n et python=3.12 -conda activate et +if ($SkipX86Windows -eq $SkipArm64Windows) { + Write-Error "Specify exactly one of -SkipArm64Windows (to build x86_64) or -SkipX86Windows (to build arm64)." + exit 1 +} -# Install CI requirements -pip install -r .ci/docker/requirements-ci.txt +if ($SkipX86Windows) { + $ArchLabel = "arm64" + py -3.12 -m venv et + .\et\Scripts\Activate.ps1 + # ARM64 prebuilt wheels are not available for some Python modules. + # To unblock the build process, only a minimal set of dependencies + # is installed via pip. `PyYAML`/`torch` for ExecuTorch's codegen, + # `requests` for download_qnn_sdk.py. + pip install pyyaml requests + pip install torch --index-url https://download.pytorch.org/whl/cpu +} else { + $ArchLabel = "x86_64" + conda create --yes --quiet -n et python=3.12 + conda activate et + # Install CI requirements + pip install -r .ci/docker/requirements-ci.txt +} # Provision the QNN SDK if ($env:QNN_SDK_ROOT -and (Test-Path -Path $env:QNN_SDK_ROOT)) { @@ -44,27 +66,26 @@ if (-not (Test-Path -Path (Join-Path $env:QNN_SDK_ROOT "include\QNN"))) { exit 1 } -# Test x86_64 Windows host build -.\backends\qualcomm\scripts\build.ps1 -SkipArm64Windows -Release +if ($SkipArm64Windows) { + .\backends\qualcomm\scripts\build.ps1 -SkipArm64Windows -Release +} else { + .\backends\qualcomm\scripts\build.ps1 -SkipX86Windows -Release +} -$x86Artifacts = @( - "build-x86_64-windows\backends\qualcomm\Release\PyQnnManagerAdaptor*.pyd", - "build-x86_64-windows\backends\qualcomm\Release\qnn_executorch_backend.dll", - "build-x86_64-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe" +$Artifacts = @( + "build-$ArchLabel-windows\backends\qualcomm\Release\qnn_executorch_backend.dll", + "build-$ArchLabel-windows\examples\qualcomm\executor_runner\Release\qnn_executor_runner.exe" ) -foreach ($artifact in $x86Artifacts) { +if ($SkipArm64Windows) { + # Only run PyQnnManagerAdaptor validation for x86_64 Windows artifacts, + # since AOT is not fully supported on native ARM64 Windows. + $Artifacts += "build-x86_64-windows\backends\qualcomm\Release\PyQnnManagerAdaptor*.pyd" +} +foreach ($artifact in $Artifacts) { if (-not (Get-ChildItem -Path $artifact -ErrorAction SilentlyContinue)) { - Write-Error "ERROR: x86_64 artifact not found: $artifact" + Write-Error "ERROR: $ArchLabel artifact not found: $artifact" exit 1 } } -# The ARM64 MSVC toolchain is currently not installed in the Windows CI -# environment. Enabling this build configuration results in build failures -# due to the missing ARM64 platform definition. -# `.\backends\qualcomm\scripts\build.ps1 -SkipX86Windows -Release` -# -# Temporarily disable this build option until ARM64 MSVC support is available -# in CI. The configuration can be re-enabled in a future update. - -Write-Host "PASSED: QNN backend Windows MSVC build completed" +Write-Host "PASSED: QNN backend Windows MSVC build ($ArchLabel) completed" diff --git a/.github/workflows/qnn-windows-msvc.yml b/.github/workflows/qnn-windows-msvc.yml index 377a2625e60..65dafc3636a 100644 --- a/.github/workflows/qnn-windows-msvc.yml +++ b/.github/workflows/qnn-windows-msvc.yml @@ -56,8 +56,8 @@ permissions: contents: read jobs: - build-qnn-windows-msvc: - name: build-qnn-windows-msvc + build-qnn-windows-x64: + name: build-qnn-windows-x64 uses: pytorch/test-infra/.github/workflows/windows_job.yml@main with: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} @@ -70,5 +70,35 @@ jobs: Set-PSDebug -Trace 1 \$ErrorActionPreference = 'Stop' \$PSNativeCommandUseErrorActionPreference = \$true - .ci/scripts/build-qnn-windows-msvc.ps1 + .ci/scripts/build-qnn-windows-msvc.ps1 -SkipArm64Windows }" + + build-qnn-windows-arm64: + name: build-qnn-windows-arm64 + runs-on: windows-11-arm + timeout-minutes: 90 + steps: + - name: Enable long paths + shell: cmd + run: | + git config --system --get core.longpaths || echo "core.longpaths is not set, setting it now" + git config --system core.longpaths true + + - name: Checkout ExecuTorch + uses: actions/checkout@v4 + with: + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + + - name: Initialize submodules + shell: pwsh + run: | + git config --global http.sslBackend openssl + git submodule update --init --recursive + + - name: Build QNN backend (arm64) + shell: pwsh + run: | + Set-PSDebug -Trace 1 + $ErrorActionPreference = 'Stop' + $PSNativeCommandUseErrorActionPreference = $true + .ci/scripts/build-qnn-windows-msvc.ps1 -SkipX86Windows From 3574bf1a3e97ab83d077073e5ea0d04c8a2e0d2b Mon Sep 17 00:00:00 2001 From: Sangwon Ha <146179778+FabulousSuperDude@users.noreply.github.com> Date: Mon, 7 Sep 2026 18:58:33 +0100 Subject: [PATCH 064/190] Arm backend: Relax Qwen3-VL 2B MXFP8 tolerance (#22590) Parameterize MXFP comparison thresholds and relax them only for the full-size 2B text model to account for accumulated quantization error. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Sangwon Ha --- .../arm/test/models/Qwen3_VL/test_qwen3_vl_model.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py b/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py index 2b83f75030a..a826c4fbdb9 100644 --- a/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py +++ b/backends/arm/test/models/Qwen3_VL/test_qwen3_vl_model.py @@ -314,6 +314,8 @@ def _test_qwen3_vl_full_models_vgf_no_quant_bf16( def _test_qwen3_vl_text_model_tosa_mxfp8_bf16( config_factory=_make_qwen3_vl_e2e_test_config, + frobenius_threshold: float = 0.1, + cosine_threshold: float = 0.98, ): # The Qwen 3 VL FP8 model only quantizes the TextModel model, inputs = TextModelWrapper.prepare_model_and_inputs(config_factory) @@ -326,8 +328,8 @@ def _test_qwen3_vl_text_model_tosa_mxfp8_bf16( aten_op=aten_op_mxfp_linear, exir_op=[], filter_fn=_is_linear, - frobenius_threshold=0.1, - cosine_threshold=0.98, + frobenius_threshold=frobenius_threshold, + cosine_threshold=cosine_threshold, mxfp_config=mxfp_config, tosa_version="1.1", tosa_extensions=["bf16", "mxfp"], @@ -404,4 +406,8 @@ def test_qwen3_vl_2b_instruct_full_models_vgf_no_quant_bf16( @pytest.mark.slow @pytest.mark.xlarge def test_qwen3_vl_2b_instruct_text_model_tosa_mxfp8_bf16(): - _test_qwen3_vl_text_model_tosa_mxfp8_bf16(_make_qwen3_vl_2b_instruct_layer_config) + _test_qwen3_vl_text_model_tosa_mxfp8_bf16( + _make_qwen3_vl_2b_instruct_layer_config, + frobenius_threshold=0.3, + cosine_threshold=0.95, + ) From 9b3f660b4a74afa87726fbbb4480196d02fcb9e9 Mon Sep 17 00:00:00 2001 From: Sangwon Ha <146179778+FabulousSuperDude@users.noreply.github.com> Date: Mon, 7 Sep 2026 19:02:04 +0100 Subject: [PATCH 065/190] Arm backend: Xfail duplicate VGF input on Darwin (#22591) MoltenVK rejects the incomplete descriptor layout generated by Model Converter for duplicated custom-shader inputs. Re-enable this test on Darwin once Model Converter preserves both descriptor bindings. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Sangwon Ha --- backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py b/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py index 0baea8b832e..c24bbf6e11f 100644 --- a/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py +++ b/backends/arm/test/runtime/test_vgf_tensor_buffer_runtime.py @@ -6,6 +6,7 @@ import sys from pathlib import Path +import pytest import torch import torch.nn.functional as F @@ -162,6 +163,11 @@ def test_two_input_add_buffer_shader_executes(tmp_path): # Covers the two-input storage-buffer shader path when both inputs are the same tensor. # Checks runtime execution matches eager output for the duplicated-input add case. +@pytest.mark.xfail( + sys.platform == "darwin", + reason="Model Converter drops duplicated custom-shader descriptor binding", + strict=True, +) @common.SkipIfNoModelConverter def test_two_input_add_buffer_shader_with_duplicated_input_executes(tmp_path): x = torch.randn(256) From b4c53c6c3e2d6ac2d6c601aab0b9cbedb3babb88 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 7 Sep 2026 12:36:27 -0700 Subject: [PATCH 066/190] Move the NXP op alias table out of the test package (#22592) The NXP backend keeps a table of short names for edge operators, for example `AddTensor = exir_ops.edge.aten.add.Tensor`. It lives in the backend's test package, but three modules that ship in the wheel read it at import time: - `backend/edge_helper.py` - `backend/node_format_inference.py` - `edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py` Two dozen production files import those modules, so the backend cannot load at all without a test package. The dependency points the wrong way: the tests should depend on the backend, not the backend on the tests. ### What this does The table has no test logic in it, so it moves to `backend/ops_aliases.py`, next to the code that reads it. That follows the shape the Arm backend already uses for shared operator constants in `backends/arm/constants.py`. The Buck target that published the table from the test package is removed, since the backend library already globs the directory it moves into. No behavior changes. The names, values and every importer stay the same, including the roughly fifty tests that use the table. ### Why the diff looks larger than it is Most of it is one line per file. Every importer switches from `nxp.tests.ops_aliases` to `nxp.backend.ops_aliases`, and because `backend` sorts before `tests`, the import sorter then moves that line to the top of the block. So a one line change shows up as several lines moving in some files. No code changed in any of them. ### Test plan Imported the moved module and every updated importer from an installed wheel, and confirmed the table exposes the same names bound to the same values as before the move. Compared the alias assignments before and after by parsing both files: 56 aliases, no names added or removed, no values changed. Checked that every NXP test target can still reach the table, through the backend library that already globs that directory, and that no reference to the old Buck target or the old import path remains anywhere in the tree. Ran the formatter and linter over all changed files. Co-authored-by: PyTorch Bot --- backends/nxp/BUCK | 1 - backends/nxp/backend/edge_helper.py | 2 +- backends/nxp/backend/node_format_inference.py | 2 +- .../nxp/{tests => backend}/ops_aliases.py | 4 ++-- ...operator_into_separate_qdq_cluster_pass.py | 2 +- backends/nxp/tests/BUCK | 11 ---------- ...add_batch_size_for_3d_input_pool_2d_ops.py | 14 ++++++------ .../generic_tests/test_convert_div_to_mul.py | 2 +- .../test_convert_scalar_to_attr.py | 2 +- .../test_decompose_split_to_slices.py | 2 +- .../tests/generic_tests/test_integration.py | 2 +- .../test_node_format_inference.py | 8 +++---- .../test_quantized_input_data.py | 2 +- backends/nxp/tests/graph_verifier.py | 10 ++++----- .../node_converter/test_abs_converter.py | 2 +- .../test_adaptive_avg_pool2d_converter.py | 8 +++---- .../test_add_tensor_converter.py | 12 +++++----- .../node_converter/test_addmm_converter.py | 14 ++++++------ .../node_converter/test_amax_converter.py | 14 ++++++------ .../node_converter/test_amin_converter.py | 14 ++++++------ .../test_avg_pool2d_converter.py | 10 ++++----- .../node_converter/test_bmm_converter.py | 6 ++++- .../node_converter/test_cat_converter.py | 12 +++++----- .../node_converter/test_clamp_converter.py | 10 ++++----- .../test_constant_pad_nd_converter.py | 2 +- .../node_converter/test_conv_converter.py | 10 ++++----- .../node_converter/test_exp_converter.py | 2 +- .../test_hardswish_converter.py | 14 ++++++------ .../node_converter/test_hardtanh_converter.py | 10 ++++----- .../test_leaky_relu_converter.py | 2 +- .../node_converter/test_log_converter.py | 2 +- .../test_max_pool_2d_converter.py | 12 +++++----- .../node_converter/test_maximum_converter.py | 12 +++++----- .../node_converter/test_mean_dim_converter.py | 14 ++++++------ .../node_converter/test_minimum_converter.py | 12 +++++----- .../node_converter/test_mm_converter.py | 2 +- .../test_mul_tensor_converter.py | 12 +++++----- .../node_converter/test_neg_converter.py | 2 +- .../node_converter/test_pad_converter.py | 2 +- .../test_permute_copy_converter.py | 12 +++++----- .../node_converter/test_prelu_converter.py | 22 +++++++++---------- .../node_converter/test_relu_converter.py | 14 ++++++------ .../node_converter/test_rsqrt_converter.py | 2 +- .../node_converter/test_sigmoid_converter.py | 2 +- .../test_slice_copy_tensor_converter.py | 10 ++++----- .../node_converter/test_softmax_converter.py | 12 +++++----- .../test_sub_tensor_converter.py | 12 +++++----- .../test_sum_dim_int_list_converter.py | 14 ++++++------ .../node_converter/test_tanh_converter.py | 2 +- .../test_upsample_bilinear2d.py | 10 ++++----- .../node_converter/test_upsample_nearest2d.py | 10 ++++----- .../test_view_copy_converter.py | 18 +++++++-------- .../test_convert_reshaping_nodes_to_view.py | 2 +- backends/nxp/tests/nsys_testing.py | 2 +- 54 files changed, 206 insertions(+), 214 deletions(-) rename backends/nxp/{tests => backend}/ops_aliases.py (94%) diff --git a/backends/nxp/BUCK b/backends/nxp/BUCK index 6dec42f04d7..81a1e04cfee 100644 --- a/backends/nxp/BUCK +++ b/backends/nxp/BUCK @@ -68,7 +68,6 @@ fbcode_target(_kind = runtime.python_library, "fbsource//third-party/pypi/neutron_converter:neutron_converter", "//caffe2:torch", "//executorch/exir:lib", - "//executorch/backends/nxp/tests:ops_aliases", ], ) diff --git a/backends/nxp/backend/edge_helper.py b/backends/nxp/backend/edge_helper.py index 408b90e264d..4708b6dfdd0 100644 --- a/backends/nxp/backend/edge_helper.py +++ b/backends/nxp/backend/edge_helper.py @@ -8,7 +8,7 @@ import torch -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddTensor, Amax, Amin, diff --git a/backends/nxp/backend/node_format_inference.py b/backends/nxp/backend/node_format_inference.py index 689de41f3e4..64595eb411f 100644 --- a/backends/nxp/backend/node_format_inference.py +++ b/backends/nxp/backend/node_format_inference.py @@ -15,7 +15,7 @@ try_get_arg, ) from executorch.backends.nxp.backend.edge_program_converter import functions_converters -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AdaptiveAvgPool2D, Amax, Amin, diff --git a/backends/nxp/tests/ops_aliases.py b/backends/nxp/backend/ops_aliases.py similarity index 94% rename from backends/nxp/tests/ops_aliases.py rename to backends/nxp/backend/ops_aliases.py index 5ccc1d67de0..a9eb04e1960 100644 --- a/backends/nxp/tests/ops_aliases.py +++ b/backends/nxp/backend/ops_aliases.py @@ -3,8 +3,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -# This file defines ops aliases for shorter and more readable test description. List is sorted alphabetically. -# When finding a missing alias, add it at the correct place. +# This file defines ops aliases for shorter and more readable descriptions of edge operators. +# List is sorted alphabetically. When finding a missing alias, add it at the correct place. import operator diff --git a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py index dafea5d259b..9ac888addb9 100644 --- a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py +++ b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py @@ -6,10 +6,10 @@ import operator import torch +from executorch.backends.nxp.backend.ops_aliases import PermuteCopy from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass from executorch.backends.nxp.neutron_partitioner import QDQClusterRecognizer -from executorch.backends.nxp.tests.ops_aliases import PermuteCopy # noinspection PyProtectedMember from executorch.exir.dialects._ops import ops as exir_ops diff --git a/backends/nxp/tests/BUCK b/backends/nxp/tests/BUCK index 7879e1e3db5..a3a5adb215c 100644 --- a/backends/nxp/tests/BUCK +++ b/backends/nxp/tests/BUCK @@ -4,17 +4,6 @@ load("@fbcode_macros//build_defs:python_pytest.bzl", "python_pytest") oncall("executorch") -fbcode_target(_kind = runtime.python_library, - name = "ops_aliases", - srcs = [ - "ops_aliases.py", - ], - deps = [ - "//caffe2:torch", - "//executorch/exir:lib", - ], -) - fbcode_target(_kind = runtime.python_library, name = "models", srcs = [ diff --git a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py index b426c34c260..679e7365577 100644 --- a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py +++ b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py @@ -15,6 +15,13 @@ from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( NeutronAtenPassManager, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AdaptiveAvgPool2D, + AvgPool2D, + GetItem, + MaxPool2DWithIndices, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -28,13 +35,6 @@ MaxPool2dModule, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AdaptiveAvgPool2D, - AvgPool2D, - GetItem, - MaxPool2DWithIndices, - ViewCopy, -) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py index 3415b79a39d..044467ca223 100644 --- a/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py +++ b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py @@ -13,6 +13,7 @@ ConvertDivToMulPass, NeutronAtenPassManager, ) +from executorch.backends.nxp.backend.ops_aliases import MulTensor from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -22,7 +23,6 @@ StaticDivLinearModel, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import MulTensor @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py index b3d68eeeb56..c68cde0c23b 100644 --- a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py +++ b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py @@ -18,6 +18,7 @@ from executorch.backends.nxp.backend.edge_helper import ( try_get_tensor_constant_from_node, ) +from executorch.backends.nxp.backend.ops_aliases import AddTensor, MulTensor, SubTensor from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -31,7 +32,6 @@ AllCloseOutputComparator, lower_run_compare, ) -from executorch.backends.nxp.tests.ops_aliases import AddTensor, MulTensor, SubTensor @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py index de4e684405c..3c3c29f95b5 100644 --- a/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py +++ b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py @@ -12,6 +12,7 @@ NeutronAtenPassManager, SplitGRUBasedOnNumLayers, ) +from executorch.backends.nxp.backend.ops_aliases import SliceCopy from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -21,7 +22,6 @@ SplitWithSize, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import SliceCopy @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_integration.py b/backends/nxp/tests/generic_tests/test_integration.py index 9916cba5bdd..f0cb25548ec 100644 --- a/backends/nxp/tests/generic_tests/test_integration.py +++ b/backends/nxp/tests/generic_tests/test_integration.py @@ -5,8 +5,8 @@ import executorch.extension.pybindings.portable_lib import executorch.kernels.quantized # noqa F401 +from executorch.backends.nxp.backend.ops_aliases import AddMM, Convolution from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.ops_aliases import AddMM, Convolution from executorch.backends.nxp.tests.use_qat import * # noqa F401 from executorch.backends.nxp.tests.executorch_pipeline import ( diff --git a/backends/nxp/tests/generic_tests/test_node_format_inference.py b/backends/nxp/tests/generic_tests/test_node_format_inference.py index 18d5f874aab..5206a029e01 100644 --- a/backends/nxp/tests/generic_tests/test_node_format_inference.py +++ b/backends/nxp/tests/generic_tests/test_node_format_inference.py @@ -13,6 +13,10 @@ NodeFormatInference, NXP_NODE_FORMAT, ) +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -21,10 +25,6 @@ MaxPool2dModule, SoftmaxModule, ) -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - MaxPool2DWithIndices, -) def test_convolution(): diff --git a/backends/nxp/tests/generic_tests/test_quantized_input_data.py b/backends/nxp/tests/generic_tests/test_quantized_input_data.py index a9f9f3e47e6..bd2a6056a0b 100644 --- a/backends/nxp/tests/generic_tests/test_quantized_input_data.py +++ b/backends/nxp/tests/generic_tests/test_quantized_input_data.py @@ -5,6 +5,7 @@ import executorch.backends.nxp.tests.nsys_testing as nsys_testing import torch +from executorch.backends.nxp.backend.ops_aliases import AvgPool2D, MulTensor from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -14,7 +15,6 @@ OUTPUTS_DIR, ReferenceModel, ) -from executorch.backends.nxp.tests.ops_aliases import AvgPool2D, MulTensor def test__single_quantized_inputs(mocker, request): diff --git a/backends/nxp/tests/graph_verifier.py b/backends/nxp/tests/graph_verifier.py index 44900b6a11b..70701bbd90f 100644 --- a/backends/nxp/tests/graph_verifier.py +++ b/backends/nxp/tests/graph_verifier.py @@ -10,17 +10,17 @@ from dataclasses import dataclass from typing import Callable, Union -from executorch.backends.nxp.neutron_partitioner import ( - NeutronPartitioner, - NXP_DELEGATION_TAG, -) -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( DequantizePerChannel, DequantizePerTensor, QuantizePerChannel, QuantizePerTensor, ) +from executorch.backends.nxp.neutron_partitioner import ( + NeutronPartitioner, + NXP_DELEGATION_TAG, +) from executorch.exir.dialects.edge._ops import EdgeOpOverload from pytest_mock import MockerFixture diff --git a/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py index d42ef4c6e7d..94a0aafc2a7 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_abs_converter.py @@ -8,12 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Abs from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import ( lower_run_compare, RandomDatasetCreator, ) -from executorch.backends.nxp.tests.ops_aliases import Abs from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py index 9646c04a3f2..65e2b82555c 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py @@ -8,6 +8,10 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AdaptiveAvgPool2D, + ExecutorchDelegateCall, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -18,10 +22,6 @@ ) from executorch.backends.nxp.tests.models import AdaptiveAvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AdaptiveAvgPool2D, - ExecutorchDelegateCall, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py index c01d0ca818d..237d15aefbd 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -21,12 +27,6 @@ ) from executorch.backends.nxp.tests.models import AddTensorModule, MaxPoolAddTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py index 1db604d5b1e..6ebb8ac128c 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py @@ -8,6 +8,13 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + ExecutorchDelegateCall, + MM, + PermuteCopy, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -15,13 +22,6 @@ from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator from executorch.backends.nxp.tests.models import AddmmModule, LinearModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddMM, - ExecutorchDelegateCall, - MM, - PermuteCopy, - ViewCopy, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py index d348c12102e..892ea70018e 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_amax_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Amax, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - Amax, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py index e2490d9c2c4..3ae13bafe40 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_amin_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Amin, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - Amin, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py index 3db1158d637..c132c9e509d 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py @@ -8,17 +8,17 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AvgPool2D, + ExecutorchDelegateCall, + ViewCopy, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import AvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AvgPool2D, - ExecutorchDelegateCall, - ViewCopy, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py index c564c024623..ee46ad94d44 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py @@ -6,6 +6,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + BMM, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import ( ViewCopy, @@ -21,7 +26,6 @@ BatchMatMulModel, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import BMM, GetItem, MaxPool2DWithIndices from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py index b28a431e3ca..30983f4a666 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + Cat, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.executorch_pipeline import ( ModelInputSpec, @@ -16,12 +22,6 @@ from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Cat, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py index b2147a0d984..04ae8757110 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_clamp_converter.py @@ -18,6 +18,11 @@ from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator as Ops, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Clamp, + ExecutorchDelegateCall, +) from executorch.backends.nxp.tests.executorch_pipeline import ( ModelInputSpec, to_quantized_edge_program, @@ -25,11 +30,6 @@ from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - Clamp, - ExecutorchDelegateCall, -) from executorch.backends.nxp.tests.use_qat import * # noqa: F403 F401 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py index b4a64447aa6..730334932d7 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py @@ -12,10 +12,10 @@ from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( ModelBuilder, ) +from executorch.backends.nxp.backend.ops_aliases import ConstantPadND, Convolution from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import PadConvModule, PadModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ConstantPadND, Convolution from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py index 3d20d38bb54..f0fb0fdb7b4 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py @@ -6,6 +6,11 @@ import numpy as np import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -16,11 +21,6 @@ lower_run_compare, ReferenceModel, ) -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - ViewCopy, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py index b304dce2c94..9f8068f67d4 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_exp_converter.py @@ -8,10 +8,10 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Exp from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Exp from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.dataset_creator import ( LinearRampDatasetCreator, diff --git a/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py index deff4e12f0c..76aaab7520e 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py @@ -8,6 +8,13 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Convolution, + Hardswish, + PermuteCopy, + ViewCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -20,13 +27,6 @@ LinearHardswishModule, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddMM, - Convolution, - Hardswish, - PermuteCopy, - ViewCopy, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py index 66a052dba4f..1199609bb17 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py @@ -18,17 +18,17 @@ from executorch.backends.nxp.backend.ir.lib.tflite.BuiltinOperator import ( BuiltinOperator as Ops, ) +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + HardTanh, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import Conv2dWithActivation, HardTanhModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - HardTanh, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py index 567cf85ebe5..5176a8c61db 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_leaky_relu_converter.py @@ -8,11 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import LeakyRelu from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import LeakyRelu from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py index 0b7fe88cffc..ebf6a8bfb9b 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_log_converter.py @@ -8,10 +8,10 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Log from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Log from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.dataset_creator import ( LinearRampDatasetCreator, diff --git a/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py index 55a47146bfc..7d3769c896d 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py @@ -8,17 +8,17 @@ # noinspection PyUnusedImports import pytest import torch - -from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( ExecutorchDelegateCall, GetItem, MaxPool2DWithIndices, ViewCopy, ) + +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py index e25ed98fb3f..349a624242a 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + Maximum, + MaxPool2DWithIndices, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -21,12 +27,6 @@ ) from executorch.backends.nxp.tests.models import MaximumModule, MaxPoolMaximumModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - Maximum, - MaxPool2DWithIndices, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py index 1674153540f..14f55de94d4 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mean_dim_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + MeanDim, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - MeanDim, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py index 9dc2b8d77d7..7c74fd4a1c9 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + Minimum, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -21,12 +27,6 @@ ) from executorch.backends.nxp.tests.models import MaxPoolMinimumModule, MinimumModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - Minimum, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py index 423999dc7ec..405e955a5ca 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py @@ -8,12 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import MM, PermuteCopy, ViewCopy from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator from executorch.backends.nxp.tests.models import LinearModule, MmModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import MM, PermuteCopy, ViewCopy from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py index 718383284be..b031ab7e47c 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + MulTensor, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -21,12 +27,6 @@ ) from executorch.backends.nxp.tests.models import MaxPoolMulTensorModule, MulTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - MulTensor, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py index 691cb3bd2ca..182e7754013 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_neg_converter.py @@ -8,6 +8,7 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Convolution, Neg from executorch.backends.nxp.tests.dataset_creator import ( LinearRampDatasetCreator, @@ -15,7 +16,6 @@ ) from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Convolution, Neg from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py index 266260f9e1f..f5c0114c81e 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py @@ -10,10 +10,10 @@ from executorch.backends.nxp.backend.ir.converter.builder.model_builder import ( ModelBuilder, ) +from executorch.backends.nxp.backend.ops_aliases import Convolution, Pad from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import PadConvModule, PadModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Convolution, Pad @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py index bdfd1e9da25..2317efd45d9 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_permute_copy_converter.py @@ -9,17 +9,17 @@ import pytest import torch from _pytest.mark import ParameterSet - -from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( ExecutorchDelegateCall, GetItem, MaxPool2DWithIndices, PermuteCopy, ) + +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py index 884e95ec20c..358320c30c7 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py @@ -11,6 +11,17 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Convolution, + ExecutorchDelegateCall, + GtScalar, + MulTensor, + PermuteCopy, + Prelu, + ViewCopy, + WhereSelf, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -25,17 +36,6 @@ ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddMM, - Convolution, - ExecutorchDelegateCall, - GtScalar, - MulTensor, - PermuteCopy, - Prelu, - ViewCopy, - WhereSelf, -) from torch.export import ExportedProgram from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program diff --git a/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py index 1f274576767..4bf9263e93f 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py @@ -7,13 +7,7 @@ import pytest import torch from executorch.backends.nxp.backend.edge_program_converter import exir_ops -from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator -from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dModule, LinearModule, ReLUModule -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddMM, Convolution, DequantizePerChannel, @@ -23,6 +17,12 @@ Relu, ViewCopy, ) +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.models import Conv2dModule, LinearModule, ReLUModule +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py index 67101410d9d..916ab5ba518 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_rsqrt_converter.py @@ -8,6 +8,7 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Rsqrt from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -15,7 +16,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Rsqrt from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py index 5c4e4f4f007..c6ea90cb5e2 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sigmoid_converter.py @@ -9,6 +9,7 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Sigmoid from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -16,7 +17,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Sigmoid from torch import nn from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py index 56d0b4bbd64..4a66c000505 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py @@ -8,6 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + SliceCopy, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -21,11 +26,6 @@ SliceTensorModule, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - SliceCopy, -) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py index 2ce0790fc98..8cec44b5274 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py @@ -6,6 +6,12 @@ import numpy as np import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, + Softmax, + ViewCopy, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -17,12 +23,6 @@ ) from executorch.backends.nxp.tests.models import SoftmaxModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - Convolution, - ExecutorchDelegateCall, - Softmax, - ViewCopy, -) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py index 1601c1e19c2..6e88bc270e4 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py @@ -8,6 +8,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + SubTensor, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -21,12 +27,6 @@ ) from executorch.backends.nxp.tests.models import MaxPoolSubTensorModule, SubTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - SubTensor, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py index 8b28142b63a..5726776c78a 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sum_dim_int_list_converter.py @@ -21,6 +21,13 @@ from executorch.backends.nxp.backend.ir.tflite_generator.builtin_options.transpose_options import ( Transpose, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + GetItem, + MaxPool2DWithIndices, + SumDimIntList, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops @@ -29,13 +36,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - GetItem, - MaxPool2DWithIndices, - SumDimIntList, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py index 51b7ee484a7..2795f607494 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py @@ -7,12 +7,12 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import Convolution, Tanh from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import Conv2dWithActivation from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import Convolution, Tanh from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py b/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py index 949f193b267..8513b87a1a3 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_upsample_bilinear2d.py @@ -8,6 +8,11 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + UpsampleBilinear2D, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program @@ -17,11 +22,6 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - UpsampleBilinear2D, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py b/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py index b3e28a7b2f8..dd3d7c9c147 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_upsample_nearest2d.py @@ -8,17 +8,17 @@ # noinspection PyUnusedImports import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + ExecutorchDelegateCall, + UpsampleNearest2D, +) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( - AddTensor, - ExecutorchDelegateCall, - UpsampleNearest2D, -) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py index 2a2d270e30a..200ffb6fc99 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py @@ -8,16 +8,8 @@ import numpy as np import pytest import torch -from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator -from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.model_output_comparator import ( - AllCloseOutputComparator, -) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare -from executorch.backends.nxp.tests.ops_aliases import ( +from executorch.backends.nxp.backend.ops_aliases import ( AddMM, AddTensor, AvgPool2D, @@ -28,6 +20,14 @@ Relu, ViewCopy, ) +from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator +from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program +from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops +from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + AllCloseOutputComparator, +) +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from torch import nn from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py b/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py index f0489b151f7..c1850f625dd 100644 --- a/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py +++ b/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py @@ -6,6 +6,7 @@ import numpy as np import pytest import torch +from executorch.backends.nxp.backend.ops_aliases import AddTensor, ViewCopy from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier from executorch.backends.nxp.tests.models import SqueezeAddModel, UnsqueezeAddModel @@ -13,7 +14,6 @@ AllCloseOutputComparator, lower_run_compare, ) -from executorch.backends.nxp.tests.ops_aliases import AddTensor, ViewCopy @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/nsys_testing.py b/backends/nxp/tests/nsys_testing.py index a8038083c37..6ec0d821bae 100644 --- a/backends/nxp/tests/nsys_testing.py +++ b/backends/nxp/tests/nsys_testing.py @@ -24,6 +24,7 @@ from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( torch_type_to_numpy_type, ) +from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.tests.config_importer import test_config from executorch.backends.nxp.tests.dataset_creator import ( @@ -44,7 +45,6 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.tests.outputs_dir_importer import outputs_dir from executorch.backends.nxp.tests.utils import save_pte_program, store_txt_input_tensor From 33fad232ef27a50746ae21324c7504c9dd1097b6 Mon Sep 17 00:00:00 2001 From: Sangwon Ha <146179778+FabulousSuperDude@users.noreply.github.com> Date: Mon, 7 Sep 2026 22:24:12 +0100 Subject: [PATCH 067/190] Arm backend: Preserve Q/DQ in rescale pass (#22559) ### Cause InsertRescaleInt32Pass runs after Q/DQ folding and assumes that quantization metadata means the operator has been fully converted to an integer path. For partially quantized binary operators, however, FoldAndAnnotateQParamsPass intentionally preserves the Q/DQ boundaries so the operator continues to execute with floating-point inputs. It still records quantization metadata for later lowering. Consequently, the rescale pass misidentifies these operators as integer operations. For add, this can fail because only one input has quantization parameters; for mul, it can insert an invalid INT32 RESCALE on a dequantized floating-point input and alter the preserved graph boundary. ### Fix Skip INT32 rescale insertion when an operator still has a direct DQ input or Q output. These nodes indicate that Q/DQ folding deliberately preserved the quantization boundary, so the operator must not be rewritten as an INT32 operation. Fully folded quantized operators continue through the existing rescale path unchanged. ### Testing Added regression coverage for partially quantized add and mul. The tests verify that the original DQ input and Q output remain connected and that no TOSA RESCALE nodes are inserted. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Sangwon Ha --- backends/arm/_passes/insert_rescales_pass.py | 6 ++ .../passes/test_insert_rescale_i32_pass.py | 56 ++++++++++++++++++- 2 files changed, 60 insertions(+), 2 deletions(-) diff --git a/backends/arm/_passes/insert_rescales_pass.py b/backends/arm/_passes/insert_rescales_pass.py index 2798d5d17f1..76c1c2995f8 100644 --- a/backends/arm/_passes/insert_rescales_pass.py +++ b/backends/arm/_passes/insert_rescales_pass.py @@ -408,6 +408,12 @@ def call(self, graph_module: GraphModule) -> PassResult: if node.op != "call_function" or node.target not in self.included_targets: continue + has_preserved_qdq = any( + input_node.target in DQ_OPS for input_node in node.all_input_nodes + ) or any(user.target in Q_OPS for user in node.users) + if has_preserved_qdq: + continue + if "input_qparams" not in node.meta or len(node.meta["input_qparams"]) == 0: continue input_qparams = node.meta["input_qparams"] diff --git a/backends/arm/test/passes/test_insert_rescale_i32_pass.py b/backends/arm/test/passes/test_insert_rescale_i32_pass.py index e685da11e05..b163db11887 100644 --- a/backends/arm/test/passes/test_insert_rescale_i32_pass.py +++ b/backends/arm/test/passes/test_insert_rescale_i32_pass.py @@ -1,16 +1,19 @@ -# Copyright 2025 Arm Limited and/or its affiliates. +# Copyright 2025-2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from typing import Tuple +from typing import Callable, Tuple +import pytest import torch from executorch.backends.arm._passes import ( FoldAndAnnotateQParamsPass, InsertRescaleInt32Pass, ) +from executorch.backends.arm.common.annotation_meta import ArmAnnotationInfo from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.exir.dialects._ops import ops as exir_ops class MultipleOpsModel(torch.nn.Module): @@ -90,6 +93,55 @@ def test_insert_rescale_int32_tosa_INT_multiple_ops(): _test_model_with_f32_data(MultipleOpsModel()) +@pytest.mark.parametrize( + "binary_target", + (exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor), +) +def test_insert_rescale_int32_preserves_partial_qdq( + binary_target: Callable[..., object], +) -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + y = graph.placeholder("y") + x_q = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + (x, 0.5, 0, -128, 127, torch.int8), + ) + x_dq = graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + (x_q, 0.5, 0, -128, 127, torch.int8), + ) + binary = graph.call_function(binary_target, (x_dq, y)) + binary.meta["custom"] = { + ArmAnnotationInfo.CUSTOM_META_KEY: ArmAnnotationInfo(quantized=True) + } + output_q = graph.call_function( + exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + (binary, 0.5, 0, -128, 127, torch.int8), + ) + output_dq = graph.call_function( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + (output_q, 0.5, 0, -128, 127, torch.int8), + ) + graph.output(output_dq) + graph_module = torch.fx.GraphModule(torch.nn.Module(), graph) + + fold_result = FoldAndAnnotateQParamsPass(preserve_partial_binary_tensor_qdq=True)( + graph_module + ) + assert fold_result is not None + rescale_result = InsertRescaleInt32Pass()(fold_result.graph_module) + assert rescale_result is not None + + assert binary.args == (x_dq, y) + assert output_q in binary.users + assert not rescale_result.graph_module.graph.find_nodes( + op="call_function", + target=exir_ops.backend.tosa.RESCALE.default, + sort=False, + ) + + def test_insert_rescale_int32_tosa_FP_dont_insert_rescales(): module = MultipleOpsModel() input_t = Tuple[torch.Tensor, torch.Tensor] From f3f0c96057ecff2873d89f0fb4349e7bcb83bf8b Mon Sep 17 00:00:00 2001 From: haowhsu-quic <111341466+haowhsu-quic@users.noreply.github.com> Date: Tue, 8 Sep 2026 13:07:28 +0800 Subject: [PATCH 068/190] Qualcomm AI Engine Direct - test framework refactor (#22072) Co-author: @winskuo-quic ### Summary - extend op / feature test for htp arch (v69~v81) - extend op / feature test for lpai / gpu backend ### Test plan ```bash pytest backends/qualcomm/tests/rework/lpai/feature/v6/test.py --device f3c0531 --soc_model SM8850 --build_folder ./build-android/ --backend lpai pytest backends/qualcomm/tests/rework/lpai/op/v6/test.py ``` ```bash pytest backends/qualcomm/tests/rework/gpu/feature/test.py --device f3c0531 --soc_model SM8650 --build_folder ./build-android/ --backend gpu ytest backends/qualcomm/tests/rework/gpu/op/test.py --device f3c0531 --soc_model SM8650 --build_folder ./build-android/ --backend gpu ``` --- .../backends/lpai/qnn_lpai_pass_manager.py | 11 +- .../lpai_partition_fallback_support.py | 9 +- backends/qualcomm/_passes/qnn_pass_manager.py | 4 +- backends/qualcomm/builders/op_batch_norm.py | 1 + backends/qualcomm/qnn_preprocess.py | 1 + .../quantizer/annotators/lpai_rules.py | 106 +- backends/qualcomm/tests/models.py | 10 + backends/qualcomm/tests/rework/conftest.py | 51 +- .../qualcomm/tests/rework/gpu/conftest.py | 37 + .../tests/rework/gpu/feature/conftest.py | 37 + .../qualcomm/tests/rework/gpu/feature/test.py | 82 + backends/qualcomm/tests/rework/gpu/op/test.py | 1279 +++++++++++++ .../tests/rework/htp/feature/v69/test.py | 107 ++ .../tests/rework/htp/feature/v73/test.py | 107 ++ .../tests/rework/htp/feature/v75/test.py | 107 ++ .../tests/rework/htp/feature/v79/test.py | 107 ++ .../tests/rework/htp/feature/v81/test.py | 107 ++ .../qualcomm/tests/rework/htp/op/v68/test.py | 69 +- .../qualcomm/tests/rework/htp/op/v69/test.py | 1408 +++++++++++++++ .../qualcomm/tests/rework/htp/op/v73/test.py | 1384 ++++++++++++++ .../qualcomm/tests/rework/htp/op/v75/test.py | 1384 ++++++++++++++ .../qualcomm/tests/rework/htp/op/v79/test.py | 1384 ++++++++++++++ .../qualcomm/tests/rework/htp/op/v81/test.py | 1384 ++++++++++++++ .../qualcomm/tests/rework/lpai/conftest.py | 97 + .../tests/rework/lpai/feature/conftest.py | 37 + .../tests/rework/lpai/feature/v6/test.py | 82 + .../qualcomm/tests/rework/lpai/op/v6/test.py | 1608 +++++++++++++++++ backends/qualcomm/tests/rework/src/feature.py | 425 +++-- backends/qualcomm/tests/rework/src/op.py | 2 +- backends/qualcomm/tests/test_qnn_delegate.py | 26 + backends/qualcomm/utils/utils.py | 5 + 31 files changed, 11136 insertions(+), 322 deletions(-) create mode 100644 backends/qualcomm/tests/rework/gpu/feature/conftest.py create mode 100644 backends/qualcomm/tests/rework/lpai/feature/conftest.py diff --git a/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py b/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py index a94766ab335..7a42b0e4a63 100644 --- a/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py +++ b/backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py @@ -51,17 +51,18 @@ def get_passes_dependency_for_capture_program(cls): { DecomposeHardsigmoid: [RemoveRedundancy], DecomposeReciprocal: [RemoveRedundancy], - LpaiPartitionFallbackSupport: [TagQuantIO], - ResolveDebugHandle: [LpaiPartitionFallbackSupport], + LpaiPartitionFallbackSupport: [TagQuantIO, ResolveDebugHandle], } ) return deps def _validate_edge_passes(self) -> None: - super()._validate_edge_passes() assert isinstance( - self.passes[-2], LpaiPartitionFallbackSupport - ), "Please ensure LpaiPartitionFallbackSupport is the last edge pass before ResolveDebugHandle." + self.passes[-2], ResolveDebugHandle + ), "Please ensure ResolveDebugHandle is the last edge pass before LpaiPartitionFallbackSupport." + assert isinstance( + self.passes[-1], LpaiPartitionFallbackSupport + ), "Please ensure LpaiPartitionFallbackSupport is the last pass." @classmethod def get_annotation_passes(cls): diff --git a/backends/qualcomm/_passes/lpai_partition_fallback_support.py b/backends/qualcomm/_passes/lpai_partition_fallback_support.py index 5983c145749..ad4a611f24b 100644 --- a/backends/qualcomm/_passes/lpai_partition_fallback_support.py +++ b/backends/qualcomm/_passes/lpai_partition_fallback_support.py @@ -254,7 +254,9 @@ def insert_partition_qdq( output_dq_node.meta[QCOM_BYPASS_NODE] = True graph_module.graph.eliminate_dead_code() - def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule): + def handle_back_to_back_nodes( + self, graph_module: torch.fx.GraphModule, unsupported_nodes: set[torch.fx.Node] + ): """ This function takes care of following cases: 1. When 2 contiguous fall back nodes ``a`` and ``b`` (both @@ -279,6 +281,7 @@ def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule): input_node for input_node in node.all_input_nodes if input_node.op == "call_function" + and input_node not in unsupported_nodes ] assert all( input_node.target in dq_ops for input_node in input_call_func_nodes @@ -327,7 +330,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: unsupported_nodes = self.get_unsupported_nodes(graph_module) for node in unsupported_nodes: self.insert_partition_qdq(graph_module, node) - self.handle_back_to_back_nodes(graph_module) + self.handle_back_to_back_nodes(graph_module, unsupported_nodes) graph_module.graph.eliminate_dead_code() graph_module.recompile() - return PassResult(graph_module, bool(unsupported_nodes)) + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 055567c2be5..2bf5e3f28b3 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -323,9 +323,7 @@ def get_passes_dependency_for_capture_program(cls): RecomposePixelUnshuffle: [RemoveRedundancy], RecomposeRmsNorm: [RemoveRedundancy], TagQuantIO: [LayoutTransform], - ResolveDebugHandle: [ - TagQuantIO - ], # IMPORTANT: Please always ensure ResolveDebugHandle is the last executed pass. + ResolveDebugHandle: [TagQuantIO], } @classmethod diff --git a/backends/qualcomm/builders/op_batch_norm.py b/backends/qualcomm/builders/op_batch_norm.py index e9675bf2397..a2623cb37b2 100644 --- a/backends/qualcomm/builders/op_batch_norm.py +++ b/backends/qualcomm/builders/op_batch_norm.py @@ -29,6 +29,7 @@ class BatchNorm(NodeVisitor): target = [ "aten._native_batch_norm_legit_no_training.default", "aten._native_batch_norm_legit.no_stats", + "aten._native_batch_norm_legit_functional.default", ] def __init__(self, *args) -> None: diff --git a/backends/qualcomm/qnn_preprocess.py b/backends/qualcomm/qnn_preprocess.py index 0d8fcd6fe3b..34d1ce05b94 100644 --- a/backends/qualcomm/qnn_preprocess.py +++ b/backends/qualcomm/qnn_preprocess.py @@ -243,6 +243,7 @@ def preprocess_multimethod( # noqa: C901 (handle_id := node.meta.get(DEBUG_HANDLE_KEY)) and QCOM_TENSOR_NAME in node.meta and len(node.meta[QCOM_TENSOR_NAME]) == 1 + and node.op == "call_function" ): debug_handle_builder.insert_delegate_mapping_entry( handles=handle_id, diff --git a/backends/qualcomm/quantizer/annotators/lpai_rules.py b/backends/qualcomm/quantizer/annotators/lpai_rules.py index 178654f18fc..8b30c9427a9 100644 --- a/backends/qualcomm/quantizer/annotators/lpai_rules.py +++ b/backends/qualcomm/quantizer/annotators/lpai_rules.py @@ -129,7 +129,7 @@ class AvgPool2d(GeneralOpDef): # TODO: Batch_norm op cannot directly map to QNN OpBatchnorm due to the number of input doesn't match. @register_annotator( - [torch.ops.aten.batch_norm.default, torch.ops.aten.instance_norm.default], + [torch.ops.aten.batch_norm.default], qnn_op=None, ) class BatchNorm(GeneralOpDef): @@ -420,7 +420,8 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: torch.ops.aten.topk.default, torch.ops.aten.sort.default, ): - out_act_quantization_spec = SharedQuantizationSpec(node.args[0]) + # assign to None since they are not supported so far + out_act_quantization_spec = None node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( output_qspec=out_act_quantization_spec, _annotated=True, @@ -807,21 +808,6 @@ class ReluMinMax(GeneralOpDef): pass -# TODO: Expand_as op cannot directly map to QNN OpTile due to the number of input doesn't match. -@register_annotator( - [ - torch.ops.aten.expand_as.default, - ], - qnn_op=None, -) -class ExpandAs(GeneralOpDef): - @staticmethod - def annotate(node: Node, quantization_config: QuantizationConfig) -> None: - annotate_in_out_obs_sharing_op(node, quantization_config) - if not _is_annotated([node]): - annotate_single_in_share_out(node, quantization_config) - - @register_annotator( [ torch.ops.aten.flatten.using_ints, @@ -854,7 +840,6 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: return act_node = node.args[0] - weight_node = node.args[2] # TODO current only support 16a16w annotate_input_qspec_map( @@ -863,94 +848,23 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None: quantization_config.input_activation, ) - annotate_input_qspec_map( - node, - weight_node, - quantization_config.input_activation, - ) + if len(node.args) > 2 and node.args[2] is not None: + weight_node = node.args[2] + annotate_input_qspec_map( + node, + weight_node, + quantization_config.input_activation, + ) nodes_to_mark_annotated = [node] annotate_output_qspec(node, quantization_config.output_activation) _mark_nodes_as_annotated(nodes_to_mark_annotated) -# TODO: There is a bug in the BackendOpInfo library, so it is bypassed now. -@register_annotator([torch.ops.aten.rsqrt.default], qnn_op=None) -class Rsqrt(GeneralOpDef): - pass - - @register_annotator([torch.ops.aten.scaled_dot_product_attention.default], qnn_op=None) class ScaledDotProductAttention(GeneralOpDef): pass -@register_annotator( - [ - torch.ops.aten.scatter.src, - torch.ops.aten.scatter.value, - torch.ops.aten.scatter_add.default, - torch.ops.aten.scatter_reduce.two, - ], - qnn_op=None, -) -class ScatterElements(GeneralOpDef): - @staticmethod - def annotate(node: Node, quantization_config: QuantizationConfig) -> None: - if _is_annotated([node]): - return - - input_act = node.args[0] - if not isinstance(input_act, Node) or not _is_float_tensor(input_act): - return - - input_qspec_map = {} - input_qspec_map[input_act] = quantization_config.input_activation - - if ( - len(node.args) > 3 - and isinstance(node.args[3], Node) - and _is_float_tensor(node.args[3]) - ): - input_qspec_map[node.args[3]] = SharedQuantizationSpec((input_act, node)) - - output_act_qspec = ( - SharedQuantizationSpec((input_act, node)) - if _is_float_tensor(node) - else None - ) - - if len(input_qspec_map) > 0 or output_act_qspec is not None: - node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=output_act_qspec, - _annotated=True, - ) - - -@register_annotator([torch.ops.aten.sort.default], QnnConstants.OpTopK.op_name) -class Sort(GeneralOpDef): - @staticmethod - def annotate(node: Node, quantization_config: QuantizationConfig) -> None: - if _is_annotated([node]): - return - - input_qspec_map = {} - input_act_qspec = quantization_config.input_activation - out_act_quantization_spec = None - if input_act_qspec is not None: - if _is_float_tensor(node.args[0]): - input_act = node.args[0] - assert isinstance(input_act, Node) - input_qspec_map[input_act] = input_act_qspec - out_act_quantization_spec = SharedQuantizationSpec((input_act, node)) - - node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation( - input_qspec_map=input_qspec_map, - output_qspec=out_act_quantization_spec, - _annotated=True, - ) - - @register_annotator( [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default], QnnConstants.OpSigmoid.op_name, diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index ad85bafec06..56b53c6f6c9 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -2943,6 +2943,16 @@ def forward(self, x): ) +class ConvRelu(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = torch.nn.ReLU() + + def forward(self, x): + return self.relu(self.conv(x)) + + class TopKandIndex(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/rework/conftest.py b/backends/qualcomm/tests/rework/conftest.py index 257166bd7c2..9730b120c16 100644 --- a/backends/qualcomm/tests/rework/conftest.py +++ b/backends/qualcomm/tests/rework/conftest.py @@ -32,6 +32,7 @@ get_qnn_context_binary_alignment, prepare_pt2e, QnnConfig, + QnnExecuTorchBackendType, QnnQuantizer, setup_common_args_and_variables, SimpleADB, @@ -269,7 +270,7 @@ def qnn_config(global_setup, request): f'invalid configuration detected, fall back to emulator workload:\n"{e}"' ) config = QnnConfig( - soc_model="unknown", build_folder="build-x86", compile_only=True + soc_model="unknown", build_folder="build-x86", enable_x86_64=True ) return config @@ -349,6 +350,7 @@ def invoke_remote( qnn_config: QnnConfig, executorch_prog: ExecutorchProgramManager, callback: callable, + inputs: Tuple[torch.Tensor] = None, ): with tempfile.TemporaryDirectory() as tmp_dir: pte_fname = f"{tmp_dir}/qnn_executorch_test.pte" @@ -363,7 +365,7 @@ def invoke_remote( pte_path=[pte_fname], workspace=f"/data/local/tmp/{device_workspace}", ) - adb.push() + adb.push(inputs=[inputs] if inputs is not None else None) callback(adb) @@ -478,7 +480,23 @@ def export_and_verify( metrics: Metrics, ): with calibrate(module, [inputs], quantizer) as exported_module: - if quantizer is not None: + fake_tensors = ( + [ + node.meta["val"] + for node in exported_module.graph.nodes + if node.op == "call_function" and "val" in node.meta + ] + if quantizer + else [] + ) + dtypes = set() + for tensor in fake_tensors: + if isinstance(tensor, (tuple, list)): + dtypes.update([n.dtype for n in tensor]) + else: + dtypes.add(tensor.dtype) + + if quantizer and {torch.float, torch.float32} & dtypes: nodes = {node.target for node in exported_module.graph.nodes} q_and_dq = { torch.ops.quantized_decomposed.quantize_per_tensor.default, @@ -505,15 +523,34 @@ def export_and_verify( ) ) execution_plan = executorch_prog.executorch_program.execution_plan[0] + + def validate(): + match qnn_config.backend: + case QnnExecuTorchBackendType.kHtpBackend: + return len(execution_plan.operators) == 0 + case QnnExecuTorchBackendType.kGpuBackend: + return len(execution_plan.operators) == 0 + case QnnExecuTorchBackendType.kLpaiBackend: + aten_op_names = { + op.name + for op in execution_plan.operators + if "quantize" not in op.name + } + return len(aten_op_names) == 0 + case _: + return True + assert all( [ - len(execution_plan.delegates) == 1, - execution_plan.delegates[0].id == "QnnBackend", - len(execution_plan.operators) == 0, + ( + len(execution_plan.delegates) == 1 + and execution_plan.delegates[0].id == "QnnBackend" + ), + validate(), ] ), EXPECT_NOT_FULLY_DELEGATED - mode = "emulator" if qnn_config.build_folder == "build-x86" else "remote" + mode = "emulator" if qnn_config.enable_x86_64 else "remote" globals()[f"verify_output_{mode}"]( module=module, inputs=inputs, diff --git a/backends/qualcomm/tests/rework/gpu/conftest.py b/backends/qualcomm/tests/rework/gpu/conftest.py index b5f86874fd4..4aaf73933cb 100644 --- a/backends/qualcomm/tests/rework/gpu/conftest.py +++ b/backends/qualcomm/tests/rework/gpu/conftest.py @@ -3,3 +3,40 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from typing import Any + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_gpu_compiler_spec, + generate_qnn_executorch_compiler_spec, + QcomChipset, +) + + +def with_gpu_context(func): + def wrapper(request, kwargs): + preserved = {k: kwargs.pop(k) for k in ["expected"]} + qnn_config = request.getfixturevalue("qnn_config") + fixtures = { + "quantizer": None, + "compile_spec": generate_qnn_executorch_compiler_spec( + soc_model=getattr(QcomChipset, qnn_config.soc_model), + backend_options=generate_gpu_compiler_spec(), + online_prepare=True, + ), + } + return func(request, fixtures | preserved) + + return wrapper + + +def enumerate_fp_dtype(metric: Any): + def wrapper(test_body): + return pytest.mark.parametrize( + "kwargs", + [pytest.param({"act": None, "expected": metric}, id="fp")], + )(test_body) + + return wrapper diff --git a/backends/qualcomm/tests/rework/gpu/feature/conftest.py b/backends/qualcomm/tests/rework/gpu/feature/conftest.py new file mode 100644 index 00000000000..5c3483a7537 --- /dev/null +++ b/backends/qualcomm/tests/rework/gpu/feature/conftest.py @@ -0,0 +1,37 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import inspect +from functools import lru_cache + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_gpu_compiler_spec, + generate_qnn_executorch_compiler_spec, +) + + +@pytest.fixture(scope="session") +def compile_specs(): + @lru_cache() + def _build(kwargs_config): + kwargs = dict(kwargs_config) + et_compile_spec_sig = set( + inspect.signature(generate_qnn_executorch_compiler_spec).parameters.keys() + ) + et_compile_spec_kwargs = { + k: kwargs[k] for k in kwargs.keys() if k in et_compile_spec_sig + } + for k in et_compile_spec_kwargs.keys(): + kwargs.pop(k) + + return generate_qnn_executorch_compiler_spec( + backend_options=generate_gpu_compiler_spec(**kwargs), + **et_compile_spec_kwargs, + ) + + return lambda kwargs_config: _build(kwargs_config) diff --git a/backends/qualcomm/tests/rework/gpu/feature/test.py b/backends/qualcomm/tests/rework/gpu/feature/test.py index b5f86874fd4..f500a0de71d 100644 --- a/backends/qualcomm/tests/rework/gpu/feature/test.py +++ b/backends/qualcomm/tests/rework/gpu/feature/test.py @@ -3,3 +3,85 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.skip(reason="multiple graphs is not supported with online-prepare") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +# GPU requires online_prepare=True +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +# SpillFill is HTP-specific (uses use_multi_contexts / SRAM spill-fill) +@pytest.mark.skip(reason="SpillFill is HTP-specific; not applicable to GPU backend") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.skip(reason="TBD on native GPU support") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 + + +# MultiGraph weight sharing is not supported on GPU +@pytest.mark.skip( + reason="Weight sharing across multiple graphs is not supported on GPU backend" +) +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/gpu/op/test.py b/backends/qualcomm/tests/rework/gpu/op/test.py index b5f86874fd4..4f64c6f547b 100644 --- a/backends/qualcomm/tests/rework/gpu/op/test.py +++ b/backends/qualcomm/tests/rework/gpu/op/test.py @@ -3,3 +3,1282 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.gpu.conftest import ( + enumerate_fp_dtype, + with_gpu_context, +) + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "expected": pytest.raises( + Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM) + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "expected": pytest.raises( + Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM) + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "expected": pytest.raises( + AssertionError, match=EXPECT_NOT_FULLY_DELEGATED + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +# 3D pooling is not supported on GPU +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +# 3D pooling is not supported on GPU +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +# 3D convolution is not supported on GPU +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": pytest.raises( + AssertionError, match=EXPECT_NOT_FULLY_DELEGATED + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +# 3D convolution is not supported on GPU +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": pytest.raises( + AssertionError, match=EXPECT_NOT_FULLY_DELEGATED + ), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype( + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)) +) +@with_gpu_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance()) +@with_gpu_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +# test_linear_block_quant has no fp variant (lpbq is quantization-specific) +@pytest.mark.skip(reason="LPBQ quantization is not applicable to GPU fp mode") +@pytest.mark.parametrize("kwargs", [pytest.param({}, id="16a4w_lpbq")]) +@with_gpu_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_gpu_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +# 3D pooling is not supported on GPU +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED)) +@with_gpu_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError)) +@with_gpu_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(pytest.raises(AssertionError)) +@with_gpu_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype( + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)) +) +@with_gpu_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_fp_dtype(Tolerance(rtol=1e-1)) +@with_gpu_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v69/test.py b/backends/qualcomm/tests/rework/htp/feature/v69/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v69/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v69/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v73/test.py b/backends/qualcomm/tests/rework/htp/feature/v73/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v73/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v73/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v75/test.py b/backends/qualcomm/tests/rework/htp/feature/v75/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v75/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v75/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v79/test.py b/backends/qualcomm/tests/rework/htp/feature/v79/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v79/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v79/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/feature/v81/test.py b/backends/qualcomm/tests/rework/htp/feature/v81/test.py index b5f86874fd4..85269adaa6a 100644 --- a/backends/qualcomm/tests/rework/htp/feature/v81/test.py +++ b/backends/qualcomm/tests/rework/htp/feature/v81/test.py @@ -3,3 +3,110 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": Tolerance()}, id="e2e"), + ], +) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"expected": nullcontext()}, id="e2e"), + ], +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v68/test.py b/backends/qualcomm/tests/rework/htp/op/v68/test.py index 03a418a41c8..420f58781d4 100644 --- a/backends/qualcomm/tests/rework/htp/op/v68/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v68/test.py @@ -15,7 +15,6 @@ CosineSimilarity, EXCEPTION_EXIR_PROGRAM, EXCEPTION_FROM_PASSES, - EXPECT_NOT_ANNOTATED, EXPECT_NOT_FULLY_DELEGATED, SkipOutputCheck, Tolerance, @@ -151,25 +150,13 @@ def test_amin(request, kwargs): AMin.test(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_any(request, kwargs): Any.test(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_arange_dtype_int(request, kwargs): Arange.test_dtype_int(request, kwargs) # noqa: F405 @@ -255,8 +242,8 @@ def test_batchnorm_2d(request, kwargs): @enumerate_activation_dtype( [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), ] ) @@ -265,13 +252,7 @@ def test_bitwise_and_numeric(request, kwargs): BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_bitwise_and_bool(request, kwargs): BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 @@ -279,8 +260,8 @@ def test_bitwise_and_bool(request, kwargs): @enumerate_activation_dtype( [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), ] ) @@ -289,13 +270,7 @@ def test_bitwise_or_numeric(request, kwargs): BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_bitwise_or_bool(request, kwargs): BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 @@ -303,8 +278,8 @@ def test_bitwise_or_bool(request, kwargs): @enumerate_activation_dtype( [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), ] ) @@ -313,13 +288,7 @@ def test_bitwise_xor_numeric(request, kwargs): BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_bitwise_xor_bool(request, kwargs): BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 @@ -818,25 +787,13 @@ def test_interpolate_nearest(request, kwargs): Interpolate.test_nearest(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(rtol=1e-1), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) @with_htp_context def test_is_inf(request, kwargs): IsInf.test(request, kwargs) # noqa: F405 -@enumerate_activation_dtype( - [ - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), - Tolerance(), - ] -) +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) @with_htp_context def test_is_nan(request, kwargs): IsNan.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v69/test.py b/backends/qualcomm/tests/rework/htp/op/v69/test.py index b5f86874fd4..3babe431a5c 100644 --- a/backends/qualcomm/tests/rework/htp/op/v69/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v69/test.py @@ -3,3 +3,1411 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 69 from ".../rework/htp/unit_test/op/v69/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; still fails on V69 +@enumerate_activation_dtype( + [ + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + Tolerance(rtol=1e-1), + ] +) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v73/test.py b/backends/qualcomm/tests/rework/htp/op/v73/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v73/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v73/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v75/test.py b/backends/qualcomm/tests/rework/htp/op/v75/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v75/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v75/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v79/test.py b/backends/qualcomm/tests/rework/htp/op/v79/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v79/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v79/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/htp/op/v81/test.py b/backends/qualcomm/tests/rework/htp/op/v81/test.py index b5f86874fd4..6362690f5fd 100644 --- a/backends/qualcomm/tests/rework/htp/op/v81/test.py +++ b/backends/qualcomm/tests/rework/htp/op/v81/test.py @@ -3,3 +3,1387 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + CosineSimilarity, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_FULLY_DELEGATED, + SkipOutputCheck, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.htp.conftest import ( + enumerate_activation_dtype, + with_htp_context, +) + +# e.g. get 73 from ".../rework/htp/unit_test/op/v73/test.py" +HTP_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_htp_context = partial(with_htp_context, hw_arch=HTP_ARCH) + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_htp_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +# bmm 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, True, Tolerance(), "8a8w_pcq"), + (16, 4, True, CosineSimilarity(0.95), "16a4w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +# LPBQ (QNN_QUANTIZATION_ENCODING_BLOCKWISE_EXPANSION) requires V69+; enabled here +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "lpbq": True, + "block_sz_map": {"conv2d": (1, 32, 1, 1)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + ], +) +@with_htp_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + # no bitwidth support for conv3d + (8, 8, True, Tolerance(), "8a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": act, "param": param, "pcq": pcq, "expected": expected}, + id=id, + ) + for act, param, pcq, expected, id in [ + (8, 8, False, Tolerance(), "8a8w_ptq"), + (16, 16, False, Tolerance(), "16a16w_ptq"), + (16, 8, True, Tolerance(), "16a8w_pcq"), + (None, None, False, Tolerance(rtol=1e-1), "fp"), + ] + ], +) +@with_htp_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + CosineSimilarity(0.95), + Tolerance(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance()]) +@with_htp_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 4, + "pcq": False, + "lpbq": True, + "block_sz_map": {"linear": (1, 32)}, + "expected": Tolerance(), + }, + id="16a4w_lpbq", + ), + ], +) +@with_htp_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 16, "param": 16, "pcq": False, "expected": Tolerance()}, + id="16a16w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 4, "pcq": True, "expected": CosineSimilarity(0.95)}, + id="16a4w_pcq", + ), + pytest.param( + {"act": 16, "param": 8, "pcq": True, "expected": Tolerance()}, + id="16a8w_pcq", + ), + pytest.param( + {"act": "fp16", "param": 8, "pcq": True, "expected": Tolerance()}, + id="fp16a8w_pcq", + ), + pytest.param( + {"act": 16, "param": 2, "pcq": True, "expected": CosineSimilarity(0.9)}, + id="16a2w_pcq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + { + "act": 16, + "param": 16, + "pcq": False, + "expected": pytest.raises(AssertionError, match=Tolerance()), + }, + id="16a16w_ptq", + ), + pytest.param( + { + "act": None, + "param": None, + "pcq": False, + "expected": Tolerance(rtol=1e-1), + }, + id="fp", + ), + ], +) +@with_htp_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + SkipOutputCheck(), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_htp_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +# sdpa 16a requires V73+; enabled here +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + Tolerance(), + Tolerance(), + pytest.raises(AssertionError), + ] +) +@with_htp_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_htp_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance(), Tolerance(), Tolerance(rtol=1e-1)]) +@with_htp_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/lpai/conftest.py b/backends/qualcomm/tests/rework/lpai/conftest.py index b5f86874fd4..9553051b08f 100644 --- a/backends/qualcomm/tests/rework/lpai/conftest.py +++ b/backends/qualcomm/tests/rework/lpai/conftest.py @@ -3,3 +3,100 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from functools import lru_cache +from typing import Any, List + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_lpai_compiler_spec, + generate_qnn_executorch_compiler_spec, + make_quantizer, + QcomChipset, + QnnExecuTorchBackendType, + QuantDtype, +) +from executorch.backends.qualcomm.serialization.qc_schema import ( + LpaiHardwareVersion, + QnnExecuTorchLpaiTargetEnv, +) + + +def with_lpai_context(func, hw_arch): + def wrapper(request, kwargs): + # extend this if necessary + preserved = {k: kwargs.pop(k) for k in ["expected"]} + callbacks_and_args = { + # extract objects from callback + "quantizers": {"arch": hw_arch} | kwargs, + "compile_specs": {"arch": hw_arch}, + } + fixtures = { + k[:-1]: request.getfixturevalue(k)(**v) + for k, v in callbacks_and_args.items() + } + return func(request, fixtures | preserved) + + return wrapper + + +def enumerate_activation_dtype(metrics: List[Any]): + def wrapper(test_body): + return pytest.mark.parametrize( + "kwargs", + [ + pytest.param({"act": act, "expected": metrics[i]}, id=id) + for i, (act, id) in enumerate( + [ + (8, "8a"), + ] + ) + ], + )(test_body) + + return wrapper + + +def _get_lpai_arch(): + # hardcoded lpai architecture with corresponding premium soc + return [ + (LpaiHardwareVersion.V6, "SM8850"), + ] + + +@pytest.fixture(scope="session") +def quantizers(): + arch_to_soc = dict(_get_lpai_arch()) + + @lru_cache() + def _build(arch, act, param, per_ch): + attr = f"use_{act}a{param}w" + if quant_dtype := getattr(QuantDtype, attr, None): + return make_quantizer( + quant_dtype=quant_dtype, + per_channel_conv=per_ch, + per_channel_linear=per_ch, + backend=QnnExecuTorchBackendType.kLpaiBackend, + soc_model=arch_to_soc[arch], + ) + + def get_quantizer(arch, act, param=None, pcq=False, **_): + param = 8 if (param is None and act is not None) else param + return _build(arch, act, param, pcq) + + return get_quantizer + + +@pytest.fixture(scope="session") +def compile_specs(): + compile_spec = { + arch: generate_qnn_executorch_compiler_spec( + soc_model=getattr(QcomChipset, soc_model), + backend_options=generate_lpai_compiler_spec( + target_env=QnnExecuTorchLpaiTargetEnv.kX86, + ), + ) + for (arch, soc_model) in _get_lpai_arch() + } + return lambda arch: compile_spec[arch] diff --git a/backends/qualcomm/tests/rework/lpai/feature/conftest.py b/backends/qualcomm/tests/rework/lpai/feature/conftest.py new file mode 100644 index 00000000000..b168d759654 --- /dev/null +++ b/backends/qualcomm/tests/rework/lpai/feature/conftest.py @@ -0,0 +1,37 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import inspect +from functools import lru_cache + +import pytest + +from executorch.backends.qualcomm.export_utils import ( + generate_lpai_compiler_spec, + generate_qnn_executorch_compiler_spec, +) + + +@pytest.fixture(scope="session") +def compile_specs(): + @lru_cache() + def _build(kwargs_config): + kwargs = dict(kwargs_config) + et_compile_spec_sig = set( + inspect.signature(generate_qnn_executorch_compiler_spec).parameters.keys() + ) + et_compile_spec_kwargs = { + k: kwargs[k] for k in kwargs.keys() if k in et_compile_spec_sig + } + for k in et_compile_spec_kwargs.keys(): + kwargs.pop(k) + + return generate_qnn_executorch_compiler_spec( + backend_options=generate_lpai_compiler_spec(**kwargs), + **et_compile_spec_kwargs, + ) + + return lambda kwargs_config: _build(kwargs_config) diff --git a/backends/qualcomm/tests/rework/lpai/feature/v6/test.py b/backends/qualcomm/tests/rework/lpai/feature/v6/test.py index b5f86874fd4..9e49d0b9ba3 100644 --- a/backends/qualcomm/tests/rework/lpai/feature/v6/test.py +++ b/backends/qualcomm/tests/rework/lpai/feature/v6/test.py @@ -3,3 +3,85 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +from contextlib import nullcontext + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import Tolerance +from executorch.backends.qualcomm.tests.rework.src.feature import * # noqa: F403 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_logging(request, kwargs): + Logging.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_multi_graph_inference(request, kwargs): + MultiGraph.test_inference(request, kwargs) # noqa: F405 + + +# LPAI forbids online_prepare=True at runtime +@pytest.mark.skip( + reason="LPAI backend only supports offline_prepare; online_prepare is forbidden" +) +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_online_prepare(request, kwargs): + OnlinePrepare.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_performance(request, kwargs): + Performance.test(request, kwargs) # noqa: F405 + + +@pytest.mark.skip(reason="TBD on native LPAI support") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_profile(request, kwargs): + Profile.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_saver(request, kwargs): + Saver.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize("kwargs", [pytest.param({"expected": Tolerance()}, id="e2e")]) +def test_shared_buffer(request, kwargs): + SharedBuffer.test(request, kwargs) # noqa: F405 + + +# SpillFill is HTP-specific (uses use_multi_contexts / SRAM spill-fill) +@pytest.mark.skip(reason="SpillFill is HTP-specific; not applicable to LPAI backend") +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_spill_fill(request, kwargs): + SpillFill.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_tensor_dump(request, kwargs): + TensorDump.test(request, kwargs) # noqa: F405 + + +# MultiGraph weight sharing requires use_weight_sharing in generate_lpai_compiler_spec (not supported) +@pytest.mark.skip( + reason="Weight sharing across multiple graphs is not supported on LPAI backend" +) +@pytest.mark.parametrize( + "kwargs", [pytest.param({"expected": nullcontext()}, id="e2e")] +) +def test_multi_graph_weight_sharing(request, kwargs): + MultiGraph.test_weight_sharing(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/lpai/op/v6/test.py b/backends/qualcomm/tests/rework/lpai/op/v6/test.py index b5f86874fd4..81eb670eebe 100644 --- a/backends/qualcomm/tests/rework/lpai/op/v6/test.py +++ b/backends/qualcomm/tests/rework/lpai/op/v6/test.py @@ -3,3 +3,1611 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + +import re +from functools import partial +from pathlib import Path + +import pytest + +from executorch.backends.qualcomm.tests.rework.conftest import ( + check_exception, + EXCEPTION_EXIR_PROGRAM, + EXCEPTION_FROM_PASSES, + EXPECT_NOT_ANNOTATED, + EXPECT_NOT_FULLY_DELEGATED, + Tolerance, +) +from executorch.backends.qualcomm.tests.rework.src.op import * # noqa: F403 +from executorch.backends.qualcomm.tests.rework.lpai.conftest import ( + enumerate_activation_dtype, + with_lpai_context, +) + + +# e.g. get 68 from ".../rework/htp/unit_test/op/v68/test.py" +LPAI_ARCH = int(re.search(r".*v([0-9]+)$", Path(__file__).parent.name).group(1)) +with_lpai_context = partial(with_lpai_context, hw_arch=LPAI_ARCH) + + +# abs not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_abs(request, kwargs): + Abs.test(request, kwargs) # noqa: F405 + + +# acos not in lpai_rules but will be decomposed into equivalent ops +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_acos(request, kwargs): + ACos.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_1d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_1d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_adaptive_avg_pool_1d(request, kwargs): + AdaptiveAvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_2d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_2d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_adaptive_avg_pool_2d(request, kwargs): + AdaptiveAvgPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_3d_unsupported_io_shape(request, kwargs): + AdaptiveAvgPool.test_3d_unsupported_io_shape(request, kwargs) # noqa: F405 + + +# adaptive_avg_pool3d not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_adaptive_avg_pool_3d(request, kwargs): + AdaptiveAvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_adaptive_max_pool_2d(request, kwargs): + AdaptiveMaxPool.test_2d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_adaptive_max_pool_2d_with_indices(request, kwargs): + AdaptiveMaxPool.test_2d_with_indices(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_add(request, kwargs): + Add.test(request, kwargs) # noqa: F405 + + +# addmm decomposes to mm+add before annotation; both in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_addmm(request, kwargs): + AddMM.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_alias(request, kwargs): + Alias.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_amax(request, kwargs): + AMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_amin(request, kwargs): + AMin.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_any(request, kwargs): + Any.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_arange_dtype_int(request, kwargs): + Arange.test_dtype_int(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_arange_dtype_float(request, kwargs): + Arange.test_dtype_float(request, kwargs) # noqa: F405 + + +# int64 cast for indices is not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_argmax(request, kwargs): + ArgMax.test(request, kwargs) # noqa: F405 + + +# int64 cast for indices is not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_argmin(request, kwargs): + ArgMin.test(request, kwargs) # noqa: F405 + + +# asin not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_asin(request, kwargs): + ASin.test(request, kwargs) # noqa: F405 + + +# some decomposed ops are not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_atan(request, kwargs): + ATan.test(request, kwargs) # noqa: F405 + + +# some decomposed ops are not supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_atan2(request, kwargs): + ATan2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_avgpool_1d(request, kwargs): + AvgPool.test_1d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_avgpool_2d(request, kwargs): + AvgPool.test_2d(request, kwargs) # noqa: F405 + + +# avg_pool3d not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_avgpool_3d(request, kwargs): + AvgPool.test_3d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_batchnorm_2d(request, kwargs): + BatchNorm2d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_and_numeric(request, kwargs): + BitwiseOp.test_and_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_and_bool(request, kwargs): + BitwiseOp.test_and_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_or_numeric(request, kwargs): + BitwiseOp.test_or_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_or_bool(request, kwargs): + BitwiseOp.test_or_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_xor_numeric(request, kwargs): + BitwiseOp.test_xor_numeric(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_bitwise_xor_bool(request, kwargs): + BitwiseOp.test_xor_bool(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_bmm(request, kwargs): + Bmm.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_cast(request, kwargs): + Cast.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_cat(request, kwargs): + Cat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_cdist(request, kwargs): + CDist.test(request, kwargs) # noqa: F405 + + +# ceil not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_ceil(request, kwargs): + Ceil.test(request, kwargs) # noqa: F405 + + +# channel_shuffle not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_channel_shuffle(request, kwargs): + ChannelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_chunk(request, kwargs): + Chunk.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clamp(request, kwargs): + Clamp.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clamp_max(request, kwargs): + ClampMax.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clamp_min(request, kwargs): + ClampMin.test(request, kwargs) # noqa: F405 + + +# clone not in lpai_rules but will be omitted +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_clone(request, kwargs): + Clone.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv1d(request, kwargs): + Conv.test_1d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv1d_transpose(request, kwargs): + Conv.test_1d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv2d(request, kwargs): + Conv.test_2d(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv2d_transpose(request, kwargs): + Conv.test_2d_transpose(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_conv2d_linear_like(request, kwargs): + Conv.test_2d_linear_like(request, kwargs) # noqa: F405 + + +# conv3d not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_conv3d(request, kwargs): + Conv.test_3d(request, kwargs) # noqa: F405 + + +# conv3d_transpose not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_conv3d_transpose(request, kwargs): + Conv.test_3d_transpose(request, kwargs) # noqa: F405 + + +# cos not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_cos(request, kwargs): + Cos.test(request, kwargs) # noqa: F405 + + +# cumsum not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_cumsum(request, kwargs): + CumSum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_div(request, kwargs): + Div.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_div_with_rounding_mode(request, kwargs): + DivWithRoundingMode.test(request, kwargs) # noqa: F405 + + +# einsum decomposes to bmm/matmul before annotation; both in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_einsum(request, kwargs): + Einsum.test(request, kwargs) # noqa: F405 + + +# elu not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_elu(request, kwargs): + Elu.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + ], +) +@with_lpai_context +def test_embedding(request, kwargs): + Embedding.test(request, kwargs) # noqa: F405 + + +# equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_equal(request, kwargs): + Equal.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_exp(request, kwargs): + Exp.test(request, kwargs) # noqa: F405 + + +# expand not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_expand(request, kwargs): + Expand.test(request, kwargs) # noqa: F405 + + +# expand_as not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_expand_as(request, kwargs): + ExpandAs.test(request, kwargs) # noqa: F405 + + +# expm1 not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_expm1(request, kwargs): + ExpM1.test(request, kwargs) # noqa: F405 + + +# fill translates to static tensor in QNN; backend-agnostic +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_fill(request, kwargs): + Fill.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_flip(request, kwargs): + Flip.test(request, kwargs) # noqa: F405 + + +# floor not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_floor(request, kwargs): + Floor.test(request, kwargs) # noqa: F405 + + +# floor_divide not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_floor_divide(request, kwargs): + FloorDivide.test(request, kwargs) # noqa: F405 + + +# fold uses col2im which is in lpai_rules (ColIm, qnn_op=None) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_fold(request, kwargs): + Fold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_lpai_context +def test_fold_unsupported_parameters(request, kwargs): + Fold.test_unsupported_parameters(request, kwargs) # noqa: F405 + + +# full/full_like translate to static tensors in QNN; backend-agnostic +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_full(request, kwargs): + Full.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_full_like(request, kwargs): + FullLike.test(request, kwargs) # noqa: F405 + + +# gather is in lpai_rules (Embedding class handles index/gather/index_select) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_gather(request, kwargs): + Gather.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_gelu(request, kwargs): + Gelu.test(request, kwargs) # noqa: F405 + + +# glu decomposes to chunk+sigmoid+mul before annotation; all in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_glu(request, kwargs): + Glu.test(request, kwargs) # noqa: F405 + + +# greater not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_greater(request, kwargs): + Greater.test_gt(request, kwargs) # noqa: F405 + + +# greater_equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_greater_equal(request, kwargs): + Greater.test_ge(request, kwargs) # noqa: F405 + + +# grid_sample not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_grid_sample_4d(request, kwargs): + GridSample.test_4d(request, kwargs) # noqa: F405 + + +# grid_sample not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_grid_sample_5d(request, kwargs): + GridSample.test_5d(request, kwargs) # noqa: F405 + + +# group_norm not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_group_norm(request, kwargs): + GroupNorm.test(request, kwargs) # noqa: F405 + + +# hardsigmoid: DecomposeHardsigmoid runs before annotation; decomposed ops in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_hardsigmoid(request, kwargs): + HardSigmoid.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_hardswish(request, kwargs): + HardSwish.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_hardtanh(request, kwargs): + HardTanh.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_index(request, kwargs): + Index.test(request, kwargs) # noqa: F405 + + +# decomposed ops might not be supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_index_copy(request, kwargs): + IndexCopy.test(request, kwargs) # noqa: F405 + + +# decomposed ops might not be supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_index_put(request, kwargs): + IndexPut.test(request, kwargs) # noqa: F405 + + +# decomposed ops might not be supported by lpai +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_index_select(request, kwargs): + IndexSelect.test(request, kwargs) # noqa: F405 + + +# instance_norm not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_instance_norm_2d(request, kwargs): + InstanceNorm2d.test(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_interpolate_bicubic(request, kwargs): + Interpolate.test_bicubic(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_interpolate_bilinear(request, kwargs): + Interpolate.test_bilinear(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_interpolate_nearest(request, kwargs): + Interpolate.test_nearest(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_is_inf(request, kwargs): + IsInf.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_is_nan(request, kwargs): + IsNan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_layer_norm(request, kwargs): + LayerNorm.test(request, kwargs) # noqa: F405 + + +# maps to prelu +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_leaky_relu(request, kwargs): + LeakyReLU.test(request, kwargs) # noqa: F405 + + +# less_equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_less_equal(request, kwargs): + LessEqual.test(request, kwargs) # noqa: F405 + + +# less_than not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_less_than(request, kwargs): + LessThan.test(request, kwargs) # noqa: F405 + + +# decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_linalg_vector_norm(request, kwargs): + LinalgVectorNorm.test(request, kwargs) # noqa: F405 + + +# LPBQ not applicable to LPAI v6 (requires HTP V69+ feature) +@pytest.mark.skip(reason="LPBQ quantization is not supported on LPAI v6") +@with_lpai_context +def test_linear_block_quant(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + pytest.param( + {"act": 8, "param": 8, "pcq": True, "expected": Tolerance()}, + id="8a8w_pcq", + ), + ], +) +@with_lpai_context +def test_linear_general(request, kwargs): + Linear.test(request, kwargs) # noqa: F405 + + +@pytest.mark.parametrize( + "kwargs", + [ + pytest.param( + {"act": 8, "param": 8, "pcq": False, "expected": Tolerance()}, + id="8a8w_ptq", + ), + ], +) +@with_lpai_context +def test_linear_non_constant_weight(request, kwargs): + LinearNonConstantWeight.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log(request, kwargs): + Log.test(request, kwargs) # noqa: F405 + + +# log10 not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log10(request, kwargs): + Log10.test(request, kwargs) # noqa: F405 + + +# log1p not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log1p(request, kwargs): + Log1p.test(request, kwargs) # noqa: F405 + + +# log2 not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log2(request, kwargs): + Log2.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_log_softmax(request, kwargs): + LogSoftmax.test(request, kwargs) # noqa: F405 + + +# logical_and not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_logical_and(request, kwargs): + LogicalAnd.test(request, kwargs) # noqa: F405 + + +# logical_not not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_logical_not(request, kwargs): + LogicalNot.test(request, kwargs) # noqa: F405 + + +# decomposed ops are not supported with invalid weight fallback triggered +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_lpai_context +def test_masked_fill(request, kwargs): + MaskedFill.test(request, kwargs) # noqa: F405 + + +# cast op for indices is not supported +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_max_dim(request, kwargs): + MaxDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_maximum(request, kwargs): + Maximum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_maxpool_2d(request, kwargs): + MaxPool2d.test(request, kwargs) # noqa: F405 + + +# max_pool3d not in lpai_rules and the decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_maxpool_3d(request, kwargs): + MaxPool3d.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_mean(request, kwargs): + Mean.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_mha(request, kwargs): + MultiheadAttention.test(request, kwargs) # noqa: F405 + + +# cast op for indices is not supported +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_min_dim(request, kwargs): + MinDim.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_minimum(request, kwargs): + Minimum.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_narrow(request, kwargs): + Narrow.test(request, kwargs) # noqa: F405 + + +# neg not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_neg(request, kwargs): + Neg.test(request, kwargs) # noqa: F405 + + +# not_equal not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_not_equal(request, kwargs): + NotEqual.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pad_constant(request, kwargs): + Pad.test_constant(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_pad_reflect(request, kwargs): + Pad.test_reflect(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_permute(request, kwargs): + Permute.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pixel_shuffle(request, kwargs): + PixelShuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pixel_unshuffle(request, kwargs): + PixelUnshuffle.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pow_scalar(request, kwargs): + PowScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_pow_tensor_scalar(request, kwargs): + PowTensorScalar.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_prelu(request, kwargs): + PReLU.test(request, kwargs) # noqa: F405 + + +# rand not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_rand(request, kwargs): + Rand.test(request, kwargs) # noqa: F405 + + +# reciprocal: DecomposeReciprocal decomposes to div(1, x); div is in lpai_rules +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reciprocal(request, kwargs): + Reciprocal.test(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_reflection_pad_1d(request, kwargs): + ReflectionPad.test_3d(request, kwargs) # noqa: F405 + + +# registered by the partitioner to not be decomposed but failed to be delegated +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_EXIR_PROGRAM)), + ] +) +@with_lpai_context +def test_reflection_pad_2d(request, kwargs): + ReflectionPad.test_4d(request, kwargs) # noqa: F405 + + +# reflection_pad3d not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reflection_pad_3d(request, kwargs): + ReflectionPad.test_5d(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_relu(request, kwargs): + Relu.test(request, kwargs) # noqa: F405 + + +# relu6 decomposes to hardtanh(0, 6); hardtanh is in lpai_rules (ReluMinMax) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_relu6(request, kwargs): + Relu6.test(request, kwargs) # noqa: F405 + + +# decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_remainder(request, kwargs): + Remainder.test(request, kwargs) # noqa: F405 + + +# repeat not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_repeat(request, kwargs): + Repeat.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_2d_to_4d_random_reshape(request, kwargs): + Reshape.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_2d_to_4d_flatten_last_two_dims(request, kwargs): + Reshape.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_5d_random_reshape(request, kwargs): + Reshape.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_reshape_5d_flatten_last_two_dims(request, kwargs): + Reshape.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_rms_norm(request, kwargs): + RmsNorm.test(request, kwargs) # noqa: F405 + + +# roll not in lpai_rules but decomposed ops are fully delegated +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_roll(request, kwargs): + Roll.test(request, kwargs) # noqa: F405 + + +# round not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_round(request, kwargs): + Round.test(request, kwargs) # noqa: F405 + + +# rsqrt not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_rsqrt(request, kwargs): + Rsqrt.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_sdpa(request, kwargs): + ScaledDotProductAttention.test(request, kwargs) # noqa: F405 + + +# scatter.src not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_scatter_src(request, kwargs): + ScatterSrc.test(request, kwargs) # noqa: F405 + + +# select_copy maps to aten.select.int which is in lpai_rules (StrideSlice) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_select_copy(request, kwargs): + SelectCopy.test(request, kwargs) # noqa: F405 + + +# select_scatter not in lpai_rules and decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_select_scatter(request, kwargs): + SelectScatter.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_sigmoid(request, kwargs): + Sigmoid.test(request, kwargs) # noqa: F405 + + +# sign not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_sign(request, kwargs): + Sign.test(request, kwargs) # noqa: F405 + + +# sin not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_sin(request, kwargs): + Sin.test(request, kwargs) # noqa: F405 + + +# slice_copy maps to aten.slice.Tensor which is in lpai_rules (StrideSlice) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_slice_copy(request, kwargs): + SliceCopy.test(request, kwargs) # noqa: F405 + + +# slice_scatter not in lpai_rules and decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_slice_scatter(request, kwargs): + SliceScatter.test(request, kwargs) # noqa: F405 + + +# scatter not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_scatter_value(request, kwargs): + ScatterValue.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_softmax(request, kwargs): + Softmax.test(request, kwargs) # noqa: F405 + + +# sort not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_sort(request, kwargs): + Sort.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_split(request, kwargs): + Split.test(request, kwargs) # noqa: F405 + + +# square is in lpai_rules (Pow class handles square.default) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_square(request, kwargs): + Square.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_squeeze(request, kwargs): + Squeeze.test(request, kwargs) # noqa: F405 + + +# stack maps to OpPack which is HTP-specific +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_stack(request, kwargs): + Stack.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_sum_int_list(request, kwargs): + SumIntList.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_swapaxes(request, kwargs): + SwapAxes.test(request, kwargs) # noqa: F405 + + +# tan not in lpai_rules and the decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_tan(request, kwargs): + Tan.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_tanh(request, kwargs): + Tanh.test(request, kwargs) # noqa: F405 + + +# threshold not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_ANNOTATED), + ] +) +@with_lpai_context +def test_threshold(request, kwargs): + Threshold.test(request, kwargs) # noqa: F405 + + +# triu not in lpai_rulesdecomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_triu(request, kwargs): + Triu.test(request, kwargs) # noqa: F405 + + +# triu not in lpai_rulesdecomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_triu_constant(request, kwargs): + Triu.test_constant(request, kwargs) # noqa: F405 + + +# trunc not in lpai_rules and decomposed ops are not fully delegated +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_trunc(request, kwargs): + Trunc.test(request, kwargs) # noqa: F405 + + +# topk not in lpai_rules, use EXPECT_NOT_FULLY_DELEGATED for there are +# other ops in the test body +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_topk(request, kwargs): + TopK.test(request, kwargs) # noqa: F405 + + +# unbind maps to OpUnpack which is HTP-specific +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_unbind(request, kwargs): + Unbind.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_unflatten(request, kwargs): + Unflatten.test(request, kwargs) # noqa: F405 + + +# unfold uses im2col which is in lpai_rules (ColIm, qnn_op=None) +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_unfold(request, kwargs): + Unfold.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype( + [ + pytest.raises(Exception, check=check_exception(EXCEPTION_FROM_PASSES)), + ] +) +@with_lpai_context +def test_unfold_unsupported(request, kwargs): + Unfold.test_unsupported(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_unsqueeze(request, kwargs): + Unsqueeze.test(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_2d_to_4d_random_reshape(request, kwargs): + View.test_2d_to_4d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_2d_to_4d_flatten_last_two_dims(request, kwargs): + View.test_2d_to_4d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_5d_random_reshape(request, kwargs): + View.test_5d_random_reshape(request, kwargs) # noqa: F405 + + +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_view_5d_flatten_last_two_dims(request, kwargs): + View.test_5d_flatten_last_two_dims(request, kwargs) # noqa: F405 + + +# where not in lpai_rules +@enumerate_activation_dtype( + [ + pytest.raises(AssertionError, match=EXPECT_NOT_FULLY_DELEGATED), + ] +) +@with_lpai_context +def test_where(request, kwargs): + Where.test(request, kwargs) # noqa: F405 + + +# var not in lpai_rules but will be decomposed +@enumerate_activation_dtype([Tolerance()]) +@with_lpai_context +def test_var(request, kwargs): + Var.test(request, kwargs) # noqa: F405 diff --git a/backends/qualcomm/tests/rework/src/feature.py b/backends/qualcomm/tests/rework/src/feature.py index 8ef6d1abec4..49499fc04e6 100644 --- a/backends/qualcomm/tests/rework/src/feature.py +++ b/backends/qualcomm/tests/rework/src/feature.py @@ -14,15 +14,24 @@ import torch +from executorch.backends.qualcomm.debugger.qcom_numerical_comparator_sample import ( + QcomCosineSimilarityComparator, +) +from executorch.backends.qualcomm.debugger.qnn_intermediate_debugger import ( + QNNIntermediateDebugger, +) from executorch.backends.qualcomm.export_utils import ( make_quantizer, QcomChipset, + QnnConfig, QnnExecuTorchBackendType, QnnExecuTorchHtpPerformanceMode, SimpleADB, to_edge_transform_and_lower_to_qnn, ) from executorch.backends.qualcomm.serialization.qc_schema import ( + QnnExecuTorchGpuPerformanceMode, + QnnExecuTorchLpaiClientPerf, QnnExecuTorchProfileLevel, ) from executorch.backends.qualcomm.tests.rework.conftest import ( @@ -51,6 +60,17 @@ def wrapper(request, kwargs): return wrapper +def get_quantizer(qnn_config: QnnConfig): + return ( + make_quantizer( + backend=qnn_config.backend, + soc_model=qnn_config.soc_model, + ) + if qnn_config.backend != QnnExecuTorchBackendType.kGpuBackend + else None + ) + + class Logging: class Model(torch.nn.Module): def __init__(self): @@ -62,6 +82,15 @@ def example_inputs(self): def forward(self, x): return torch.nn.ReLU()(x) + @staticmethod + def _get_log_pattern(backend): + return { + QnnExecuTorchBackendType.kHtpBackend: "QnnDsp ", + QnnExecuTorchBackendType.kGpuBackend: "OpenCL", + # looks like no special keyword appears + QnnExecuTorchBackendType.kLpaiBackend: "", + }[backend] + @staticmethod def _test(qnn_config, compile_specs, expected, aot): def callback(adb: SimpleADB, pattern): @@ -78,9 +107,7 @@ def verify(log): model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, @@ -91,28 +118,28 @@ def verify(log): invoke_remote( qnn_config=qnn_config, executorch_prog=executorch_prog_mgr, - callback=partial(callback, pattern="QnnDsp "), + callback=partial( + callback, + pattern=Logging._get_log_pattern(qnn_config.backend), + ), ) @staticmethod @unpack_fixtures def test(subtests, qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { QnnExecuTorchBackendType.kHtpBackend: [ - compile_specs(tuple(d.items())) - for d in [ - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": True, - "use_fp16": False, - }, - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": False, - "use_fp16": False, - }, - ] + {"soc_model": soc_model, "debug": True, "use_fp16": False}, + {"soc_model": soc_model, "debug": False, "use_fp16": False}, + ], + QnnExecuTorchBackendType.kGpuBackend: [ + {"soc_model": soc_model, "debug": True, "online_prepare": True}, + {"soc_model": soc_model, "debug": False, "online_prepare": True}, + ], + QnnExecuTorchBackendType.kLpaiBackend: [ + {"soc_model": soc_model, "debug": True}, + {"soc_model": soc_model, "debug": False}, ], } @@ -120,7 +147,9 @@ def test(subtests, qnn_config, compile_specs, expected): with subtests.test(msg=config): __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend][i], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend][i].items()) + ), expected=expected, aot=config == "compile_time_option", ) @@ -152,7 +181,7 @@ def compile(models, compile_specs): with calibrate( models[i], [inputs], - make_quantizer(soc_model=qnn_config.soc_model), + get_quantizer(qnn_config), ) as model: modules_dict[graph_name] = model sample_inputs_dict[graph_name] = inputs @@ -221,21 +250,24 @@ def test_weight_sharing(qnn_config, compile_specs, expected): @staticmethod @unpack_fixtures def test_inference(qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: {"soc_model": soc_model}, } __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), expected=expected, ) @@ -254,27 +286,28 @@ def forward(self, x): @staticmethod @unpack_fixtures def test(qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "online_prepare": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "online_prepare": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "online_prepare": True, + }, } module = __class__.Model() - qnn_config.online_prepare = True export_and_verify( module=module, inputs=module.example_inputs(), qnn_config=qnn_config, - quantizer=make_quantizer(soc_model=qnn_config.soc_model), - compile_specs=backend_compile_specs[qnn_config.backend], + quantizer=get_quantizer(qnn_config), + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), metrics=expected, ) @@ -304,53 +337,82 @@ def verify(log): adb.extra_cmds += "" if aot else " --htp_performance_mode 6" adb.execute(output_callback=verify) + # TODO: extend performance check for following backends + def callback_gpu(adb: SimpleADB): + adb.execute() + + def callback_lpai(adb: SimpleADB): + adb.execute() + with expected: # model declaration model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, inputs=inputs, compiler_specs=compile_specs, ).to_executorch() - # verifier per backend dispatcher = { QnnExecuTorchBackendType.kHtpBackend: callback_htp, + QnnExecuTorchBackendType.kGpuBackend: callback_gpu, + QnnExecuTorchBackendType.kLpaiBackend: callback_lpai, } # remote testing invoke_remote( qnn_config=qnn_config, executorch_prog=executorch_prog_mgr, - callback=partial(dispatcher[qnn_config.backend], voltage=80), + callback=( + partial(dispatcher[qnn_config.backend], voltage=80) + if qnn_config.backend == QnnExecuTorchBackendType.kHtpBackend + else dispatcher[qnn_config.backend] + ), ) @staticmethod @unpack_fixtures def test(subtests, qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { QnnExecuTorchBackendType.kHtpBackend: [ - compile_specs(tuple(d.items())) - for d in [ - # compile_time option - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": True, - "use_fp16": False, - "htp_performance_mode": QnnExecuTorchHtpPerformanceMode.kHtpHighPowerSaver, - }, - # runtime_option (performance mode defaults to kHtpBurst) - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "debug": True, - "use_fp16": False, - }, - ] + # compile_time option + { + "soc_model": soc_model, + "debug": True, + "use_fp16": False, + "htp_performance_mode": QnnExecuTorchHtpPerformanceMode.kHtpHighPowerSaver, + }, + # runtime_option (performance mode defaults to kHtpBurst) + {"soc_model": soc_model, "debug": True, "use_fp16": False}, + ], + QnnExecuTorchBackendType.kGpuBackend: [ + # compile_time option: set low perf hint to GPU + { + "soc_model": soc_model, + "online_prepare": True, + "performance_mode": QnnExecuTorchGpuPerformanceMode.kGpuPerfHintLow, + }, + # runtime_option: scaffold — GPUruntime perf hint not yet wired in C++ + # TODO: extend GPU runtime to accept dynamic performance settings + { + "soc_model": soc_model, + "debug": True, + "online_prepare": True, + }, + ], + QnnExecuTorchBackendType.kLpaiBackend: [ + { + "soc_model": soc_model, + "fps": 30, + "ftrt_ratio": 10, + "client_perf_type": QnnExecuTorchLpaiClientPerf.kRealTime, + }, + # runtime_option: scaffold — LPAI runtime perf hint not yet wired in C++ + # TODO: extend LPAI runtime to accept dynamic performance settings + {"soc_model": soc_model, "debug": True}, ], } @@ -358,7 +420,9 @@ def test(subtests, qnn_config, compile_specs, expected): with subtests.test(msg=config): __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend][i], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend][i].items()) + ), expected=expected, aot=config == "compile_time_option", ) @@ -409,9 +473,7 @@ def callback(adb: SimpleADB, executorch_prog_mgr, expected_profile_events): model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, @@ -426,30 +488,43 @@ def callback(adb: SimpleADB, executorch_prog_mgr, expected_profile_events): callback=partial( callback, executorch_prog_mgr=executorch_prog_mgr, - expected_profile_events=20, + expected_profile_events=2, ), ) @staticmethod @unpack_fixtures def test(subtests, qnn_config, compile_specs, expected): - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { QnnExecuTorchBackendType.kHtpBackend: [ - compile_specs(tuple(d.items())) - for d in [ - # compile_time option - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, - "use_fp16": False, - }, - # runtime_option - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "use_fp16": False, - }, - ] + # compile_time option + { + "soc_model": soc_model, + "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, + "use_fp16": False, + }, + # runtime_option + {"soc_model": soc_model, "use_fp16": False}, + ], + QnnExecuTorchBackendType.kGpuBackend: [ + # compile_time option + { + "soc_model": soc_model, + "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, + "online_prepare": True, + }, + # runtime_option + {"soc_model": soc_model, "online_prepare": True}, + ], + QnnExecuTorchBackendType.kLpaiBackend: [ + # compile_time option + { + "soc_model": soc_model, + "profile_level": QnnExecuTorchProfileLevel.kProfileDetailed, + }, + # runtime_option + {"soc_model": soc_model}, ], } @@ -457,7 +532,9 @@ def test(subtests, qnn_config, compile_specs, expected): with subtests.test(msg=config): __class__._test( qnn_config=qnn_config, - compile_specs=backend_compile_specs[qnn_config.backend][i], + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend][i].items()) + ), expected=expected, aot=config == "compile_time_option", ) @@ -482,17 +559,23 @@ def test(qnn_config, compile_specs, expected): option_to_flatbuffer, ) - # extend this for other backends + # saver=True is a top-level QnnExecuTorchOptions field; works across backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "saver": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "saver": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "saver": True, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: { + "soc_model": soc_model, + "saver": True, + }, } with expected: @@ -500,13 +583,13 @@ def test(qnn_config, compile_specs, expected): model = __class__.Model() inputs = model.example_inputs() # perform ptq - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering with tempfile.TemporaryDirectory() as tmp_dir: # hack saver output folder - cs = backend_compile_specs[qnn_config.backend] + cs = compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ) option = flatbuffer_to_option(cs[0].value) option.saver_output_dir = f"{tmp_dir}/saver_output" cs[0].value = option_to_flatbuffer(option) @@ -539,17 +622,23 @@ def forward(self, x): @staticmethod @unpack_fixtures def test(qnn_config, compile_specs, expected): - # extend this for other backends + # shared_buffer=True is a top-level QnnExecuTorchOptions field; works across backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "shared_buffer": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "shared_buffer": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "shared_buffer": True, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: { + "soc_model": soc_model, + "shared_buffer": True, + }, } module = __class__.Model() @@ -558,8 +647,10 @@ def test(qnn_config, compile_specs, expected): module=module, inputs=module.example_inputs(), qnn_config=qnn_config, - quantizer=make_quantizer(soc_model=qnn_config.soc_model), - compile_specs=backend_compile_specs[qnn_config.backend], + quantizer=get_quantizer(qnn_config), + compile_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), metrics=expected, ) @@ -605,9 +696,7 @@ def test(qnn_config, compile_specs, expected): # perform ptq model = __class__.Model() inputs = model.example_inputs() - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering edge_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, @@ -621,22 +710,23 @@ def test(qnn_config, compile_specs, expected): class TensorDump: + # Simple Conv2d+ReLU model that is supported by all backends class Model(torch.nn.Module): def __init__(self): super().__init__() - self.idx_source = torch.rand(10, 3) + self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1) + self.relu = torch.nn.ReLU() def example_inputs(self): - return (torch.randn(3, 10),) + return (torch.randn(1, 3, 8, 8),) def forward(self, x): - a, b = torch.topk(x, 3) - return a + self.idx_source[b] + return self.relu(self.conv(x)) @staticmethod @unpack_fixtures def test(qnn_config, compile_specs, expected): - def callback(adb: SimpleADB, expected_intermediate_events): + def callback(adb: SimpleADB, debugger, expected_compared_events): with tempfile.TemporaryDirectory() as tmp_dir: etdump_path = f"{tmp_dir}/etdump.etdp" debug_output_path = f"{tmp_dir}/debug_output.bin" @@ -644,49 +734,82 @@ def callback(adb: SimpleADB, expected_intermediate_events): adb.pull_debug_output( etdump_path=etdump_path, debug_buffer_path=debug_output_path ) - inspector = Inspector( - etdump_path=etdump_path, debug_buffer_path=debug_output_path + debugger.setup_inspector( + etdump_path=etdump_path, + debug_buffer_path=debug_output_path, ) - for event_block in inspector.event_blocks: - if event_block.name == "Execute": - assert ( - len(event_block.events) == expected_intermediate_events - ), ( - f"unexpected number of intermediate events, expecting " - f"{expected_intermediate_events}, but has {len(event_block.events)} events.", - ) + comparator = debugger.create_comparator(QcomCosineSimilarityComparator) + numeric_results = debugger.inspector.calculate_numeric_gap( + distance=comparator, + reference_graph=debugger.reference_graph_name, + ) + numeric_results = numeric_results.set_index("runtime_debug_handle") + assert len(numeric_results) == expected_compared_events, ( + f"unexpected number of compared events, expecting " + f"{expected_compared_events}, but has {len(numeric_results)} events." + ) + for _, row in numeric_results.iterrows(): + assert comparator.is_valid_score(row.gap[0]), ( + f"Node {row.aot_ops} is failing " + f"{comparator.metric_name()} test, {row.gap[0]} is lower " + f"than {comparator.threshold}." + ) - # extend this for other backends + soc_model = getattr(QcomChipset, qnn_config.soc_model) + # dump_intermediate_outputs=True is a top-level QnnExecuTorchOptions field; works across backends backend_compile_specs = { - QnnExecuTorchBackendType.kHtpBackend: compile_specs( - tuple( - { - "soc_model": getattr(QcomChipset, qnn_config.soc_model), - "dump_intermediate_outputs": True, - "use_fp16": False, - }.items() - ) - ), + QnnExecuTorchBackendType.kHtpBackend: { + "soc_model": soc_model, + "dump_intermediate_outputs": True, + "use_fp16": False, + }, + QnnExecuTorchBackendType.kGpuBackend: { + "soc_model": soc_model, + "dump_intermediate_outputs": True, + "online_prepare": True, + }, + QnnExecuTorchBackendType.kLpaiBackend: { + "soc_model": soc_model, + "dump_intermediate_outputs": True, + }, } with expected: # perform ptq model = __class__.Model() inputs = model.example_inputs() - with calibrate( - model, [inputs], make_quantizer(soc_model=qnn_config.soc_model) - ) as model: + with calibrate(model, [inputs], get_quantizer(qnn_config)) as model: # start lowering executorch_prog_mgr = to_edge_transform_and_lower_to_qnn( module=model, inputs=inputs, - compiler_specs=backend_compile_specs[qnn_config.backend], + compiler_specs=compile_specs( + tuple(backend_compile_specs[qnn_config.backend].items()) + ), generate_etrecord=True, ).to_executorch() - # remote testing - qnn_config.dump_intermediate_outputs = True - invoke_remote( - qnn_config=qnn_config, - executorch_prog=executorch_prog_mgr, - callback=partial(callback, expected_intermediate_events=9), - ) + + with tempfile.TemporaryDirectory() as etrecord_dir: + etrecord_path = f"{etrecord_dir}/etrecord.bin" + etrecord = executorch_prog_mgr.get_etrecord() + debugger = QNNIntermediateDebugger(inputs) + debugger.set_etrecord_file_path(etrecord_path) + debugger.set_edge_ep( + edge_ep=etrecord.graph_map[debugger.reference_graph_name] + ) + etrecord.update_representative_inputs(debugger.sample_input) + etrecord.save(etrecord_path) + + # remote testing + qnn_config.dump_intermediate_outputs = True + invoke_remote( + qnn_config=qnn_config, + executorch_prog=executorch_prog_mgr, + inputs=inputs, + # conv + relu = 2 intermediate outputs + callback=partial( + callback, + debugger=debugger, + expected_compared_events=2, + ), + ) diff --git a/backends/qualcomm/tests/rework/src/op.py b/backends/qualcomm/tests/rework/src/op.py index 963285725cb..a782235732f 100644 --- a/backends/qualcomm/tests/rework/src/op.py +++ b/backends/qualcomm/tests/rework/src/op.py @@ -3198,7 +3198,7 @@ def forward(self, x): @unpack_fixtures def test(subtests, qnn_config, quantizer, compile_spec, expected): inputs = (torch.randn(1, 4, 8, 8),) - dims = [-1, 1, 2] + dims = [-1, 3] for dim in dims: with subtests.test(msg=f"dim:{dim}"): with expected as metrics: diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 1139085ec31..3c05a816c6d 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -7787,6 +7787,32 @@ def test_qnn_backend_dump_intermediate_outputs_simple_model(self): expected_compared_events=expected_compared_events, ) + def test_qnn_backend_dump_intermediate_outputs_conv_relu(self): + match get_backend_type(self.backend): + case QnnExecuTorchBackendType.kHtpBackend: + backend_options = generate_htp_compiler_spec(use_fp16=False) + case QnnExecuTorchBackendType.kLpaiBackend: + backend_options = generate_lpai_compiler_spec( + target_env=self.get_lpai_target_env() + ) + case _: + raise ValueError("Backend is not implemented yet") + TestQNN.compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=self.chipset_table[TestQNN.soc_model], + backend_options=backend_options, + dump_intermediate_outputs=True, + ) + sample_input = (torch.randn(1, 3, 8, 8),) + module = ConvRelu() # noqa: F405 + module = self.get_qdq_module(module, sample_input) + + self.lower_module_and_test_output( + module, + sample_input, + expected_partitions=1, + expected_compared_events=2, + ) + def test_qnn_backend_dump_intermediate_outputs_topk(self): torch.manual_seed(8) backend_options = generate_htp_compiler_spec(use_fp16=False) diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 75ef9c1f128..2caa8905fda 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -34,6 +34,7 @@ QnnExecuTorchBackendOptions, QnnExecuTorchBackendType, QnnExecuTorchGpuBackendOptions, + QnnExecuTorchGpuPerformanceMode, QnnExecuTorchGpuPrecision, QnnExecuTorchHtpBackendOptions, QnnExecuTorchHtpPerformanceMode, @@ -1018,6 +1019,7 @@ def draw_graph(title, path, graph_module: torch.fx.GraphModule, format=DrawForma def generate_gpu_compiler_spec( + performance_mode: QnnExecuTorchGpuPerformanceMode = QnnExecuTorchGpuPerformanceMode.kGpuPerfHintHigh, precision: QnnExecuTorchGpuPrecision = QnnExecuTorchGpuPrecision.kGpuPrecisionUserProvided, use_memory_optimizations: bool = True, use_node_optimizations: bool = True, @@ -1028,6 +1030,8 @@ def generate_gpu_compiler_spec( Helper function generating backend options for QNN HTP Args: + performance_mode: + kGpuPerfHintHigh / kGpuPerfHintNormal / kGpuPerfHintLow precision: kGpuPrecisionFp32 - Sets the precision mode to floating point 32-bit (FP32). kGpuPrecisionFp16 - Sets the precision mode to floating point 16-bit (FP16). @@ -1046,6 +1050,7 @@ def generate_gpu_compiler_spec( """ # TODO: enable performance hint mechanism in runtime and make this as an option gpu_options = QnnExecuTorchGpuBackendOptions() + gpu_options.performance_mode = performance_mode gpu_options.precision = precision gpu_options.use_memory_optimizations = use_memory_optimizations gpu_options.use_node_optimizations = use_node_optimizations From c9eee945a0113c0cbb19ba17dec8be2b44acdb9f Mon Sep 17 00:00:00 2001 From: qti-horodnic Date: Mon, 7 Sep 2026 22:10:58 -0700 Subject: [PATCH 069/190] Qualcomm AI Engine Direct - Enabling Support for Qualcomm Chipsets for Snapdragon 7+ Gen 3 (#22542) ### Summary Adding SoC support for SM7675 (Snapdragon 7+ Gen 3) and SM8635 (Snapdragon 8s Gen 3), both HTP V73. Scope of Support: - Quantized HTP: supported*. Quantized HTP validated across 8a8w/16a4w/16a2w, per-channel, block-wise, and QAT paths, plus the full quantized utils suite. See test cases in the Test Plan section. I did not run the whole quantized ops/model suite, which is why support is denoted with a (*). - FP16: not supported. QNN rejects these SoCs for FP16 (`"The SocModel doesn't support FP16"`), so FP16 lowering does not produce a delegated graph. Note it also does not degrade gracefully today: models containing ops on the partitioner's non-decompose list (e.g. `linear`, `layer_norm`) fail to export rather than falling back to CPU. Note that these changes need the fix added in https://github.com/pytorch/executorch/pull/22543 for the tests to pass. ### Test plan ``` python backends/qualcomm/tests/test_qnn_delegate.py \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_conv2d \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_linear \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_layer_norm \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_conv2d \ -v --soc_model SM8635 --host aisw-local-vm3 --device 4603496b \ --build_folder build-android python backends/qualcomm/tests/test_qnn_delegate.py \ TestQNNQuantizedOperator.test_qnn_backend_sort \ TestQNNQuantizedOperator.test_qnn_backend_conv2d \ TestQNNQuantizedOperator.test_qnn_backend_linear \ TestQNNQuantizedOperator.test_qnn_backend_layer_norm \ TestQNNQuantizedOperator.test_qnn_backend_element_wise_add \ -v --soc_model SM8635 --host aisw-local-vm3 --device 4603496b \ --build_folder build-android python backends/qualcomm/tests/test_qnn_delegate.py \ TestQNNQuantizedOperator.test_qnn_backend_16a2w_conv2d \ TestQNNQuantizedOperator.test_qnn_backend_16a2w_linear \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_per_channel_linear \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_per_channel_linear_with_bias \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_conv2d_qat \ TestQNNQuantizedOperator.test_qnn_backend_16a4w_block_conv2d_qat \ -v --soc_model SM8635 --host aisw-local-vm3 --device 4603496b \ --build_folder build-android python backends/qualcomm/tests/test_qnn_delegate.py TestQNNQuantizedUtils -v --soc_model SM8635 --host aisw-local-vm3 --device 4603496b --build_folder build-android ``` --- backends/qualcomm/serialization/qc_compiler_spec.fbs | 2 ++ backends/qualcomm/serialization/qc_schema.py | 4 ++++ backends/qualcomm/utils/utils.py | 6 ++++++ 3 files changed, 12 insertions(+) diff --git a/backends/qualcomm/serialization/qc_compiler_spec.fbs b/backends/qualcomm/serialization/qc_compiler_spec.fbs index 4d3d59c4fbe..100404329af 100644 --- a/backends/qualcomm/serialization/qc_compiler_spec.fbs +++ b/backends/qualcomm/serialization/qc_compiler_spec.fbs @@ -48,10 +48,12 @@ enum QcomChipset: int { SA8295 = 39, SA8797 = 72, SC8380XP = 60, + SM7675 = 70, SM8350 = 30, SM8450 = 36, SM8475 = 42, SM8550 = 43, + SM8635 = 68, SM8650 = 57, SM8750 = 69, SM8850 = 87, diff --git a/backends/qualcomm/serialization/qc_schema.py b/backends/qualcomm/serialization/qc_schema.py index b3992bc97f1..692c3f3f217 100644 --- a/backends/qualcomm/serialization/qc_schema.py +++ b/backends/qualcomm/serialization/qc_schema.py @@ -55,10 +55,12 @@ class QcomChipset(IntEnum): SA8295 = 39 # v68 SA8797 = 72 # v81 SC8380XP = 60 # v73 + SM7675 = 70 # v73 SM8350 = 30 # v68 SM8450 = 36 # v69 SM8475 = 42 # v69 SM8550 = 43 # v73 + SM8635 = 68 # v73 SM8650 = 57 # v75 SM8750 = 69 # v79 SM8850 = 87 # v81 @@ -86,11 +88,13 @@ class SocInfo: QcomChipset.SA8295: SocInfo(QcomChipset.SA8295, HtpInfo(HtpArch.V68, 8)), QcomChipset.SA8797: SocInfo(QcomChipset.SA8797, HtpInfo(HtpArch.V81, 16)), QcomChipset.SC8380XP: SocInfo(QcomChipset.SC8380XP, HtpInfo(HtpArch.V73, 8)), + QcomChipset.SM7675: SocInfo(QcomChipset.SM7675, HtpInfo(HtpArch.V73, 4)), QcomChipset.SM8350: SocInfo(QcomChipset.SM8350, HtpInfo(HtpArch.V68, 4)), QcomChipset.SM8450: SocInfo(QcomChipset.SM8450, HtpInfo(HtpArch.V69, 8)), QcomChipset.SM8475: SocInfo(QcomChipset.SM8475, HtpInfo(HtpArch.V69, 8)), QcomChipset.SM8550: SocInfo(QcomChipset.SM8550, HtpInfo(HtpArch.V73, 8)), QcomChipset.SA8255: SocInfo(QcomChipset.SA8255, HtpInfo(HtpArch.V73, 8)), + QcomChipset.SM8635: SocInfo(QcomChipset.SM8635, HtpInfo(HtpArch.V73, 4)), QcomChipset.SM8650: SocInfo(QcomChipset.SM8650, HtpInfo(HtpArch.V75, 8)), QcomChipset.SM8750: SocInfo(QcomChipset.SM8750, HtpInfo(HtpArch.V79, 8)), QcomChipset.SM8850: SocInfo( diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 2caa8905fda..f78d47525ad 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -1199,9 +1199,11 @@ def generate_qnn_executorch_compiler_spec( # noqa: C901 Args: soc_model: The SoC you plan to run the compiled model. Please check QcomChipset for supported SoC. + SM7675(Snapdragon 7+ Gen 3) SM8450 (Snapdragon 8 Gen 1) SM8475(Snapdragon 8 Gen 1+) SM8550(Snapdragon 8 Gen 2) + SM8635(Snapdragon 8s Gen 3) SM8650(Snapdragon 8 Gen 3) SM8750(Snapdragon 8 Elite) SM8850(Snapdragon 8 Elite Gen 5) @@ -1332,7 +1334,9 @@ def get_soc_to_htp_arch_map(): "SM8450": HtpArch.V69, "SM8475": HtpArch.V69, "SM8550": HtpArch.V73, + "SM7675": HtpArch.V73, "SA8255": HtpArch.V73, + "SM8635": HtpArch.V73, "SM8650": HtpArch.V75, "SM8750": HtpArch.V79, "SM8850": HtpArch.V81, @@ -1362,11 +1366,13 @@ def get_soc_to_chipset_map(): "SA8295": QcomChipset.SA8295, "SA8797": QcomChipset.SA8797, "SC8380XP": QcomChipset.SC8380XP, + "SM7675": QcomChipset.SM7675, "SM8350": QcomChipset.SM8350, "SM8450": QcomChipset.SM8450, "SM8475": QcomChipset.SM8475, "SM8550": QcomChipset.SM8550, "SA8255": QcomChipset.SA8255, + "SM8635": QcomChipset.SM8635, "SM8650": QcomChipset.SM8650, "SM8750": QcomChipset.SM8750, "SM8850": QcomChipset.SM8850, From f115e69452991f6c90d78b93d2c6ca2372e4879e Mon Sep 17 00:00:00 2001 From: Vaclav Novak Date: Tue, 8 Sep 2026 14:14:26 +0200 Subject: [PATCH 070/190] NXP backend: fixed bug in NeutronBackend.cpp (#22435) ### Summary There is a bug that causes incorrect handling of an operator with one input channel. This PR fixes it. ### Test plan tests can be manually run using `pytest -c /dev/null backends/nxp/tests/` cc @robert-kalmar @JakeStevens @digantdesai @rascani @MartinPavella @jirioc --- backends/nxp/runtime/NeutronBackend.cpp | 33 +++++++++++++++++++++---- 1 file changed, 28 insertions(+), 5 deletions(-) diff --git a/backends/nxp/runtime/NeutronBackend.cpp b/backends/nxp/runtime/NeutronBackend.cpp index 4c75f76712a..88945c7428a 100644 --- a/backends/nxp/runtime/NeutronBackend.cpp +++ b/backends/nxp/runtime/NeutronBackend.cpp @@ -442,8 +442,20 @@ class NeutronBackend final : public PyTorchBackendInterface { auto arg = args[cfg->inputMap[i]]->toTensor(); auto dim_order = arg.dim_order().data(); - if (cfg->inputTranspositionFlags[i] && - multipleChannelsPresent(arg.sizes())) { + if (cfg->inputTranspositionFlags[i]) { + if (!multipleChannelsPresent(arg.sizes())) { + // The input has only 1 channel, so NCHW and NHWC data is equivalent + // and no transposition is needed. + if (!is_channels_last_dim_order(dim_order, arg.dim()) && + !is_contiguous_dim_order(dim_order, arg.dim())) { + ET_LOG(Error, "Input %d uses unsupported dim-order.", i); + print_dim_order(dim_order, arg.dim()); + return Error::InvalidProgram; + } + + cfg->dcfg.inputs[i] = arg.const_data_ptr(); + continue; + } // The input must be transposed. if (arg.sizes().size() < 3) { ET_LOG(Error, "Unable to transpose 1D and 2D input to channel last"); @@ -496,10 +508,21 @@ class NeutronBackend final : public PyTorchBackendInterface { auto arg = args[cfg->numInputArgs + cfg->outputMap[i]]->toTensor(); auto dim_order = arg.dim_order().data(); - if (cfg->outputTranspositionFlags[i] && - multipleChannelsPresent(arg.sizes())) { - // The output will have to be transposed. + if (cfg->outputTranspositionFlags[i]) { + if (!multipleChannelsPresent(arg.sizes())) { + // The output has only 1 channel, so NCHW and NHWC data is equivalent + // and no transposition is needed. + if (!is_channels_last_dim_order(dim_order, arg.dim()) && + !is_contiguous_dim_order(dim_order, arg.dim())) { + ET_LOG(Error, "Output %d uses unsupported dim-order.", i); + print_dim_order(dim_order, arg.dim()); + return Error::InvalidProgram; + } + cfg->dcfg.outputs[i] = arg.mutable_data_ptr(); + continue; + } + // The output will have to be transposed. if (is_channels_last_dim_order(dim_order, arg.dim())) { // The tensor will already be correctly permuted. No transposition // needed. From 978e5844f79dbccb3e3274644842f2f5a47a2eff Mon Sep 17 00:00:00 2001 From: Irina Korchakova Date: Tue, 8 Sep 2026 15:16:20 +0200 Subject: [PATCH 071/190] Use neutron kernel names in profiling table (#21614) ### Summary Replace neutron kernel kinds with names in profiling table if possible: 1. Add driver version into etdump as metadata for the last traced kernel (which corresponds to time needed for profiling info print) in runtime. 2. Check if eIQ SDK installed on host PC corresponds to the driver version from runtime. - If versions match -> get list of kernel names from convertor (compiler) and replace kernel numbers with kernels name. - If versions doesn't match -> use kernel numbers in result table. ### Test plan Test included cc @robert-kalmar @JakeStevens @digantdesai @rascani --- backends/nxp/runtime/NeutronBackend.cpp | 13 +- backends/nxp/runtime/NeutronDriver.h | 13 ++ .../nxp/tests/generic_tests/test_profiling.py | 50 +++++-- backends/nxp/tests/profiling_utils.py | 139 ++++++++++++++++++ examples/nxp/analyzing_with_inspector.py | 23 ++- 5 files changed, 223 insertions(+), 15 deletions(-) create mode 100644 backends/nxp/tests/profiling_utils.py diff --git a/backends/nxp/runtime/NeutronBackend.cpp b/backends/nxp/runtime/NeutronBackend.cpp index 88945c7428a..1effb21a536 100644 --- a/backends/nxp/runtime/NeutronBackend.cpp +++ b/backends/nxp/runtime/NeutronBackend.cpp @@ -659,6 +659,14 @@ class NeutronBackend final : public PyTorchBackendInterface { index++; } } + // The neutronGetSdkVersion() function is available starting with Neutron + // Software 3.2.1. The code below is not backward compatible with earlier + // Neutron Software versions. + NeutronSdkVersion neutron_sdk_version = neutronGetSdkVersion(); + uint16_t neutron_sdk_version_uint16 = + static_cast(neutron_sdk_version.major << 8) | + static_cast(neutron_sdk_version.minor << 4) | + static_cast(neutron_sdk_version.patch); event_tracer_log_profiling_delegate( tracer, nullptr, @@ -666,9 +674,8 @@ class NeutronBackend final : public PyTorchBackendInterface { neutron_events[events_num - 1].startEvent.time, neutron_events[events_num - 1].stopEvent.time + stop_ticks - start_ticks, - static_cast( - &neutron_events[events_num - 1].startEvent.functionCode), - sizeof(uint8_t)); + static_cast(&neutron_sdk_version_uint16), + sizeof(uint16_t)); } #endif diff --git a/backends/nxp/runtime/NeutronDriver.h b/backends/nxp/runtime/NeutronDriver.h index 5c47bd74eab..9280cd6adba 100644 --- a/backends/nxp/runtime/NeutronDriver.h +++ b/backends/nxp/runtime/NeutronDriver.h @@ -124,6 +124,16 @@ typedef struct { void (*wait)(uint32_t channel); } NeutronConfig; +/// This structure contains semantic version of the Neutron SDK +/// (major.minor.patch) and the SHA version (string and uint32_t). +typedef struct { + uint32_t major; + uint32_t minor; + uint32_t patch; + const char* hashString; + uint32_t hashUint32; +} NeutronSdkVersion; + /* Invalid handle, returned by neutronModelPrepare() if an error occurred. */ #define NEUTRON_INVALID_HANDLE NULL @@ -224,6 +234,9 @@ NeutronError neutronSetConfig(NeutronConfig* config); /// - Used to get NeutronContext size. size_t neutronGetModelContextSize(); +/// - Used to get Neutron SDK version. +NeutronSdkVersion neutronGetSdkVersion(); + /// - Allocates size bytes and returns a pointer to the allocated memory. /// The returned pointer address will be a multiple of the alignment. /// Returns NULL on failure. diff --git a/backends/nxp/tests/generic_tests/test_profiling.py b/backends/nxp/tests/generic_tests/test_profiling.py index cd90bdd345b..7e147dc989c 100644 --- a/backends/nxp/tests/generic_tests/test_profiling.py +++ b/backends/nxp/tests/generic_tests/test_profiling.py @@ -22,6 +22,11 @@ OUTPUTS_DIR, ) +from executorch.backends.nxp.tests.profiling_utils import ( + get_neutron_compiler_version, + get_neutron_driver_version, + get_neutron_kernel_kinds, +) from executorch.devtools.inspector._inspector import Inspector from executorch.examples.models.mlperf_tiny import ( DeepAutoEncoder, @@ -65,12 +70,15 @@ def inspector_check(test_name: str) -> None: 5. The profiling dump event does not have associated op types. """ + # Global mapping of Neutron kernel IDs to names used by the delegate metadata parser. + kernel_kinds = {} + def parse_delegate_metadata( delegate_metadatas: list[bytes], ) -> Union[list[str], dict[str, Any]]: """Metadata parser for Neutron Backend metadata. - The parser is a callable that deserializes the data and returns neutron kernel number. + The parser deserializes delegate metadata and converts kernel IDs into human-readable kernel names when available. The deserialized data is then added back to the corresponding event in the event block for user consumption. """ @@ -81,7 +89,13 @@ def parse_delegate_metadata( if function_code == 0: metadata_list.append("Profiling dump") else: - metadata_list.append("Neutron kernel " + str(function_code)) + metadata_list.append( + kernel_kinds.get( + function_code, "Neutron kernel " + str(function_code) + ) + ) + elif len(metadata_bytes) == 2: + metadata_list.append("Profiling dump") else: metadata_list.append("Invalid metadata size") return metadata_list @@ -96,6 +110,16 @@ def parse_delegate_metadata( file_path ), f"Required profiling file does not exist: {file_path}" + # Validate driver/compiler version compatibility and load kernel names + # used to decode delegate metadata. + driver_version = get_neutron_driver_version(etdump_path) + compiler_version = get_neutron_compiler_version() + if driver_version: + assert ( + driver_version == compiler_version + ), "Driver and compiler versions do not match" + kernel_kinds = get_neutron_kernel_kinds() + # Create Inspector and parse profiling data. try: inspector = Inspector( @@ -123,18 +147,22 @@ def parse_delegate_metadata( assert numeric_events, "No numeric delegate profiling events found" - # All delegate events except the last one should describe - # individual Neutron kernels. + # All numeric delegate events except the last contain either + # resolved kernel names or fallback "Neutron kernel " metadata. for event in numeric_events[:-1]: - metadata = str(event.delegate_debug_metadatas) - - assert "Neutron kernel" in metadata, ( - f"Event {event.name}: expected 'Neutron kernel', " f"got {metadata}" - ) + metadata = event.delegate_debug_metadatas + if kernel_kinds: + assert "Neutron kernel" not in metadata, ( + f"Event {event.name}: expected kernel kind, " f"got {metadata}" + ) + else: + assert "Neutron kernel" in metadata, ( + f"Event {event.name}: expected 'Neutron kernel', " f"got {metadata}" + ) # The final numeric event should represent the profiling dump. profiling_dump_event = numeric_events[-1] - profiling_metadata = str(profiling_dump_event.delegate_debug_metadatas) + profiling_metadata = profiling_dump_event.delegate_debug_metadatas assert "Profiling dump" in profiling_metadata, ( f"Event {profiling_dump_event.name}: " @@ -142,7 +170,7 @@ def parse_delegate_metadata( ) # Profiling dump event is expected to have no associated operators. - assert profiling_dump_event.op_types == [], ( + assert not profiling_dump_event.op_types, ( f"Event {profiling_dump_event.name}: expected empty op_types, " f"got {profiling_dump_event.op_types}" ) diff --git a/backends/nxp/tests/profiling_utils.py b/backends/nxp/tests/profiling_utils.py new file mode 100644 index 00000000000..f43cb5455cd --- /dev/null +++ b/backends/nxp/tests/profiling_utils.py @@ -0,0 +1,139 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +import re +import subprocess + +from executorch.devtools.etdump.serialize import deserialize_from_etdump_flatcc + + +def get_neutron_driver_version(etdump_path: str) -> str: + """ + Extract the Neutron Driver version from an ETDump file. + + The Neutron Driver version is stored in the metadata of the last Neutron + delegate event. This event is emitted when the profiling dump is generated. + The version is encoded as a 16-bit value in little-endian format: + - 4 bits - major version + - 4 bits - minor version + - 4 bits - patch version + - 4 bits - reserved + + :param etdump_path: Path to the ETDump binary file. + :return: Neutron Driver version string (e.g. "1.2.3") if successfully decoded, + otherwise empty string. Errors are logged instead of raised. + """ + + try: + with open(etdump_path, "rb") as f: + data = f.read() + etdump = deserialize_from_etdump_flatcc(data) + except Exception as e: + logging.exception("Failed to load ETDump: %s", e) + return "" + + events = [] + try: + for run in etdump.run_data: + for event in run.events: + profile_event = getattr(event, "profile_event", None) + if ( + profile_event is not None + and getattr(profile_event, "delegate_debug_id_int", 0) > 0 + ): + events.append(event) + except Exception as e: + logging.exception("Failed while processing events: %s", e) + return "" + + try: + metadata = events[-1].profile_event.delegate_debug_metadata + if not metadata or len(metadata) < 2: + logging.error("Invalid delegate_debug_metadata") + return "" + + major, minor, patch = [ + (int.from_bytes(metadata, "little") >> shift) & 0xF for shift in (8, 4, 0) + ] + return f"{major}.{minor}.{patch}" + + except Exception as e: + logging.exception("Failed to extract version from metadata: %s", e) + return "" + + +def get_neutron_compiler_version() -> str: + """ + Get the Neutron Compiler version reported by the neutron_compiler tool. + + Executes `neutron_compiler --version` and returns the version as + {major}.{minor}.{patch} string. + + :return: The version string returned by neutron_compiler, or empty string if + the command fails, times out, or the executable is not available. + Errors are logged instead of being raised. + """ + + try: + # Use neutron_compiler because neutron_converter executable is unavailable since NS 3.2.1. + proc = subprocess.Popen( + ["neutron_compiler", "--version"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + stdout, stderr = proc.communicate(timeout=10) + if proc.returncode != 0: + logging.error( + "Failed to get compiler version: %s", + stderr.strip(), + ) + return "" + version_match = re.search(r"version\s(\d+\.\d+\.\d+)", stdout) + if version_match: + return version_match.group(1) + else: + logging.exception("Unexpected error while getting neutron compiler version") + return "" + except Exception: + logging.exception("Error while getting neutron compiler version") + return "" + + +def get_neutron_kernel_kinds(target: str = "imxrt700") -> dict[int, str]: + """ + Retrieve kernel kinds supported by neutron_compiler for the specified target. + + Executes the neutron_compiler command with the --show-kernel-kinds option, + parses its output, and returns a dictionary mapping kernel IDs to kernel + names. + + :param target: Target platform for which kernel kinds should be queried. + Defaults to "imxrt700". + :return: Returns empty dict if neutron_compiler exits with an error. + Otherwise, a dictionary where: + - key: kernel ID (int) + - value: kernel name (str) + """ + + # Use neutron_compiler because neutron_converter executable is unavailable since NS 3.2.1. + proc = subprocess.Popen( + ["neutron_compiler", "--target", target, "--show-kernel-kinds"], + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + ) + stdout, stderr = proc.communicate(timeout=10) + if proc.returncode != 0: + logging.error( + "Failed to get kernrl kinds from neutron_compiler: %s", + stderr.strip(), + ) + return {} + return { + int(op_id): name + for op_id, name in re.findall(r"\[\s*(\d+)\s*\]\s+(.+)", stdout) + } diff --git a/examples/nxp/analyzing_with_inspector.py b/examples/nxp/analyzing_with_inspector.py index b339af79d6e..e5f6d97a8c1 100644 --- a/examples/nxp/analyzing_with_inspector.py +++ b/examples/nxp/analyzing_with_inspector.py @@ -7,8 +7,17 @@ from typing import Any, Union +from executorch.backends.nxp.tests.profiling_utils import ( + get_neutron_compiler_version, + get_neutron_driver_version, + get_neutron_kernel_kinds, +) + from executorch.devtools import Inspector +# Global mapping of Neutron kernel IDs to names used by the delegate metadata parser. +kernel_kinds = {} + def parse_delegate_metadata( delegate_metadatas: list[bytes], @@ -26,7 +35,13 @@ def parse_delegate_metadata( if function_code == 0: metadata_list.append("Profiling dump") else: - metadata_list.append("Neutron kernel " + str(function_code)) + metadata_list.append( + kernel_kinds.get( + function_code, "Neutron kernel " + str(function_code) + ) + ) + elif len(metadata_bytes) == 2: + metadata_list.append("Profiling dump") else: metadata_list.append("Invalid metadata size") return metadata_list @@ -37,6 +52,12 @@ def parse_delegate_metadata( try: etrecord_path = "etrecord/etrecord.bin" etdump_path = "etdump/trace.etdump" + + driver_version = get_neutron_driver_version(etdump_path) + compiler_version = get_neutron_compiler_version() + if driver_version and driver_version == compiler_version: + kernel_kinds = get_neutron_kernel_kinds() + inspector = Inspector( etdump_path=etdump_path, etrecord=etrecord_path, From 7dc8641963f23b6129b97c42b204a5355fdb3bc2 Mon Sep 17 00:00:00 2001 From: Longfang Date: Tue, 8 Sep 2026 08:55:16 -0700 Subject: [PATCH 072/190] Report whether XNNPACK packed weights fell back to heap (#22413) Differential Revision: D118206001 Pull Request resolved: https://github.com/pytorch/executorch/pull/22413 --- backends/xnnpack/runtime/XNNPACKBackend.cpp | 12 + backends/xnnpack/runtime/XNNPACKBackend.h | 113 ++++++ backends/xnnpack/runtime/XNNWeightsCache.cpp | 153 ++++++++- backends/xnnpack/runtime/XNNWeightsCache.h | 53 +++ .../runtime/XNNWeightsCacheManager.cpp | 93 +++++ .../xnnpack/runtime/XNNWeightsCacheManager.h | 16 + .../test/runtime/test_weight_cache.cpp | 45 +++ .../test_xnn_weights_cache_manager.cpp | 325 +++++++++++++++++- 8 files changed, 784 insertions(+), 26 deletions(-) diff --git a/backends/xnnpack/runtime/XNNPACKBackend.cpp b/backends/xnnpack/runtime/XNNPACKBackend.cpp index b9b4e82f6ca..a76a0832def 100644 --- a/backends/xnnpack/runtime/XNNPACKBackend.cpp +++ b/backends/xnnpack/runtime/XNNPACKBackend.cpp @@ -255,6 +255,12 @@ class XnnpackBackend final return first_err; } + public: + /** See xnnpack::get_packed_cache_report(). */ + xnnpack::PackedCacheReport packed_cache_report() const { + return options_.weights_cache_manager().report(); + } + private: mutable xnnpack::XnnpackBackendOptions options_; @@ -272,5 +278,11 @@ Backend backend{xnnpack::xnnpack_backend_key, &backend_instance}; static auto success_with_compiler = register_backend(backend); } // namespace +namespace xnnpack { +PackedCacheReport get_packed_cache_report() { + return backend_instance.packed_cache_report(); +} +} // namespace xnnpack + } // namespace backends } // namespace executorch diff --git a/backends/xnnpack/runtime/XNNPACKBackend.h b/backends/xnnpack/runtime/XNNPACKBackend.h index 1053a206360..df83a500489 100644 --- a/backends/xnnpack/runtime/XNNPACKBackend.h +++ b/backends/xnnpack/runtime/XNNPACKBackend.h @@ -1,5 +1,11 @@ #pragma once +#include +#include +#include +#include +#include + namespace executorch::backends::xnnpack { /// The key for the backend. This is used to register the backend, check /// availability, and get/set options. @@ -61,4 +67,111 @@ enum class WorkspaceSharingMode { // maximum enum value. Count, }; + +/// Outcome of opening the packed-weight cache file. +enum class PackedCacheState : int32_t { + /// No cache path configured — the caller never opted in. + Disabled = 0, + /// The cache file opened. Does NOT imply zero heap; see PackedCacheStats. + FileBacked = 1, + /// A path was configured but the file could not be used. + HeapFallback = 2, +}; + +/** Why an individual allocation was served from heap. */ +enum class PackedCacheHeapReason : int32_t { + None = 0, + /// The instance has no cache path: it never opted into file backing, so + /// heap is the intended behaviour rather than a fallback. Bucketed + /// separately and excluded from heap_bytes — a process that mixes an + /// opted-in model with a non-opted-in one would otherwise report the + /// latter's packed weights as if the former had fallen back. + NotOptedIn = 1, + /// Unnamed constant — can never be reloaded by name. By design. + UnnamedConstant = 2, + /// Incidental re-pack after a successful load. By design *if* the loaded + /// cache is complete; a large volume here means it was not. + RepackAfterLoad = 3, + /// No usable file descriptor at allocation time. + NoFileBacking = 4, + /// ftruncate() to extend the file failed. + GrowFailed = 5, + /// mmap() of the grown region failed. + MmapFailed = 6, + /// Not a reason; bounds the per-reason counters. Matches the + /// WorkspaceSharingMode convention in this header. + Count, +}; + +/** Which step failed when a configured path still ended up on heap. */ +enum class PackedCacheFailure : int32_t { + None = 0, + OpenFailed = 1, + TruncateFailed = 2, + GrowFailed = 3, + MmapFailed = 4, +}; + +/** + * Per-cache counters. `heap_bytes` against `mapped_bytes` is the signal; + * `state` alone calls a partially-loaded cache healthy. + */ +struct PackedCacheStats { + PackedCacheState state{PackedCacheState::Disabled}; + PackedCacheFailure failure{PackedCacheFailure::None}; + int32_t last_errno{0}; + /// Cache file size as of the last successful save. + int64_t file_bytes{0}; + /// Packed bytes served from heap when the file was supposed to serve them. + /// Excludes NotOptedIn, so this is only ever "bytes that should have been + /// file-backed and were not". + int64_t heap_bytes{0}; + /// Packed bytes served from the mmap'd file (clean, file-backed). + int64_t mapped_bytes{0}; + /// Reason accounting for the largest share of heap_bytes. On the aggregate + /// this is the argmax over per-reason bytes summed across caches, not the + /// local reason of whichever cache happened to allocate the most. + PackedCacheHeapReason heap_reason{PackedCacheHeapReason::None}; + /// Heap bytes split by reason, so callers can sum per reason rather than + /// per cache. Index with PackedCacheHeapReason. Excludes nothing — the + /// NotOptedIn slot is populated here but omitted from `heap_bytes`. + std::array(PackedCacheHeapReason::Count)> + heap_bytes_by_reason{}; +}; + +/** One live cache instance and its own counters. */ +struct PackedCacheEntry { + /// Cache file path. Empty for the shared heap-only instance handed to + /// callers that never configured one. + std::string path; + PackedCacheStats stats; +}; + +/** + * Aggregate plus the per-instance breakdown behind it. + * + * Both come from one pass, so the summary and the detail always describe the + * same instant. The breakdown exists because the aggregate alone cannot be + * attributed: a process running several models folds them into one number, so + * a fallback in one model is indistinguishable from a fallback in another. + * The manager already keys caches by path — this stops discarding that. + * + * Takes no per-instance lock: the counters are atomics, so this never waits + * on a model compile. + */ +struct PackedCacheReport { + /// Summed counters. `failure` / `last_errno` are left unset here; read them + /// from the `dominant_fallback` entry so they stay tied to one cache. + PackedCacheStats aggregate; + /// Sorted by path, so repeated calls agree regardless of map iteration + /// order. + std::vector per_cache; + /// Index into `per_cache` of the cache that best explains a fallback: the + /// largest heap contributor, or if none allocated, the first cache in + /// HeapFallback. -1 when nothing fell back. + int32_t dominant_fallback{-1}; +}; + +PackedCacheReport get_packed_cache_report(); + } // namespace executorch::backends::xnnpack diff --git a/backends/xnnpack/runtime/XNNWeightsCache.cpp b/backends/xnnpack/runtime/XNNWeightsCache.cpp index 34479c1c369..3b6b6d310e9 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCache.cpp @@ -80,14 +80,28 @@ static T read_le(const uint8_t* src) { // Open the cache file and take an advisory exclusive lock. Returns the // fd, or -1 if open/flock failed (logs the failure). The caller decides // how to recover (typically: skip the mmap path for this init). -static int open_locked(const std::string& path, int flags) { +// out_errno receives the errno of whichever call failed. Reading errno at the +// call site does not work: the flock path closes the fd first, and close() (or +// ET_LOG) can overwrite it. +static int open_locked(const std::string& path, int flags, int* out_errno) { + if (out_errno != nullptr) { + *out_errno = 0; + } int fd = open(path.c_str(), flags, 0600); if (fd < 0) { - ET_LOG(Error, "open(%s) failed (errno=%d)", path.c_str(), errno); + const int err = errno; + if (out_errno != nullptr) { + *out_errno = err; + } + ET_LOG(Error, "open(%s) failed (errno=%d)", path.c_str(), err); return -1; } if (flock(fd, LOCK_EX | LOCK_NB) != 0) { - ET_LOG(Error, "flock(%s) failed (errno=%d)", path.c_str(), errno); + const int err = errno; + if (out_errno != nullptr) { + *out_errno = err; + } + ET_LOG(Error, "flock(%s) failed (errno=%d)", path.c_str(), err); close(fd); return -1; } @@ -128,6 +142,72 @@ void XNNWeightsCache::reset_for_fresh_write() { } #endif +void XNNWeightsCache::record_cache_failure( + PackedCacheFailure failure, + int err) noexcept { + state_.store( + static_cast(PackedCacheState::HeapFallback), + std::memory_order_relaxed); + failure_.store(static_cast(failure), std::memory_order_relaxed); + last_errno_.store(err, std::memory_order_relaxed); +} + +PackedCacheStats XNNWeightsCache::stats() const noexcept { + PackedCacheStats out; + out.state = + static_cast(state_.load(std::memory_order_relaxed)); + out.failure = + static_cast(failure_.load(std::memory_order_relaxed)); + out.last_errno = last_errno_.load(std::memory_order_relaxed); + out.file_bytes = file_bytes_.load(std::memory_order_relaxed); + out.mapped_bytes = mapped_bytes_.load(std::memory_order_relaxed); + int64_t worst = 0; + for (size_t i = 0; i < heap_bytes_by_reason_.size(); ++i) { + const int64_t bytes = + heap_bytes_by_reason_[i].load(std::memory_order_relaxed); + out.heap_bytes_by_reason[i] = bytes; + if (i == static_cast(PackedCacheHeapReason::NotOptedIn)) { + continue; // intended heap use, not a fallback + } + out.heap_bytes += bytes; + if (bytes > worst) { + worst = bytes; + out.heap_reason = static_cast(i); + } + } + return out; +} + +void XNNWeightsCache::record_heap_alloc( + size_t n, + PackedCacheHeapReason reason) noexcept { + // Re-bucket every reason to NotOptedIn when no path was configured. Such an + // instance is the shared heap-only cache handed to callers that never asked + // for file backing; counting its bytes as a fallback would inflate the + // metric for whichever model in the process *did* opt in. + // packed_cache_path_ is set once before the instance is published and never + // mutated, so this read needs no synchronization. + const PackedCacheHeapReason bucket = + packed_cache_path_.empty() ? PackedCacheHeapReason::NotOptedIn : reason; + heap_bytes_by_reason_[static_cast(bucket)].fetch_add( + static_cast(n), std::memory_order_relaxed); +} + +void XNNWeightsCache::record_mapped_alloc(size_t n) noexcept { + mapped_bytes_.fetch_add(static_cast(n), std::memory_order_relaxed); +} + +void XNNWeightsCache::mark_cache_file_backed() noexcept { + // Only ever upgrades Disabled -> FileBacked. A fallback already recorded + // describes memory the process is carrying, so a later success must not + // mask it. + int32_t expected = static_cast(PackedCacheState::Disabled); + state_.compare_exchange_strong( + expected, + static_cast(PackedCacheState::FileBacked), + std::memory_order_relaxed); +} + Error XNNWeightsCache::initialize_for_runtime( MemoryAllocator* runtime_allocator, const NamedDataMap* named_data_map) { @@ -147,7 +227,13 @@ Error XNNWeightsCache::initialize_for_runtime( // where fresh-write→save→re-init re-enters load_packed_cache and // double-mmaps the same file. if (!name_to_packed_data_metadata_.empty()) { - packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR); + int open_errno = 0; + packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR, &open_errno); + if (packed_file_fd_ < 0) { + record_cache_failure(PackedCacheFailure::OpenFailed, open_errno); + } else { + mark_cache_file_backed(); + } return Error::Ok; } @@ -160,27 +246,47 @@ Error XNNWeightsCache::initialize_for_runtime( "Loaded packed weight cache: %s (%zu entries)", packed_cache_path_.c_str(), name_to_packed_data_metadata_.size()); - packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR); + int open_errno = 0; + packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR, &open_errno); + // Loaded entries are already mmap'd, so reads stay file-backed even if the + // write fd could not be reopened. Record the errno anyway: without it a + // partial cache silently re-packs to heap every launch with no reason. + if (packed_file_fd_ < 0) { + record_cache_failure(PackedCacheFailure::OpenFailed, open_errno); + } + mark_cache_file_backed(); return Error::Ok; } // Fresh write. Skip O_TRUNC in open_locked so a concurrent holder's // mmap stays valid; truncate explicitly only after we hold the lock. - packed_file_fd_ = open_locked(packed_cache_path_, O_RDWR | O_CREAT); + int create_errno = 0; + packed_file_fd_ = + open_locked(packed_cache_path_, O_RDWR | O_CREAT, &create_errno); if (packed_file_fd_ < 0) { + const int err = create_errno; + ET_LOG( + Error, + "open(O_RDWR|O_CREAT) failed for %s (errno=%d); heap fallback this init", + packed_cache_path_.c_str(), + err); + record_cache_failure(PackedCacheFailure::OpenFailed, err); return Error::Ok; } if (ftruncate(packed_file_fd_, 0) != 0) { + const int err = errno; ET_LOG( Error, "ftruncate(0) failed for %s (errno=%d); heap fallback this init", packed_cache_path_.c_str(), - errno); + err); + record_cache_failure(PackedCacheFailure::TruncateFailed, err); close(packed_file_fd_); packed_file_fd_ = -1; return Error::Ok; } reset_for_fresh_write(); + mark_cache_file_backed(); ET_LOG( Info, "Opened packed weight file for writing: %s", @@ -394,6 +500,10 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { // instead of re-packing into heap (dirty memory) every time. if (context->last_lookup_unnamed_ || (context->loaded_from_disk_ && !seed_mismatch_repack)) { + context->record_heap_alloc( + n, + context->last_lookup_unnamed_ ? PackedCacheHeapReason::UnnamedConstant + : PackedCacheHeapReason::RepackAfterLoad); return context->reserve_space_heap(n); } if (context->packed_file_fd_ >= 0) { @@ -403,13 +513,16 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { size_t map_size = (n + page_size - 1) & ~(page_size - 1); if (ftruncate(context->packed_file_fd_, file_offset + map_size) != 0) { + const int err = errno; ET_LOG( Error, "reserve_space ftruncate to %zu failed (errno=%d)", file_offset + map_size, - errno); + err); + context->record_cache_failure(PackedCacheFailure::GrowFailed, err); close(context->packed_file_fd_); context->packed_file_fd_ = -1; + context->record_heap_alloc(n, PackedCacheHeapReason::GrowFailed); return context->reserve_space_heap(n); } @@ -421,13 +534,16 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { context->packed_file_fd_, file_offset); if (ptr == MAP_FAILED) { + const int err = errno; ET_LOG( Error, "reserve_space mmap %zu bytes failed (errno=%d)", map_size, - errno); + err); + context->record_cache_failure(PackedCacheFailure::MmapFailed, err); close(context->packed_file_fd_); context->packed_file_fd_ = -1; + context->record_heap_alloc(n, PackedCacheHeapReason::MmapFailed); return context->reserve_space_heap(n); } @@ -439,12 +555,16 @@ void* XNNWeightsCache::reserve_space(XNNWeightsCache* context, size_t n) { kPackedAllocationAlignment); context->packed_file_used_ = file_offset + map_size; + // n, not map_size: the heap side records the raw request too, and the + // heap:mapped ratio is only meaningful if both measure the same thing. + context->record_mapped_alloc(n); context->file_ptr_to_region_index_[ptr] = context->mmap_regions_.size(); context->mmap_regions_.push_back({ptr, map_size}); context->ptr_to_file_offset_[ptr] = file_offset; return ptr; } #endif + context->record_heap_alloc(n, PackedCacheHeapReason::NoFileBacking); return context->reserve_space_heap(n); } @@ -609,6 +729,8 @@ Error XNNWeightsCache::save_packed_index() { // trailer drops the old entry. Monitoring file_bytes over time tells // us when GC or a size cap is needed. const size_t file_bytes = index_start + buf.size(); + file_bytes_.store( + static_cast(file_bytes), std::memory_order_relaxed); ET_LOG( Info, "Saved packed weight index: %u entries at offset %zu, file_bytes=%zu", @@ -699,6 +821,9 @@ bool XNNWeightsCache::load_packed_cache() { } mmap_regions_.push_back({map, file_size}); + // Bytes actually referenced by the index. Less than file_size whenever an + // earlier run re-packed a name and orphaned its old bytes. + size_t loaded_bytes = 0; const uint8_t* cursor = static_cast(map) + index_start; const uint8_t* end = static_cast(map) + index_region_end; @@ -766,6 +891,7 @@ bool XNNWeightsCache::load_packed_cache() { meta.in_current_runtime = false; meta.from_load = true; meta.seed = seed; + loaded_bytes += static_cast(data_size); name_to_packed_data_metadata_[name] = meta; } @@ -783,6 +909,15 @@ bool XNNWeightsCache::load_packed_cache() { mmap_regions_at_last_save_ = mmap_regions_.size(); mmap_regions_synced_ = mmap_regions_.size(); loaded_from_disk_ = true; + // Success path only: the truncated-entry branch above munmaps and rolls + // back, so counting at the mmap call would over-report. + // + // loaded_bytes, not file_size. The file is append-only, so a same-name + // re-pack leaves the old bytes behind; file_size counts those orphans and + // would inflate mapped_bytes against heap_bytes. Without this a warm launch + // reports heap=0/mapped=0/file=0, identical to the feature being off. + record_mapped_alloc(loaded_bytes); + file_bytes_.store(static_cast(file_size), std::memory_order_relaxed); return true; #else return false; diff --git a/backends/xnnpack/runtime/XNNWeightsCache.h b/backends/xnnpack/runtime/XNNWeightsCache.h index f584199e307..8ac023f63c4 100644 --- a/backends/xnnpack/runtime/XNNWeightsCache.h +++ b/backends/xnnpack/runtime/XNNWeightsCache.h @@ -10,10 +10,13 @@ #include +#include #include #include #include #include +#include +#include #include #include #include @@ -52,6 +55,13 @@ struct PackedDataMeta { uint32_t seed{0}; }; +// Telemetry types live in XNNPACKBackend.h — hosts read them without pulling +// in xnnpack.h through this header. +using xnnpack::PackedCacheFailure; +using xnnpack::PackedCacheHeapReason; +using xnnpack::PackedCacheState; +using xnnpack::PackedCacheStats; + class XNNWeightsCache { public: XNNWeightsCache(); @@ -162,7 +172,50 @@ class XNNWeightsCache { return instance_mutex_; } + /** + * Outcome of the file-backed path for this instance. HeapFallback is + * sticky: once an init has been served from heap the instance keeps + * reporting it, because that is the memory the process is actually + * carrying for the rest of its life. + */ + PackedCacheStats stats() const noexcept; + private: + /** Record a fallback. Overwrites any previous failure for this instance. */ + void record_cache_failure(PackedCacheFailure failure, int err) noexcept; + /** Note a working file-backed path; never downgrades a recorded fallback. */ + void mark_cache_file_backed() noexcept; + /** Attribute `n` packed bytes to heap under `reason`. */ + void record_heap_alloc(size_t n, PackedCacheHeapReason reason) noexcept; + /** Attribute `n` packed bytes to the mmap'd file. */ + void record_mapped_alloc(size_t n) noexcept; + + // Telemetry counters. Written from the XNNPACK callbacks (which run under + // the caller-held instance mutex) and read by hosts through + // XNNWeightsCacheManager::aggregate_stats() with no lock at all — atomics, + // not the mutex, are what make that read safe. The mutex is held across the + // whole of xnn_create_runtime, so a telemetry read that waited on it could + // stall an inference thread for the length of a model compile. + // + // relaxed ordering throughout: these are independent accumulators, and a + // reader that observes one field slightly ahead of another still gets a + // usable picture. There is no invariant spanning them. + // + // Cumulative for the instance's lifetime — delete_packed_data and + // full_unload do not decrement. Decrementing would need a ptr -> reason map + // kept alive purely for telemetry, and hosts sample right after a load or a + // generate, before anything is released, so the two agree in practice. + // Read them as "bytes this cache ever packed", not current residency. + std::atomic state_{static_cast(PackedCacheState::Disabled)}; + std::atomic failure_{static_cast(PackedCacheFailure::None)}; + std::atomic last_errno_{0}; + std::atomic file_bytes_{0}; + std::atomic mapped_bytes_{0}; + std::array< + std::atomic, + static_cast(PackedCacheHeapReason::Count)> + heap_bytes_by_reason_{}; + static constexpr uint32_t kCacheMagic = 0x58505743; // "XPWC" // Bump when the on-disk layout (footer or per-entry record) changes. // v2: per-entry seed added — old v1 files don't carry seeds and would diff --git a/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp b/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp index 0f122aa8ab0..116f7d9f5a0 100644 --- a/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp +++ b/backends/xnnpack/runtime/XNNWeightsCacheManager.cpp @@ -10,6 +10,7 @@ #include +#include #include #include @@ -77,6 +78,98 @@ Error XNNWeightsCacheManager::save_all() { return first_err; } +xnnpack::PackedCacheReport XNNWeightsCacheManager::report() const { + // Snapshot path + instance under the owning mutexes, then read the counters + // without XNNWeightsCache::mutex(). That mutex is held across the whole of + // xnn_create_runtime, so waiting on it here would let a telemetry read stall + // an inference thread for the length of a model compile. + std::vector< + std::pair>> + live; + { + std::scoped_lock lock(meta_mutex_); + live.reserve(caches_.size()); + for (const auto& entry : caches_) { + if (auto cache = entry.second.lock()) { + live.emplace_back(entry.first, std::move(cache)); + } + } + } + { + std::scoped_lock lock(empty_path_mutex_); + if (auto cache = empty_path_cache_.lock()) { + live.emplace_back(std::string{}, std::move(cache)); + } + } + // caches_ is unordered; sort so the report and the index into it are stable + // across calls. + std::sort(live.begin(), live.end(), [](const auto& a, const auto& b) { + return a.first < b.first; + }); + + xnnpack::PackedCacheReport out; + out.per_cache.reserve(live.size()); + for (const auto& [path, cache] : live) { + out.per_cache.push_back(xnnpack::PackedCacheEntry{path, cache->stats()}); + } + + int64_t best_heap = -1; + int32_t first_fallback = -1; + for (size_t i = 0; i < out.per_cache.size(); ++i) { + const auto& s = out.per_cache[i].stats; + auto& agg = out.aggregate; + agg.file_bytes += s.file_bytes; + agg.heap_bytes += s.heap_bytes; + agg.mapped_bytes += s.mapped_bytes; + for (size_t r = 0; r < s.heap_bytes_by_reason.size(); ++r) { + agg.heap_bytes_by_reason[r] += s.heap_bytes_by_reason[r]; + } + // A fallback anywhere is the reportable outcome: if any cache on this + // process went to heap, the process is carrying that memory. + if (s.state == delegate::PackedCacheState::HeapFallback) { + agg.state = s.state; + if (first_fallback < 0) { + first_fallback = static_cast(i); + } + } else if ( + s.state == delegate::PackedCacheState::FileBacked && + agg.state == delegate::PackedCacheState::Disabled) { + agg.state = s.state; + } + if (s.heap_bytes > best_heap) { + best_heap = s.heap_bytes; + if (s.heap_bytes > 0) { + out.dominant_fallback = static_cast(i); + } + } + } + // A cache can fail before it ever allocates, so fall back to the first + // cache in HeapFallback rather than reporting no explanation at all. + if (out.dominant_fallback < 0) { + out.dominant_fallback = first_fallback; + } + + // Argmax over per-reason totals summed across caches — not the local reason + // of whichever cache allocated the most, which can disagree with the global + // picture when one cache mixes reasons. + int64_t worst = 0; + for (size_t r = 0; r < out.aggregate.heap_bytes_by_reason.size(); ++r) { + if (r == static_cast(delegate::PackedCacheHeapReason::NotOptedIn)) { + continue; // intended heap use, not a fallback + } + if (out.aggregate.heap_bytes_by_reason[r] > worst) { + worst = out.aggregate.heap_bytes_by_reason[r]; + out.aggregate.heap_reason = + static_cast(r); + } + } + return out; +} + +delegate::PackedCacheStats XNNWeightsCacheManager::aggregate_stats() const { + return report().aggregate; +} + size_t XNNWeightsCacheManager::live_count() const { std::scoped_lock lock(meta_mutex_); size_t count = 0; diff --git a/backends/xnnpack/runtime/XNNWeightsCacheManager.h b/backends/xnnpack/runtime/XNNWeightsCacheManager.h index c35285b6337..39e5c4b711a 100644 --- a/backends/xnnpack/runtime/XNNWeightsCacheManager.h +++ b/backends/xnnpack/runtime/XNNWeightsCacheManager.h @@ -54,6 +54,22 @@ class XNNWeightsCacheManager { * expired weak_ptrs. */ runtime::Error save_all(); + /** + * Worst outcome across live caches, for host telemetry. A HeapFallback + * anywhere wins over a FileBacked elsewhere: if any cache on this process + * went to heap, the process is carrying that memory. `file_bytes` sums + * across live instances. + */ + delegate::PackedCacheStats aggregate_stats() const; + + /** + * Aggregate plus the per-instance breakdown, from a single pass so the two + * always agree. Callers that need to attribute a fallback to a specific + * model use the breakdown; the aggregate answers "is this process carrying + * heap memory at all". + */ + xnnpack::PackedCacheReport report() const; + /** Test-only: count of live (non-expired) entries. */ size_t live_count() const; diff --git a/backends/xnnpack/test/runtime/test_weight_cache.cpp b/backends/xnnpack/test/runtime/test_weight_cache.cpp index d2c079c057a..384edbaec8c 100644 --- a/backends/xnnpack/test/runtime/test_weight_cache.cpp +++ b/backends/xnnpack/test/runtime/test_weight_cache.cpp @@ -11,13 +11,19 @@ #include #include #include + #include #include #include #include +#include using namespace ::testing; +using executorch::backends::xnnpack::get_packed_cache_report; +using executorch::backends::xnnpack::packed_cache_path_option_key; +using executorch::backends::xnnpack::PackedCacheHeapReason; +using executorch::backends::xnnpack::save_weight_cache_on_disk_option_key; using executorch::backends::xnnpack::weight_cache_option_key; using executorch::backends::xnnpack::workspace_sharing_mode_option_key; using executorch::backends::xnnpack::WorkspaceSharingMode; @@ -148,3 +154,42 @@ TEST(RuntimeSpec, OverridesGlobalWeightCache) { get_option(xnnpack_backend_key, read_option); ASSERT_EQ(std::get(read_option.value), true); } + +TEST(PackedCacheStats, GlobalAccessorReachesTheBackendSingleton) { + executorch::runtime::runtime_init(); + + // The wiring under test is get_packed_cache_report() -> the registered + // backend instance -> XnnpackBackendOptions -> XNNWeightsCacheManager. + // Hosts call only this entry point, and nothing else in the suite exercises + // it. Absolute values depend on what else has run in the process, so this + // asserts reachability and invariants rather than specific counts. + const auto report = get_packed_cache_report(); + const auto& stats = report.aggregate; + + EXPECT_GE(stats.heap_bytes, 0); + EXPECT_GE(stats.mapped_bytes, 0); + EXPECT_GE(stats.file_bytes, 0); + EXPECT_LT( + static_cast(stats.heap_reason), + static_cast(PackedCacheHeapReason::Count)); + EXPECT_NE(stats.heap_reason, PackedCacheHeapReason::NotOptedIn) + << "NotOptedIn is excluded from heap_bytes and must never be reported"; + + // The breakdown must be able to attribute the aggregate: several models + // share a process, and a single folded number cannot say which one fell + // back. + for (const auto& entry : report.per_cache) { + EXPECT_GE(entry.stats.heap_bytes, 0); + EXPECT_GE(entry.stats.mapped_bytes, 0); + } +} + +// NOTE: the warm path — load_packed_cache() succeeding and its mapped bytes +// being counted — is deliberately NOT covered here. Producing a loadable +// cache file needs a model with packed weights and a populated index; the +// models wired into this target (ModuleAddLarge / ModuleSubLarge) are +// elementwise and pack nothing, so a cold run writes a zero-entry trailer +// that load_packed_cache correctly rejects. Covering it needs ModuleLinear +// plus its external .ptd and a NamedDataMap, as test_xnn_data_separation +// does. Until then the warm path is verified on device by the +// PackedWeights log line reporting non-zero mapped/cache_file. diff --git a/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp b/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp index 06bc74211ad..c2764776f6f 100644 --- a/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp +++ b/backends/xnnpack/test/runtime/test_xnn_weights_cache_manager.cpp @@ -14,8 +14,11 @@ #include #include +#include #include +#include #include +#include #include #include @@ -31,14 +34,34 @@ class XNNWeightsCacheManagerTest : public ::testing::Test { manager_ = std::make_unique(); } + void TearDown() override { + for (const auto& path : temp_paths_) { + std::remove(path.c_str()); + } + } + + // Unique per test and per process. A leftover file from an earlier run + // flips initialize_for_runtime between the load and fresh-create branches, + // and two concurrent runs would race on the same path. + std::string TempPath(const char* tag) { + const auto* info = ::testing::UnitTest::GetInstance()->current_test_info(); + auto path = std::string(::testing::TempDir()) + "xnnwc_" + info->name() + + "_" + tag + "_" + std::to_string(static_cast(::getpid())) + + ".bin"; + std::remove(path.c_str()); + temp_paths_.push_back(path); + return path; + } + std::unique_ptr manager_; + std::vector temp_paths_; }; // --- Core dedup semantics --- TEST_F(XNNWeightsCacheManagerTest, SamePathReturnsSameInstance) { - auto a = manager_->get_or_create("/tmp/test_cache_same.bin"); - auto b = manager_->get_or_create("/tmp/test_cache_same.bin"); + auto a = manager_->get_or_create(TempPath("same")); + auto b = manager_->get_or_create(TempPath("same")); ASSERT_TRUE(a.ok()); ASSERT_TRUE(b.ok()); EXPECT_EQ(a.get().get(), b.get().get()) @@ -46,8 +69,8 @@ TEST_F(XNNWeightsCacheManagerTest, SamePathReturnsSameInstance) { } TEST_F(XNNWeightsCacheManagerTest, DifferentPathsReturnDifferentInstances) { - auto a = manager_->get_or_create("/tmp/test_cache_a.bin"); - auto b = manager_->get_or_create("/tmp/test_cache_b.bin"); + auto a = manager_->get_or_create(TempPath("a")); + auto b = manager_->get_or_create(TempPath("b")); ASSERT_TRUE(a.ok()); ASSERT_TRUE(b.ok()); EXPECT_NE(a.get().get(), b.get().get()) @@ -85,7 +108,7 @@ TEST_F(XNNWeightsCacheManagerTest, EmptyPathRecreatedAfterAllRefsDrop) { TEST_F(XNNWeightsCacheManagerTest, EmptyPathDoesNotShareWithMmapPath) { auto empty = manager_->get_or_create(""); - auto mmap = manager_->get_or_create("/tmp/test_cache_isolation.bin"); + auto mmap = manager_->get_or_create(TempPath("isolation")); ASSERT_TRUE(empty.ok()); ASSERT_TRUE(mmap.ok()); // Empty-path cache stays separate from any mmap-path cache — @@ -100,7 +123,7 @@ TEST_F(XNNWeightsCacheManagerTest, EmptyPathDoesNotShareWithMmapPath) { TEST_F(XNNWeightsCacheManagerTest, ExpiredEntryDoesNotLeak) { { - auto a = manager_->get_or_create("/tmp/test_cache_expire.bin"); + auto a = manager_->get_or_create(TempPath("expire")); ASSERT_TRUE(a.ok()); EXPECT_EQ(manager_->live_count(), 1u); } @@ -112,13 +135,13 @@ TEST_F(XNNWeightsCacheManagerTest, ExpiredEntryDoesNotLeak) { TEST_F(XNNWeightsCacheManagerTest, ExpiredEntryRecreatedOnNextCall) { void* first_addr = nullptr; { - auto a = manager_->get_or_create("/tmp/test_cache_recreate.bin"); + auto a = manager_->get_or_create(TempPath("recreate")); ASSERT_TRUE(a.ok()); first_addr = a.get().get(); } // Address re-use is allowed but not required; the only guarantee is // that we get a usable instance, not a dangling shared_ptr. - auto b = manager_->get_or_create("/tmp/test_cache_recreate.bin"); + auto b = manager_->get_or_create(TempPath("recreate")); ASSERT_TRUE(b.ok()); ASSERT_NE(b.get(), nullptr); // Live count should be 1 again — the stale entry was erased and @@ -136,15 +159,18 @@ TEST_F(XNNWeightsCacheManagerTest, ConcurrentSamePathSameInstance) { std::vector threads; threads.reserve(kThreads); std::atomic ready{0}; + // Resolve the path up front: TempPath() appends to temp_paths_, which is + // not safe to call from the racing threads. + const std::string race_path = TempPath("race"); for (int i = 0; i < kThreads; ++i) { - threads.emplace_back([this, &results, &ready, i] { + threads.emplace_back([this, &results, &ready, &race_path, i] { // Spin to maximize the chance of true concurrent entry into // get_or_create. ready.fetch_add(1, std::memory_order_acq_rel); while (ready.load(std::memory_order_acquire) < kThreads) { std::this_thread::yield(); } - auto r = manager_->get_or_create("/tmp/test_cache_race.bin"); + auto r = manager_->get_or_create(race_path); ASSERT_TRUE(r.ok()); results[i] = r.get(); }); @@ -169,10 +195,14 @@ TEST_F(XNNWeightsCacheManagerTest, ConcurrentDifferentPathsIndependent) { std::vector> results(kThreads); std::vector threads; threads.reserve(kThreads); + std::vector paths; + paths.reserve(kThreads); for (int i = 0; i < kThreads; ++i) { - threads.emplace_back([this, &results, i] { - std::string path = "/tmp/test_cache_diff_" + std::to_string(i) + ".bin"; - auto r = manager_->get_or_create(path); + paths.push_back(TempPath(("diff_" + std::to_string(i)).c_str())); + } + for (int i = 0; i < kThreads; ++i) { + threads.emplace_back([this, &results, &paths, i] { + auto r = manager_->get_or_create(paths[i]); ASSERT_TRUE(r.ok()); results[i] = r.get(); }); @@ -195,8 +225,8 @@ TEST_F(XNNWeightsCacheManagerTest, SaveAllNoLiveInstancesIsOk) { } TEST_F(XNNWeightsCacheManagerTest, SaveAllWalksLiveCaches) { - auto a = manager_->get_or_create("/tmp/test_cache_save_a.bin"); - auto b = manager_->get_or_create("/tmp/test_cache_save_b.bin"); + auto a = manager_->get_or_create(TempPath("save_a")); + auto b = manager_->get_or_create(TempPath("save_b")); ASSERT_TRUE(a.ok()); ASSERT_TRUE(b.ok()); EXPECT_EQ(manager_->live_count(), 2u); @@ -208,7 +238,7 @@ TEST_F(XNNWeightsCacheManagerTest, SaveAllWalksLiveCaches) { TEST_F(XNNWeightsCacheManagerTest, SaveAllSkipsExpiredEntries) { { - auto a = manager_->get_or_create("/tmp/test_cache_save_expired.bin"); + auto a = manager_->get_or_create(TempPath("save_expired")); ASSERT_TRUE(a.ok()); } // The entry's weak_ptr is now expired. save_all must not crash on @@ -220,7 +250,268 @@ TEST_F(XNNWeightsCacheManagerTest, SaveAllSkipsExpiredEntries) { // --- Path is set on the instance before publishing --- TEST_F(XNNWeightsCacheManagerTest, NonEmptyPathRegistersInMap) { - auto a = manager_->get_or_create("/tmp/test_cache_register.bin"); + auto a = manager_->get_or_create(TempPath("register")); ASSERT_TRUE(a.ok()); EXPECT_EQ(manager_->live_count(), 1u); } + +// --- Packed-cache telemetry (host-visible fallback reporting) --- + +TEST_F(XNNWeightsCacheManagerTest, StatsDisabledWhenNoCacheEverUsed) { + const auto stats = manager_->aggregate_stats(); + EXPECT_EQ( + stats.state, + executorch::backends::xnnpack::delegate::PackedCacheState::Disabled); + EXPECT_EQ(stats.last_errno, 0); + EXPECT_EQ(stats.file_bytes, 0); + EXPECT_EQ(stats.heap_bytes, 0); + EXPECT_EQ(stats.mapped_bytes, 0); +} + +TEST_F(XNNWeightsCacheManagerTest, StatsReportOpenFailureWithErrno) { + // A path whose parent directory does not exist: open(O_RDWR|O_CREAT) fails + // with ENOENT, which is the same branch a full disk takes with ENOSPC. + auto cache = manager_->get_or_create("/nonexistent_dir_xnnwc/cache.bin"); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok) + << "a fallback must stay non-fatal"; + } + + const auto report = manager_->report(); + EXPECT_EQ( + report.aggregate.state, + executorch::backends::xnnpack::delegate::PackedCacheState::HeapFallback) + << "an unusable path must be reported as a heap fallback, not silently"; + + // failure/errno are deliberately absent from the aggregate: they belong to + // one cache. dominant_fallback names which one, even though this cache + // never allocated (open failed before any pack). + ASSERT_GE(report.dominant_fallback, 0); + const auto& dominant = + report.per_cache[static_cast(report.dominant_fallback)].stats; + EXPECT_EQ( + dominant.failure, + executorch::backends::xnnpack::delegate::PackedCacheFailure::OpenFailed); + EXPECT_NE(dominant.last_errno, 0) + << "errno is what distinguishes ENOSPC from a path problem"; + EXPECT_EQ( + report.aggregate.failure, + executorch::backends::xnnpack::delegate::PackedCacheFailure::None) + << "the aggregate must not adopt one cache's failure"; +} + +TEST_F(XNNWeightsCacheManagerTest, HeapReasonIsGlobalArgmaxNotPerCache) { + // Two caches whose local dominant reasons disagree with the global one. + // Cache A: a failed grow plus a smaller unnamed pack. Cache B: unnamed only, + // larger in total than A's grow contribution. Summing per cache would report + // A's reason; summing per reason reports UnnamedConstant, which is correct. + auto a = manager_->get_or_create(TempPath("argmax_a")); + auto b = manager_->get_or_create(TempPath("argmax_b")); + ASSERT_TRUE(a.ok()); + ASSERT_TRUE(b.ok()); + + const auto unnamed_pack = [](XNNWeightsCache* c, size_t n) { + auto* provider = c->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, n), nullptr); + }; + + { + std::lock_guard lock(a.get()->mutex()); + ASSERT_EQ(a.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + unnamed_pack(a.get().get(), 8192); + } + { + std::lock_guard lock(b.get()->mutex()); + ASSERT_EQ(b.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + unnamed_pack(b.get().get(), 16384); + unnamed_pack(b.get().get(), 16384); + } + + const auto report = manager_->report(); + EXPECT_EQ( + report.aggregate.heap_reason, + executorch::backends::xnnpack::delegate::PackedCacheHeapReason:: + UnnamedConstant); + const auto unnamed_idx = + static_cast(executorch::backends::xnnpack::delegate:: + PackedCacheHeapReason::UnnamedConstant); + EXPECT_EQ( + report.aggregate.heap_bytes_by_reason[unnamed_idx], + report.aggregate.heap_bytes) + << "per-reason totals must sum to the same heap_bytes"; +} + +TEST_F(XNNWeightsCacheManagerTest, PerCacheIsSortedByPathForStableIndices) { + // dominant_fallback is an index into per_cache, so the order must not + // depend on unordered_map iteration. + auto z = manager_->get_or_create(TempPath("zzz")); + auto a = manager_->get_or_create(TempPath("aaa")); + ASSERT_TRUE(z.ok()); + ASSERT_TRUE(a.ok()); + + const auto report = manager_->report(); + ASSERT_GE(report.per_cache.size(), 2u); + for (size_t i = 1; i < report.per_cache.size(); ++i) { + EXPECT_LE(report.per_cache[i - 1].path, report.per_cache[i].path); + } +} + +TEST_F(XNNWeightsCacheManagerTest, HeapFallbackWinsOverFileBackedInAggregate) { + auto bad = manager_->get_or_create("/nonexistent_dir_xnnwc/cache.bin"); + auto good = manager_->get_or_create(TempPath("stats_ok")); + ASSERT_TRUE(bad.ok()); + ASSERT_TRUE(good.ok()); + { + std::lock_guard lock(good.get()->mutex()); + ASSERT_EQ(good.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + } + { + std::lock_guard lock(bad.get()->mutex()); + ASSERT_EQ(bad.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + } + + EXPECT_EQ( + manager_->aggregate_stats().state, + executorch::backends::xnnpack::delegate::PackedCacheState::HeapFallback) + << "if any live cache fell back, the process is carrying that memory"; +} + +// A binary "did the file open" flag is not enough: a cache can load +// successfully and still serve most of its packed bytes from heap. These +// cover the byte accounting that distinguishes the two. + +TEST_F(XNNWeightsCacheManagerTest, MappedBytesCountedWhenFileBacked) { + auto cache = manager_->get_or_create(TempPath("bytes_mapped")); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache.get()->get(); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_GT(stats.mapped_bytes, 0); + EXPECT_EQ(stats.heap_bytes, 0) << "the healthy case must report zero heap"; +} + +TEST_F(XNNWeightsCacheManagerTest, HeapBytesAttributedToUnnamedConstant) { + auto cache = manager_->get_or_create(TempPath("bytes_unnamed")); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache.get()->get(); + // A look_up whose kernel pointer was never named marks the next + // reserve_space as an unnamed constant, which routes to heap. + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_GT(stats.heap_bytes, 0) << "heap bytes must be counted, not hidden"; + EXPECT_EQ( + stats.heap_reason, + executorch::backends::xnnpack::delegate::PackedCacheHeapReason:: + UnnamedConstant); +} + +TEST_F(XNNWeightsCacheManagerTest, FileBackedStateDoesNotImplyZeroHeap) { + // The case a state flag alone reports as healthy: the file opened fine, so + // state is FileBacked, yet packed bytes still went to heap. Only the byte + // split makes that visible. + auto cache = manager_->get_or_create(TempPath("bytes_split")); + ASSERT_TRUE(cache.ok()); + { + std::lock_guard lock(cache.get()->mutex()); + ASSERT_EQ(cache.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache.get()->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_EQ( + stats.state, + executorch::backends::xnnpack::delegate::PackedCacheState::FileBacked) + << "the file opened, so state alone looks healthy"; + EXPECT_GT(stats.heap_bytes, 0) + << "but heap bytes are non-zero — this is what state alone hides"; +} + +TEST_F(XNNWeightsCacheManagerTest, AggregateStatsTakesNoInstanceLock) { + // aggregate_stats() must not wait on XNNWeightsCache::mutex(): that mutex is + // held across all of xnn_create_runtime, so a telemetry read that blocked on + // it would stall inference for the length of a model compile. + auto cache = manager_->get_or_create(TempPath("nolock")); + ASSERT_TRUE(cache.ok()); + std::lock_guard held(cache.get()->mutex()); + const auto stats = manager_->aggregate_stats(); // must not deadlock + EXPECT_EQ(stats.heap_bytes, 0); +} + +TEST_F(XNNWeightsCacheManagerTest, EmptyPathHeapIsNotCountedAsFallback) { + // The shared heap-only instance handed to callers that never configured a + // path. Its heap use is intended, so it must not inflate heap_bytes for a + // model in the same process that did opt in. + auto opted_out = manager_->get_or_create(""); + ASSERT_TRUE(opted_out.ok()); + { + std::lock_guard lock(opted_out.get()->mutex()); + ASSERT_EQ( + opted_out.get()->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = opted_out.get()->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 4096), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_EQ(stats.heap_bytes, 0) + << "a cache that never opted into file backing is not a fallback"; +} + +TEST_F(XNNWeightsCacheManagerTest, OptedInHeapStillCountedAlongsideOptedOut) { + // Both kinds live at once: only the opted-in instance's heap bytes count. + auto opted_out = manager_->get_or_create(""); + auto opted_in = manager_->get_or_create(TempPath("mixed")); + ASSERT_TRUE(opted_out.ok()); + ASSERT_TRUE(opted_in.ok()); + for (auto* cache : {opted_out.get().get(), opted_in.get().get()}) { + std::lock_guard lock(cache->mutex()); + ASSERT_EQ(cache->initialize_for_runtime(nullptr, nullptr), Error::Ok); + auto* provider = cache->get(); + int dummy = 0; + xnn_weights_cache_look_up_key key{}; + key.kernel = &dummy; + key.bias = nullptr; + provider->look_up(provider->context, &key); + ASSERT_NE(provider->reserve_space(provider->context, 8192), nullptr); + } + + const auto stats = manager_->aggregate_stats(); + EXPECT_GT(stats.heap_bytes, 0) << "the opted-in instance's heap must count"; + EXPECT_EQ( + stats.heap_reason, + executorch::backends::xnnpack::delegate::PackedCacheHeapReason:: + UnnamedConstant) + << "and NotOptedIn must never win the argmax"; +} From 8dc72ce36bd4ee04531b1226c662f6183250952f Mon Sep 17 00:00:00 2001 From: Jiri Ocenasek Date: Tue, 1 Sep 2026 09:19:16 +0200 Subject: [PATCH 073/190] NXP backend: switch to neutron compiler --- backends/nxp/README.md | 6 +- ...add_batch_size_for_3d_input_pool_2d_ops.py | 2 +- .../ops_converters/mean_dim_converter.py | 2 +- .../ops_converters/permute_copy_converter.py | 2 +- .../upsample_bilinear2d_converter.py | 2 +- .../upsample_nearest2d_converter.py | 2 +- ...manager.py => neutron_compiler_manager.py} | 72 ++++++++++++------- backends/nxp/backend/neutron_map.py | 22 +++--- backends/nxp/backend/neutron_target_spec.py | 12 ++-- backends/nxp/nxp_backend.py | 14 ++-- backends/nxp/runtime/NeutronDriver.h | 6 +- .../test_context_sensitive_delegation.py | 4 +- ...er.py => test_neutron_compiler_manager.py} | 12 ++-- .../node_converter/test_cat_converter.py | 2 +- .../backends/nxp/nxp-kernel-selection.md | 8 +-- docs/source/backends/nxp/nxp-overview.md | 2 +- docs/source/backends/nxp/nxp-partitioner.rst | 6 +- docs/source/backends/nxp/nxp-profiling.md | 12 ++-- docs/source/backends/nxp/nxp-quantization.md | 2 +- .../nxp/tutorials/nxp-basic-tutorial.md | 4 +- examples/nxp/aot_neutron_compile.py | 2 +- 21 files changed, 109 insertions(+), 87 deletions(-) rename backends/nxp/backend/{neutron_converter_manager.py => neutron_compiler_manager.py} (68%) rename backends/nxp/tests/generic_tests/{test_neutron_converter_manager.py => test_neutron_compiler_manager.py} (82%) diff --git a/backends/nxp/README.md b/backends/nxp/README.md index 4188dd8f810..05204fd7c78 100644 --- a/backends/nxp/README.md +++ b/backends/nxp/README.md @@ -33,9 +33,9 @@ The eIQ Neutron NPU Backend should be considered as prototype quality at this mo improvements. NXP and the ExecuTorch community is actively developing this codebase. ## Neutron Backend implementation and SW architecture -Neutron Backend uses the eIQ Neutron Converter as ML compiler to compile the delegated subgraph to Neutron microcode. -The Neutron Converter accepts the ML model in LiteRT format, for the **eIQ Neutron N3** class therefore the Neutron Backend -uses the LiteRT flatbuffers format as IR between the ExecuTorch and Neutron Converter ML compiler. +Neutron Backend uses the eIQ Neutron Compiler as ML compiler to compile the delegated subgraph to Neutron microcode. +The Neutron Compiler accepts the ML model in LiteRT format, for the **eIQ Neutron N3** class therefore the Neutron Backend +uses the LiteRT flatbuffers format as IR between the ExecuTorch and Neutron Compiler ML compiler. ## Layout * `backend/ir/` - TFLite/LiteRT based IR to represent the Edge Subgraph, taken from onnx2tflite code base and extended to diff --git a/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py index 7435b3b6969..cbd000befbb 100644 --- a/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py +++ b/backends/nxp/aten_passes/add_batch_size_for_3d_input_pool_2d_ops.py @@ -13,7 +13,7 @@ class AddBatchSizeFor3DInputPool2DOps(PassBase): """Adds batch size dimension for aten.adaptive_avg_pool2d.default, aten.avg_pool2d.default - and aten.max_pool2d.default ops with 3D input, as the Neutron Converter is unable to convert these ops with 3D input. + and aten.max_pool2d.default ops with 3D input, as the Neutron Compiler is unable to compile these ops with 3D input. │ ┌──────▼──────┐ diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py index 4d03e5e97b7..b9c5a11ca35 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/mean_dim_converter.py @@ -41,7 +41,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if is_alone_in_partition and keepdim and all(input_shape[d] == 1 for d in dim): - # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the + # The operator is a no-op, so the Neutron Compiler will skip it. If it's the only node in the # partition, the graph would end up empty. return False diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py index 3e4908c2211..6f0eb6ad757 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/permute_copy_converter.py @@ -400,7 +400,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if has_static_input and is_alone_in_partition: # Transpose with a static input is a no-op on Neutron. If it was the only operator in the partition, - # Neutron Converter would produce and empty graph, so delegation is prohibited. + # Neutron Compiler would produce and empty graph, so delegation is prohibited. return False return True diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py index 2f0126e9aae..544638f5b5b 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_bilinear2d_converter.py @@ -39,7 +39,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if is_alone_in_partition and input_shape == output_shape: - # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the + # The operator is a no-op, so the Neutron Compiler will skip it. If it's the only node in the # partition, the graph would end up empty. return False diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py index a3c8db14f51..cdf8ad7d46e 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/upsample_nearest2d_converter.py @@ -41,7 +41,7 @@ def supports_partitioning_result( is_alone_in_partition = cls.is_node_alone_in_partition(node, partition_list) if is_alone_in_partition and h_scale == w_scale == 1: - # The operator is a no-op, so the Neutron Converter will skip it. If it's the only node in the + # The operator is a no-op, so the Neutron Compiler will skip it. If it's the only node in the # partition, the graph would end up empty. return False diff --git a/backends/nxp/backend/neutron_converter_manager.py b/backends/nxp/backend/neutron_compiler_manager.py similarity index 68% rename from backends/nxp/backend/neutron_converter_manager.py rename to backends/nxp/backend/neutron_compiler_manager.py index 92b4e25a5de..f7ca416cb0c 100644 --- a/backends/nxp/backend/neutron_converter_manager.py +++ b/backends/nxp/backend/neutron_compiler_manager.py @@ -6,19 +6,35 @@ import logging import multiprocessing import os +import warnings try: - from eiq_neutron_sdk import neutron_converter, neutron_library_utils + from eiq_neutron_sdk import neutron_compiler, neutron_library_utils + + _USING_NEUTRON_COMPILER = True except ImportError: - raise RuntimeError( - "eIQ Neutron SDK not found. To install it, run 'examples/nxp/setup.sh'." - ) + try: + from eiq_neutron_sdk import ( + neutron_converter as neutron_compiler, + neutron_library_utils, + ) + + _USING_NEUTRON_COMPILER = False + warnings.warn( + "The support for eIQ Neutron SDK <= 3.2.2 will be removed in future releases.", + DeprecationWarning, + stacklevel=2, + ) + except ImportError: + raise RuntimeError( + "eIQ Neutron SDK not found. To install it, run 'examples/nxp/setup.sh'." + ) def _build_compilation_context(compilation_opts): """Build a CompilationContext from a plain dict of options.""" - cctx = neutron_converter.CompilationContext() - cctx.targetOpts = neutron_converter.getNeutronTarget(compilation_opts["target"]) + cctx = neutron_compiler.CompilationContext() + cctx.targetOpts = neutron_compiler.getNeutronTarget(compilation_opts["target"]) cctx.compilationOpts.minNumOpsPerGraph = compilation_opts["minNumOpsPerGraph"] cctx.compilationOpts.excludeGraphPasses = compilation_opts["excludeGraphPasses"] cctx.compilationOpts.fetchConstantsToSRAM = compilation_opts["fetchConstantsToSRAM"] @@ -37,19 +53,22 @@ def _build_compilation_context(compilation_opts): return cctx -def convert_unsafe(tflite_model, compilation_opts, queue): +def compile_unsafe(tflite_model, compilation_opts, queue): """ - Run neutron_converter on given tflite_model with the provided compilation options. + Run neutron_compiler on given tflite_model with the provided compilation options. This routine is supposed to run in a separate process. - If properly finished, the output queue contains the converted model, - otherwise the neutron_converter exits and the output queue is empty. + If properly finished, the output queue contains the compiled model, + otherwise the neutron_compiler exits and the output queue is empty. """ cctx = _build_compilation_context(compilation_opts) - model_converted = neutron_converter.convertModel(list(tflite_model), cctx) - queue.put(model_converted) + if _USING_NEUTRON_COMPILER: + model_compiled = neutron_compiler.compileModel(list(tflite_model), cctx) + else: + model_compiled = neutron_compiler.convertModel(list(tflite_model), cctx) + queue.put(model_compiled) -class NeutronConverterManager: +class NeutronCompilerManager: """ Manager for conversion of TFLite model in flatbuffers format into TFLite model that contains NeutronGraph nodes. @@ -69,8 +88,8 @@ def _rename_partition_kernel_selection_file(delegation_tag): except OSError: logging.error("Failed to rename partition kernel selection file.") - def get_converter(self): - return neutron_converter + def get_compiler(self): + return neutron_compiler def get_library_utils(self): return neutron_library_utils @@ -84,7 +103,7 @@ def verify_target(self, target: str): f"Target `{target}` is not a valid target. Must be one of `{valid_targets}`." ) - def convert( + def compile( self, tflite_model: bytes, target: str, @@ -93,9 +112,9 @@ def convert( use_profiling: bool = False, ) -> bytes: """ - Call Neutron Converter. + Call Neutron Compiler. - :param tflite_model: A generic TFLite model to be converted. + :param tflite_model: A generic TFLite model to be compiled. :param target: The target platform. :param delegation_tag: The delegation tag of model partition. :param fetch_constants_to_sram: Add microcode that fetches weights from external memory. @@ -104,7 +123,7 @@ def convert( :return: TFLite model with Neutron microcode as bytes. """ - # Neutron converter crashes if we provide invalid target -> verify. + # Neutron compiler crashes if we provide invalid target -> verify. self.verify_target(target) compilation_opts = { @@ -124,7 +143,7 @@ def convert( queue = multiprocessing.Manager().Queue() process = multiprocessing.Process( - target=convert_unsafe, + target=compile_unsafe, args=(tflite_model, compilation_opts, queue), ) process.start() @@ -132,20 +151,23 @@ def convert( if queue.empty(): # signals the unsafe task did not run till the end raise RuntimeError( - f"Neutron converter module terminated unexpectedly with exit code {process.exitcode}" + f"Neutron compiler module terminated unexpectedly with exit code {process.exitcode}" ) - model_converted = queue.get() + model_compiled = queue.get() process.close() except (EOFError, OSError, TypeError) as e: # Multiprocessing failed (likely due to environment restrictions) # Fall back to direct execution logging.warning( - f"Multiprocessing not available ({e}), running neutron converter directly" + f"Multiprocessing not available ({e}), running neutron compiler directly" ) cctx = _build_compilation_context(compilation_opts) - model_converted = neutron_converter.convertModel(list(tflite_model), cctx) + if _USING_NEUTRON_COMPILER: + model_compiled = neutron_compiler.compileModel(list(tflite_model), cctx) + else: + model_compiled = neutron_compiler.convertModel(list(tflite_model), cctx) if self.dump_kernel_selection_code: self._rename_partition_kernel_selection_file(delegation_tag) - return bytes(model_converted) + return bytes(model_compiled) diff --git a/backends/nxp/backend/neutron_map.py b/backends/nxp/backend/neutron_map.py index da497565726..a5becafc4f0 100644 --- a/backends/nxp/backend/neutron_map.py +++ b/backends/nxp/backend/neutron_map.py @@ -91,15 +91,15 @@ def get_tensors_name(tensors: str) -> list[str]: class NeutronMap: - """Mapping between Neutron, TFLite, and Edge operators based on the Neutron converter log. + """Mapping between Neutron, TFLite, and Edge operators based on the Neutron compiler log. - Parses the Neutron converter log to extract information about TFLite nodes and Neutron subgraphs. + Parses the Neutron compiler log to extract information about TFLite nodes and Neutron subgraphs. Maps TFLite operators to corresponding Neutron operators. Maps Edge operators to Neutron operators via the Edge-to-TFLite mapping. Attributes: - tflite_nodes (list[Node]): TFLite node information extracted from the converter log. - neutron_subgraphs (list[SubgraphInfo]): Neutron subgraph information extracted from the converter log. + tflite_nodes (list[Node]): TFLite node information extracted from the compiler log. + neutron_subgraphs (list[SubgraphInfo]): Neutron subgraph information extracted from the compiler log. neutron_graphs (list[int]): Indices of final Neutron graphs derived from neutron_subgraphs. edge_to_tflite_map (dict[int, tuple[int, ...]]): Mapping from Edge operators to TFLite operators. edge_to_neutron_map (dict[int, tuple[int, ...]]): Mapping from Edge operators to Neutron operators. @@ -118,12 +118,12 @@ class NeutronMap: tflite_to_neutron_map: dict[int, tuple[int, ...]] def __init__( - self, neutron_converter_log: str, edge_to_tflite_map: dict[int, tuple[int, ...]] + self, neutron_compiler_log: str, edge_to_tflite_map: dict[int, tuple[int, ...]] ) -> None: - """Initialize neutron map from neutron converter log. + """Initialize neutron map from neutron compiler log. - :param neutron_converter_log: neutron converter log obtained during model conversion. It should contain - original tflite graph and neutron graph dump. To add these dumps to converter log the dumpAfterImport and + :param neutron_compiler_log: neutron compiler log obtained during model compilation. It should contain + original tflite graph and neutron graph dump. To add these dumps to compiler log the dumpAfterImport and dumpAfterGenerate flags have to be set to "console". """ super().__init__() @@ -134,12 +134,12 @@ def __init__( self.tflite_to_neutron_map = {} self.edge_to_neutron_map = {} self.neutron_kernels_num = 0 - self._split_profiling_log(neutron_converter_log) + self._split_profiling_log(neutron_compiler_log) def _split_profiling_log(self, log: str) -> None: """Process profiling log to split it into original TFLite and converted Neutron nodes. - :param log: Neutron converter log obtained during model conversion, containing the original + :param log: Neutron compiler log obtained during model compilation, containing the original TFLite graph and Neutron graph dump. :return: None. Sets class attributes tflite_nodes and neutron_subgraphs with node information. """ @@ -175,7 +175,7 @@ def _split_profiling_log(self, log: str) -> None: def _get_neutron_subgraphs(self, graph_dump: str) -> list[SubgraphInfo]: """Parse Neutron graph dump and extract subgraph information. - :param graph_dump: String containing the Neutron graph dump from the converter log. + :param graph_dump: String containing the Neutron graph dump from the compiler log. :return: List of SubgraphInfo objects containing subgraph metadata and operator nodes. """ diff --git a/backends/nxp/backend/neutron_target_spec.py b/backends/nxp/backend/neutron_target_spec.py index 5a75caf9a75..a51d437f060 100644 --- a/backends/nxp/backend/neutron_target_spec.py +++ b/backends/nxp/backend/neutron_target_spec.py @@ -8,8 +8,8 @@ from enum import Enum import torch -from executorch.backends.nxp.backend.neutron_converter_manager import ( - NeutronConverterManager, +from executorch.backends.nxp.backend.neutron_compiler_manager import ( + NeutronCompilerManager, ) from executorch.exir.dialects._ops import ops as exir_ops from torch.fx import Node @@ -98,10 +98,10 @@ class NeutronTargetSpec: def __init__(self, target: str): - converter_manager = NeutronConverterManager() - converter_manager.verify_target(target) - neutron_converter = converter_manager.get_converter() - self.neutron_target = neutron_converter.getNeutronTarget(target) + compiler_manager = NeutronCompilerManager() + compiler_manager.verify_target(target) + neutron_compiler = compiler_manager.get_compiler() + self.neutron_target = neutron_compiler.getNeutronTarget(target) if self.is_subsystem(): raise ValueError( diff --git a/backends/nxp/nxp_backend.py b/backends/nxp/nxp_backend.py index 2f4bb07316f..16430afada5 100644 --- a/backends/nxp/nxp_backend.py +++ b/backends/nxp/nxp_backend.py @@ -25,8 +25,8 @@ EdgeProgramToIRConverter, ) from executorch.backends.nxp.backend.ir.conversion_config import ConversionConfig -from executorch.backends.nxp.backend.neutron_converter_manager import ( - NeutronConverterManager, +from executorch.backends.nxp.backend.neutron_compiler_manager import ( + NeutronCompilerManager, ) from executorch.backends.nxp.backend.neutron_map import NeutronMap @@ -86,10 +86,10 @@ def neutron_compile_spec( :param use_neutron_for_format_conversion: If True, the EdgeProgramToIRConverter will insert `Transpose` ops to ensure that the IO matches the executorch partition, which will be delegated to Neutron. - :param fetch_constants_to_sram: If True, the Neutron Converter will insert microinstructions to prefetch weights + :param fetch_constants_to_sram: If True, the Neutron Compiler will insert microinstructions to prefetch weights from FLASH to SRAM. This should be used when the whole model does not fit into SRAM. - :param dump_kernel_selection_code: Whether Neutron converter dumps kernel selection code. - :param use_profiling: If true Neutron Converter will enable profiling for neutron delegated model + :param dump_kernel_selection_code: Whether Neutron Compiler dumps kernel selection code. + :param use_profiling: If true Neutron Compiler will enable profiling for neutron delegated model :return: self for method chaining """ @@ -282,9 +282,9 @@ def preprocess( # noqa C901 ) with capture_fd_output() as tmp: - neutron_model = NeutronConverterManager( + neutron_model = NeutronCompilerManager( dump_kernel_selection_code - ).convert( + ).compile( tflite_model, target, delegation_tag, diff --git a/backends/nxp/runtime/NeutronDriver.h b/backends/nxp/runtime/NeutronDriver.h index 9280cd6adba..0d29f8b172b 100644 --- a/backends/nxp/runtime/NeutronDriver.h +++ b/backends/nxp/runtime/NeutronDriver.h @@ -42,7 +42,7 @@ typedef void* NeutronModelHandle; typedef struct { /// Neutron microcode buffer address. - /// The Neutron microcode is generated by the Neutron converter tool. + /// The Neutron microcode is generated by the Neutron compiler tool. /// The microcode buffer, 16 bytes aligned, is allocated and initialized by /// the application or ML framework. The microcode buffer is passed by /// reference to the Neutron firmware. The microcode buffer is specific for a @@ -50,7 +50,7 @@ typedef struct { const void* microcode; /// Neutron weights buffer address. - /// The Neutron weights is generated by the Neutron converter tool. + /// The Neutron weights is generated by the Neutron compiler tool. /// The weights buffer, 16 bytes aligned, is allocated and initialized by the /// application or ML framework. The weights buffer address is passed by /// reference to the Neutron-firmware. The weights buffer is specific for a @@ -58,7 +58,7 @@ typedef struct { const void* weights; /// Neutron kernels buffer address. - /// The Neutron kernels are generated by the Neutron converter tool. + /// The Neutron kernels are generated by the Neutron compiler tool. /// The kernels buffer, 16 bytes aligned, is allocated and initialized by the /// application or ML framework. The kernels buffer address is passed by /// reference to the Neutron-firmware. The kernels buffer is specific for a diff --git a/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py index 1b1aaed897e..677312f9483 100644 --- a/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py +++ b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py @@ -122,10 +122,10 @@ def test_noop_partitions__concatenate_one_tensor_and_add_zeros(): @pytest.mark.xfail( strict=True, - reason="Neutron Converter currently supports these 2 noops in sequence.", + reason="Neutron Compiler currently supports these 2 noops in sequence.", ) def test_noop_partitions__concatenate_one_tensor_and_add_zeros__forced_delegation(): - # When the noop `Concatenate` and noop `Add` are in sequence, Neutron Converter supports them. This edge case is + # When the noop `Concatenate` and noop `Add` are in sequence, Neutron Compiler supports them. This edge case is # not reflected in our logic. But as this edge case is extremely rare (and even if it ever happened in a real # model, the consequences would be minimal), fixing it is not a priority. diff --git a/backends/nxp/tests/generic_tests/test_neutron_converter_manager.py b/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py similarity index 82% rename from backends/nxp/tests/generic_tests/test_neutron_converter_manager.py rename to backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py index 8bd3446da7a..fab2c8afc06 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_converter_manager.py +++ b/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py @@ -6,8 +6,8 @@ import multiprocessing import pickle -from executorch.backends.nxp.backend.neutron_converter_manager import ( - NeutronConverterManager, +from executorch.backends.nxp.backend.neutron_compiler_manager import ( + NeutronCompilerManager, ) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.models import LinearModule @@ -17,23 +17,23 @@ def test_conv2d_neutron_conversion__prefetching(mocker): model = LinearModule(True) input_shape = (1, 1, 32, 32) - converter_spy = mocker.spy(NeutronConverterManager, "convert") + compiler_spy = mocker.spy(NeutronCompilerManager, "compile") _ = to_quantized_edge_program( model, input_shape, fetch_constants_to_sram=True ).exported_program() - neutron_model_prefetch = converter_spy.spy_return + neutron_model_prefetch = compiler_spy.spy_return _ = to_quantized_edge_program( model, input_shape, fetch_constants_to_sram=False ).exported_program() - neutron_model_regular = converter_spy.spy_return + neutron_model_regular = compiler_spy.spy_return assert len(neutron_model_prefetch) != len( neutron_model_regular ), "The weight prefetching flag does not make a difference!" -def test_convert_unsafe_args_are_picklable(mocker): +def test_compile_unsafe_args_are_picklable(mocker): """Verify that all args passed to `multiprocessing.Process` are picklable. The subprocess uses forkserver/spawn in some environments, which requires diff --git a/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py index 30983f4a666..ac541ed0b77 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_cat_converter.py @@ -136,7 +136,7 @@ def test__different_shapes__channels_first(self, mocker, request, dim, num_input lower_run_compare(model, input_shapes, graph_verifier, request) def test__single_input__alone_in_partition__not_delegated(self): - # The operator is a noop, and there is no other op in the model. The Neutron Converter would produce an empty + # The operator is a noop, and there is no other op in the model. The Neutron Compiler would produce an empty # graph, so the `cat` is not delegated. input_shape = [ModelInputSpec((2, 3, 5))] model = CatModule(1) diff --git a/docs/source/backends/nxp/nxp-kernel-selection.md b/docs/source/backends/nxp/nxp-kernel-selection.md index 307f06d1d02..4cd3bfc33f7 100644 --- a/docs/source/backends/nxp/nxp-kernel-selection.md +++ b/docs/source/backends/nxp/nxp-kernel-selection.md @@ -1,12 +1,12 @@ # NXP eIQ Neutron Kernel Selective Kernel Registration The NXP ExecuTorch backend supports selective Neutron kernel registration for `Neutron-C` targets, which reduces the -size of the Neutron Firmware. During the backend's conversion to the Neutron representation by the Neutron Converter, +size of the Neutron Firmware. During the backend's conversion to the Neutron representation by the Neutron Compiler, microcode for the Neutron accelerator is generated. The microcode consists of kernel calls executed by the Neutron Driver. The code for kernel call functions is distributed in the Neutron Firmware. -The `eiq_neutron_sdk.neutron_converter` optionally generates a `*_kernel_selection.c` file, registering +The `eiq_neutron_sdk.neutron_compiler` optionally generates a `*_kernel_selection.c` file, registering only kernels that are required for a particular model or, in the case of ExecuTorch, a delegated subgraph. This `*_kernel_selection.c`, when used during application linking, takes precedence over the default list of registered kernels in the Neutron Firmware, and allows the linker to include only the necessary Neutron kernels. @@ -21,7 +21,7 @@ final application with unused code. In memory-constrained environments, you can deployed models. This way you can reduce the size of the final application by linking only selected kernels, used in one or more models. -The feature works as follows: The Neutron Converter with the appropriate flag exports a kernel selection file for each +The feature works as follows: The Neutron Compiler with the appropriate flag exports a kernel selection file for each converted subgraph, the kernel selection files are then merged and ready to be included in the MCUXpresso SDK to use for a selection-only build. @@ -52,7 +52,7 @@ python -m eiq_neutron_sdk.neutron_library_utils.merge_kernel_selection_code \ -output-file merged_kernel_selection.c ``` -Each particular model must be converted by the same Neutron converter version, so the `*_kernel_selection.c` files +Each particular model must be compiled by the same Neutron Compiler version, so the `*_kernel_selection.c` files share the same version. ## MCUXpresso SDK build with kernel selection diff --git a/docs/source/backends/nxp/nxp-overview.md b/docs/source/backends/nxp/nxp-overview.md index 604b480e614..87249faa05c 100644 --- a/docs/source/backends/nxp/nxp-overview.md +++ b/docs/source/backends/nxp/nxp-overview.md @@ -46,7 +46,7 @@ For a quick overview how to convert a custom PyTorch model, take a look at our [ An example runtime application using the eIQ NSYS (eIQ Neutron Simulator) is available [examples/nxp/executor_runner](https://github.com/pytorch/executorch/blob/main/examples/nxp/executor_runner/), described in the tutorial [Getting started with eIQ Neutron NPU ExecuTorch backend](tutorials/nxp-basic-tutorial.md) -To learn how to run the converted model on the NXP hardware, use one of our example projects on using ExecuTorch runtime from MCUXpresso IDE example projects list. +To learn how to run the delegated model on the NXP hardware, use one of our example projects on using ExecuTorch runtime from MCUXpresso IDE example projects list. For more finegrained tutorial, visit [this manual page](https://mcuxpresso.nxp.com/mcuxsdk/latest/html/middleware/eiq/executorch/docs/nxp/topics/example_applications.html). For guideline how to update the eIQ Neutron Runtime on MCUXpresso SDK, follow the instructions from the eIQ Neutron SDK package `docs/NeutronSDKUserGuide.md` available diff --git a/docs/source/backends/nxp/nxp-partitioner.rst b/docs/source/backends/nxp/nxp-partitioner.rst index 4ddc38fb2db..b24f4dde78e 100644 --- a/docs/source/backends/nxp/nxp-partitioner.rst +++ b/docs/source/backends/nxp/nxp-partitioner.rst @@ -27,9 +27,9 @@ Following fields can be set: * `extra_flags` - Extra flags for the Neutron compiler. * `operators_not_to_delegate` - List of operators that will not be delegated. * `use_neutron_for_format_conversion` - If True, let the eIQ Neutron NPU to handle conversion between channel-first (NCHW) and channel-last (NHWC) data formats. That is the Neutron backend will insert `Transpose` ops to ensure that the IO matches the executorch partition, which will be delegated to Neutron. -* `fetch_constants_to_sram` - If True, the Neutron Converter will insert microinstructions to prefetch weights from FLASH to SRAM. This should be used when the whole model does not fit into SRAM on Neutron-C devices, like i.MX RT700. -* `dump_kernel_selection_code` - Whether Neutron converter dumps kernel selection code, which is used by the selective kernel registration, see :doc:`Neutron Firmware Kernel Selection support `. -* `use_profiling` - If true Neutron Converter will enable profiling for neutron delegated model. +* `fetch_constants_to_sram` - If True, the Neutron Compiler will insert microinstructions to prefetch weights from FLASH to SRAM. This should be used when the whole model does not fit into SRAM on Neutron-C devices, like i.MX RT700. +* `dump_kernel_selection_code` - Whether Neutron Compiler dumps kernel selection code, which is used by the selective kernel registration, see :doc:`Neutron Firmware Kernel Selection support `. +* `use_profiling` - If true Neutron Compiler will enable profiling for neutron delegated model. ------------------------- Custom Delegation Options diff --git a/docs/source/backends/nxp/nxp-profiling.md b/docs/source/backends/nxp/nxp-profiling.md index 17e352e479d..1b8f3b49035 100644 --- a/docs/source/backends/nxp/nxp-profiling.md +++ b/docs/source/backends/nxp/nxp-profiling.md @@ -7,16 +7,16 @@ to provide visibility into delegated operator execution time. There are three steps required to obtain profiling results for an NXP‑delegated model: -* Convert the model with profiling support enabled. +* Delegate the model with profiling support enabled. * Generate the artifacts consumed by the Developer Tools (`ETRecord`, `ETDump`). * Create and run the Inspector class to consume these artifacts and print the results. --- -## Convert a model with the profiling support +## Delegate a model with the profiling support Profiling data is generated only for a **profilable** model. -To convert a model with profiling enabled, the `--use-profiling` flag must be set. +To delegate a model with profiling enabled, the `--use-profiling` flag must be set. See the `aot_neutron_compile.py` example and its [README](https://github.com/pytorch/executorch/blob/main/examples/nxp/README.md) @@ -89,7 +89,7 @@ A full implementation is available in [aot_neutron_compile.py](https://github.com/pytorch/executorch/blob/main/examples/nxp/aot_neutron_compile.py). The `--use_profiling` flag is used to create a **profilable** model and the corresponding `ETRecord` file -(see [Convert a model with profiling support](#convert-a-model-with-profiling-support) for the full command). +(see [Delegate a model with the profiling support](#delegate-a-model-with-the-profiling-support) for the full command). --- @@ -102,7 +102,7 @@ The next step is to generate an `ETDump`. An `ETDump` contains runtime data coll To generate an `ETDump`, ensure that the ExecuTorch runtime library is integrated with the Developer Tools and built with the `ET_EVENT_TRACER_ENABLED` flag enabled. -Only models converted with profiling support will produce an `ETDump` containing execution times for all Neutron +Only models delegated with profiling support will produce an `ETDump` containing execution times for all Neutron operators. Otherwise, the dump will include only the final delegate execution time. Neutron software provides a profiling mechanism that logs individual operator execution times to a dedicated runtime @@ -176,7 +176,7 @@ The [Inspector](https://docs.pytorch.org/executorch/1.0/model-inspector.html) AP contents of `ETRecord` and `ETDump`, enabling developers to gain insights into model architecture and performance statistics. -`ETRecord` is an optional argument used to obtain a mapping between the original model and the converted Neutron model. +`ETRecord` is an optional argument used to obtain a mapping between the original model and the delegated Neutron model. An `ETDump` generated on the board contains metadata for each Neutron operator, including its unique identifier. To visualize this metadata in the Inspector results table, set the `include_delegate_debug_data = True` argument. diff --git a/docs/source/backends/nxp/nxp-quantization.md b/docs/source/backends/nxp/nxp-quantization.md index 61cd00632df..ef038e47f44 100644 --- a/docs/source/backends/nxp/nxp-quantization.md +++ b/docs/source/backends/nxp/nxp-quantization.md @@ -248,7 +248,7 @@ Moving from PTQ to QAT check-list: #### Known limitations of QAT In the current ExecuTorch/TorchAO implementation, there is an issue when quantizing biasless convolutions during QAT. -The pipeline produces a non‑quantized empty bias, which causes the Neutron Converter to fail. +The pipeline produces a non‑quantized empty bias, which causes the Neutron Compiler to fail. To mitigate this issue, use the `QuantizeFusedConvBnBiasAtenPass` post‑quantization: ```python diff --git a/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md b/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md index b2e07bb7c1d..264a3fc2dfc 100644 --- a/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md +++ b/docs/source/backends/nxp/tutorials/nxp-basic-tutorial.md @@ -13,14 +13,14 @@ You need to install the ExecuTorch. Please follow the tutorial to install the Ex In addition to this, you will need to install the eIQ Neutron Simulator, called NSYS, -and the Neutron Converter for generating the byte-code for the eIQ Neutron NPU, +and the Neutron Compiler for generating the byte-code for the eIQ Neutron NPU, during the model conversion in ExecuTorch AoT flow. To install the eIQ Neutron dependencies, run: ```bash examples/nxp/setup.sh ``` This will install: -* Neutron Converter, for converting the Neutron IR to Neutron byte-code +* Neutron Compiler, for compiling the Neutron IR to Neutron byte-code * eIQ Neutron SDK, containing the eIQ Neutron runtimes (driver and firmware) for various NXP SoC and simulator * eIQ NSYS, the Neutron behavioral simulator diff --git a/examples/nxp/aot_neutron_compile.py b/examples/nxp/aot_neutron_compile.py index 697953f7946..7b7eeff8d19 100644 --- a/examples/nxp/aot_neutron_compile.py +++ b/examples/nxp/aot_neutron_compile.py @@ -245,7 +245,7 @@ def _get_arg_parser(): required=False, default=False, action="store_true", - help="During conversion to Neutron microcode by Neutron Converter, a kernel selection file will be dumped in " + help="During compilation to Neutron microcode by Neutron Compiler, a kernel selection file will be dumped in " "the working directory. This file can be used for reduction of Neutron Firmware size in the built app." "See `docs/source/backends/nxp/nxp-kernel-selection.md` for details.", ) From 162a6ac11cfe0118107a916a53672d8624f02676 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Tue, 8 Sep 2026 11:28:17 -0700 Subject: [PATCH 074/190] Restore pull request cancellation on the untagged workflows (#22610) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit #21622 added `github.ref_type == 'branch' && github.sha` to the concurrency key of all 29 workflows, to stop ciflow tag re-pushes flooding the fleet. That term is also truthy on `pull_request`, where `github.sha` is the merge commit and changes with every push, so pushes stopped cancelling superseded runs — before #21622 a branch would show seven of eight runs cancelled, after it three full runs overlap and all complete. These five workflows have no tag trigger, so there is nothing for the guard to protect and the key goes back to pytorch's `pull.yml` form, which they now match exactly; `lint.yml` was costing the most, running on every push to every pull request. `github.run_id` on the dispatch term comes from the same key and fixes `check-labels.yml`, where two dispatches for different pull requests currently cancel each other. The 27 tag-exposed workflows are left alone, since #21622 is doing its job on their tag path. Authored with Claude Code. --- .github/workflows/build-cmsis-pack.yml | 2 +- .github/workflows/build-presets.yml | 2 +- .github/workflows/check-c10-sync.yml | 2 +- .github/workflows/check-labels.yml | 2 +- .github/workflows/lint.yml | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) diff --git a/.github/workflows/build-cmsis-pack.yml b/.github/workflows/build-cmsis-pack.yml index c9c40d670f1..9974ceea4ee 100644 --- a/.github/workflows/build-cmsis-pack.yml +++ b/.github/workflows/build-cmsis-pack.yml @@ -40,7 +40,7 @@ on: type: string concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: diff --git a/.github/workflows/build-presets.yml b/.github/workflows/build-presets.yml index 89d36cd6b0a..48c745121be 100644 --- a/.github/workflows/build-presets.yml +++ b/.github/workflows/build-presets.yml @@ -9,7 +9,7 @@ on: workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: diff --git a/.github/workflows/check-c10-sync.yml b/.github/workflows/check-c10-sync.yml index 73a4837adf6..db62e15cc4b 100644 --- a/.github/workflows/check-c10-sync.yml +++ b/.github/workflows/check-c10-sync.yml @@ -8,7 +8,7 @@ on: - runtime/core/portable_type/c10/** concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: diff --git a/.github/workflows/check-labels.yml b/.github/workflows/check-labels.yml index ebaa38cf0bd..b103f734a20 100644 --- a/.github/workflows/check-labels.yml +++ b/.github/workflows/check-labels.yml @@ -25,7 +25,7 @@ on: required: true concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }} cancel-in-progress: true jobs: diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index ea9fa1dfbe1..9b217d00292 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -11,7 +11,7 @@ on: workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: From f2dc5e3a7a9980b8d088cd88270844addb386294 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Tue, 8 Sep 2026 11:31:48 -0700 Subject: [PATCH 075/190] Restore pull request cancellation in pull.yml (#22609) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two things, one enabling the other. **`ciflow/trunk/*` on `pull.yml` is a mistake.** It arrived in #18172, which moved the ARM Cortex-M size test from `trunk.yml` to `pull.yml` and carried trunk's tag trigger across with the job. `ciflow/trunk` should run `trunk.yml`; `pull.yml` already runs on every pull request via `pull_request`, so those tag runs were duplicates — on #22549 the tag run had 120 jobs against the pull request run's 121. `trunk.yml` keeps the trigger. **With no tag push reaching this workflow, the concurrency key can go back to pytorch's.** #21622 added `github.ref_type == 'branch' && github.sha` to stop ciflow tag re-pushes flooding the fleet, but that term is truthy on `pull_request` too, where `github.sha` is the merge commit and changes with every push — so pushes stopped cancelling superseded runs. Before #21622: eight consecutive runs on one branch, seven cancelled. After: three overlapping full runs on `bump-torch-pin-2.14`, none cancelled. The key now matches pytorch's `pull.yml` exactly, `github.run_id` included so two manual dispatches don't cancel each other. Authored with Claude Code. --- .github/workflows/pull.yml | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 30d3849a25e..2522cca3113 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -6,12 +6,10 @@ on: branches: - main - release/* - tags: - - ciflow/trunk/* workflow_dispatch: concurrency: - group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}-${{ github.event_name == 'workflow_dispatch' && github.run_id }}-${{ github.event_name == 'schedule' }} cancel-in-progress: true jobs: From faf73b9859e05a1c86f8fbd5f99e437ff88a792c Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 8 Sep 2026 12:19:33 -0700 Subject: [PATCH 076/190] Fix C++ standard selection for ExecuTorch tests (#22615) ### Summary Resolve ATen dependencies returned through selects so affected targets compile with C++20. Keep embedded tests on C++17 while overriding Apple tests to C++20 because generated XPlugins sources use `std::span` and `std::ranges`. Test Plan: - Built representative Apple and fbcode test targets - Passed ATen detection unit tests - Passed lint Authored with assistance from Codex. --- .../executorch/build/runtime_wrapper.bzl | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/shim_et/xplat/executorch/build/runtime_wrapper.bzl b/shim_et/xplat/executorch/build/runtime_wrapper.bzl index 85c390fe6c2..0a2162996a7 100644 --- a/shim_et/xplat/executorch/build/runtime_wrapper.bzl +++ b/shim_et/xplat/executorch/build/runtime_wrapper.bzl @@ -158,6 +158,14 @@ def _is_aten_target(kwargs): "libtorch_python", "torch-core-cpp", ] + aten_resolved_external_deps = [ + "c10", + "libtorch", + "libtorch_python", + "torch-core-cpp", + ] + # The ATen-flavored gtest and gmock names resolve to the same internal + # labels as their ordinary variants, so only their short names are unique. for key in ["external_deps", "exported_external_deps"]: for dep in kwargs.get(key) or []: if dep in aten_external_deps: @@ -166,12 +174,17 @@ def _is_aten_target(kwargs): # A target can also name one of those through external_dep_location, which # hands back the resolved label and puts it in an ordinary dep list. aten_targets = [] - for name in aten_external_deps: + + def _note_aten_targets(targets): + for target in targets: + if target not in aten_targets: + aten_targets.append(target) + return targets + + for name in aten_resolved_external_deps: resolved = env.resolve_external_dep(name) if resolved != env.EXTERNAL_DEP_FALLTHROUGH: - for target in resolved: - if target not in aten_targets: - aten_targets.append(target) + selects.apply(obj = resolved, function = _note_aten_targets) # A dep list can be a select(), so collect through selects.apply rather than # walking it. The lists it holds are the same shape either way. @@ -199,7 +212,8 @@ def _patch_test_compiler_flags(kwargs, aten_mode = False): kwargs["compiler_flags"] = [] # A test that compiles against ATen needs C++20, which PyTorch's headers - # require. Every other test stays at C++17, which the embedded builds use. + # require. Other tests stay at C++17 for embedded builds, but Apple plugin + # generation also requires C++20. name = kwargs.get("name", "") is_aten_test = ( aten_mode or @@ -214,6 +228,11 @@ def _patch_test_compiler_flags(kwargs, aten_mode = False): kwargs["compiler_flags"] += [ "-std=c++17", ] + if env.is_xplat(): + kwargs["fbobjc_compiler_flags"] = kwargs.get( + "fbobjc_compiler_flags", + [], + ) + ["-std=c++20"] # Relaxing some constraints for tests kwargs["compiler_flags"] += [ From 98bcd679f6bb9e4185035575386bffa9c69338e5 Mon Sep 17 00:00:00 2001 From: Kiymet Akdemir <54183514+kiymetakdemir@users.noreply.github.com> Date: Tue, 8 Sep 2026 12:33:58 -0700 Subject: [PATCH 077/190] Add an executor for batched generation (#22534) ### Summary Implements the Executor seam against an ExecuTorch program. Backend-neutral: it reaches the cache through CacheBuilderRegistry and BatchControl, with the backend id read from the program and the cache kind and cache-key option as configuration, so any backend registering a batch-capable cache is served. A session is one cache sequence, and a batch is one forward carrying every input's tokens end to end on a single axis, with the cache's mask keeping the sequences apart. SessionIds are monotonic and never reissued, mapped onto cache sequence ids that do recycle; the cache frees an id once a sequence's last cell goes, which the seam forbids for a SessionId. Read build_step first; it holds the reasoning. It flattens a batch into tokens and positions in one pass so entry i of each names the same token, which is how the cache pairs them when placing cells and building the mask. Sequence lengths come from the cache and do not move as inputs are laid down, so a per-sequence cursor carries the batch's own writes. A batch wider than the method was traced at is split into slices rather than refused, so a scheduler need not know the width; preferred_batch_tokens() reports it for one that wants to pack a single pass anyway. Slices run in order, so each attends the cells its predecessors wrote. --- CMakeLists.txt | 4 +- extension/llm/batching/CMakeLists.txt | 20 +- extension/llm/batching/module_executor.cpp | 519 +++++++++++++++++++++ extension/llm/batching/module_executor.h | 132 ++++++ 4 files changed, 673 insertions(+), 2 deletions(-) create mode 100644 extension/llm/batching/module_executor.cpp create mode 100644 extension/llm/batching/module_executor.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 93e02c521eb..bf1a34e6f47 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -997,7 +997,9 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache) list(APPEND _executorch_extensions extension_llm_cache) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/batching) - list(APPEND _executorch_extensions extension_llm_batching) + list(APPEND _executorch_extensions extension_llm_batching + extension_llm_batching_module + ) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/extension/llm/batching/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt index 2175964352b..fac3cfff5f0 100644 --- a/extension/llm/batching/CMakeLists.txt +++ b/extension/llm/batching/CMakeLists.txt @@ -9,6 +9,10 @@ # scheduler and the executor seam are header-only and free of ExecuTorch runtime # types; the runner owns a thread, so this is a static library rather than an # INTERFACE target. +# +# extension_llm_batching_module is a separate target because it implements that +# seam against a program and a KV cache, and so carries the runtime types the +# seam itself is kept clear of. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) @@ -26,8 +30,22 @@ target_compile_options(extension_llm_batching PUBLIC ${_common_compile_options}) find_package(Threads REQUIRED) target_link_libraries(extension_llm_batching PUBLIC Threads::Threads) +add_library(extension_llm_batching_module module_executor.cpp) +target_link_libraries( + extension_llm_batching_module + PUBLIC extension_llm_batching extension_llm_cache extension_module + extension_tensor + PRIVATE extension_llm_sampler +) +target_include_directories( + extension_llm_batching_module PUBLIC ${_common_include_directories} +) +target_compile_options( + extension_llm_batching_module PUBLIC ${_common_compile_options} +) + install( - TARGETS extension_llm_batching + TARGETS extension_llm_batching extension_llm_batching_module EXPORT ExecuTorchTargets DESTINATION ${CMAKE_INSTALL_LIBDIR} INCLUDES diff --git a/extension/llm/batching/module_executor.cpp b/extension/llm/batching/module_executor.cpp new file mode 100644 index 00000000000..fb9ad912b92 --- /dev/null +++ b/extension/llm/batching/module_executor.cpp @@ -0,0 +1,519 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +using ::executorch::extension::make_tensor_ptr; +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +namespace { + +// Constant methods carry no delegate, so the layout reads with the program +// loaded and the method not. Sizing is the caller's and is left unset. +Result config_from_program(Module& module) { + const auto read_int = [&module](const char* name) -> std::optional { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isInt()) { + return std::nullopt; + } + return r->at(0).toInt(); + }; + const auto read_ints = + [&module](const char* name) -> std::optional> { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isTensor()) { + return std::nullopt; + } + const auto t = r->at(0).toTensor(); + if (t.scalar_type() != ::executorch::aten::ScalarType::Int) { + return std::nullopt; + } + const int32_t* p = t.const_data_ptr(); + return std::vector(p, p + t.numel()); + }; + + const auto n_caches = read_int("get_n_caches"); + const auto kv_heads = read_ints("get_kv_heads"); + const auto head_dims = read_ints("get_head_dims"); + const auto windows = read_ints("get_windows"); + ET_CHECK_OR_RETURN_ERROR( + n_caches && kv_heads && head_dims && windows, + InvalidArgument, + "ModuleExecutor: the program publishes no KV layout"); + const auto n = static_cast(*n_caches); + ET_CHECK_OR_RETURN_ERROR( + kv_heads->size() == n && head_dims->size() == n && windows->size() == n, + InvalidArgument, + "ModuleExecutor: the published KV layout names %zu caches inconsistently", + n); + + cache::CacheConfig cfg{}; + cfg.n_layers = static_cast(n); + cfg.layers.reserve(n); + for (size_t l = 0; l < n; ++l) { + cache::LayerConfig lc{}; + lc.n_kv_heads = (*kv_heads)[l]; + lc.head_dim = (*head_dims)[l]; + lc.policy = (*windows)[l] > 0 + ? cache::LayerPolicy{cache::LayerPolicy::Kind::Ring, (*windows)[l]} + : cache::LayerPolicy{cache::LayerPolicy::Kind::Flat, 0}; + cfg.layers.push_back(lc); + } + return cfg; +} + +// The backend-load option the delegate resolves the cache through. +constexpr char kCacheKeyOption[] = "cache_key"; + +std::uint64_t nondeterministic_seed() { + std::random_device device; + return device(); +} + +// One forward's inputs, flattened across the batch. Entry i of `tokens` and of +// `positions` names the same token, which is how the cache pairs them. +struct Step { + // Signed to match the model's token input, not Token. + std::vector tokens; + std::vector positions; + // The sequence each token belongs to, declared to the cache one slice at a + // time so a declaration always matches the forward that places it. + std::vector seq_ids; + // Per input: the logits row it draws from, or -1 when its prediction is + // discarded. An input of any width contributes one, since only its last row + // predicts a token the session does not hold. + std::vector logit_indices; +}; + +// Flatten the batch and truncate whatever it reopens; execute() declares each +// slice to the cache as it runs it. A per-sequence cursor carries the batch's +// own writes, so consecutive chunks of one prompt abut and only the first can +// reopen committed ground. Every input is checked before any is truncated, so +// a refusal leaves the cache untouched. +Result build_step( + cache::BatchControl& ctl, + const BatchInput& batch, + const std::unordered_map& sessions, + int max_session_tokens) { + Step step; + const std::size_t total = batch.size(); + step.tokens.reserve(total); + step.positions.reserve(total); + step.logit_indices.reserve(batch.inputs.size()); + + step.seq_ids.reserve(total); + // Truncations the batch asks for, held until every input has been checked. + std::vector> rewinds; + // Where each sequence stands mid-batch: the cache still reports what it held + // before the step, so the batch's own writes live here. + std::unordered_map cursor; + + for (const Input& input : batch.inputs) { + const auto seq_it = sessions.find(input.sid); + if (seq_it == sessions.end()) { + ET_LOG(Error, "build_step: session %" PRId64 " is not open", input.sid); + return Error::InvalidArgument; + } + const std::int32_t seq_id = seq_it->second.seq_id; + if (input.size == 0 || !input.tokens || + input.offset + input.size > input.tokens->size()) { + ET_LOG( + Error, + "build_step: session %" PRId64 " gave a slice its tokens do not hold", + input.sid); + return Error::InvalidArgument; + } + + const std::int64_t start = static_cast(input.position) + + static_cast(input.offset); + const auto [cursor_it, first_for_seq] = + cursor.try_emplace(seq_id, ctl.next_pos(seq_id)); + int& at = cursor_it->second; + if (start > at) { + // Positions nothing attended, and nothing later reaches back to fill. + ET_LOG( + Error, + "build_step: session %" PRId64 " starts at %" PRId64 + " over a sequence holding %d", + input.sid, + start, + at); + return Error::InvalidArgument; + } + if (start < at) { + if (!first_for_seq) { + // Its predecessor in this batch has already been laid down, so a + // rewind now would truncate committed cells for a step whose + // positions repeat and cannot be placed. + ET_LOG( + Error, + "build_step: session %" PRId64 " overlaps its earlier input", + input.sid); + return Error::InvalidArgument; + } + if (start == 0) { + // Emptying a sequence hands its id back, and the step names it. + ET_LOG( + Error, + "build_step: session %" PRId64 " reopens from the start", + input.sid); + return Error::InvalidArgument; + } + rewinds.emplace_back(seq_id, static_cast(start)); + at = static_cast(start); + } + + const std::int64_t end = start + static_cast(input.size); + if (end > max_session_tokens) { + ET_LOG( + Error, + "build_step: session %" PRId64 " reaches %" PRId64 " of %d cells", + input.sid, + end, + max_session_tokens); + return Error::OutOfResources; + } + + const Token* slice = input.tokens->data() + input.offset; + for (std::size_t k = 0; k < input.size; ++k) { + step.tokens.push_back(static_cast(slice[k])); + } + for (std::size_t k = 0; k < input.size; ++k) { + step.positions.push_back(start + static_cast(k)); + } + step.seq_ids.insert(step.seq_ids.end(), input.size, seq_id); + at = static_cast(end); + step.logit_indices.push_back( + input.produce_output ? static_cast(step.tokens.size()) - 1 : -1); + } + + for (const auto& [seq_id, from] : rewinds) { + if (!ctl.seq_rm(seq_id, from, std::nullopt)) { + ET_LOG(Error, "build_step: sequence %d would not truncate", seq_id); + return Error::Internal; + } + } + return step; +} + +} // namespace + +ModuleExecutor::ModuleExecutor( + std::unique_ptr module, + std::shared_ptr cache, + std::unique_ptr session, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::string method, + std::int32_t vocab_size, + int max_step_tokens) + : session_(std::move(session)), + cache_(std::move(cache)), + module_(std::move(module)), + ctl_(cache_->as_batch_control()), + max_sessions_(max_sessions), + max_session_tokens_(max_session_tokens), + backend_id_(std::move(backend_id)), + method_(std::move(method)), + vocab_size_(vocab_size), + max_step_tokens_(max_step_tokens) {} + +ModuleExecutor::~ModuleExecutor() = default; + +std::unique_ptr ModuleExecutor::create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity, + std::string cache_kind, + std::string method) { + if (module == nullptr) { + ET_LOG(Error, "ModuleExecutor: no program"); + return nullptr; + } + if (max_sessions <= 0 || max_session_tokens <= 0) { + ET_LOG(Error, "ModuleExecutor: session limits must be positive"); + return nullptr; + } + if (module->load() != Error::Ok) { // a no-op once the caller has loaded it + ET_LOG(Error, "ModuleExecutor: the program did not load"); + return nullptr; + } + + auto cfg = config_from_program(*module); + if (!cfg.ok()) { + return nullptr; + } + cfg->capacity = max_sessions * max_session_tokens; + cfg->kv_dtype = kv_dtype; + if (initial_capacity >= 0) { + cfg->initial_capacity = initial_capacity; + } + if (!cache::valid(*cfg)) { + ET_LOG(Error, "ModuleExecutor: the program's layout is unusable"); + return nullptr; + } + + const auto meta = module->method_meta(method); + if (!meta.ok()) { + ET_LOG(Error, "ModuleExecutor: %s has no metadata", method.c_str()); + return nullptr; + } + + std::string backend_id; + for (std::size_t i = 0; i < meta->num_backends(); ++i) { + const auto name = meta->get_backend_name(i); + if (!name.ok()) { + ET_LOG( + Error, "ModuleExecutor: %s has an unnamed delegate", method.c_str()); + return nullptr; + } + if (backend_id.empty()) { + backend_id = name.get(); + } else if (backend_id != name.get()) { + ET_LOG( + Error, + "ModuleExecutor: %s spans more than one backend, so which holds the " + "cache is ambiguous", + method.c_str()); + return nullptr; + } + } + if (backend_id.empty()) { + ET_LOG(Error, "ModuleExecutor: %s delegates to nothing", method.c_str()); + return nullptr; + } + + auto built = + cache::CacheBuilderRegistry::global().build(backend_id, cache_kind, *cfg); + if (!built.ok()) { + ET_LOG( + Error, + "ModuleExecutor: backend %s registers no %s cache", + backend_id.c_str(), + cache_kind.c_str()); + return nullptr; + } + std::shared_ptr cache = built.get(); + if (cache->as_batch_control() == nullptr) { + ET_LOG(Error, "ModuleExecutor: the cache carries no sequence identity"); + return nullptr; + } + + if (meta->num_outputs() == 0) { + ET_LOG(Error, "ModuleExecutor: %s publishes no outputs", method.c_str()); + return nullptr; + } + const auto logits_info = meta->output_tensor_meta(0); + if (!logits_info.ok() || logits_info->sizes().empty()) { + ET_LOG(Error, "ModuleExecutor: %s has no logits shape", method.c_str()); + return nullptr; + } + const auto logits_sizes = logits_info->sizes(); + + const auto tokens_info = meta->input_tensor_meta(0); + if (!tokens_info.ok() || tokens_info->sizes().empty()) { + ET_LOG( + Error, "ModuleExecutor: %s has no token input shape", method.c_str()); + return nullptr; + } + const auto tokens_sizes = tokens_info->sizes(); + + auto session = + std::make_unique(cache::make_unique_key(), cache); + + return std::unique_ptr(new ModuleExecutor( + std::move(module), + std::move(cache), + std::move(session), + max_sessions, + max_session_tokens, + std::move(backend_id), + std::move(method), + logits_sizes[logits_sizes.size() - 1], + tokens_sizes[tokens_sizes.size() - 1])); +} + +bool ModuleExecutor::initialize() { + // The delegate resolves the cache from this key while the method loads. + char key[::executorch::runtime::kMaxOptionKeyLength] = {}; + std::memcpy(key, kCacheKeyOption, sizeof(kCacheKeyOption) - 1); + ::executorch::runtime::BackendOptions<1> options; + ::executorch::runtime::LoadBackendOptionsMap options_map; + if (options.set_option(key, session_->key().c_str()) != Error::Ok || + options_map.set_options(backend_id_.c_str(), options.view()) != + Error::Ok) { + ET_LOG(Error, "ModuleExecutor: could not name the cache to the backend"); + return false; + } + if (module_->load_method( + method_, + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + &options_map) != Error::Ok) { + ET_LOG(Error, "ModuleExecutor: could not load %s", method_.c_str()); + return false; + } + return true; +} + +std::optional ModuleExecutor::open_session() { + if (static_cast(sessions_.size()) >= max_sessions_) { + return std::nullopt; + } + const std::optional seq_id = ctl_->seq_new(); + if (!seq_id) { + return std::nullopt; + } + const SessionId session = next_session_++; + sessions_.emplace(session, SessionInfo{*seq_id, nullptr}); + return session; +} + +void ModuleExecutor::close_session(SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + // Frees the cells and hands the sequence id back. The session id is not. + ctl_->seq_rm(it->second.seq_id, 0, std::nullopt); + sessions_.erase(it); +} + +void ModuleExecutor::set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + // One sampler per generation, carrying its own generator state from here on. + it->second.sampler = std::make_unique( + vocab_size_, + params.temperature, + params.top_p, + seed.value_or(nondeterministic_seed())); + it->second.sampler->set_topk(params.top_k); +} + +bool ModuleExecutor::execute(const BatchInput& batch, BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + + const Result step = + build_step(*ctl_, batch, sessions_, max_session_tokens_); + if (!step.ok()) { + return false; + } + + // A batch wider than the method was traced at runs as several forwards. They + // go in order, so a slice attends the cells its predecessors wrote, and each + // input's logits row falls in exactly one of them. + const int total = static_cast(step->tokens.size()); + for (int off = 0; off < total; off += max_step_tokens_) { + const int n = std::min(max_step_tokens_, total - off); + // Placement checks the forward's token count against the declaration, so + // each slice declares its own. + if (!ctl_->declare_step(std::vector( + step->seq_ids.begin() + off, step->seq_ids.begin() + off + n))) { + ET_LOG(Error, "ModuleExecutor: the cache refused a slice of %d", n); + return false; + } + auto tokens = make_tensor_ptr( + {1, n}, + std::vector( + step->tokens.begin() + off, step->tokens.begin() + off + n)); + auto positions = make_tensor_ptr( + {n}, + std::vector( + step->positions.begin() + off, step->positions.begin() + off + n)); + auto result = module_->execute(method_, {tokens, positions}); + if (!result.ok()) { + ET_LOG( + Error, + "ModuleExecutor: %s failed with 0x%x", + method_.c_str(), + static_cast(result.error())); + return false; + } + if (result->empty() || !result->at(0).isTensor()) { + ET_LOG(Error, "ModuleExecutor: %s returned no logits", method_.c_str()); + return false; + } + // Non-const: the sampler reduces each row in place. Each is read once. + auto logits = result->at(0).toTensor(); + + for (std::size_t i = 0; i < batch.inputs.size(); ++i) { + const int row = step->logit_indices[i]; + if (row < off || row >= off + n) { + continue; // another slice's row, or a chunk whose prediction is dropped + } + const SessionId session = batch.inputs[i].sid; + const std::optional token = sample_row(logits, row - off, session); + if (!token) { + return false; + } + out.outputs[i] = Output{session, {*token}}; + } + } + return true; +} + +std::optional ModuleExecutor::sample_row( + ::executorch::aten::Tensor& logits, + int row, + SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end() || it->second.sampler == nullptr) { + ET_LOG( + Error, + "ModuleExecutor: session %" PRId64 " has no sampling policy", + session); + return std::nullopt; + } + if (row >= logits.numel() / vocab_size_) { + ET_LOG(Error, "ModuleExecutor: logits hold no row %d", row); + return std::nullopt; + } + // A one-row view over the model's own output: sample_from_logits reduces in + // place and reads the last dimension. + auto one_row = make_tensor_ptr( + {vocab_size_}, + static_cast(logits.mutable_data_ptr()) + + static_cast(row) * vocab_size_ * + ::executorch::runtime::elementSize(logits.scalar_type()), + logits.scalar_type()); + return static_cast(sample_from_logits(*one_row, *it->second.sampler)); +} + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/module_executor.h b/extension/llm/batching/module_executor.h new file mode 100644 index 00000000000..9f64d10ff11 --- /dev/null +++ b/extension/llm/batching/module_executor.h @@ -0,0 +1,132 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// An Executor that runs an ExecuTorch Module over a registered KV cache, built +// from the layout the program publishes. A session is one cache sequence, a +// batch is one forward carrying every input's tokens on a single axis, and the +// cache's mask keeps the sequences apart. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { + +class Sampler; + +namespace batching { + +namespace cache = ::executorch::extension::llm::cache; + +// A session's cache sequence and the sampler its generation draws from. +struct SessionInfo { + std::int32_t seq_id; + std::unique_ptr sampler; +}; + +class ET_EXPERIMENTAL ModuleExecutor : public Executor { + public: + ~ModuleExecutor() override; + + // Builds the cache from the layout `module` publishes and pairs it with the + // backend, which is read from the program -- so the method's attention must + // be delegated to just one. The method itself loads in initialize(); the + // program must be loaded and its method must not be, since the delegate + // resolves the cache while that load runs. + // + // Capacity is `max_sessions` x `max_session_tokens` cells exactly, and + // open_session() holds the count, so exhaustion is unreachable rather than + // handled. `kv_dtype` is the ET ScalarType K/V is stored in; a negative + // `initial_capacity` leaves the pools to grow from their own default. + // `cache_kind` must name a builder that carries batch control -- a cache + // serving one sequence cannot back a batch of them. + // + // nullptr = unusable limits, no published KV layout, a method spanning + // several backends, or no such cache for the backend it names. A method that + // will not load is reported by initialize(). + static std::unique_ptr create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity = -1, + std::string cache_kind = "cell", + std::string method = "forward"); + + // The widest step this method takes, from the shape its token input was + // traced at. A wider batch is sliced; a narrower one leaves the forward + // partly unused. + std::size_t preferred_batch_tokens() const override { + return static_cast(max_step_tokens_); + } + + // Loads the method here so the delegate that resolves the cache binds on the + // thread that runs it. + bool initialize() override; + + std::optional open_session() override; + void close_session(SessionId session) override; + void set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) override; + bool execute(const BatchInput& batch, BatchOutput& out) override; + + private: + ModuleExecutor( + std::unique_ptr module, + std::shared_ptr cache, + std::unique_ptr session, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::string method, + std::int32_t vocab_size, + int max_step_tokens); + + // Draw the token an input produced from its row of `logits`, which the + // session's sampler consumes in place. + std::optional + sample_row(::executorch::aten::Tensor& logits, int row, SessionId session); + + // Ordered so the module dies first, releasing the delegate that resolved the + // cache before the registry entry naming it goes. + std::unique_ptr session_; + std::shared_ptr cache_; + std::unique_ptr module_; + cache::BatchControl* ctl_; + int max_sessions_; + int max_session_tokens_; + std::string backend_id_; + std::string method_; + // The method's logits width, so a sampler can be built by its policy. + std::int32_t vocab_size_; + int max_step_tokens_; + + SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids + std::unordered_map sessions_; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch From ee2cd8ac9c0aa7226a5f604ea0b877197fb3a258 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 8 Sep 2026 13:05:53 -0700 Subject: [PATCH 078/190] Cortex-M: preserve unsupported transpose ranks (#22571) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Keep int8 `aten.permute_copy` nodes with ranks outside 1–4 on the portable operator. The Cortex-M transpose kernel only supports ranks 1–4, so lowering a rank-5 permutation produced a program that failed when loaded by the runtime. The existing supported-rank FVP coverage remains unchanged. A new dialect test verifies that a rank-5 int8 permutation stays as `aten.permute_copy`. ### Test plan `source examples/arm/arm-scratch/setup_path.sh && python -m pytest --config-file=backends/arm/test/pytest.ini backends/cortex_m/test/ops/test_transpose.py` `lintrunner -m origin/main` AI-assisted: Codex. --- .../cortex_m/passes/aten_to_cortex_m_pass.py | 2 +- backends/cortex_m/test/ops/test_transpose.py | 19 +++++++++++++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/backends/cortex_m/passes/aten_to_cortex_m_pass.py b/backends/cortex_m/passes/aten_to_cortex_m_pass.py index 7f0cd2dd434..5799896e17c 100644 --- a/backends/cortex_m/passes/aten_to_cortex_m_pass.py +++ b/backends/cortex_m/passes/aten_to_cortex_m_pass.py @@ -1259,7 +1259,7 @@ def _get_permute_replacement( ) -> DialectNodeSpec | None: del dialect_pass input_tensor = _get_input_tensor_data(node) - if input_tensor.dtype != torch.int8: + if input_tensor.dtype != torch.int8 or not 1 <= input_tensor.dim() <= 4: return None return _transpose_spec(node, input_tensor) diff --git a/backends/cortex_m/test/ops/test_transpose.py b/backends/cortex_m/test/ops/test_transpose.py index 2e5f5112bd9..26a026358df 100644 --- a/backends/cortex_m/test/ops/test_transpose.py +++ b/backends/cortex_m/test/ops/test_transpose.py @@ -24,6 +24,12 @@ "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, } +RANK5_OPS_AFTER_PASSES = { + "executorch_exir_dialects_edge__ops_cortex_m_quantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_aten_permute_copy_default": 1, + "executorch_exir_dialects_edge__ops_cortex_m_dequantize_per_tensor_default": 1, +} + class CortexMPermute(torch.nn.Module): ops_before_transforms = OPS_BEFORE_PASSES @@ -98,6 +104,19 @@ def test_dialect_transpose(test_case, cortex_m_target): ) +def test_dialect_rank5_permute_stays_portable(cortex_m_target): + tester = CortexMTester( + CortexMPermute((0, 2, 1, 4, 3)), + (ramp_tensor(-1.0, 1.0, (1, 2, 3, 4, 5)),), + target_config=cortex_m_target, + ) + tester.test_dialect( + OPS_BEFORE_PASSES, + RANK5_OPS_AFTER_PASSES, + qtol=1, + ) + + @parametrize("test_case", test_cases) def test_implementation_transpose(test_case, cortex_m_target): tester = CortexMTester( From aca0acec292299ad7149614a10292a3b7cc8507b Mon Sep 17 00:00:00 2001 From: Jacob Stevens Date: Tue, 8 Sep 2026 17:49:00 -0400 Subject: [PATCH 079/190] Add SmolLM2 360M export support (#22459) Differential Revision: D118479813 Pull Request resolved: https://github.com/pytorch/executorch/pull/22459 --- examples/models/llama/export_llama_lib.py | 4 +++- examples/models/smollm2/360M_config.json | 16 ++++++++++++++++ examples/models/smollm2/BUCK | 1 + extension/llm/export/config/llm_config.py | 1 + 4 files changed, 21 insertions(+), 1 deletion(-) create mode 100644 examples/models/smollm2/360M_config.json diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index ffe9a89a570..d982eff7468 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -115,6 +115,7 @@ "qwen3_5_4b", "phi_4_mini", "smollm2", + "smollm2_360m", "lfm2_350m", # hybrid "lfm2_700m", # hybrid "lfm2_1_2b", # hybrid @@ -128,6 +129,7 @@ "qwen2_5_coder_32b": "Qwen/Qwen2.5-Coder-32B-Instruct", "phi_4_mini": "microsoft/Phi-4-mini-instruct", "smollm2": "HuggingFaceTB/SmolLM2-135M", + "smollm2_360m": "HuggingFaceTB/SmolLM2-360M", "qwen3_0_6b": "Qwen/Qwen3-0.6B", "qwen3_1_7b": "Qwen/Qwen3-1.7B", "qwen3_4b": "Qwen/Qwen3-4B", @@ -712,7 +714,7 @@ def export_llama( # noqa: C901 from executorch.examples.models.qwen3 import convert_weights elif model_name == "phi_4_mini": from executorch.examples.models.phi_4_mini import convert_weights - elif model_name == "smollm2": + elif model_name in ("smollm2", "smollm2_360m"): from executorch.examples.models.smollm2 import convert_weights elif model_name.startswith("lfm2"): from executorch.examples.models.lfm2 import convert_weights diff --git a/examples/models/smollm2/360M_config.json b/examples/models/smollm2/360M_config.json new file mode 100644 index 00000000000..332c4c12337 --- /dev/null +++ b/examples/models/smollm2/360M_config.json @@ -0,0 +1,16 @@ +{ + "dim": 960, + "ffn_dim_multiplier": 1, + "hidden_dim": 2560, + "n_heads": 15, + "n_kv_heads": 5, + "n_layers": 32, + "norm_eps": 1e-05, + "rope_theta": 100000.0, + "use_scaled_rope": false, + "vocab_size": 49152, + "use_hf_rope": false, + "attention_qkv_bias": false, + "bos_idx": 0, + "eos_idx": 0 +} diff --git a/examples/models/smollm2/BUCK b/examples/models/smollm2/BUCK index 6d81065373b..45173b5cbf5 100644 --- a/examples/models/smollm2/BUCK +++ b/examples/models/smollm2/BUCK @@ -14,6 +14,7 @@ fbcode_target(_kind = runtime.python_library, base_module = "executorch.examples.models.smollm2", resources = { "135M_config.json": "135M_config.json", + "360M_config.json": "360M_config.json", }, deps = [ "//caffe2:torch", diff --git a/extension/llm/export/config/llm_config.py b/extension/llm/export/config/llm_config.py index acdc9771141..3eb8b8a18f6 100644 --- a/extension/llm/export/config/llm_config.py +++ b/extension/llm/export/config/llm_config.py @@ -50,6 +50,7 @@ class ModelType(str, Enum): qwen3_5_4b = "qwen3_5_4b" phi_4_mini = "phi_4_mini" smollm2 = "smollm2" + smollm2_360m = "smollm2_360m" lfm2_350m = "lfm2_350m" lfm2_700m = "lfm2_700m" lfm2_1_2b = "lfm2_1_2b" From a519efa64e9f5aea9b449fb21a0abbfd59761396 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Tue, 8 Sep 2026 14:59:46 -0700 Subject: [PATCH 080/190] Move pull.yml to linux_job_v3 (#22247) Third of six splitting up #22107. Stacked on #22246. 45 call sites in one file, all mechanical: the v2 -> v3 rename, EC2 runner labels swapped for their OSDC equivalents, `use-custom-docker-registry` dropped since v3 ignores it, and the ECR image spelled out against the hash `_docker-image.yml` resolves. | EC2 | OSDC | |---|---| | `linux.2xlarge`, `linux.2xlarge.memory` | `mt-l-x86iavx512-8-64` | | `linux.4xlarge.memory` | `mt-l-x86iavx512-16-128` | | `linux.24xlarge` | `mt-l-x86iavx512-94-192` | | `linux.24xlarge.memory` | `mt-l-x86iavx512-94-768` | | `linux.arm64.2xlarge` | `mt-l-arm64g4-16-62` | | `linux.g5.4xlarge.nvidia.gpu` | `mt-l-x86aavx2-29-113-a10g` | Two jobs needed more than the rename: - `test-qnn-delegate-linux` goes to `mt-l-x86iamx-8-64` instead. The vendor letter names the ISA, not the silicon, so `x86iavx512` is `r7a`, i.e. AMD EPYC. The QNN backend's `disable_mkldnn_on_amd()` only runs on an AMD host, where it sets `torch.backends.mkldnn.enabled` outside a `flags()` context and raises once the tests have frozen the flags, failing all 1040. `mt-l-x86iamx-8-64` is `r7i` at the same 8 vCPU / 64Gi, so the host stays Intel as it was on `c5`. Worth fixing in the backend separately for anyone who does run it on AMD. - The pytest-xdist worker cap, since `auto` and `logical` size themselves from the machine's cores rather than the pod's limit and get OOM-killed. Authored with Claude Code. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- .github/workflows/pull.yml | 361 +++++++++++++++++++++---------------- 1 file changed, 206 insertions(+), 155 deletions(-) diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 2522cca3113..d17be95b3aa 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -13,6 +13,10 @@ concurrency: cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + # Emits the list of changed files for the current PR or push commit. # On PR: PR diff. On push: diff against `github.event.before`. # On events without a diff base (workflow_dispatch, tag creation, @@ -33,8 +37,9 @@ jobs: uses: ./.github/workflows/_ci-run-decision.yml test-qnn-wheel-packages-linux: + needs: docker-image name: test-qnn-wheel-packages-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -43,8 +48,8 @@ jobs: matrix: python-version: [ "3.10", "3.11", "3.12", "3.13" ] with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -77,7 +82,7 @@ jobs: contents: read test-minimal-wheel-linux: - needs: changed-files + needs: [docker-image, changed-files] if: | github.event_name != 'pull_request' || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') || @@ -90,13 +95,13 @@ jobs: contains(needs.changed-files.outputs.changed-files, 'setup.py') || contains(needs.changed-files.outputs.changed-files, 'tools/cmake/') name: test-minimal-wheel-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -107,16 +112,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh test-setup-linux-gcc: + needs: docker-image name: test-setup-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -132,8 +138,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "add" "${BUILD_TOOL}" "portable" test-models-linux-basic: + needs: docker-image name: test-models-linux-basic - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -142,23 +149,23 @@ jobs: model: [mv3, vit] backend: [portable, xnnpack-quantization-delegation] build-tool: [cmake, buck2] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 # TODO: Need to figure out why buck2 doesnt work on Graviton instances. - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 build-tool: buck2 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -176,8 +183,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-models-linux: + needs: docker-image name: test-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -185,33 +193,33 @@ jobs: matrix: model: [linear, add, add_mul, ic3, mv2, resnet18, resnet50, mobilebert, emformer_transcribe] backend: [portable, xnnpack-quantization-delegation] - runner: [linux.2xlarge] + runner: [mt-l-x86iavx512-8-64] include: - model: ic4 backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: ic4 backend: xnnpack-quantization-delegation - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: emformer_join backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: emformer_join backend: xnnpack-quantization-delegation - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: phi_4_mini backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: llama3_2_vision_encoder backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: w2l backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -229,16 +237,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-parakeet-xnnpack-linux: + needs: docker-image name: test-parakeet-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.4xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -262,16 +271,17 @@ jobs: echo "::endgroup::" test-voxtral-realtime-xnnpack-linux: + needs: docker-image name: test-voxtral-realtime-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.4xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -298,9 +308,10 @@ jobs: echo "::endgroup::" test-llama-runner-linux: + needs: docker-image # Test Both linux x86 and linux aarch64 name: test-llama-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -308,25 +319,25 @@ jobs: matrix: dtype: [fp32] mode: [xnnpack+custom+qe,xnnpack+custom+quantize_kv,xnnpack+quantize_kv] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] include: - dtype: bf16 mode: custom - runner: linux.2xlarge + runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-clang12 # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -349,16 +360,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -dtype "${DTYPE}" -mode "${MODE}" -upload "${ARTIFACTS_DIR_NAME}" test-llama-runner-linux-android: + needs: docker-image name: test-llama-runner-linux-android - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -374,16 +386,17 @@ jobs: bash .ci/scripts/build_llama_android.sh "${BUILD_TOOL}" test-custom-ops-linux: + needs: docker-image name: test-custom-ops-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -398,16 +411,17 @@ jobs: PYTHON_EXECUTABLE=python bash examples/portable/custom_ops/test_custom_ops.sh "${BUILD_TOOL}" test-selective-build-linux: + needs: docker-image name: test-selective-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -422,9 +436,10 @@ jobs: PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" test-multimodal-linux: + needs: docker-image if: ${{ !github.event.pull_request.head.repo.fork }} name: test-multimodal-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -435,8 +450,8 @@ jobs: model: ["gemma3-4b"] # llava gives segfault so not covering. with: secrets-env: EXECUTORCH_HF_TOKEN - runner: linux.24xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-768 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -461,16 +476,17 @@ jobs: echo "::endgroup::" test-moshi-linux: + needs: docker-image name: test-moshi-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -493,16 +509,17 @@ jobs: python -m unittest examples.models.moshi.mimi.test_mimi test-quantized-aot-lib-linux: + needs: docker-image name: test-quantized-aot-lib-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -516,16 +533,17 @@ jobs: PYTHON_EXECUTABLE=python bash examples/xnnpack/quantization/test_quantize.sh "${BUILD_TOOL}" mv2 test-binary-size-linux-gcc: + needs: docker-image name: test-binary-size-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc9-nopytorch + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc9-nopytorch-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -559,16 +577,17 @@ jobs: fi test-binary-size-linux: + needs: docker-image name: test-binary-size-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -603,8 +622,9 @@ jobs: fi test-arm-cortex-m-size-test: + needs: docker-image name: test-arm-cortex-m-size-test - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -613,8 +633,8 @@ jobs: os: [bare_metal, zephyr-preset] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -696,14 +716,15 @@ jobs: fi test-mcu-cortex-m-backend: + needs: docker-image name: test-mcu-cortex-m-backend - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -763,16 +784,17 @@ jobs: docker-image: ci-image:executorch-ubuntu-22.04-clang12 test-qnn-buck-build-linux: + needs: docker-image name: test-qnn-buck-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -797,8 +819,9 @@ jobs: buck2 build //backends/qualcomm/... test-arm-backend-no-driver: + needs: docker-image name: test-arm-backend-no-driver - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -811,8 +834,8 @@ jobs: - test_arm_backend: test_run_tosa fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -832,14 +855,15 @@ jobs: backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-public-api-backward-compatibility: + needs: docker-image name: test-arm-backend-public-api-backward-compatibility - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-24.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -859,8 +883,9 @@ jobs: python backends/arm/test/public_api_bc/run_public_api_bc_scenarios.py test-llama-runner-qnn-linux: + needs: docker-image name: test-llama-runner-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -871,8 +896,8 @@ jobs: mode: [qnn] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -898,8 +923,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -mode "${MODE}" -dtype "${DTYPE}" -pt2e_quantize "${PT2E_QUANTIZE}" test-static-llama-qnn-linux: + needs: docker-image name: test-static-llama-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -908,8 +934,8 @@ jobs: task: [stories_110m, stories_260k_bc] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -932,8 +958,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} test-sqnr-static-llm-qnn-linux: + needs: docker-image name: test-sqnr-static-llm-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -942,8 +969,8 @@ jobs: task: [smollm2_135m] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -966,8 +993,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} sqnr test-qnn-models-linux: + needs: docker-image name: test-qnn-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -976,8 +1004,8 @@ jobs: model: [mv2, mv3, dl3] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -991,14 +1019,15 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn" test-qnn-direct-build-linux: + needs: docker-image name: test-qnn-direct-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1021,22 +1050,23 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true - # No runner-linux, so this takes the memory-optimized default. The suite - # runs one export worker per core, so what matters is memory per core, not - # core count: the previous instance gave each worker about 4 GiB and the - # job was killed. The memory-optimized default gives each worker more. + # No runner-linux, so this takes _test_backend.yml's default. The suite + # runs one export worker per core, so what matters is memory per core: + # the instance this used to run on gave each worker about 4 GiB and the + # job was killed. The default label is a little under 8 GiB per core. test-qnn-passes-linux: + needs: docker-image name: test-qnn-passes-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1065,16 +1095,21 @@ jobs: pytest -xvs backends/qualcomm/tests/test_import_side_effects.py test-qnn-delegate-linux: + needs: docker-image name: test-qnn-delegate-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + # Intel, unlike the avx512 labels: those are r7a, i.e. AMD EPYC, and the + # QNN backend disables MKLDNN on an AMD host through a non-bracketed + # torch.backends mutation that raises once the tests have frozen the + # flags. Same 8 vCPU / 64Gi, on r7i. + runner: mt-l-x86iamx-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1103,16 +1138,17 @@ jobs: -k "TestQNNFloatingPointOperator or TestQNNQuantizedOperator" test-phi-3-mini-runner-linux: + needs: docker-image name: test-phi-3-mini-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1133,16 +1169,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_phi_3_mini.sh Release test-qnn-python-imports-linux: + needs: docker-image name: test-qnn-python-imports-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 15 @@ -1181,16 +1218,17 @@ jobs: --module-prefix executorch.examples.qualcomm test-eval_llama-wikitext-linux: + needs: docker-image name: test-eval_llama-wikitext-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1210,7 +1248,7 @@ jobs: # TODO(larryliu0820): Fix this issue before reenabling it: https://gist.github.com/larryliu0820/7377ecd0d79dbc06076cec8d9f2b85d2 # test-eval_llama-mmlu-linux: # name: test-eval_llama-mmlu-linux - # uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + # uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main # permissions: # id-token: write # contents: read @@ -1236,16 +1274,17 @@ jobs: # PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh test-llama_runner_eager-linux: + needs: docker-image name: test-llama_runner_eager-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1263,16 +1302,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama_runner_eager.sh test-lora-linux: + needs: docker-image name: test-lora-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1290,16 +1330,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora.sh test-lora-multimethod-linux: + needs: docker-image name: test-lora-multimethod-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1317,16 +1358,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora_multimethod.sh test-mediatek-models-linux: + needs: docker-image name: test-mediatek-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-mediatek-sdk + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-mediatek-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1344,16 +1386,17 @@ jobs: # placeholder for mediatek to add more tests test-openvino-linux: + needs: docker-image name: test-openvino-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1366,16 +1409,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_openvino.sh test-build-wasm-linux: + needs: docker-image name: test-build-wasm-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1394,8 +1438,9 @@ jobs: PYTHON_EXECUTABLE=python bash examples/wasm/test_build_wasm.sh unittest-wasm-bindings: + needs: docker-image name: unittest-wasm-bindings - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -1404,8 +1449,8 @@ jobs: enable-etdump: ['', '--enable-etdump'] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1440,13 +1485,14 @@ jobs: pnpm test unittest-nxp-neutron: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 150 @@ -1484,18 +1530,19 @@ jobs: bash backends/nxp/run_unittests.sh test-samsung-quantmodels-linux: + needs: docker-image name: test-samsung-quantmodels-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -1522,18 +1569,19 @@ jobs: done test-samsung-models-linux: + needs: docker-image name: test-samsung-models-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 360 @@ -1564,14 +1612,15 @@ jobs: python -m unittest discover -s backends/samsung/test/models -p "test_*.py" test-vulkan-models-linux: + needs: docker-image name: test-vulkan-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1605,14 +1654,15 @@ jobs: done test-vulkan-operators-linux: + needs: docker-image name: test-vulkan-operators-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1697,14 +1747,15 @@ jobs: echo "::endgroup::" nxp-build-test: + needs: docker-image name: nxp-build-test - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 From dacc1d0d2dbaef882dc0decf3bd285c613136736 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:09:50 -0700 Subject: [PATCH 081/190] Add text stream (#22263) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds TextStream to extension/llm/runner, which assembles a stream of token ids into text a caller can hand straight to a UI, a socket, or a JSON encoder. A byte-level tokenizer can emit a token carrying only part of a UTF-8 character, so decoding tokens one at a time yields pieces that are not valid UTF-8 even though their concatenation is — survivable when printing to a terminal, broken for anything that treats a piece as a standalone string. TextStream owns that state: append() decodes each token against the previous one, emits only the characters completed so far, and holds the incomplete tail until its remaining bytes arrive, while flush() releases whatever is still held once generation ends. A tokenizer failure makes the stream sticky-failed rather than leaving the caller to guess which tokens reached the sink. Stop-string filtering deliberately stays a separate stage applied to the assembled text via the existing stop_safe_prefix_len, so this class does one job. Two supporting changes ride along: batching::Token becomes std::uint64_t to match the tokenizer's decode signature, and FakeExecutor::stop_token moves from a -1 sentinel to std::optional now that the type is unsigned. Covered by 12 new test cases in test_text_stream.cpp. --- extension/llm/runner/test/CMakeLists.txt | 1 + .../llm/runner/test/test_text_stream.cpp | 318 ++++++++++++++++++ extension/llm/runner/text_stream.h | 128 +++++++ extension/llm/runner/util.h | 22 +- 4 files changed, 467 insertions(+), 2 deletions(-) create mode 100644 extension/llm/runner/test/test_text_stream.cpp create mode 100644 extension/llm/runner/text_stream.h diff --git a/extension/llm/runner/test/CMakeLists.txt b/extension/llm/runner/test/CMakeLists.txt index 81b69c0ab9a..942cd8d75ec 100644 --- a/extension/llm/runner/test/CMakeLists.txt +++ b/extension/llm/runner/test/CMakeLists.txt @@ -23,6 +23,7 @@ set(_test_srcs test_text_prefiller.cpp test_text_decoder_runner.cpp test_multimodal_input.cpp + test_text_stream.cpp test_util.cpp test_wav_loader.cpp ) diff --git a/extension/llm/runner/test/test_text_stream.cpp b/extension/llm/runner/test/test_text_stream.cpp new file mode 100644 index 00000000000..8cddcdb88b5 --- /dev/null +++ b/extension/llm/runner/test/test_text_stream.cpp @@ -0,0 +1,318 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include + +#include + +using executorch::extension::llm::TextStream; +using executorch::runtime::Error; + +namespace { + +// Maps each token to the bytes it contributes. `pair_pieces` overrides the +// piece for a (previous, token) pair, so a test can observe that the preceding +// token reaches the tokenizer at all. +class FakeTokenizer : public tokenizers::Tokenizer { + public: + std::map pieces; + std::map, std::string> pair_pieces; + // Tokens the tokenizer refuses to decode. + std::vector rejected; + + tokenizers::Error load(const std::string&) override { + initialized_ = true; + return tokenizers::Error::Ok; + } + + tokenizers::Result decode( + uint64_t previous, + uint64_t token, + bool /*skip_special_tokens*/ = false) const override { + for (uint64_t bad : rejected) { + if (token == bad) { + return tokenizers::Error::Internal; + } + } + auto pair = pair_pieces.find({previous, token}); + if (pair != pair_pieces.end()) { + return pair->second; + } + auto single = pieces.find(token); + if (single != pieces.end()) { + return single->second; + } + return std::string(); + } + + tokenizers::Result> + encode(const std::string&, int8_t, int8_t) const override { + return std::vector{}; + } + tokenizers::Result id_to_piece(uint64_t) const override { + return std::string(); + } + tokenizers::Result piece_to_id(const std::string&) const override { + return uint64_t{0}; + } +}; + +// Collects what the stream emitted, both as separate pieces and joined. +struct Sink { + std::vector pieces; + std::string joined; + + void operator()(const std::string& piece) { + pieces.push_back(piece); + joined += piece; + } +}; + +// A stream writing into `sink`. +TextStream +stream_into(const FakeTokenizer& tokenizer, Sink& sink, uint64_t previous = 0) { + return TextStream( + tokenizer, [&sink](const std::string& piece) { sink(piece); }, previous); +} + +// The three bytes of U+4E16 (CJK), which a byte-level tokenizer can split. +constexpr const char kCjkByte0[] = "\xE4"; +constexpr const char kCjkByte1[] = "\xB8"; +constexpr const char kCjkByte2[] = "\x96"; +constexpr const char kCjk[] = "\xE4\xB8\x96"; + +} // namespace + +TEST(TextStreamTest, EmitsWholeCharactersImmediately) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "Hello"}, {2, " world"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + EXPECT_EQ(stream.append(std::vector{1, 2}), Error::Ok); + EXPECT_EQ(sink.joined, "Hello world"); + EXPECT_EQ(sink.pieces.size(), 2u); + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, NeverEmitsAnEmptyPiece) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, ""}, {2, "x"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(std::vector{1, 2}), Error::Ok); + EXPECT_EQ(sink.pieces, (std::vector{"x"})); +} + +// The reason this class exists: a character split across tokens must not reach +// the sink in pieces, or a consumer treating each piece as a string breaks. +TEST(TextStreamTest, HoldsBackAPartialCharacterUntilItCompletes) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, kCjkByte0}, {2, kCjkByte1}, {3, kCjkByte2}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + EXPECT_TRUE(sink.pieces.empty()) << "one third of a character is not text"; + EXPECT_TRUE(stream.has_pending()); + + ASSERT_EQ(stream.append(2u), Error::Ok); + EXPECT_TRUE(sink.pieces.empty()); + + ASSERT_EQ(stream.append(3u), Error::Ok); + EXPECT_EQ(sink.joined, kCjk); + EXPECT_EQ(sink.pieces.size(), 1u) << "the character arrives whole, once"; + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, EmitsTheCompletePrefixAndKeepsTheRest) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, std::string("ab") + kCjkByte0}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + EXPECT_EQ(sink.joined, "ab") << "the finished characters go now"; + EXPECT_TRUE(stream.has_pending()) << "the split character waits"; +} + +TEST(TextStreamTest, FlushReleasesAnUnfinishedCharacter) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, kCjkByte0}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + ASSERT_TRUE(sink.pieces.empty()); + + stream.flush(); + EXPECT_EQ(sink.joined, kCjkByte0) + << "a generation that ends mid-character must not swallow the bytes"; + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, FlushIsIdempotentAndSilentWhenEmpty) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "done"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + stream.flush(); + stream.flush(); + EXPECT_EQ(sink.pieces, (std::vector{"done"})); +} + +// Only SentencePiece reads the preceding token, but the stream must still +// forward it and advance it, or that tokenizer would strip the wrong space. +TEST(TextStreamTest, ForwardsThePrecedingTokenAndAdvancesIt) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{7, "?"}, {8, "!"}}; + tokenizer.pair_pieces = {{{5, 7}, " seeded"}, {{7, 8}, " advanced"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink, /*previous=*/5); + + ASSERT_EQ(stream.append(7u), Error::Ok); + EXPECT_EQ(sink.joined, " seeded") << "the constructor seed reaches decode"; + + ASSERT_EQ(stream.append(8u), Error::Ok); + EXPECT_EQ(sink.joined, " seeded advanced") + << "the previous token becomes the one just decoded"; +} + +TEST(TextStreamTest, ATokenizerErrorFailsTheStreamForGood) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "a"}, {3, "c"}}; + tokenizer.rejected = {2}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + EXPECT_EQ(stream.append(1u), Error::Ok); + EXPECT_EQ(stream.append(2u), Error::InvalidArgument); + EXPECT_TRUE(stream.failed()); + EXPECT_EQ(stream.append(3u), Error::InvalidState) + << "a failed stream must not resume and emit text out of order"; + EXPECT_EQ(sink.joined, "a"); +} + +// The batch stops where it broke rather than skipping past the bad token. +TEST(TextStreamTest, ABatchStopsAtTheTokenThatFailed) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "a"}, {3, "c"}}; + tokenizer.rejected = {2}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + EXPECT_EQ( + stream.append(std::vector{1, 2, 3}), Error::InvalidArgument); + EXPECT_EQ(sink.joined, "a") << "nothing after the failure is emitted"; + EXPECT_TRUE(stream.failed()); +} + +// A speculative executor hands back several tokens at once, so a batch must +// read the same as the tokens arriving one by one. +TEST(TextStreamTest, ABatchMatchesTokenByTokenDelivery) { + FakeTokenizer tokenizer; + tokenizer.pieces = { + {1, "He"}, {2, kCjkByte0}, {3, kCjkByte1}, {4, kCjkByte2}}; + Sink batched; + Sink one_at_a_time; + + TextStream a = stream_into(tokenizer, batched); + ASSERT_EQ(a.append(std::vector{1, 2, 3, 4}), Error::Ok); + + TextStream b = stream_into(tokenizer, one_at_a_time); + for (uint64_t token : {1u, 2u, 3u, 4u}) { + ASSERT_EQ(b.append(token), Error::Ok); + } + + EXPECT_EQ(batched.joined, one_at_a_time.joined); + EXPECT_EQ(batched.joined, std::string("He") + kCjk); +} + +TEST(TextStreamTest, AnInvalidLeadByteIsEmittedRatherThanStallingOutput) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "\xFF"}, {2, "ok"}}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(std::vector{1, 2}), Error::Ok); + EXPECT_EQ( + sink.joined, + "\xFF" + "ok") + << "a byte that can never start a character must not hold up the stream"; + EXPECT_FALSE(stream.has_pending()); +} + +TEST(TextStreamTest, ToleratesAnAbsentSink) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "x"}}; + TextStream stream(tokenizer, nullptr); + + EXPECT_EQ(stream.append(1u), Error::Ok); + stream.flush(); +} + +// A four-byte codepoint takes the len == 4 branch, which the three-byte cases +// above leave untested, and is the widest split a byte-level tokenizer can +// make. +TEST(TextStreamTest, HoldsBackAFourByteCharacterUntilItCompletes) { + FakeTokenizer tokenizer; + tokenizer.pieces = { + {1, "\xF0"}, {2, "\x9F"}, {3, "\x98"}, {4, "\x80"}}; // U+1F600 + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + for (uint64_t id : {1u, 2u, 3u}) { + ASSERT_EQ(stream.append(id), Error::Ok); + EXPECT_TRUE(sink.joined.empty()) << "emitted before the character finished"; + EXPECT_TRUE(stream.has_pending()); + } + ASSERT_EQ(stream.append(4u), Error::Ok); + EXPECT_EQ(sink.joined, "\xF0\x9F\x98\x80"); + EXPECT_FALSE(stream.has_pending()); +} + +// A lead byte promises continuation bytes that never arrive. Holding them would +// stall the stream for good, so they go out as-is: the sink can see invalid +// UTF-8 here, which is the documented trade for never blocking. +TEST(TextStreamTest, AMalformedSequenceIsEmittedRatherThanHeldForever) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, "\xE4\x41"}}; // 3-byte lead, then ASCII 'A' + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + EXPECT_EQ(sink.joined, "\xE4\x41"); + EXPECT_FALSE(stream.has_pending()) << "a malformed tail must not be held"; +} + +// The bytes held when a decode fails are not lost: the stream is sticky-failed, +// but flush() still surrenders what it was holding. +TEST(TextStreamTest, FlushAfterAFailureStillReleasesTheHeldBytes) { + FakeTokenizer tokenizer; + tokenizer.pieces = {{1, kCjkByte0}}; + tokenizer.rejected = {2}; + Sink sink; + TextStream stream = stream_into(tokenizer, sink); + + ASSERT_EQ(stream.append(1u), Error::Ok); + ASSERT_TRUE(stream.has_pending()); + EXPECT_NE(stream.append(2u), Error::Ok); + EXPECT_TRUE(stream.has_pending()) << "a failure keeps what was held"; + + stream.flush(); + EXPECT_EQ(sink.joined, kCjkByte0); + EXPECT_FALSE(stream.has_pending()); +} diff --git a/extension/llm/runner/text_stream.h b/extension/llm/runner/text_stream.h new file mode 100644 index 00000000000..8a58b28ea2c --- /dev/null +++ b/extension/llm/runner/text_stream.h @@ -0,0 +1,128 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +// Assembles a stream of token ids into text a caller can hand straight to a +// UI, a socket, or a JSON encoder. +// +// A byte-level tokenizer can emit a token that is only part of a character, so +// decoding each token on its own yields pieces that are not valid UTF-8 even +// though their concatenation is. Printing to a terminal survives that; +// anything that treats a piece as a standalone string does not. Holding the +// incomplete tail back until it completes is the state this exists to own, +// because a character split across tokens outlives any one delivery of them. +// +// Stop strings are a separate stage. Assemble here, then filter the text with +// stop_safe_prefix_len before it reaches the sink. + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +class ET_EXPERIMENTAL TextStream { + public: + // Never called with an empty string. + using Sink = std::function; + + // The tokenizer is borrowed and must outlive the stream. + // + // `previous` is the token the next one follows, normally the last token of + // the prompt. Only a SentencePiece tokenizer reads it, to drop the leading + // space of the first token after BOS; every BPETokenizerBase discards it. + // The default is for BPE, which cannot observe the value. A SentencePiece + // caller must pass the real previous token: 0 is a valid id there, so + // leaving it defaulted silently mis-handles the space on the first token. + TextStream( + const tokenizers::Tokenizer& tokenizer, + Sink on_text, + uint64_t previous = 0) + : tokenizer_(tokenizer), + on_text_(std::move(on_text)), + previous_(previous) {} + + // Emits every character the token completed. A token that only extends an + // unfinished character emits nothing and is not lost. + // + // On a tokenizer error the stream stops emitting and stays failed, rather + // than leaving the caller to guess which tokens reached the sink. + ::executorch::runtime::Error append(uint64_t token) { + if (failed_) { + return ::executorch::runtime::Error::InvalidState; + } + const tokenizers::Result piece = + tokenizer_.decode(previous_, token); + if (!piece.ok()) { + failed_ = true; + return ::executorch::runtime::Error::InvalidArgument; + } + previous_ = token; + pending_ += *piece; + emit_(utf8_complete_prefix_len(pending_)); + return ::executorch::runtime::Error::Ok; + } + + // Stops at the first token that fails. + ::executorch::runtime::Error append(const std::vector& tokens) { + for (uint64_t token : tokens) { + const ::executorch::runtime::Error error = append(token); + if (error != ::executorch::runtime::Error::Ok) { + return error; + } + } + return ::executorch::runtime::Error::Ok; + } + + // Emits whatever is held back, including a trailing character the tokens + // never finished. Call it once the generation has ended, or those bytes are + // dropped. Idempotent. + void flush() { + emit_(pending_.size()); + } + + // A character whose bytes have not all arrived yet. + bool has_pending() const { + return !pending_.empty(); + } + + bool failed() const { + return failed_; + } + + private: + void emit_(size_t length) { + if (length == 0) { + return; + } + std::string ready = pending_.substr(0, length); + pending_.erase(0, length); + if (on_text_) { + on_text_(ready); + } + } + + const tokenizers::Tokenizer& tokenizer_; + Sink on_text_; + uint64_t previous_; + std::string pending_; + bool failed_ = false; +}; + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/util.h b/extension/llm/runner/util.h index f5bcb945dbf..b45f1b51f67 100644 --- a/extension/llm/runner/util.h +++ b/extension/llm/runner/util.h @@ -87,8 +87,13 @@ ET_EXPERIMENTAL void inline safe_printf(const char* piece) { // UTF-8 multi-byte sequence. A byte-level tokenizer can emit a token that is // only part of a character (e.g. one byte of a 3-byte CJK codepoint or emoji), // so a caller streaming text must hold the incomplete tail until it completes -// rather than decode the partial bytes. An invalid lead byte counts as length 1 -// (emitted, so the caller can replace it) rather than stalling output. +// rather than decode the partial bytes. +// +// Malformed input is emitted rather than held, so a bad byte never stalls the +// stream: an invalid lead byte counts as length 1, and so does a valid lead +// whose continuation bytes are not 0x80-0xBF. Both reach the caller as single +// bytes it can replace, instead of being run together into a sequence that +// merely looks multi-byte. Only a *well-formed* truncated tail is held back. ET_EXPERIMENTAL size_t inline utf8_complete_prefix_len(const std::string& s) { size_t i = 0; const size_t n = s.size(); @@ -106,6 +111,19 @@ ET_EXPERIMENTAL size_t inline utf8_complete_prefix_len(const std::string& s) { } else { len = 1; // invalid lead byte; emit it and let the caller replace it } + // A lead byte only promises a length; the bytes after it have to agree. + // Checking the ones already present means a malformed sequence degrades to + // the invalid-lead path above rather than being emitted whole, and a + // truncated *valid* sequence is still held for the bytes that finish it. + if (len > 1) { + const size_t seen = (n - i) < len ? (n - i) : len; + for (size_t k = 1; k < seen; ++k) { + if ((static_cast(s[i + k]) & 0xC0) != 0x80) { + len = 1; + break; + } + } + } if (i + len > n) { break; // incomplete trailing sequence: hold it for more bytes } From 595e8c148f3528e5ea171ead7b949e0915a1eae7 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 8 Sep 2026 15:19:18 -0700 Subject: [PATCH 082/190] Let Buck tests inherit the platform C++ standard (#22621) ### Summary Let `cxx_test` targets inherit the C++ language standard selected by their platform toolchain. Forcing C++17 on tests can break builds against Folly headers that require C++20; the wrapper should not override the toolchain's choice. ATen detection and standard selection for other target types remain unchanged. The macro tests now cover dependencies resolved through selects, shared googletest aliases, and test compiler flag patching without an injected language standard. The port maps the mirrored internal files to their OSS paths and applies Python formatting plus `__slots__` to satisfy the OSS linter. Authored with AI assistance using Codex. ### Test plan All 10 macro tests pass. Black, Flake8, and whitespace checks pass: ```bash python3 .ci/scripts/tests/test_is_aten_target.py -v black --check .ci/scripts/tests/test_is_aten_target.py flake8 .ci/scripts/tests/test_is_aten_target.py git diff --check ``` Also checked that the previous macro injects a standard for ATen tests, and that the updated macro initializes missing flags and preserves an explicitly supplied standard. The internal Buck integration target from the original diff was not rerun in this OSS worktree. Full lintrunner was unavailable locally. --- .ci/scripts/tests/test_is_aten_target.py | 74 +++++++++++++++++-- .../executorch/build/runtime_wrapper.bzl | 29 +------- 2 files changed, 68 insertions(+), 35 deletions(-) diff --git a/.ci/scripts/tests/test_is_aten_target.py b/.ci/scripts/tests/test_is_aten_target.py index b8debfddcc9..8a468b2c77b 100644 --- a/.ci/scripts/tests/test_is_aten_target.py +++ b/.ci/scripts/tests/test_is_aten_target.py @@ -27,20 +27,33 @@ REPO_ROOT / "shim_et" / "xplat" / "executorch" / "build" / "runtime_wrapper.bzl" ) -# What the open source dependency map resolves these names to. Kept here rather than -# imported so the test states the mapping it depends on. + +class _Select: + __slots__ = ("values",) + + def __init__(self, values: dict[str, list[str]]) -> None: + self.values = values + + +# Representative dependency-map results, including platform-specific values and +# aliases shared by ATen and non-ATen dependency names. RESOLVED = { - "c10": ["//third-party:libtorch"], + "c10": _Select( + { + "ovr_config//os:android": ["fbsource//xplat/caffe2/c10:c10"], + "DEFAULT": ["fbsource//xplat/caffe2/c10:c10_ovrsource"], + } + ), "libtorch": ["//third-party:libtorch"], "libtorch_python": ["//third-party:libtorch_python"], "torch-core-cpp": ["//third-party:libtorch"], - "gtest_aten": ["//third-party:gtest_aten"], - "gmock_aten": ["//third-party:gmock_aten"], + "gtest_aten": ["fbsource//third-party/googletest:gtest"], + "gmock_aten": ["fbsource//third-party/googletest:gmock"], } FALLTHROUGH = "@fallthrough@" -def _load_is_aten_target(): +def _load_macro_function(name: str): """Execute the real macro text, with the little of Starlark it uses shimmed.""" text = MACROS.read_text() start = text.index("def _has_pytorch_dep") @@ -53,6 +66,8 @@ def _load_is_aten_target(): def _apply(obj, function): """Stand-in for selects.apply: run over each list the object holds.""" + if isinstance(obj, _Select): + return _Select({key: function(value) for key, value in obj.values.items()}) if isinstance(obj, dict): return {key: function(value) for key, value in obj.items()} return function(obj) @@ -64,7 +79,11 @@ def _apply(obj, function): "selects": types.SimpleNamespace(apply=_apply), } exec(compile(text[start:end], str(MACROS), "exec"), namespace) - return namespace["_is_aten_target"] + return namespace[name] + + +def _load_is_aten_target(): + return _load_macro_function("_is_aten_target") class TestIsAtenTarget(unittest.TestCase): @@ -93,6 +112,16 @@ def test_resolved_label_in_exported_deps(self) -> None: ) ) + def test_resolved_label_returned_inside_a_select(self) -> None: + for dep in [ + "fbsource//xplat/caffe2/c10:c10", + "fbsource//xplat/caffe2/c10:c10_ovrsource", + ]: + with self.subTest(dep=dep): + self.assertTrue( + self.is_aten_target({"name": "some_lib", "deps": [dep]}) + ) + def test_short_name_in_external_deps(self) -> None: for name in RESOLVED: with self.subTest(name=name): @@ -101,7 +130,6 @@ def test_short_name_in_external_deps(self) -> None: ) def test_plain_target_is_not_aten(self) -> None: - """The embedded builds rely on these staying at the older standard.""" self.assertFalse( self.is_aten_target( { @@ -114,6 +142,16 @@ def test_plain_target_is_not_aten(self) -> None: ) ) + def test_plain_gtest_target_is_not_aten(self) -> None: + self.assertFalse( + self.is_aten_target( + { + "name": "some_test", + "deps": ["fbsource//third-party/googletest:gtest"], + } + ) + ) + def test_executorch_label_alone_is_not_aten(self) -> None: """Every label under the project contains the word torch.""" self.assertFalse( @@ -147,5 +185,25 @@ def test_select_without_aten_is_not_aten(self) -> None: ) +class TestPatchTestCompilerFlags(unittest.TestCase): + def test_inherits_platform_standard(self) -> None: + patch_test_compiler_flags = _load_macro_function("_patch_test_compiler_flags") + for name in ["some_test", "some_aten_test"]: + with self.subTest(name=name): + kwargs = { + "name": name, + "compiler_flags": ["-DTEST"], + "fbobjc_compiler_flags": ["-DAPPLE_TEST"], + } + + result = patch_test_compiler_flags(kwargs) + + self.assertFalse( + any(flag.startswith("-std=") for flag in result["compiler_flags"]) + ) + self.assertEqual(["-DAPPLE_TEST"], result["fbobjc_compiler_flags"]) + self.assertIn("-Wno-error", result["compiler_flags"]) + + if __name__ == "__main__": unittest.main() diff --git a/shim_et/xplat/executorch/build/runtime_wrapper.bzl b/shim_et/xplat/executorch/build/runtime_wrapper.bzl index 0a2162996a7..5d525b73c00 100644 --- a/shim_et/xplat/executorch/build/runtime_wrapper.bzl +++ b/shim_et/xplat/executorch/build/runtime_wrapper.bzl @@ -207,33 +207,10 @@ def _is_aten_target(kwargs): return True return False -def _patch_test_compiler_flags(kwargs, aten_mode = False): +def _patch_test_compiler_flags(kwargs): if "compiler_flags" not in kwargs: kwargs["compiler_flags"] = [] - # A test that compiles against ATen needs C++20, which PyTorch's headers - # require. Other tests stay at C++17 for embedded builds, but Apple plugin - # generation also requires C++20. - name = kwargs.get("name", "") - is_aten_test = ( - aten_mode or - "_aten" in name or - "aten_" in name - ) - if is_aten_test: - kwargs["compiler_flags"] += [ - "-std=c++20", - ] - else: - kwargs["compiler_flags"] += [ - "-std=c++17", - ] - if env.is_xplat(): - kwargs["fbobjc_compiler_flags"] = kwargs.get( - "fbobjc_compiler_flags", - [], - ) + ["-std=c++20"] - # Relaxing some constraints for tests kwargs["compiler_flags"] += [ "-Wno-missing-prototypes", @@ -391,12 +368,10 @@ def _cxx_test(*args, **kwargs): kwargs["deps"] = [] kwargs["deps"].append("//executorch/test/utils:utils") - # Before _patch_kwargs_cxx, which consumes external_deps. - aten_mode = _is_aten_target(kwargs) _patch_kwargs_cxx(kwargs) env.patch_headers(kwargs) _patch_build_mode_flags(kwargs) - _patch_test_compiler_flags(kwargs, aten_mode) + _patch_test_compiler_flags(kwargs) env.patch_platform_build_mode_flags(kwargs) env.cxx_test(*args, **kwargs) From e2d5b6536f15f0873d16a384313b1a7c94903cbf Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Tue, 8 Sep 2026 15:27:34 -0700 Subject: [PATCH 083/190] Make MLX XCFramework patches upstream-friendly (#22537) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Refactor the MLX XCFramework integration to reduce ExecuTorch-specific changes in the vendored MLX source and make the remaining Metal SDK patch suitable for upstreaming. - Update the per-platform Metal SDK patch to use standard CMake inputs: - `CMAKE_SYSTEM_NAME` - `CMAKE_OSX_SYSROOT` - `CMAKE_OSX_DEPLOYMENT_TARGET` - Remove the patch that customized MLX’s SwiftPM metallib name. - Resolve ExecuTorch’s per-platform metallib from its SwiftPM resource bundle and configure it through MLX’s public `metal::set_metallib_path()` API. - Add `EXECUTORCH_MLX_SWIFTPM_RESOURCES`, defaulting to `OFF`, so ordinary macOS and wheel builds do not compile the Objective-C++ resolver or link Foundation. - Enable the resolver explicitly for Apple framework builds. - Add focused resolver tests and run them in MLX CI. --- .github/workflows/mlx.yml | 6 +- CMakePresets.json | 22 +++ backends/mlx/CMakeLists.txt | 61 +++---- backends/mlx/patches/apply.sh | 0 .../patches/mlx_metal_sdk_per_platform.patch | 63 +++++--- .../patches/mlx_swiftpm_metallib_name.patch | 35 ---- backends/mlx/runtime/MLXBackend.cpp | 32 ++++ backends/mlx/runtime/SwiftPMMetallibPath.h | 24 +++ backends/mlx/runtime/SwiftPMMetallibPath.mm | 152 ++++++++++++++++++ backends/mlx/test/CMakeLists.txt | 26 +++ backends/mlx/test/mlx_metallib_path_test.mm | 123 ++++++++++++++ docs/source/using-executorch-ios.md | 2 +- scripts/build_apple_frameworks.sh | 2 +- tools/cmake/preset/default.cmake | 4 + 14 files changed, 453 insertions(+), 99 deletions(-) mode change 100644 => 100755 backends/mlx/patches/apply.sh delete mode 100644 backends/mlx/patches/mlx_swiftpm_metallib_name.patch create mode 100644 backends/mlx/runtime/SwiftPMMetallibPath.h create mode 100644 backends/mlx/runtime/SwiftPMMetallibPath.mm create mode 100644 backends/mlx/test/mlx_metallib_path_test.mm diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 13a7f9d35b9..7977b234589 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -68,7 +68,7 @@ jobs: echo "::endgroup::" echo "::group::Build test runners" - ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_mutable_state_test mlx_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 )) + ${CONDA_RUN} cmake --build cmake-out --target op_test_runner multi_thread_test_runner mlx_metallib_path_test mlx_mutable_state_test mlx_sequence_cache_test mlx_cell_cache_test -j$(( $(sysctl -n hw.ncpu) - 1 )) echo "::endgroup::" echo "::group::Check MLX artifact sizes" @@ -102,6 +102,10 @@ jobs: exit 1 fi + echo "::group::Run SwiftPM metallib path unit test" + ./cmake-out/backends/mlx/test/mlx_metallib_path_test + echo "::endgroup::" + echo "::group::Run mutable-state (multi-session) unit test" ./cmake-out/backends/mlx/test/mlx_mutable_state_test echo "::endgroup::" diff --git a/CMakePresets.json b/CMakePresets.json index 34a35123ef6..4989990a962 100644 --- a/CMakePresets.json +++ b/CMakePresets.json @@ -90,6 +90,28 @@ "rhs": "Darwin" } }, + { + "name": "apple-framework-resources", + "hidden": true, + "cacheVariables": { + "EXECUTORCH_MLX_SWIFTPM_RESOURCES": "ON" + } + }, + { + "name": "apple-framework-macos", + "displayName": "Build ExecuTorch frameworks for macOS", + "inherits": ["macos", "apple-framework-resources"] + }, + { + "name": "apple-framework-ios", + "displayName": "Build ExecuTorch frameworks for iOS", + "inherits": ["ios", "apple-framework-resources"] + }, + { + "name": "apple-framework-ios-simulator", + "displayName": "Build ExecuTorch frameworks for iOS Simulator", + "inherits": ["ios-simulator", "apple-framework-resources"] + }, { "name": "linux", "displayName": "Build ExecuTorch for Linux", diff --git a/backends/mlx/CMakeLists.txt b/backends/mlx/CMakeLists.txt index 75bfc1f6ed6..b58c82b93cc 100644 --- a/backends/mlx/CMakeLists.txt +++ b/backends/mlx/CMakeLists.txt @@ -187,34 +187,8 @@ message( # each patch file under patches/ for its rationale. set(_mlx_patches ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_sdk_per_platform.patch - ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_swiftpm_metallib_name.patch ${CMAKE_CURRENT_SOURCE_DIR}/patches/mlx_metal_remove_addrspace_compat.patch ) -# In a framework build the delegate is static, so MLX cannot find a colocated -# metallib and instead loads one from a SwiftPM resource bundle. SWIFTPM_BUNDLE -# is the bundle name and MLX_SWIFTPM_METALLIB_NAME the per-slice file inside it -# (one per slice, since a slice's metallib does not load on another). PLATFORM -# is set only for the Apple presets; a plain wheel build leaves it empty and -# keeps the colocated path. -set(_mlx_extra_cxx_flags "") -set(_mlx_cxx_flags_arg "") -if(PLATFORM) - if(PLATFORM STREQUAL "OS64") - set(_mlx_metallib_slice "ios") - elseif(PLATFORM STREQUAL "SIMULATORARM64") - set(_mlx_metallib_slice "ios-simulator") - else() - set(_mlx_metallib_slice "macos") - endif() - set(_mlx_extra_cxx_flags - "-DSWIFTPM_BUNDLE=\\\"executorch_backend_mlx_resources\\\" -DMLX_SWIFTPM_METALLIB_NAME=\\\"mlx-${_mlx_metallib_slice}\\\"" - ) - # Only override the sub-build's CXX flags when there is something to add. An - # empty -DCMAKE_CXX_FLAGS= on the command line beats the environment, so - # passing it unconditionally would silently drop a wheel build's CXXFLAGS for - # MLX only. - set(_mlx_cxx_flags_arg "-DCMAKE_CXX_FLAGS=${_mlx_extra_cxx_flags}") -endif() # Prefer the preset's per-slice DEPLOYMENT_TARGET; fall back to the toolchain # value. @@ -264,7 +238,6 @@ ExternalProject_Add( CMAKE_GENERATOR "Unix Makefiles" CMAKE_ARGS "-DCMAKE_BUILD_TYPE=${_mlx_build_type}" -DCMAKE_CXX_STANDARD=${CMAKE_CXX_STANDARD} - ${_mlx_cxx_flags_arg} # Feed the preset's per-slice DEPLOYMENT_TARGET in as # CMAKE_OSX_DEPLOYMENT_TARGET: the ios.toolchain does not reliably # carry it to this sub-build, and the shader-flag patch reads @@ -309,14 +282,12 @@ ExternalProject_Add( # ExternalProject stamps the patch step and BUILD_ALWAYS does not re-run it, so # a reused build directory whose MLX source was reset (patches reverted) would -# recompile an unpatched MLX and silently drop the iOS Metal SDK selection and -# the SwiftPM metallib name. Re-apply the patches on every build; apply.sh is -# idempotent (it reverse-checks each patch, skipping those already applied). -# DEPENDERS configure, not build: mlx_metal_sdk_per_platform.patch edits the -# sub-project's own CMake to pick the Metal SDK from PLATFORM, so it must land -# before the sub-configure runs, or the shaders are built against the macOS SDK -# for every slice. That does mean the sub-configure re-runs on each build, which -# is the price of the ordering. +# recompile an unpatched MLX and silently drop the iOS Metal SDK selection. +# Re-apply the patches on every build; apply.sh is idempotent (it reverse-checks +# each patch, skipping those already applied). DEPENDERS configure, not build: +# mlx_metal_sdk_per_platform.patch edits the sub-project's own CMake to select +# the Metal SDK, so it must land before the sub-configure runs. That means the +# sub-configure re-runs on each build, which is the price of the ordering. ExternalProject_Add_Step( mlx_external reapply_patches COMMAND bash ${CMAKE_CURRENT_SOURCE_DIR}/patches/apply.sh ${MLX_SOURCE_DIR} @@ -366,6 +337,21 @@ set(_mlx_backend__srcs ${CMAKE_CURRENT_SOURCE_DIR}/runtime/MLXBackend.cpp ${CMAKE_CURRENT_SOURCE_DIR}/runtime/mlx_mutable_state.cpp ) +if(EXECUTORCH_MLX_SWIFTPM_RESOURCES) + if(NOT APPLE) + message( + FATAL_ERROR "EXECUTORCH_MLX_SWIFTPM_RESOURCES requires an Apple target" + ) + endif() + enable_language(OBJCXX) + list(APPEND _mlx_backend__srcs + ${CMAKE_CURRENT_SOURCE_DIR}/runtime/SwiftPMMetallibPath.mm + ) + set_source_files_properties( + ${CMAKE_CURRENT_SOURCE_DIR}/runtime/SwiftPMMetallibPath.mm + PROPERTIES COMPILE_FLAGS "-fobjc-arc" + ) +endif() # Build the delegate as a shared library for the wheel so a C++ consumer can # link it, and keep it static everywhere else so no other build changes. @@ -424,6 +410,11 @@ target_include_directories( target_link_libraries( mlxdelegate PRIVATE mlx_schema extension_llm_cache $ ) +if(EXECUTORCH_MLX_SWIFTPM_RESOURCES) + target_compile_definitions( + mlxdelegate PRIVATE EXECUTORCH_MLX_SWIFTPM_RESOURCES + ) +endif() if(EXECUTORCH_BUILD_SHARED) set_target_properties( mlxdelegate PROPERTIES OUTPUT_NAME executorch_backend_mlx diff --git a/backends/mlx/patches/apply.sh b/backends/mlx/patches/apply.sh old mode 100644 new mode 100755 diff --git a/backends/mlx/patches/mlx_metal_sdk_per_platform.patch b/backends/mlx/patches/mlx_metal_sdk_per_platform.patch index e8983d61844..595b452feca 100644 --- a/backends/mlx/patches/mlx_metal_sdk_per_platform.patch +++ b/backends/mlx/patches/mlx_metal_sdk_per_platform.patch @@ -1,45 +1,56 @@ -Select the Metal SDK per target platform. +Select the Metal SDK from the CMake target platform. -MLX compiles and links its Metal shader library with a hardcoded -`xcrun -sdk macosx metal` and a hardcoded `-mmacosx-version-min` flag. ExecuTorch -builds MLX for iOS device, iOS simulator, and macOS from one source tree, so the -hardcoded macOS SDK produces a metallib built for the wrong platform on the iOS -and simulator slices. +MLX supports both macOS and iOS CMake targets, but its Metal compile and link +commands use the macOS SDK and deployment flag unconditionally. That produces a +metallib for the wrong platform when cross-compiling for an iOS device or +simulator. -Derive the Metal SDK and the deployment-version flag from PLATFORM (which -ExecuTorch already passes to this build): iphoneos for OS64, iphonesimulator for -SIMULATORARM64, macosx otherwise. This mirrors how the rest of the Apple build -selects its SDK. +Use CMAKE_SYSTEM_NAME to distinguish macOS from iOS and CMAKE_OSX_SYSROOT to +distinguish iOS device and simulator SDKs. CMAKE_OSX_SYSROOT may be either a +short SDK name or an absolute, versioned SDK path, so match it case-insensitively. +Fail ambiguous or unsupported configurations instead of silently producing a +metallib for the wrong platform. -Upstream candidate; carried locally until MLX selects the Metal SDK by platform. +Upstream candidate; carried locally until MLX selects the Metal SDK from its +CMake target platform. diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt -index edc169ee..94154f80 100644 +index ecabbda5..d0b86d3b 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt -@@ -9,6 +9,21 @@ set(BASE_HEADERS +@@ -9,6 +9,31 @@ set(BASE_HEADERS logging.h utils.h) -+# The Metal SDK and deployment flag must follow the target platform. ExecuTorch -+# builds this for iOS device, iOS simulator, and macOS from one source tree and -+# passes PLATFORM in for each. Without this the shaders are always built against -+# the macOS SDK, so the iOS and simulator metallibs are wrong for their slice. -+if(PLATFORM STREQUAL "OS64") -+ set(MLX_METAL_SDK iphoneos) -+ set(MLX_METAL_VERSION_MIN_FLAG "-mios-version-min") -+elseif(PLATFORM STREQUAL "SIMULATORARM64") -+ set(MLX_METAL_SDK iphonesimulator) -+ set(MLX_METAL_VERSION_MIN_FLAG "-mios-simulator-version-min") -+else() ++# Match the Metal SDK and deployment flag to the CMake target platform. ++if(CMAKE_SYSTEM_NAME STREQUAL "Darwin") + set(MLX_METAL_SDK macosx) + set(MLX_METAL_VERSION_MIN_FLAG "-mmacosx-version-min") ++elseif(CMAKE_SYSTEM_NAME STREQUAL "iOS") ++ string(TOLOWER "${CMAKE_OSX_SYSROOT}" _mlx_osx_sysroot) ++ if(_mlx_osx_sysroot MATCHES "iphonesimulator") ++ set(MLX_METAL_SDK iphonesimulator) ++ set(MLX_METAL_VERSION_MIN_FLAG "-mios-simulator-version-min") ++ elseif(_mlx_osx_sysroot MATCHES "iphoneos") ++ set(MLX_METAL_SDK iphoneos) ++ set(MLX_METAL_VERSION_MIN_FLAG "-mios-version-min") ++ else() ++ message( ++ FATAL_ERROR ++ "Unable to select the Metal SDK for iOS from CMAKE_OSX_SYSROOT='${CMAKE_OSX_SYSROOT}'" ++ ) ++ endif() ++else() ++ message( ++ FATAL_ERROR ++ "MLX Metal supports only macOS and iOS, got CMAKE_SYSTEM_NAME='${CMAKE_SYSTEM_NAME}'" ++ ) +endif() + function(build_kernel_base TARGET SRCFILE DEPS) set(METAL_FLAGS -x -@@ -26,10 +41,10 @@ function(build_kernel_base TARGET SRCFILE DEPS) +@@ -27,10 +52,10 @@ function(build_kernel_base TARGET SRCFILE DEPS) endif() if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") set(METAL_FLAGS ${METAL_FLAGS} @@ -52,7 +63,7 @@ index edc169ee..94154f80 100644 -I${PROJECT_SOURCE_DIR} -o ${TARGET}.air DEPENDS ${SRCFILE} ${DEPS} ${BASE_HEADERS} OUTPUT ${TARGET}.air -@@ -180,12 +195,12 @@ endif() +@@ -188,12 +213,12 @@ endif() set(METAL_LINK_FLAGS) if(NOT CMAKE_OSX_DEPLOYMENT_TARGET STREQUAL "") diff --git a/backends/mlx/patches/mlx_swiftpm_metallib_name.patch b/backends/mlx/patches/mlx_swiftpm_metallib_name.patch deleted file mode 100644 index 6f08cb6578c..00000000000 --- a/backends/mlx/patches/mlx_swiftpm_metallib_name.patch +++ /dev/null @@ -1,35 +0,0 @@ -Make the SwiftPM metallib name a build-time define. - -MLX loads its Metal library from a SwiftPM resource bundle by a fixed name, -"default". ExecuTorch ships one resource bundle that holds a separate metallib for -each Apple platform slice (device, simulator, macOS), because a metallib built for -one slice does not load on another. Each slice's binary therefore has to ask for -its own file. - -Read the name from MLX_SWIFTPM_METALLIB_NAME when defined, falling back to -"default" so a plain MLX build is unchanged. - -Upstream candidate; carried locally until MLX supports a per-slice bundle name. - -diff --git a/mlx/backend/metal/device.cpp b/mlx/backend/metal/device.cpp -index 29eecf55..61a1180f 100644 ---- a/mlx/backend/metal/device.cpp -+++ b/mlx/backend/metal/device.cpp -@@ -217,8 +217,15 @@ MTL::Library* load_default_library(MTL::Device* device) { - return lib; - } - -- // Then try default.metallib in a SwiftPM bundle if we have one -- std::tie(lib, error[2]) = load_swiftpm_library(device, "default"); -+ // Then try the metallib in a SwiftPM bundle if we have one. The name is a -+ // build-time define because ExecuTorch ships one bundle holding a metallib per -+ // platform slice, so each slice's binary must ask for its own file rather than a -+ // single shared "default". Falls back to "default" for a plain MLX build. -+#ifndef MLX_SWIFTPM_METALLIB_NAME -+#define MLX_SWIFTPM_METALLIB_NAME "default" -+#endif -+ std::tie(lib, error[2]) = -+ load_swiftpm_library(device, MLX_SWIFTPM_METALLIB_NAME); - if (lib) { - return lib; - } diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index 615c4f1c6ce..e16e46ddca6 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/backends/mlx/runtime/MLXBackend.cpp @@ -14,6 +14,10 @@ #include "MLXSequenceCache.h" #include "mlx_mutable_state.h" +#ifdef EXECUTORCH_MLX_SWIFTPM_RESOURCES +#include "SwiftPMMetallibPath.h" +#endif + #include #include @@ -214,6 +218,23 @@ static std::mutex& mlx_global_mutex() { return m; } +#ifdef EXECUTORCH_MLX_SWIFTPM_RESOURCES +// Must be called while holding mlx_global_mutex() and before MLX initializes +// its Metal device. An application-provided path always takes precedence. +static bool configure_metallib_path_locked() { + if (!::mlx::core::metal::get_metallib_path().empty()) { + return true; + } + + const auto path = resolve_swiftpm_metallib_path(); + if (!path.has_value()) { + return false; + } + ::mlx::core::metal::set_metallib_path(*path); + return true; +} +#endif + class MLXBackend final : public ::executorch::runtime::BackendInterface { public: ~MLXBackend() override = default; @@ -236,6 +257,17 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { FreeableBuffer* processed, ArrayRef compile_specs) const override { std::lock_guard lock(mlx_global_mutex()); +#ifdef EXECUTORCH_MLX_SWIFTPM_RESOURCES + if (!configure_metallib_path_locked()) { + ET_LOG( + Error, + "Failed to find the MLX metallib in the SwiftPM resource bundle"); + if (processed != nullptr) { + processed->Free(); + } + return Error::NotFound; + } +#endif auto* handle = context.get_runtime_allocator()->allocateInstance(); if (handle == nullptr) { diff --git a/backends/mlx/runtime/SwiftPMMetallibPath.h b/backends/mlx/runtime/SwiftPMMetallibPath.h new file mode 100644 index 00000000000..8bc838cf93d --- /dev/null +++ b/backends/mlx/runtime/SwiftPMMetallibPath.h @@ -0,0 +1,24 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include + +namespace executorch::backends::mlx { + +// Finds the current platform's metallib in the named SwiftPM resource bundle. +// Each path may be a containing directory, a code bundle, or the resource +// bundle itself. Exposed for focused path tests. +std::optional find_swiftpm_metallib_path( + const std::vector& container_paths); + +// Resolves the current platform's metallib from loaded Apple bundles. +std::optional resolve_swiftpm_metallib_path(); + +} // namespace executorch::backends::mlx diff --git a/backends/mlx/runtime/SwiftPMMetallibPath.mm b/backends/mlx/runtime/SwiftPMMetallibPath.mm new file mode 100644 index 00000000000..977a35b4eee --- /dev/null +++ b/backends/mlx/runtime/SwiftPMMetallibPath.mm @@ -0,0 +1,152 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#import +#import + +#include "SwiftPMMetallibPath.h" + +#include +#include + +namespace executorch::backends::mlx { +namespace { + +constexpr const char* kResourceBundleName = + "executorch_backend_mlx_resources.bundle"; + +const char* metallib_filename() { +#if TARGET_OS_SIMULATOR + return "mlx-ios-simulator.metallib"; +#elif TARGET_OS_IOS + return "mlx-ios.metallib"; +#elif TARGET_OS_OSX + return "mlx-macos.metallib"; +#else + return nullptr; +#endif +} + +std::optional regular_file_path(NSURL* url) { + if (url == nil || !url.fileURL || url.path == nil) { + return std::nullopt; + } + + std::error_code error; + const std::filesystem::path path(url.fileSystemRepresentation); + if (!std::filesystem::is_regular_file(path, error)) { + return std::nullopt; + } + return path.string(); +} + +std::optional find_in_resource_bundle(NSURL* bundle_url) { + if (bundle_url == nil || !bundle_url.fileURL) { + return std::nullopt; + } + + NSBundle* bundle = [NSBundle bundleWithURL:bundle_url]; + if (bundle == nil) { + return std::nullopt; + } + + NSString* filename = [NSString stringWithUTF8String:metallib_filename()]; + if (filename == nil) { + return std::nullopt; + } + + if (auto path = regular_file_path( + [bundle URLForResource:filename.stringByDeletingPathExtension + withExtension:filename.pathExtension])) { + return path; + } + + // SwiftPM's native build system can emit a flat resource bundle. Check the + // bundle root explicitly in addition to Foundation's platform resource URL. + return regular_file_path([bundle_url URLByAppendingPathComponent:filename]); +} + +std::optional find_from_container(NSURL* container_url) { + if (container_url == nil || !container_url.fileURL) { + return std::nullopt; + } + + NSString* resource_bundle_name = + [NSString stringWithUTF8String:kResourceBundleName]; + if ([container_url.lastPathComponent isEqualToString:resource_bundle_name]) { + return find_in_resource_bundle(container_url); + } + + NSBundle* container_bundle = [NSBundle bundleWithURL:container_url]; + if (container_bundle != nil) { + NSURL* resource_bundle_url = [container_bundle + URLForResource:resource_bundle_name.stringByDeletingPathExtension + withExtension:resource_bundle_name.pathExtension]; + if (auto path = find_in_resource_bundle(resource_bundle_url)) { + return path; + } + } + + return find_in_resource_bundle( + [container_url URLByAppendingPathComponent:resource_bundle_name]); +} + +void append_path(NSMutableOrderedSet* paths, NSURL* url) { + if (url != nil && url.fileURL && url.path != nil) { + [paths addObject:url.path]; + } +} + +} // namespace + +std::optional find_swiftpm_metallib_path( + const std::vector& container_paths) { + if (metallib_filename() == nullptr) { + return std::nullopt; + } + + @autoreleasepool { + for (const auto& container_path : container_paths) { + NSString* path = [NSString stringWithUTF8String:container_path.c_str()]; + if (path == nil) { + continue; + } + if (auto metallib_path = + find_from_container([NSURL fileURLWithPath:path])) { + return metallib_path; + } + } + } + + return std::nullopt; +} + +std::optional resolve_swiftpm_metallib_path() { + @autoreleasepool { + NSMutableOrderedSet* paths = [NSMutableOrderedSet orderedSet]; + NSBundle* main_bundle = NSBundle.mainBundle; + append_path(paths, main_bundle.bundleURL); + append_path(paths, main_bundle.resourceURL); + + for (NSBundle* bundle in NSBundle.allBundles) { + append_path(paths, bundle.bundleURL); + append_path(paths, bundle.resourceURL); + } + for (NSBundle* framework in NSBundle.allFrameworks) { + append_path(paths, framework.bundleURL); + append_path(paths, framework.resourceURL); + } + + std::vector container_paths; + container_paths.reserve(paths.count); + for (NSString* path in paths) { + container_paths.emplace_back(path.fileSystemRepresentation); + } + return find_swiftpm_metallib_path(container_paths); + } +} + +} // namespace executorch::backends::mlx diff --git a/backends/mlx/test/CMakeLists.txt b/backends/mlx/test/CMakeLists.txt index 377818ec75a..44693ad8c66 100644 --- a/backends/mlx/test/CMakeLists.txt +++ b/backends/mlx/test/CMakeLists.txt @@ -53,6 +53,32 @@ add_dependencies(op_test_runner strict_compile_test) # Multi-threaded inference test include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) +if(APPLE) + et_cxx_test( + mlx_metallib_path_test SOURCES + ${CMAKE_CURRENT_LIST_DIR}/mlx_metallib_path_test.mm + ${CMAKE_CURRENT_LIST_DIR}/../runtime/SwiftPMMetallibPath.mm + ) + target_include_directories( + mlx_metallib_path_test PRIVATE ${CMAKE_CURRENT_LIST_DIR}/../runtime + ) + set_source_files_properties( + ${CMAKE_CURRENT_LIST_DIR}/mlx_metallib_path_test.mm + ${CMAKE_CURRENT_LIST_DIR}/../runtime/SwiftPMMetallibPath.mm + PROPERTIES COMPILE_FLAGS "-fobjc-arc" + ) + if(EXECUTORCH_MLX_ENABLE_SANITIZERS) + target_compile_options( + mlx_metallib_path_test PRIVATE -fsanitize=address,undefined + -fno-omit-frame-pointer + ) + target_link_options( + mlx_metallib_path_test PRIVATE ${_mlx_sanitizer_link_options} + ) + endif() + target_link_libraries(mlx_metallib_path_test ${FOUNDATION_FRAMEWORK}) +endif() + et_cxx_test( multi_thread_test_runner SOURCES diff --git a/backends/mlx/test/mlx_metallib_path_test.mm b/backends/mlx/test/mlx_metallib_path_test.mm new file mode 100644 index 00000000000..17f15298636 --- /dev/null +++ b/backends/mlx/test/mlx_metallib_path_test.mm @@ -0,0 +1,123 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#import +#import + +#include "SwiftPMMetallibPath.h" + +#include + +#include +#include +#include +#include + +namespace executorch::backends::mlx { +namespace { + +const char* expected_metallib_filename() { +#if TARGET_OS_SIMULATOR + return "mlx-ios-simulator.metallib"; +#elif TARGET_OS_IOS + return "mlx-ios.metallib"; +#else + return "mlx-macos.metallib"; +#endif +} + +const char* wrong_metallib_filename() { +#if TARGET_OS_OSX + return "mlx-ios.metallib"; +#else + return "mlx-macos.metallib"; +#endif +} + +class MLXMetallibPathTest : public ::testing::Test { + protected: + void SetUp() override { + NSString* name = [NSString + stringWithFormat:@"executorch_mlx_metallib_path_%@", + NSUUID.UUID.UUIDString]; + root_ = std::filesystem::temp_directory_path() / + std::string(name.fileSystemRepresentation); + ASSERT_TRUE(std::filesystem::create_directories(root_)); + } + + void TearDown() override { + std::error_code error; + std::filesystem::remove_all(root_, error); + } + + std::filesystem::path create_bundle(bool deep) { + const auto bundle = root_ / "executorch_backend_mlx_resources.bundle"; + const auto contents = deep ? bundle / "Contents" : bundle; + const auto resources = deep ? contents / "Resources" : bundle; + EXPECT_TRUE(std::filesystem::create_directories(resources)); + + std::ofstream(contents / "Info.plist") + << "\n" + << "\n" + << "" + << "CFBundleIdentifier" + << "org.pytorch.executorch.mlx-test-resources" + << "CFBundlePackageTypeBNDL" + << "\n"; + return resources; + } + + std::filesystem::path root_; +}; + +TEST_F(MLXMetallibPathTest, MissingBundleReturnsNoPath) { + EXPECT_FALSE(find_swiftpm_metallib_path({root_.string()}).has_value()); +} + +TEST_F(MLXMetallibPathTest, ProcessWithoutSwiftPMBundleReturnsNoPath) { + EXPECT_FALSE(resolve_swiftpm_metallib_path().has_value()); +} + +TEST_F(MLXMetallibPathTest, FindsFlatBundleResource) { + const auto resources = create_bundle(/*deep=*/false); + const auto metallib = resources / expected_metallib_filename(); + std::ofstream(metallib) << "fixture"; + + EXPECT_EQ(find_swiftpm_metallib_path({root_.string()}), metallib.string()); +} + +TEST_F(MLXMetallibPathTest, FindsMacOSDeepBundleResource) { + const auto resources = create_bundle(/*deep=*/true); + const auto metallib = resources / expected_metallib_filename(); + std::ofstream(metallib) << "fixture"; + + EXPECT_EQ(find_swiftpm_metallib_path({root_.string()}), metallib.string()); + EXPECT_EQ( + find_swiftpm_metallib_path( + {(root_ / "executorch_backend_mlx_resources.bundle").string()}), + metallib.string()); +} + +TEST_F(MLXMetallibPathTest, SelectsCurrentPlatformSlice) { + const auto resources = create_bundle(/*deep=*/true); + const auto expected = resources / expected_metallib_filename(); + const auto wrong = resources / wrong_metallib_filename(); + std::ofstream(expected) << "expected"; + std::ofstream(wrong) << "wrong"; + + EXPECT_EQ(find_swiftpm_metallib_path({root_.string()}), expected.string()); +} + +TEST_F(MLXMetallibPathTest, IgnoresWrongPlatformSlice) { + const auto resources = create_bundle(/*deep=*/true); + std::ofstream(resources / wrong_metallib_filename()) << "fixture"; + + EXPECT_FALSE(find_swiftpm_metallib_path({root_.string()}).has_value()); +} + +} // namespace +} // namespace executorch::backends::mlx diff --git a/docs/source/using-executorch-ios.md b/docs/source/using-executorch-ios.md index 7053c28fd76..c4dc29e198e 100644 --- a/docs/source/using-executorch-ios.md +++ b/docs/source/using-executorch-ios.md @@ -171,7 +171,7 @@ OTHER_LDFLAGS = $(inherited) \ **Note:** In the example above, we link against the Debug version of the ExecuTorch runtime (`libexecutorch_debug`) to preserve the logs. Normally, that does not impact the performance too much. Nevertheless, remember to link against the release version of the runtime (`libexecutorch`) for the best performance and no logs. -**Note:** The MLX backend loads its Metal kernels at runtime from a per-slice metallib inside a resource bundle named `executorch_backend_mlx_resources`, not from the frameworks in `cmake-out`. The build stages the correctly named files (`mlx-ios.metallib`, `mlx-ios-simulator.metallib`, `mlx-macos.metallib`) under `.Package.swift/backend_mlx_resources/`. If you integrate MLX from a source build, ship those files in a bundle of that name for the slices you use, or MLX links and registers but has no kernels to run. +**Note:** The MLX backend loads its Metal kernels at runtime from a per-slice metallib inside a resource bundle named `executorch_backend_mlx_resources`, not from the frameworks in `cmake-out`. The Apple framework presets enable this resource lookup when MLX is available, and the framework build stages the correctly named files (`mlx-ios.metallib`, `mlx-ios-simulator.metallib`, `mlx-macos.metallib`) under `.Package.swift/backend_mlx_resources/`. Generic Apple presets retain MLX's native colocated-metallib lookup. If you use a custom CMake configuration for SwiftPM packaging, enable `EXECUTORCH_MLX_SWIFTPM_RESOURCES` and ship the matching slice in a bundle of that name. You can assign such a config file to your target in Xcode: diff --git a/scripts/build_apple_frameworks.sh b/scripts/build_apple_frameworks.sh index bc201858b30..0fe7b27e7c1 100755 --- a/scripts/build_apple_frameworks.sh +++ b/scripts/build_apple_frameworks.sh @@ -8,7 +8,7 @@ set -euxo pipefail MODES=() -PRESETS=("ios" "ios-simulator" "macos") +PRESETS=("apple-framework-ios" "apple-framework-ios-simulator" "apple-framework-macos") # To support backwards compatibility, we want to retain the same output directory. PRESETS_RELATIVE_OUT_DIR=("ios" "simulator" "macos") diff --git a/tools/cmake/preset/default.cmake b/tools/cmake/preset/default.cmake index f0b1285209e..89e679aec7c 100644 --- a/tools/cmake/preset/default.cmake +++ b/tools/cmake/preset/default.cmake @@ -128,6 +128,10 @@ define_overridable_option( EXECUTORCH_BUILD_EXTENSION_APPLE "Build the Apple extension" BOOL OFF ) define_overridable_option(EXECUTORCH_BUILD_MLX "Build the MLX backend" BOOL OFF) +define_overridable_option( + EXECUTORCH_MLX_SWIFTPM_RESOURCES + "Load the MLX metallib from the ExecuTorch SwiftPM resource bundle" BOOL OFF +) define_overridable_option( EXECUTORCH_BUILD_NEURON "Build the backends/mediatek directory" BOOL OFF ) From 7828708b8429b5056f76374e37098663c8192889 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Tue, 8 Sep 2026 16:50:16 -0700 Subject: [PATCH 084/190] merge pybind test into e2e Differential Revision: D119238464 Pull Request resolved: https://github.com/pytorch/executorch/pull/22623 --- .ci/scripts/tests/test_cuda_workflow.py | 84 +++++++++++++++ .github/workflows/cuda.yml | 135 +++++++----------------- 2 files changed, 124 insertions(+), 95 deletions(-) create mode 100644 .ci/scripts/tests/test_cuda_workflow.py diff --git a/.ci/scripts/tests/test_cuda_workflow.py b/.ci/scripts/tests/test_cuda_workflow.py new file mode 100644 index 00000000000..7eeb5e1e393 --- /dev/null +++ b/.ci/scripts/tests/test_cuda_workflow.py @@ -0,0 +1,84 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from pathlib import Path + +import yaml + + +ROOT = Path(__file__).resolve().parents[3] +WORKFLOW = yaml.safe_load((ROOT / ".github" / "workflows" / "cuda.yml").read_text()) + + +def _all_keys(value): + if isinstance(value, dict): + for key, child in value.items(): + yield key + yield from _all_keys(child) + elif isinstance(value, list): + for child in value: + yield from _all_keys(child) + + +def _model_quant(entry): + return (entry["model"]["repo"], entry["model"]["name"], entry["quant"]) + + +class CudaWorkflowTest(unittest.TestCase): + def test_pybind_runs_inline_for_the_expected_matrix_cells(self): + job = WORKFLOW["jobs"]["test-model-cuda-e2e"] + matrix = job["strategy"]["matrix"] + pybind_rows = [row for row in matrix["include"] if "pybind_model" in row] + + actual = { + (*_model_quant(row), row["pybind_model"], row["pybind_quantized"]) + for row in pybind_rows + } + expected = { + ( + "google", + "gemma-3-4b-it", + "quantized-int4-tile-packed", + "gemma3-4b", + True, + ), + ( + "Qwen", + "Qwen3-0.6B", + "non-quantized", + "qwen3-0.6b", + False, + ), + ( + "Qwen", + "Qwen3-0.6B", + "quantized-int4-tile-packed", + "qwen3-0.6b", + True, + ), + } + self.assertEqual(expected, actual) + + excluded = {_model_quant(row) for row in matrix["exclude"]} + active = { + (model["repo"], model["name"], quant) + for model in matrix["model"] + for quant in matrix["quant"] + } - excluded + self.assertTrue({_model_quant(row) for row in pybind_rows} <= active) + + script = job["with"]["script"] + self.assertIn('if [ -n "${{ matrix.pybind_model }}" ]', script) + self.assertIn("test_huggingface_optimum_model.py", script) + self.assertIn("--run_only", script) + self.assertGreaterEqual(script.count('"${MODEL_DIR}"'), 2) + + def test_model_e2e_does_not_transfer_artifacts(self): + self.assertNotIn("test-cuda-pybind", WORKFLOW["jobs"]) + keys = set(_all_keys(WORKFLOW["jobs"]["test-model-cuda-e2e"])) + self.assertNotIn("upload-artifact", keys) + self.assertNotIn("download-artifact", keys) diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 21c0fe6e844..553e6ac721b 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -504,6 +504,25 @@ jobs: repo: "openai" name: "whisper-large-v3-turbo" quant: "non-quantized" + include: + - model: + repo: "google" + name: "gemma-3-4b-it" + quant: "quantized-int4-tile-packed" + pybind_model: "gemma3-4b" + pybind_quantized: true + - model: + repo: "Qwen" + name: "Qwen3-0.6B" + quant: "non-quantized" + pybind_model: "qwen3-0.6b" + pybind_quantized: false + - model: + repo: "Qwen" + name: "Qwen3-0.6B" + quant: "quantized-int4-tile-packed" + pybind_model: "qwen3-0.6b" + pybind_quantized: true with: timeout: 240 secrets-env: EXECUTORCH_HF_TOKEN @@ -512,7 +531,6 @@ jobs: gpu-arch-version: "13.0" use-custom-docker-registry: false submodules: recursive - upload-artifact: ${{ matrix.model.repo }}-${{ matrix.model.name }}-cuda-${{ matrix.quant }} ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux @@ -547,10 +565,30 @@ jobs: echo "::endgroup::" fi + MODEL_DIR="$(mktemp -d "${RUNNER_TEMP:-/tmp}/cuda_model_XXXXXX")" RUN_EXPORT=1 source .ci/scripts/test_model_e2e.sh cuda \ "${{ matrix.model.repo }}/${{ matrix.model.name }}" \ "${{ matrix.quant }}" \ - "${RUNNER_ARTIFACT_DIR}" + "${MODEL_DIR}" + + if [ -n "${{ matrix.pybind_model }}" ]; then + echo "::group::Run CUDA model with pybind" + conda install -y -c conda-forge 'libstdcxx-ng>=12' + export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + strings /opt/conda/lib/libstdc++.so.6 | grep GLIBCXX_3.4.29 + + PYBIND_ARGS=() + if [ "${{ matrix.pybind_quantized }}" = "true" ]; then + PYBIND_ARGS+=(--quantize) + fi + python .ci/scripts/test_huggingface_optimum_model.py \ + --model "${{ matrix.pybind_model }}" \ + --recipe cuda \ + --model_dir "${MODEL_DIR}" \ + --run_only \ + "${PYBIND_ARGS[@]}" + echo "::endgroup::" + fi test-muse-glimmer-cuda-e2e: name: test-muse-glimmer-cuda-e2e-${{ matrix.variant }}-${{ matrix.mode }} @@ -621,96 +659,3 @@ jobs: "${{ matrix.variant }}" \ "${RUNNER_TEMP}/muse_glimmer" \ "${{ matrix.mode }}" - - test-cuda-pybind: - name: test-cuda-pybind - # This job downloads models exported by test-model-cuda-e2e and runs them using pybind. - # Explicitly check the producer job so a skipped run (fork PR, - # non-sampled push, or no path match) auto-skips this job too. - needs: [changed-files, test-model-cuda-e2e, run-decision] - if: | - needs.test-model-cuda-e2e.result == 'success' && - ( - contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || - contains(needs.changed-files.outputs.changed-files, 'backends/aoti') || - contains(needs.changed-files.outputs.changed-files, '.github/workflows/cuda.yml') || - contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test-cuda-build.sh') || - contains(needs.changed-files.outputs.changed-files, '.ci/scripts/export_model_artifact.sh') || - contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_model_e2e.sh') || - needs.run-decision.outputs.is-full-run == 'true' - ) - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main - permissions: - id-token: write - contents: read - secrets: inherit - strategy: - fail-fast: false - matrix: - include: - - model: "gemma3-4b" - quantize: "--quantize" - artifact: "google-gemma-3-4b-it-cuda-quantized-int4-tile-packed" - - model: "qwen3-0.6b" - quantize: "" - artifact: "Qwen-Qwen3-0.6B-cuda-non-quantized" - - model: "qwen3-0.6b" - quantize: "--quantize" - artifact: "Qwen-Qwen3-0.6B-cuda-quantized-int4-tile-packed" - with: - timeout: 120 - secrets-env: EXECUTORCH_HF_TOKEN - download-artifact: ${{ matrix.artifact }} - runner: linux.g5.4xlarge.nvidia.gpu - gpu-arch-type: cuda - gpu-arch-version: "13.0" - use-custom-docker-registry: false - submodules: recursive - ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} - script: | - set -eux - - # OSDC mounts HF_HOME read-only at /mnt/hf_cache; redirect to a writable dir - # (RUNNER_TEMP, or /tmp when RUNNER_TEMP isn't writable inside the container). - export HF_HOME="${RUNNER_TEMP:-/tmp}/hf_cache" - mkdir -p "${HF_HOME}" 2>/dev/null || export HF_HOME=/tmp/hf_cache - mkdir -p "${HF_HOME}" - - echo "::group::Setup ExecuTorch" - # Disable MKL to avoid duplicate target error when conda has multiple MKL installations - export USE_MKL=OFF - ./install_executorch.sh - echo "::endgroup::" - - echo "::group::Fix libstdc++ GLIBCXX version" - # The embedded .so files in the CUDA blob require GLIBCXX_3.4.29 - # which the default conda libstdc++ doesn't have. Install a newer - # libstdc++ from conda-forge and use it via LD_PRELOAD. - conda install -y -c conda-forge 'libstdcxx-ng>=12' - export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH - # Verify the new libstdc++ has GLIBCXX_3.4.29 - strings /opt/conda/lib/libstdc++.so.6 | grep GLIBCXX_3.4.29 || { - echo "Error: GLIBCXX_3.4.29 not found in /opt/conda/lib/libstdc++.so.6" - exit 1 - } - echo "::endgroup::" - - echo "::group::Setup Huggingface" - pip install -U "huggingface_hub[cli]>=1.2.1,<2.0" - export HF_TOKEN="$(printf '%s' "$SECRET_EXECUTORCH_HF_TOKEN" | tr -d '\r\n')" - echo "::endgroup::" - - echo "::group::Install optimum-executorch" - OPTIMUM_ET_VERSION=$(cat .ci/docker/ci_commit_pins/optimum-executorch.txt) - pip install "optimum~=2.0.0" "transformers==5.0.0rc1" - pip install --no-deps git+https://github.com/huggingface/optimum-executorch.git@${OPTIMUM_ET_VERSION} - echo "::endgroup::" - - echo "::group::Test CUDA Model: ${{ matrix.model }} ${{ matrix.quantize }}" - python .ci/scripts/test_huggingface_optimum_model.py \ - --model ${{ matrix.model }} \ - --recipe cuda \ - --model_dir "${RUNNER_ARTIFACT_DIR}" \ - --run_only \ - ${{ matrix.quantize }} - echo "::endgroup::" From 622b260382f46083f9fcf25bf8cb7e6ad9a12e25 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Tue, 8 Sep 2026 17:19:36 -0700 Subject: [PATCH 085/190] Cache cleanup (#22525) Refactors the off-graph KV cache API around clearer ownership and extension boundaries: * Replaces per-interface lookup and MLX dynamic_cast with one extensible, no-RTTI named-face mechanism. * Renames and tightens the factory/registry APIs, including generated installation keys, RAII publication, checked builder registration, and explicit cache kinds. * Moves planner-specific APIs out of the neutral cache interface, removes dead cache_et.h, and updates MLX integration and callers. * Adds documentation and tests covering face recovery, registration failures, registry lifetime, and shared ownership. --------- Co-authored-by: kiymetakdemir --- backends/mlx/examples/llm/run_llm_hf.cpp | 675 +++++++++++---------- backends/mlx/runtime/MLXBackend.cpp | 44 +- backends/mlx/runtime/MLXCache.h | 6 +- backends/mlx/runtime/MLXCellCache.h | 8 + backends/mlx/runtime/MLXSequenceCache.h | 8 + backends/mlx/runtime/backend_options.h | 8 - backends/mlx/test/mlx_cell_cache_test.cpp | 14 +- backends/mlx/test/op_test_runner.cpp | 26 +- extension/llm/batching/module_executor.cpp | 137 ++--- extension/llm/batching/module_executor.h | 40 +- extension/llm/cache/README.md | 201 ++++++ extension/llm/cache/cache.h | 131 ++-- extension/llm/cache/cache_et.h | 64 -- extension/llm/cache/cache_registry.cpp | 74 ++- extension/llm/cache/cache_registry.h | 108 ++-- extension/llm/cache/cell_cache.h | 27 +- extension/llm/cache/sequence_cache.h | 72 ++- extension/llm/cache/test/cache_test.cpp | 212 +++++-- 18 files changed, 1122 insertions(+), 733 deletions(-) create mode 100644 extension/llm/cache/README.md delete mode 100644 extension/llm/cache/cache_et.h diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp index 37ae7bc8921..1872b9a8a33 100644 --- a/backends/mlx/examples/llm/run_llm_hf.cpp +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -181,21 +181,21 @@ std::optional> const_ints(Module& module, const char* name) { } // Fill in the cache geometry the export published: get_n_caches, then one -// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat), -// plus get_prefill_chunk_size, which the export validates against the sliding -// window and which becomes max_write -- the largest step the cache may see. +// entry per cache in get_kv_heads / get_head_dims / get_windows (0 = flat). // Capacity and dtype stay with the flags. False means this is not an off-graph // model. -bool read_kv_layout(Module& module, cache::CacheConfig& cfg) { +bool read_kv_layout( + Module& module, + int prefill_chunk, + cache::CacheConfig& cfg) { const auto n_caches = const_int(module, "get_n_caches"); const auto kv_heads = const_ints(module, "get_kv_heads"); const auto head_dims = const_ints(module, "get_head_dims"); const auto windows = const_ints(module, "get_windows"); - const auto chunk = const_int(module, "get_prefill_chunk_size"); - if (!n_caches || !kv_heads || !head_dims || !windows || !chunk) { + if (!n_caches || !kv_heads || !head_dims || !windows) { return false; } - cfg.max_write = static_cast(*chunk); + cfg.max_write = prefill_chunk; const size_t n = static_cast(*n_caches); if (kv_heads->size() != n || head_dims->size() != n || windows->size() != n) { return false; @@ -351,19 +351,10 @@ int main(int argc, char** argv) { return 1; } - // Off-graph models (update_and_attend) need a cache bound via cache_key; - // in-graph models (mlx::kv_cache_update) don't -- omit --kv-max-capacity - // for those. session/options are outer-scoped: session must outlive the - // Module (it keeps the cache in the registry) and mlx_opts must outlive - // load_method() (the map holds a view into it). - std::optional session; + // Outer-scoped because mlx_opts must outlive load_method(): the map holds + // a view into it. ::executorch::runtime::BackendOptions<1> mlx_opts; ::executorch::runtime::LoadBackendOptionsMap options_map; - const bool off_graph = kv_capacity > 0; - // Tokens per prefill step, from the .pte. 0 means one step: an - // in-graph model publishes no chunk and has no ring to bound. - int prefill_chunk = 0; - // Load the program but not forward: the cache must exist before forward's // backend init reads its key, and the layout it needs is published by // constant methods in the same file. @@ -373,337 +364,383 @@ int main(int argc, char** argv) { std::cerr << "Failed to load " << pte << std::endl; return 1; } - - if (off_graph) { - cache::CacheConfig cfg{}; - cfg.capacity = kv_capacity; - cfg.kv_dtype = storage_dtype(kv_dtype); - if (cfg.kv_dtype < 0) { - std::cerr << "Invalid --kv-storage-dtype: " << kv_dtype - << " (bf16|fp16|fp32)" << std::endl; - return 1; - } - if (!read_kv_layout(module, cfg)) { - std::cerr << "No KV cache layout in " << pte - << "; re-export with --use-offgraph-cache" << std::endl; - return 1; - } - if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) { - std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl; - return 1; - } - if (!cache::valid(cfg)) { - std::cerr << "Invalid cache config" << std::endl; + const auto published_prefill_chunk = + const_int(module, "get_prefill_chunk_size"); + if (!published_prefill_chunk || *published_prefill_chunk <= 0 || + *published_prefill_chunk > std::numeric_limits::max()) { + std::cerr << "Invalid or missing get_prefill_chunk_size in " << pte + << std::endl; + return 1; + } + const int prefill_chunk = static_cast(*published_prefill_chunk); + + // Everything past load_method is identical for both model kinds; only + // setup differs. ctl is null for an in-graph model, which owns its cache + // inside the graph and exposes no control face. + auto run = + [&](cache::SequenceControl* ctl, + const ::executorch::runtime::LoadBackendOptionsMap* load_opts, + int run_prefill_chunk) -> int { + if (module.load_method( + "forward", + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + load_opts) != Error::Ok) { + std::cerr << "Failed to load forward" << std::endl; return 1; } - if (initial_capacity >= 0) { - cfg.initial_capacity = initial_capacity; - } - auto built = cache::CacheBuilderRegistry::global().build( - ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); - if (!built.ok()) { - std::cerr << "Failed to build cache: " - << static_cast(built.error()) << std::endl; + // Timings reported at the end, in the shared runner's format. + ::executorch::extension::llm::Stats stats; + stats.model_load_start_ms = load_start_ms; + stats.model_load_end_ms = ::executorch::extension::llm::time_in_ms(); + + // Weights-only baseline, so the deltas below isolate the cache. + const double mem_at_load = ::mlx::core::get_active_memory() / 1048576.0; + std::cout << "[mem] after load : " << mem_at_load << " MiB" + << std::endl; + + // Encode. HFTokenizer maps special-token markers in the string to their + // ids, so the template's <|...|> tokens encode correctly; it already + // carries <|begin_of_text|>, so pass bos=0 to avoid a doubled BOS. + std::string enc_input; + if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { + std::cerr << "Unknown --chat template: " << chat + << " (expected llama3, gemma, gemma4, or 0)" << std::endl; return 1; } - prefill_chunk = cfg.max_write ? *cfg.max_write : 0; - session.emplace(cache::make_unique_key(), built.get()); - - print_cache_summary(cfg); - if (mlx_opts.set_option( - ::executorch::backends::mlx::kCacheKeyKey, - session->key().c_str()) != Error::Ok || - options_map.set_options( - ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != - Error::Ok) { - std::cerr << "Failed to set cache_key option" << std::endl; + // The template carries its own BOS, so only a raw prompt asks for one. + const int8_t bos = chat == "0" ? 1 : 0; + auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); + if (!enc.ok()) { + std::cerr << "Encode failed" << std::endl; return 1; } - } - - if (module.load_method( - "forward", - /*planned_memory=*/nullptr, - /*event_tracer=*/nullptr, - off_graph ? &options_map : nullptr) != Error::Ok) { - std::cerr << "Failed to load forward" << std::endl; - return 1; - } - // Timings reported at the end, in the shared runner's format. - ::executorch::extension::llm::Stats stats; - stats.model_load_start_ms = load_start_ms; - stats.model_load_end_ms = ::executorch::extension::llm::time_in_ms(); - - // Weights-only baseline, so the deltas below isolate the cache. - const double mem_at_load = ::mlx::core::get_active_memory() / 1048576.0; - std::cout << "[mem] after load : " << mem_at_load << " MiB" << std::endl; - - // Encode. HFTokenizer maps special-token markers in the string to their - // ids, so the template's <|...|> tokens encode correctly; it already - // carries <|begin_of_text|>, so pass bos=0 to avoid a doubled BOS. - std::string enc_input; - if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { - std::cerr << "Unknown --chat template: " << chat - << " (expected llama3, gemma, gemma4, or 0)" << std::endl; - return 1; - } - // The template carries its own BOS, so only a raw prompt asks for one. - const int8_t bos = chat == "0" ? 1 : 0; - auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); - if (!enc.ok()) { - std::cerr << "Encode failed" << std::endl; - return 1; - } - std::vector tokens = std::move(*enc); - const int prompt_len = static_cast(tokens.size()); - - // End-of-text from the model's metadata when it publishes any, else the - // tokenizer's. The turn-end token is ours: it depends on --chat, which the - // .pte knows nothing about. - std::unordered_set stop_ids = - ::executorch::extension::llm::get_eos_ids(tokenizer.get(), &module); - std::optional turn_end_id; - if (chat != "0") { - const char* turn_end = chat == "llama3" ? "<|eot_id|>" - : chat == "gemma4" ? "" - : ""; - if (auto eot = tokenizer->piece_to_id(turn_end); eot.ok()) { - turn_end_id = static_cast(*eot); - stop_ids.insert(*eot); - } - } - auto is_stop = [&](int64_t t) { - for (uint64_t s : stop_ids) { - if (t == static_cast(s)) { - return true; + std::vector tokens = std::move(*enc); + const int prompt_len = static_cast(tokens.size()); + + // End-of-text from the model's metadata when it publishes any, else the + // tokenizer's. The turn-end token is ours: it depends on --chat, which + // the .pte knows nothing about. + std::unordered_set stop_ids = + ::executorch::extension::llm::get_eos_ids(tokenizer.get(), &module); + std::optional turn_end_id; + if (chat != "0") { + const char* turn_end = chat == "llama3" ? "<|eot_id|>" + : chat == "gemma4" ? "" + : ""; + if (auto eot = tokenizer->piece_to_id(turn_end); eot.ok()) { + turn_end_id = static_cast(*eot); + stop_ids.insert(*eot); } } - return false; - }; - - // One Sampler for the whole run, as the shared runner does: constructing - // one per token would reseed its RNG from the wall clock every time. Built - // on first use because the vocab size comes from the logits -- this export - // publishes no get_vocab_size. - std::optional<::executorch::extension::llm::Sampler> sampler; - - auto step = [&](const std::vector& ids, - const std::vector& pos) { - auto in = - make_tensor_ptr({1, (int)ids.size()}, std::vector(ids)); - auto cp = make_tensor_ptr({(int)pos.size()}, std::vector(pos)); - auto out = module.execute("forward", {in, cp}); - if (!out.ok()) { - throw std::runtime_error("execute failed"); - } - const auto& logits = out->at(0).toTensor(); - if (!sampler) { - sampler.emplace( - static_cast(logits.size(logits.dim() - 1)), temperature); - } - stats.on_sampling_begin(); - const int32_t tok = - ::executorch::extension::llm::sample_from_logits(logits, *sampler); - stats.on_sampling_end(); - return static_cast(tok); - }; - - // Prefill in chunks, so a ring layer holds window + chunk - 1 slots rather - // than growing with the prompt. Only the last chunk's token is kept; the - // earlier ones exist to place their K/V in the cache. - auto prefill = [&](const std::vector& ids, - const std::vector& pos) { - const size_t step_size = - prefill_chunk > 0 ? static_cast(prefill_chunk) : ids.size(); - int64_t next = 0; - for (size_t off = 0; off < ids.size(); off += step_size) { - const size_t n = std::min(step_size, ids.size() - off); - next = step( - {ids.begin() + off, ids.begin() + off + n}, - {pos.begin() + off, pos.begin() + off + n}); - } - return next; - }; - - // Multi-turn: history stays in the cache, so each turn only prefills its - // own tokens at the running position. /reset and /undo drive the cache's - // control face directly -- off-graph only, since an in-graph cache gives - // the runner no handle to its state. - if (interactive) { - if (!off_graph) { - std::cerr << "--interactive requires --kv-max-capacity\n"; - return 1; - } - auto* control = session->control(); - std::cout << "Multi-turn chat. /reset clears, /undo drops the last turn, " - "/undo N drops N tokens, /quit exits.\n"; - int64_t position = 0; - int64_t turn_start = 0; // position this turn began at, for /undo - std::string line; - while (std::cout << "\n> " && std::getline(std::cin, line)) { - if (line == "/quit") { - break; + auto is_stop = [&](int64_t t) { + for (uint64_t s : stop_ids) { + if (t == static_cast(s)) { + return true; + } + } + return false; + }; + + // One Sampler for the whole run, as the shared runner does: constructing + // one per token would reseed its RNG from the wall clock every time. + // Built on first use because the vocab size comes from the logits -- this + // export publishes no get_vocab_size. + std::optional<::executorch::extension::llm::Sampler> sampler; + + auto step = [&](const std::vector& ids, + const std::vector& pos) { + auto in = + make_tensor_ptr({1, (int)ids.size()}, std::vector(ids)); + auto cp = make_tensor_ptr({(int)pos.size()}, std::vector(pos)); + auto out = module.execute("forward", {in, cp}); + if (!out.ok()) { + throw std::runtime_error("execute failed"); + } + const auto& logits = out->at(0).toTensor(); + if (!sampler) { + sampler.emplace( + static_cast(logits.size(logits.dim() - 1)), temperature); + } + stats.on_sampling_begin(); + const int32_t tok = + ::executorch::extension::llm::sample_from_logits(logits, *sampler); + stats.on_sampling_end(); + return static_cast(tok); + }; + + // Prefill in chunks, so a ring layer holds window + chunk - 1 slots + // rather than growing with the prompt. Only the last chunk's token is + // kept; the earlier ones exist to place their K/V in the cache. + auto prefill = [&](const std::vector& ids, + const std::vector& pos) { + const size_t step_size = static_cast(run_prefill_chunk); + int64_t next = 0; + for (size_t off = 0; off < ids.size(); off += step_size) { + const size_t n = std::min(step_size, ids.size() - off); + next = step( + {ids.begin() + off, ids.begin() + off + n}, + {pos.begin() + off, pos.begin() + off + n}); } - if (line == "/reset") { - control->clear(); - position = turn_start = 0; - std::cout << "[cleared]\n"; - continue; + return next; + }; + + // Multi-turn: history stays in the cache, so each turn only prefills its + // own tokens at the running position. /reset and /undo drive the cache's + // control face directly -- off-graph only, since an in-graph cache gives + // the runner no handle to its state. + if (interactive) { + if (ctl == nullptr) { + std::cerr << "--interactive requires --kv-max-capacity\n"; + return 1; } - if (line == "/undo" || line.rfind("/undo ", 0) == 0) { - // Bare /undo drops the last turn; /undo N drops N tokens. - int64_t target = turn_start; - if (line.size() > 6) { - try { - const int64_t n = std::stoll(line.substr(6)); - target = n >= position ? 0 : position - n; - } catch (const std::exception&) { - std::cout << "[usage: /undo [n_tokens]]\n"; - continue; + auto* control = ctl; + std::cout + << "Multi-turn chat. /reset clears, /undo drops the last turn, " + "/undo N drops N tokens, /quit exits.\n"; + int64_t position = 0; + int64_t turn_start = 0; // position this turn began at, for /undo + std::string line; + while (std::cout << "\n> " && std::getline(std::cin, line)) { + if (line == "/quit") { + break; + } + if (line == "/reset") { + control->clear(); + position = turn_start = 0; + std::cout << "[cleared]\n"; + continue; + } + if (line == "/undo" || line.rfind("/undo ", 0) == 0) { + // Bare /undo drops the last turn; /undo N drops N tokens. + int64_t target = turn_start; + if (line.size() > 6) { + try { + const int64_t n = std::stoll(line.substr(6)); + target = n >= position ? 0 : position - n; + } catch (const std::exception&) { + std::cout << "[usage: /undo [n_tokens]]\n"; + continue; + } + } + if (control->rewind(static_cast(target))) { + position = target; + turn_start = std::min(turn_start, position); + std::cout << "[rewound to " << position << "]\n"; + } else { + // A sliding-window layer has physically dropped those cells. + std::cout << "[cannot rewind to " << target << "]\n"; } + continue; } - if (control->rewind(static_cast(target))) { - position = target; - turn_start = std::min(turn_start, position); - std::cout << "[rewound to " << position << "]\n"; - } else { - // A sliding-window layer has physically dropped those cells. - std::cout << "[cannot rewind to " << target << "]\n"; + if (line.empty()) { + continue; } - continue; - } - if (line.empty()) { - continue; + + std::string turn; + wrap_turn(chat, line, /*with_bos=*/position == 0, turn); + auto te = tokenizer->encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); + if (!te.ok()) { + std::cerr << "Encode failed\n"; + continue; + } + const int n = static_cast(te->size()); + // Admit the turn if its prompt plus one token fits; reserving the + // whole max_new budget up front would report "full" with most of the + // cache still free. Generation is then clamped to the room that + // remains. + if (!control->can_extend(n + 1)) { + std::cout << "[cache full: " << position << "/" + << control->capacity() << ", turn " << n << " tokens" + << (control->can_extend(1) ? "" : ", length at capacity") + << ", use /reset]\n"; + continue; + } + const int budget = std::min( + max_new, control->capacity() - static_cast(position) - n); + + turn_start = position; + std::vector tin(te->begin(), te->end()), tpos; + for (int i = 0; i < n; ++i) { + tpos.push_back(position + i); + } + int64_t next = prefill(tin, tpos); + position += n; + + uint64_t prev = te->back(); + for (int i = 0; i < budget && !is_stop(next); ++i) { + if (auto piece = + tokenizer->decode(prev, static_cast(next)); + piece.ok()) { + std::cout << *piece << std::flush; + } + prev = static_cast(next); + next = step({next}, {position}); + ++position; + } + // The turn-end token stops generation, so it is neither printed nor + // fed back -- but the next turn opens without closing this one, and + // an unterminated assistant turn compounds over a session. Commit it, + // at the cost of one extra step per turn. + if (turn_end_id && next == *turn_end_id && control->can_extend(1)) { + step({next}, {position}); + ++position; + } + std::cout << "\n[" << position << "/" << control->capacity() + << " tokens" + << (budget < max_new ? ", generation capped by capacity" + : "") + << "]\n"; } + return 0; + } - std::string turn; - wrap_turn(chat, line, /*with_bos=*/position == 0, turn); - auto te = tokenizer->encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); - if (!te.ok()) { - std::cerr << "Encode failed\n"; - continue; + std::vector ids(tokens.begin(), tokens.end()), prefill_pos; + for (int i = 0; i < prompt_len; ++i) { + prefill_pos.push_back(i); + } + auto ms = [](auto a, auto b) { + return std::chrono::duration(b - a).count(); + }; + // Sequence length against the configured ceiling, with what MLX actually + // holds for it. Pools start at initial_capacity and grow by doubling, so + // the bytes lag the token count in steps; bf16 storage (kv_dtype 15) + // halves them vs fp32 (6). + auto print_footprint = [&](const char* when, int len) { + if (ctl == nullptr) { + return; } - const int n = static_cast(te->size()); - // Admit the turn if its prompt plus one token fits; reserving the whole - // max_new budget up front would report "full" with most of the cache - // still free. Generation is then clamped to the room that remains. - if (!control->can_extend(n + 1)) { - std::cout << "[cache full: " << position << "/" << control->capacity() - << ", turn " << n << " tokens" - << (control->can_extend(1) ? "" : ", length at capacity") - << ", use /reset]\n"; - continue; + const int cap = ctl->capacity(); + const double pct = cap > 0 ? 100.0 * len / cap : 0.0; + std::cout << "[cache] " << when << ": " << len << " / " << cap + << " tokens (" << pct << "%)" << std::endl; + const double mem = ::mlx::core::get_active_memory() / 1048576.0; + std::cout << "[mem] " << when << ": " << mem << " MiB (+" + << (mem - mem_at_load) << " MiB since load)" << std::endl; + }; + + // One optional warmup run to absorb JIT and pool growth, then one + // measured run, as the shared LLM runners do. Repeats belong in a harness + // that restarts the process: clear() rewinds the sequence but leaves the + // pools at their grown size, so an in-process repeat cannot see + // reallocation. + for (int iter = 0; iter < (warmup ? 2 : 1); ++iter) { + const bool measured = !warmup || iter == 1; + if (iter > 0 && ctl != nullptr) { + ctl->clear(); } - const int budget = std::min( - max_new, control->capacity() - static_cast(position) - n); - - turn_start = position; - std::vector tin(te->begin(), te->end()), tpos; - for (int i = 0; i < n; ++i) { - tpos.push_back(position + i); + stats.inference_start_ms = ::executorch::extension::llm::time_in_ms(); + int64_t next = prefill(ids, prefill_pos); + stats.prompt_eval_end_ms = ::executorch::extension::llm::time_in_ms(); + // prefill returns the first generated token, so TTFT ends with prefill + stats.first_token_ms = stats.prompt_eval_end_ms; + if (measured) { + std::cout << "\n"; + print_footprint("after prefill", prompt_len); + std::cout << "\n"; // blank line before the streamed generation } - int64_t next = prefill(tin, tpos); - position += n; - - uint64_t prev = te->back(); - for (int i = 0; i < budget && !is_stop(next); ++i) { - if (auto piece = tokenizer->decode(prev, static_cast(next)); - piece.ok()) { - std::cout << *piece << std::flush; + + uint64_t prev = tokens.back(); + int generated = 0; + for (int i = 0; i < max_new; ++i) { + if (is_stop(next)) { + break; + } + if (measured) { + if (auto piece = + tokenizer->decode(prev, static_cast(next)); + piece.ok()) { + ::executorch::extension::llm::safe_printf(piece->c_str()); + fflush(stdout); + } } prev = static_cast(next); - next = step({next}, {position}); - ++position; + ++generated; + next = step({next}, {prompt_len + i}); } - // The turn-end token stops generation, so it is neither printed nor - // fed back -- but the next turn opens without closing this one, and an - // unterminated assistant turn compounds over a session. Commit it, at - // the cost of one extra step per turn. - if (turn_end_id && next == *turn_end_id && control->can_extend(1)) { - step({next}, {position}); - ++position; + stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); + if (measured) { + std::cout << "\n\n"; // close the generation line + blank separator + // trailing space aligns the colon with the "after prefill" line above + print_footprint("after decode ", prompt_len + generated); + stats.num_prompt_tokens = prompt_len; + stats.num_generated_tokens = generated; } - std::cout << "\n[" << position << "/" << control->capacity() - << " tokens" - << (budget < max_new ? ", generation capped by capacity" : "") - << "]\n"; } + std::cout << std::endl; + ::executorch::extension::llm::print_report(stats); return 0; + }; + + // An in-graph model (mlx::kv_cache_update) binds no cache: nothing to + // build, no key to hand the delegate, and so no registry entry to guard. + if (kv_capacity <= 0) { + return run( + /*ctl=*/nullptr, + /*load_opts=*/nullptr, + /*run_prefill_chunk=*/prefill_chunk); } - std::vector ids(tokens.begin(), tokens.end()), prefill_pos; - for (int i = 0; i < prompt_len; ++i) { - prefill_pos.push_back(i); + cache::CacheConfig cfg{}; + cfg.capacity = kv_capacity; + cfg.kv_dtype = storage_dtype(kv_dtype); + if (cfg.kv_dtype < 0) { + std::cerr << "Invalid --kv-storage-dtype: " << kv_dtype + << " (bf16|fp16|fp32)" << std::endl; + return 1; + } + if (!read_kv_layout(module, prefill_chunk, cfg)) { + std::cerr << "No KV cache layout in " << pte + << "; re-export with --use-offgraph-cache" << std::endl; + return 1; + } + if (!kv_windows.empty() && !apply_window_override(kv_windows, cfg)) { + std::cerr << "Invalid --kv-windows: " << kv_windows << std::endl; + return 1; + } + if (!cache::valid(cfg)) { + std::cerr << "Invalid cache config" << std::endl; + return 1; + } + if (initial_capacity >= 0) { + cfg.initial_capacity = initial_capacity; } - auto ms = [](auto a, auto b) { - return std::chrono::duration(b - a).count(); - }; - // Sequence length against the configured ceiling, with what MLX actually - // holds for it. Pools start at initial_capacity and grow by doubling, so - // the bytes lag the token count in steps; bf16 storage (kv_dtype 15) halves - // them vs fp32 (6). - auto print_footprint = [&](const char* when, int len) { - if (!session) { - return; - } - const int cap = session->control()->capacity(); - const double pct = cap > 0 ? 100.0 * len / cap : 0.0; - std::cout << "[cache] " << when << ": " << len << " / " << cap - << " tokens (" << pct << "%)" << std::endl; - const double mem = ::mlx::core::get_active_memory() / 1048576.0; - std::cout << "[mem] " << when << ": " << mem << " MiB (+" - << (mem - mem_at_load) << " MiB since load)" << std::endl; - }; - // One optional warmup run to absorb JIT and pool growth, then one measured - // run, as the shared LLM runners do. Repeats belong in a harness that - // restarts the process: clear() rewinds the sequence but leaves the pools - // at their grown size, so an in-process repeat cannot see reallocation. - for (int iter = 0; iter < (warmup ? 2 : 1); ++iter) { - const bool measured = !warmup || iter == 1; - if (iter > 0 && off_graph) { - session->control()->clear(); - } - stats.inference_start_ms = ::executorch::extension::llm::time_in_ms(); - int64_t next = prefill(ids, prefill_pos); - stats.prompt_eval_end_ms = ::executorch::extension::llm::time_in_ms(); - // prefill returns the first generated token, so TTFT ends with prefill - stats.first_token_ms = stats.prompt_eval_end_ms; - if (measured) { - std::cout << "\n"; - print_footprint("after prefill", prompt_len); - std::cout << "\n"; // blank line before the streamed generation - } + const char* const cache_kind = cache::kind::kSingle; + auto built = cache::CacheFactory::global().build( + ::executorch::backends::mlx::kMLXBackendId, cache_kind, cfg); + if (!built.ok()) { + std::cerr << "Failed to build cache: " << static_cast(built.error()) + << std::endl; + return 1; + } + const std::shared_ptr kv = built.get(); + + // Published for the delegate to find by key, and erased when this scope + // exits. That is after run() returns, so the entry is still there for the + // load_method() inside it. + const cache::InstallGuard guard{kv}; + + print_cache_summary(cfg); + if (guard.set_option(mlx_opts) != Error::Ok || + options_map.set_options( + ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != + Error::Ok) { + std::cerr << "Failed to set cache_key option" << std::endl; + return 1; + } - uint64_t prev = tokens.back(); - int generated = 0; - for (int i = 0; i < max_new; ++i) { - if (is_stop(next)) { - break; - } - if (measured) { - if (auto piece = tokenizer->decode(prev, static_cast(next)); - piece.ok()) { - ::executorch::extension::llm::safe_printf(piece->c_str()); - fflush(stdout); - } - } - prev = static_cast(next); - ++generated; - next = step({next}, {prompt_len + i}); - } - stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); - if (measured) { - std::cout << "\n\n"; // close the generation line + blank separator - // trailing space aligns the colon with the "after prefill" line above - print_footprint("after decode ", prompt_len + generated); - stats.num_prompt_tokens = prompt_len; - stats.num_generated_tokens = generated; - } + // Checked here so a null ctl inside run() can only mean "in-graph model". + // A cache kind that offers BatchControl instead would otherwise be run as + // if it had no cache at all, with a key published and options set. + auto* ctl = kv->as(); + if (ctl == nullptr) { + std::cerr << "Cache kind '" << cache_kind + << "' offers no single-sequence control face" << std::endl; + return 1; } - std::cout << std::endl; - ::executorch::extension::llm::print_report(stats); - return 0; + + return run(ctl, &options_map, *cfg.max_write); } catch (const std::exception& e) { std::cerr << "Error: " << e.what() << std::endl; return 1; diff --git a/backends/mlx/runtime/MLXBackend.cpp b/backends/mlx/runtime/MLXBackend.cpp index e16e46ddca6..f5a79a9402d 100644 --- a/backends/mlx/runtime/MLXBackend.cpp +++ b/backends/mlx/runtime/MLXBackend.cpp @@ -25,6 +25,7 @@ #include #include #include +#include #include @@ -190,8 +191,8 @@ struct MLXHandle { // Keep-alive for the off-graph KV cache bound in init(). state.cache is a // non-owning view of the same object, so the cache must outlive the handle - // even if the runner's session is torn down first. - std::shared_ptr<::executorch::extension::llm::cache::CacheBase> cache_shared; + // even if the runner drops its InstallGuard first. + std::shared_ptr<::executorch::extension::llm::cache::Cache> cache_shared; // Keep the constant buffers alive for zero-copy constants // Each FreeableBuffer must outlive the MLX arrays that reference it @@ -371,7 +372,8 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { // Bind the off-graph KV cache, if the runner installed one under a key it // passed as a runtime spec. Bound before the init chain runs so an // update_and_attend node there sees the same cache execute() will. - if (auto spec = context.get_runtime_spec(kCacheKeyKey); + if (auto spec = + context.get_runtime_spec(cache::kCacheKeyOption); spec.ok() && spec.get() != nullptr && *spec.get() != '\0') { const char* cache_key = spec.get(); handle->cache_shared = @@ -381,12 +383,10 @@ class MLXBackend final : public ::executorch::runtime::BackendInterface { std::string("init: cache_key '") + cache_key + "' is not installed in the CacheRegistry"); } - // Cross-cast from the neutral ownership anchor to this backend's - // tensor-typed op face; the two are deliberately unrelated bases (see - // MLXCache.h), so nullptr here means the key names another backend's - // cache. - handle->state.cache = - dynamic_cast(handle->cache_shared.get()); + // Ask the neutral ownership anchor for this backend's tensor-typed op + // face. It is named by MLXCache itself rather than by cache.h, so + // nullptr here means the key names another backend's cache. + handle->state.cache = handle->cache_shared->as(); if (handle->state.cache == nullptr) { throw std::runtime_error( std::string("init: cache under key '") + cache_key + @@ -603,18 +603,30 @@ static auto success_with_compiler = register_backend(backend); // Cache kind is named by the builder tag rather than an enum on the config: a // runner asks the registry for (backend_id, kind) and gets back a neutral -// CacheBase it installs under a cache_key. Adding a kind is a new builder here. +// Cache it installs under a cache_key. Adding a kind is a new builder here. const int cache_builders_registered = [] { - cache::CacheBuilderRegistry::global().register_builder( - kMLXBackendId, "seq", [](const cache::CacheConfig& cfg) { - return std::shared_ptr( + const Error single = cache::CacheFactory::global().register_builder( + kMLXBackendId, cache::kind::kSingle, [](const cache::CacheConfig& cfg) { + return std::shared_ptr( std::make_shared(cfg)); }); - cache::CacheBuilderRegistry::global().register_builder( - kMLXBackendId, "cell", [](const cache::CacheConfig& cfg) { - return std::shared_ptr( + ET_CHECK_MSG( + single == Error::Ok, + "Failed to register cache builder for %s:%s", + kMLXBackendId, + cache::kind::kSingle); + const Error batched_cell = cache::CacheFactory::global().register_builder( + kMLXBackendId, + cache::kind::kBatchedCell, + [](const cache::CacheConfig& cfg) { + return std::shared_ptr( std::make_shared(cfg)); }); + ET_CHECK_MSG( + batched_cell == Error::Ok, + "Failed to register cache builder for %s:%s", + kMLXBackendId, + cache::kind::kBatchedCell); return 0; }(); } // namespace diff --git a/backends/mlx/runtime/MLXCache.h b/backends/mlx/runtime/MLXCache.h index 523d3054712..40da8e94110 100644 --- a/backends/mlx/runtime/MLXCache.h +++ b/backends/mlx/runtime/MLXCache.h @@ -30,12 +30,16 @@ struct AttendSpec { }; // Tensor-typed op face of the off-graph KV cache, kept separate from the -// neutral CacheBase (which is tensor-free) so a cache can expose both without a +// neutral Cache (which is tensor-free) so a cache can expose both without a // diamond. ExecutionState holds one; nothing assigns it yet -- the registry // that owns the cache and hands this pointer to the executor lands in a // follow-up, until which exec_update_and_attend is unreachable. class MLXCache { public: + // Named here, not in cache.h: a backend face is tensor-typed and the + // neutral header cannot know about it. + static constexpr const char* kFaceName = "mlx.MLXCache"; + virtual ~MLXCache() = default; // Write this step's K/V for `layer` at `positions`, one host int per query diff --git a/backends/mlx/runtime/MLXCellCache.h b/backends/mlx/runtime/MLXCellCache.h index 566fc6f7bfa..91719fa0526 100644 --- a/backends/mlx/runtime/MLXCellCache.h +++ b/backends/mlx/runtime/MLXCellCache.h @@ -84,6 +84,14 @@ class MLXCellCache : public cache::CellCache, public MLXCache { mask(*step)}; } + protected: + void* face(cache::FaceId id) override { + if (void* p = cache::CellCache::face(id)) { + return p; + } + return cache::expose(this, id); + } + private: // The step's bits as SDPA wants them: [1, 1, length, read_len], one row per // query token. diff --git a/backends/mlx/runtime/MLXSequenceCache.h b/backends/mlx/runtime/MLXSequenceCache.h index 79beeeb846e..3680b16087c 100644 --- a/backends/mlx/runtime/MLXSequenceCache.h +++ b/backends/mlx/runtime/MLXSequenceCache.h @@ -125,6 +125,14 @@ class MLXSequenceCache : public cache::SequenceCache, public MLXCache { return AttendSpec{K, V, AttendSpec::Mask::Causal, std::nullopt}; } + protected: + void* face(cache::FaceId id) override { + if (void* p = cache::SequenceCache::face(id)) { + return p; + } + return cache::expose(this, id); + } + private: // A sequence cache holds one run of one sequence, so the step is described by // where it starts; the remaining positions carry no information beyond diff --git a/backends/mlx/runtime/backend_options.h b/backends/mlx/runtime/backend_options.h index af9a993ce7e..2f96ce93525 100644 --- a/backends/mlx/runtime/backend_options.h +++ b/backends/mlx/runtime/backend_options.h @@ -42,14 +42,6 @@ inline constexpr char kClearCacheIntervalKey[] = "clear_cache_interval"; // errors otherwise). Saves one full mutable-buffer (KV-cache) copy per handle. inline constexpr char kSkipMutableBufferInitKey[] = "skip_mutable_buffer_init"; -// Per-model runtime-spec key (string). Names the off-graph KV cache this handle -// binds to: the runner creates the cache, installs it in the process-global -// CacheRegistry under this key, and the delegate looks it up in init(). The -// DelegateHandle is opaque to the host, so the key is the only rendezvous -// channel. Unset means no cache, and any update_and_attend node then fails at -// execute() rather than silently attending nothing. -inline constexpr char kCacheKeyKey[] = "cache_key"; - } // namespace mlx } // namespace backends } // namespace executorch diff --git a/backends/mlx/test/mlx_cell_cache_test.cpp b/backends/mlx/test/mlx_cell_cache_test.cpp index a93590b453e..6e68ba79ef5 100644 --- a/backends/mlx/test/mlx_cell_cache_test.cpp +++ b/backends/mlx/test/mlx_cell_cache_test.cpp @@ -268,12 +268,16 @@ TEST_F(MLXCellCacheTest, InvalidConfigThrows) { // A runner reaches a layout by (backend_id, kind), so the builder registration // is as much a part of the layout as the class. TEST_F(MLXCellCacheTest, RegistryBuildsCellLayout) { - auto built = cache::CacheBuilderRegistry::global().build( - kMLXBackendId, "cell", flat_config(32, 1, H, D, kHalf)); + auto built = cache::CacheFactory::global().build( + kMLXBackendId, + cache::kind::kBatchedCell, + flat_config(32, 1, H, D, kHalf)); ASSERT_TRUE(built.ok()); - const std::shared_ptr& c = *built; - EXPECT_NE(c->as_batch_control(), nullptr); - EXPECT_EQ(c->as_control(), nullptr); + const std::shared_ptr& c = *built; + EXPECT_NE(c->as(), nullptr); + EXPECT_NE(c->as(), nullptr) << "the backend face comes back too"; + // A cell layout is multi-sequence, so it offers no single-sequence face. + EXPECT_EQ(c->as(), nullptr); } } // namespace diff --git a/backends/mlx/test/op_test_runner.cpp b/backends/mlx/test/op_test_runner.cpp index 53291e41064..12eb7235ce5 100644 --- a/backends/mlx/test/op_test_runner.cpp +++ b/backends/mlx/test/op_test_runner.cpp @@ -300,38 +300,33 @@ int main(int argc, char* argv[]) { namespace cache = ::executorch::extension::llm::cache; - // Build and install the off-graph KV cache before the Module, so the - // registry entry exists by the time the delegate's init() looks it up. - // Declared here so the session outlives the module. - std::optional cache_session; + // Publish the off-graph KV cache until the delegate resolves its key while + // loading the method. + std::optional cache_install_guard; if (!kv_cache_spec.empty()) { cache::CacheConfig cfg{}; if (!parse_kv_cache_spec(kv_cache_spec, cfg)) { std::cerr << "Invalid --kv-cache spec: " << kv_cache_spec << std::endl; return 1; } - auto built = cache::CacheBuilderRegistry::global().build( - ::executorch::backends::mlx::kMLXBackendId, "seq", cfg); + auto built = cache::CacheFactory::global().build( + ::executorch::backends::mlx::kMLXBackendId, + cache::kind::kSingle, + cfg); if (!built.ok()) { std::cerr << "Failed to build KV cache: " << static_cast(built.error()) << std::endl; return 1; } - cache_session.emplace(cache::make_unique_key(), built.get()); - if (verbose) { - std::cout << "Installed KV cache under key " << cache_session->key() - << std::endl; - } + cache_install_guard.emplace(built.get()); } Module module(pte_path); Error load_error = Error::Ok; - if (cache_session) { + if (cache_install_guard) { ::executorch::runtime::BackendOptions<1> mlx_opts; ::executorch::runtime::LoadBackendOptionsMap options_map; - if (mlx_opts.set_option( - ::executorch::backends::mlx::kCacheKeyKey, - cache_session->key().c_str()) != Error::Ok || + if (cache_install_guard->set_option(mlx_opts) != Error::Ok || options_map.set_options( ::executorch::backends::mlx::kMLXBackendId, mlx_opts.view()) != Error::Ok) { @@ -358,6 +353,7 @@ int main(int argc, char* argv[]) { << static_cast(load_method_error) << std::endl; return 1; } + cache_install_guard.reset(); if (verbose) { std::cout << "Reading inputs from: " << input_path << std::endl; diff --git a/extension/llm/batching/module_executor.cpp b/extension/llm/batching/module_executor.cpp index fb9ad912b92..522a928e508 100644 --- a/extension/llm/batching/module_executor.cpp +++ b/extension/llm/batching/module_executor.cpp @@ -10,7 +10,7 @@ #include #include -#include +#include #include #include @@ -19,6 +19,7 @@ #include #include #include +#include #include namespace executorch { @@ -86,39 +87,20 @@ Result config_from_program(Module& module) { return cfg; } -// The backend-load option the delegate resolves the cache through. -constexpr char kCacheKeyOption[] = "cache_key"; - std::uint64_t nondeterministic_seed() { std::random_device device; return device(); } -// One forward's inputs, flattened across the batch. Entry i of `tokens` and of -// `positions` names the same token, which is how the cache pairs them. -struct Step { - // Signed to match the model's token input, not Token. - std::vector tokens; - std::vector positions; - // The sequence each token belongs to, declared to the cache one slice at a - // time so a declaration always matches the forward that places it. - std::vector seq_ids; - // Per input: the logits row it draws from, or -1 when its prediction is - // discarded. An input of any width contributes one, since only its last row - // predicts a token the session does not hold. - std::vector logit_indices; -}; - -// Flatten the batch and truncate whatever it reopens; execute() declares each -// slice to the cache as it runs it. A per-sequence cursor carries the batch's -// own writes, so consecutive chunks of one prompt abut and only the first can -// reopen committed ground. Every input is checked before any is truncated, so -// a refusal leaves the cache untouched. -Result build_step( - cache::BatchControl& ctl, - const BatchInput& batch, - const std::unordered_map& sessions, - int max_session_tokens) { +} // namespace + +Result ModuleExecutor::build_step( + const BatchInput& batch) { + // Flatten the batch and truncate whatever it reopens; execute() declares each + // slice to the cache as it runs it. A per-sequence cursor carries the batch's + // own writes, so consecutive chunks of one prompt abut and only the first can + // reopen committed ground. Every input is checked before any is truncated, so + // a refusal leaves the cache untouched. Step step; const std::size_t total = batch.size(); step.tokens.reserve(total); @@ -133,14 +115,15 @@ Result build_step( std::unordered_map cursor; for (const Input& input : batch.inputs) { - const auto seq_it = sessions.find(input.sid); - if (seq_it == sessions.end()) { + const auto seq_it = sessions_.find(input.sid); + if (seq_it == sessions_.end()) { ET_LOG(Error, "build_step: session %" PRId64 " is not open", input.sid); return Error::InvalidArgument; } const std::int32_t seq_id = seq_it->second.seq_id; if (input.size == 0 || !input.tokens || - input.offset + input.size > input.tokens->size()) { + input.offset > input.tokens->size() || + input.size > input.tokens->size() - input.offset) { ET_LOG( Error, "build_step: session %" PRId64 " gave a slice its tokens do not hold", @@ -151,7 +134,7 @@ Result build_step( const std::int64_t start = static_cast(input.position) + static_cast(input.offset); const auto [cursor_it, first_for_seq] = - cursor.try_emplace(seq_id, ctl.next_pos(seq_id)); + cursor.try_emplace(seq_id, ctl_->next_pos(seq_id)); int& at = cursor_it->second; if (start > at) { // Positions nothing attended, and nothing later reaches back to fill. @@ -188,13 +171,13 @@ Result build_step( } const std::int64_t end = start + static_cast(input.size); - if (end > max_session_tokens) { + if (end > max_session_tokens_) { ET_LOG( Error, "build_step: session %" PRId64 " reaches %" PRId64 " of %d cells", input.sid, end, - max_session_tokens); + max_session_tokens_); return Error::OutOfResources; } @@ -212,7 +195,7 @@ Result build_step( } for (const auto& [seq_id, from] : rewinds) { - if (!ctl.seq_rm(seq_id, from, std::nullopt)) { + if (!ctl_->seq_rm(seq_id, from, std::nullopt)) { ET_LOG(Error, "build_step: sequence %d would not truncate", seq_id); return Error::Internal; } @@ -220,22 +203,18 @@ Result build_step( return step; } -} // namespace - ModuleExecutor::ModuleExecutor( std::unique_ptr module, - std::shared_ptr cache, - std::unique_ptr session, + std::shared_ptr cache, int max_sessions, int max_session_tokens, std::string backend_id, std::string method, std::int32_t vocab_size, int max_step_tokens) - : session_(std::move(session)), - cache_(std::move(cache)), + : install_guard_(cache), module_(std::move(module)), - ctl_(cache_->as_batch_control()), + ctl_(cache->as()), max_sessions_(max_sessions), max_session_tokens_(max_session_tokens), backend_id_(std::move(backend_id)), @@ -245,7 +224,7 @@ ModuleExecutor::ModuleExecutor( ModuleExecutor::~ModuleExecutor() = default; -std::unique_ptr ModuleExecutor::create( +Result> ModuleExecutor::create( std::unique_ptr module, int max_sessions, int max_session_tokens, @@ -255,20 +234,26 @@ std::unique_ptr ModuleExecutor::create( std::string method) { if (module == nullptr) { ET_LOG(Error, "ModuleExecutor: no program"); - return nullptr; + return Error::InvalidArgument; } if (max_sessions <= 0 || max_session_tokens <= 0) { ET_LOG(Error, "ModuleExecutor: session limits must be positive"); - return nullptr; + return Error::InvalidArgument; } - if (module->load() != Error::Ok) { // a no-op once the caller has loaded it + const Error load_error = + module->load(); // no-op once the caller has loaded it + if (load_error != Error::Ok) { ET_LOG(Error, "ModuleExecutor: the program did not load"); - return nullptr; + return load_error; } auto cfg = config_from_program(*module); if (!cfg.ok()) { - return nullptr; + return cfg.error(); + } + if (max_sessions > std::numeric_limits::max() / max_session_tokens) { + ET_LOG(Error, "ModuleExecutor: total cache capacity exceeds int range"); + return Error::InvalidArgument; } cfg->capacity = max_sessions * max_session_tokens; cfg->kv_dtype = kv_dtype; @@ -277,13 +262,13 @@ std::unique_ptr ModuleExecutor::create( } if (!cache::valid(*cfg)) { ET_LOG(Error, "ModuleExecutor: the program's layout is unusable"); - return nullptr; + return Error::InvalidProgram; } const auto meta = module->method_meta(method); if (!meta.ok()) { ET_LOG(Error, "ModuleExecutor: %s has no metadata", method.c_str()); - return nullptr; + return meta.error(); } std::string backend_id; @@ -292,7 +277,7 @@ std::unique_ptr ModuleExecutor::create( if (!name.ok()) { ET_LOG( Error, "ModuleExecutor: %s has an unnamed delegate", method.c_str()); - return nullptr; + return name.error(); } if (backend_id.empty()) { backend_id = name.get(); @@ -302,56 +287,63 @@ std::unique_ptr ModuleExecutor::create( "ModuleExecutor: %s spans more than one backend, so which holds the " "cache is ambiguous", method.c_str()); - return nullptr; + return Error::InvalidProgram; } } if (backend_id.empty()) { ET_LOG(Error, "ModuleExecutor: %s delegates to nothing", method.c_str()); - return nullptr; + return Error::InvalidProgram; } auto built = - cache::CacheBuilderRegistry::global().build(backend_id, cache_kind, *cfg); + cache::CacheFactory::global().build(backend_id, cache_kind, *cfg); if (!built.ok()) { ET_LOG( Error, "ModuleExecutor: backend %s registers no %s cache", backend_id.c_str(), cache_kind.c_str()); - return nullptr; + return built.error(); } - std::shared_ptr cache = built.get(); - if (cache->as_batch_control() == nullptr) { + std::shared_ptr cache = built.get(); + if (cache->as() == nullptr) { ET_LOG(Error, "ModuleExecutor: the cache carries no sequence identity"); - return nullptr; + return Error::InvalidType; } if (meta->num_outputs() == 0) { ET_LOG(Error, "ModuleExecutor: %s publishes no outputs", method.c_str()); - return nullptr; + return Error::InvalidProgram; } const auto logits_info = meta->output_tensor_meta(0); - if (!logits_info.ok() || logits_info->sizes().empty()) { + if (!logits_info.ok()) { + ET_LOG(Error, "ModuleExecutor: %s has no logits metadata", method.c_str()); + return logits_info.error(); + } + if (logits_info->sizes().empty()) { ET_LOG(Error, "ModuleExecutor: %s has no logits shape", method.c_str()); - return nullptr; + return Error::InvalidProgram; } const auto logits_sizes = logits_info->sizes(); const auto tokens_info = meta->input_tensor_meta(0); - if (!tokens_info.ok() || tokens_info->sizes().empty()) { + if (!tokens_info.ok()) { + ET_LOG( + Error, + "ModuleExecutor: %s has no token input metadata", + method.c_str()); + return tokens_info.error(); + } + if (tokens_info->sizes().empty()) { ET_LOG( Error, "ModuleExecutor: %s has no token input shape", method.c_str()); - return nullptr; + return Error::InvalidProgram; } const auto tokens_sizes = tokens_info->sizes(); - auto session = - std::make_unique(cache::make_unique_key(), cache); - return std::unique_ptr(new ModuleExecutor( std::move(module), std::move(cache), - std::move(session), max_sessions, max_session_tokens, std::move(backend_id), @@ -362,11 +354,9 @@ std::unique_ptr ModuleExecutor::create( bool ModuleExecutor::initialize() { // The delegate resolves the cache from this key while the method loads. - char key[::executorch::runtime::kMaxOptionKeyLength] = {}; - std::memcpy(key, kCacheKeyOption, sizeof(kCacheKeyOption) - 1); ::executorch::runtime::BackendOptions<1> options; ::executorch::runtime::LoadBackendOptionsMap options_map; - if (options.set_option(key, session_->key().c_str()) != Error::Ok || + if (install_guard_.set_option(options) != Error::Ok || options_map.set_options(backend_id_.c_str(), options.view()) != Error::Ok) { ET_LOG(Error, "ModuleExecutor: could not name the cache to the backend"); @@ -392,7 +382,7 @@ std::optional ModuleExecutor::open_session() { return std::nullopt; } const SessionId session = next_session_++; - sessions_.emplace(session, SessionInfo{*seq_id, nullptr}); + sessions_.emplace(session, SessionState{*seq_id, nullptr}); return session; } @@ -427,8 +417,7 @@ bool ModuleExecutor::execute(const BatchInput& batch, BatchOutput& out) { out.outputs.clear(); out.outputs.resize(batch.inputs.size()); - const Result step = - build_step(*ctl_, batch, sessions_, max_session_tokens_); + const Result step = build_step(batch); if (!step.ok()) { return false; } diff --git a/extension/llm/batching/module_executor.h b/extension/llm/batching/module_executor.h index 9f64d10ff11..ec249b3ecfa 100644 --- a/extension/llm/batching/module_executor.h +++ b/extension/llm/batching/module_executor.h @@ -37,12 +37,6 @@ namespace batching { namespace cache = ::executorch::extension::llm::cache; -// A session's cache sequence and the sampler its generation draws from. -struct SessionInfo { - std::int32_t seq_id; - std::unique_ptr sampler; -}; - class ET_EXPERIMENTAL ModuleExecutor : public Executor { public: ~ModuleExecutor() override; @@ -60,16 +54,16 @@ class ET_EXPERIMENTAL ModuleExecutor : public Executor { // `cache_kind` must name a builder that carries batch control -- a cache // serving one sequence cannot back a batch of them. // - // nullptr = unusable limits, no published KV layout, a method spanning - // several backends, or no such cache for the backend it names. A method that - // will not load is reported by initialize(). - static std::unique_ptr create( + // Returns an error for unusable limits, no published KV layout, a method + // spanning several backends, or no such cache for the backend it names. A + // method that will not load is reported by initialize(). + static ::executorch::runtime::Result> create( std::unique_ptr module, int max_sessions, int max_session_tokens, int kv_dtype, int initial_capacity = -1, - std::string cache_kind = "cell", + std::string cache_kind = cache::kind::kBatchedCell, std::string method = "forward"); // The widest step this method takes, from the shape its token input was @@ -92,10 +86,23 @@ class ET_EXPERIMENTAL ModuleExecutor : public Executor { bool execute(const BatchInput& batch, BatchOutput& out) override; private: + struct SessionState { + std::int32_t seq_id; + std::unique_ptr sampler; + }; + + struct Step { + std::vector tokens; + std::vector positions; + std::vector seq_ids; + std::vector logit_indices; + }; + + ::executorch::runtime::Result build_step(const BatchInput& batch); + ModuleExecutor( std::unique_ptr module, - std::shared_ptr cache, - std::unique_ptr session, + std::shared_ptr cache, int max_sessions, int max_session_tokens, std::string backend_id, @@ -110,10 +117,9 @@ class ET_EXPERIMENTAL ModuleExecutor : public Executor { // Ordered so the module dies first, releasing the delegate that resolved the // cache before the registry entry naming it goes. - std::unique_ptr session_; - std::shared_ptr cache_; + cache::InstallGuard install_guard_; std::unique_ptr module_; - cache::BatchControl* ctl_; + cache::BatchControl* const ctl_; int max_sessions_; int max_session_tokens_; std::string backend_id_; @@ -123,7 +129,7 @@ class ET_EXPERIMENTAL ModuleExecutor : public Executor { int max_step_tokens_; SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids - std::unordered_map sessions_; + std::unordered_map sessions_; }; } // namespace batching diff --git a/extension/llm/cache/README.md b/extension/llm/cache/README.md new file mode 100644 index 00000000000..0673956a077 --- /dev/null +++ b/extension/llm/cache/README.md @@ -0,0 +1,201 @@ +# Off-graph KV cache + +A KV cache that lives outside the exported graph. The runner creates it, the +backend writes into it during a forward, and neither holds a pointer to the +other. They meet through a string key. + +Keeping the cache out of the graph lets the runner do things the graph cannot +express: rewind a turn, clear between prompts, or hand one pool of memory to +several concurrent sequences. + +## Who uses what + +Three audiences touch this directory, and they need almost disjoint parts of +it. + +### The control plane + +The runner, or a batch executor. It decides *what* to cache and *when* to +discard it, and it runs between forwards, never during one. + +It uses `CacheFactory` to build a cache, `InstallGuard` to publish it, and one +runner-facing face for the rest of the session: + +```cpp +auto built = CacheFactory::global().build(kMLXBackendId, kind::kSingle, cfg); +if (!built.ok()) { return built.error(); } + +const std::shared_ptr kv = built.get(); +const InstallGuard guard{kv}; // published while in scope +guard.set_option(mlx_opts); // hand key to backend + +auto* ctl = kv->as(); +ctl->can_extend(n); ctl->rewind(len); ctl->clear(); +``` + +It includes `cache.h` and `cache_registry.h`. It never includes +`sequence_cache.h` or `cell_cache.h`, and never calls a planner face. It does +not know how bytes are arranged, only how much room is left and how to give +some back. + +### The backend + +The byte layer inside the delegate. It owns the actual tensors and runs during +a forward. + +At init it resolves the key and asks for its own face: + +```cpp +handle->cache_shared = CacheRegistry::global().get(cache_key); +handle->state.cache = handle->cache_shared->as(); +``` + +During a forward it asks a planner face where the bytes go: + +```cpp +auto plan = planner->plan(layer, position, T); // integers, no tensors +// ... write K/V into those rows, attend over those runs ... +planner->commit(*plan); +``` + +It includes the layout headers, because it subclasses them to attach its own +tensor storage. It is the only caller of `plan()`, `commit()`, and `place_step()`. + +### The cache implementer + +Someone adding a layout or a backend face. Subclass a neutral layout, add +whatever face your backend needs, and list them: + +```cpp +class MLXCellCache : public cache::CellCache, public MLXCache { + void* face(cache::FaceId id) override { + if (void* p = cache::CellCache::face(id)) { return p; } + return cache::expose(this, id); + } +}; +``` + +Then register a builder so the control plane can ask for it by name. Registration +is insertion-only: an empty builder or a duplicate `(backend_id, kind)` returns +`Error::InvalidArgument`, leaving any existing builder unchanged. + +## Faces + +A cache is owned as a `Cache*` and asked for the interface you want: + +```cpp +auto* ctl = cache->as(); // null if this cache does not offer it +``` + +`Cache` has one virtual method. Each face declares its own name, and an +implementation lists the faces it offers through `expose`. + +| | single sequence | pooled cells | +| ---------------- | ----------------- | -------------- | +| control plane | `SequenceControl` | `BatchControl` | +| backend | `SequencePlanner` | `CellStepper` | +| backend-specific | each backend names its own, such as `MLXCache` || + +Control-plane faces live in `cache.h`. A runner calls `as()` +on something it got from the registry, so it must see the face without choosing +a layout. Backend faces live with their layout in `sequence_cache.h` or +`cell_cache.h`, because only a byte layer calls them and it already includes +that header to construct the cache. + +**No RTTI.** The core avoids `dynamic_cast` so it can build with `-fno-rtti` +under `EXECUTORCH_OPTIMIZE_SIZE`. The `static_cast` inside `expose` also +applies the pointer adjustment a face at a non-zero offset needs, refuses to +compile if the type is not really a base, and is bound to its own name, so the +two cannot be mismatched. Because `as()` names `T::kFaceName`, asking for a +type that is not a face fails to compile instead of returning null. + +**Names, not an enum.** The set of faces is open. A backend adds one without +this directory learning about it: `MLXCache` declares its name in `MLXCache.h` +and `cache.h` never sees it. Names compare by pointer first and fall back to +`strcmp`, which covers a cache built in one shared object and queried from +another. + +A face name is a global ABI identifier. It must be non-null, remain stable, and +identify exactly one C++ interface across the core, every backend, and every +shared object. Reusing a name for an unrelated or incompatible interface makes +the erased pointer cast invalid. The raw lookup hook is protected; consumers use +`as()`, and each concrete cache must explicitly implement the faces it offers. + +## Layouts + +**`SequenceCache`** holds one sequence with a single logical length for the +whole model. Each layer is flat, keeping all history, or ring, sliding a +window, so a model that mixes both stays coherent. Offers `SequenceControl` and +`SequencePlanner`. + +**`CellCache`** holds many sequences over a shared pool of per-token cells. A +cell is freed once no sequence owns it. Offers `BatchControl` and `CellStepper`. + +## Cache kinds + +| constant | value | layout | +| -------------------- | --------------- | --------------------------------- | +| `kind::kSingle` | `single` | one sequence, per-layer runs | +| `kind::kBatchedCell` | `batched-cell` | many sequences over a shared pool | + +Kinds are strings so a backend can register a layout this directory has never +heard of. The constants name the kinds it does know about. Use them: a typo in +a literal is a runtime `NotFound`, while a typo in a constant does not compile. + +## Lifetimes + +`InstallGuard` is the only way to publish. `CacheRegistry::install` is private, +so an entry cannot outlive its owner and a second caller cannot clobber it. + +Three lifetimes overlap: + +- The **registry entry** must exist across every `load_method()` that resolves + the key. +- The **guard** controls that discoverability and may be destroyed after the + final such initialization. +- The **cache** may outlive the entry and guard. Each backend that resolved the + key holds its own `shared_ptr`. + +Destroying the guard unpublishes the key without invalidating an already +resolved cache. A later `load_method()` using that key fails, so the guard must +remain alive for as long as new delegates may still need to resolve it. + +## Two layers + +`cache.h`, `sequence_cache.h`, and `cell_cache.{h,cpp}` include nothing but the +C++ standard library. No tensors and no ExecuTorch. They describe where bytes +go using integers: which physical rows a step writes, which it reads, what the +mask should be. Failures come back as `bool` and `std::optional`. + +`cache_registry.{h,cpp}` is ExecuTorch-specific. It uses `Result`, `Error`, and +`ET_LOG`, but the stronger tie is its reason for existing. `DelegateHandle` is +opaque and backend options carry only strings, so a runner cannot pass the +backend a pointer. Publishing under a generated key works around that. Give a +framework where the cache can be handed to the op directly, and this layer +disappears. + +> The build does not yet honour this split. One `extension_llm_cache` target +> compiles both halves and links `executorch_core`, so the neutral core cannot +> currently be built without ExecuTorch. + +## Files + +``` +cache.h faces, the face mechanism, config neutral +sequence_cache.h SequenceCache, SequencePlanner, flat/ring neutral +cell_cache.{h,cpp} CellCache, CellStepper, the cell pool neutral +cache_registry.{h,cpp} CacheRegistry, CacheFactory, InstallGuard ExecuTorch +``` + +## Known gaps + +**`CacheConfig` fields do not mean the same thing to every layout.** `capacity` +is a position ceiling for `kSingle` and a count of pool slots for +`kBatchedCell`. `max_write` is read only by ring layers, which in turn ignore +`initial_capacity`. Splitting it into model-dictated shape and per-kind options +is the intended fix. + +**Registration relies on static-initializer side effects.** A builder is +registered only if the linker pulls its object file in. Whole-archive linking +of backends covers this today. If that changes, the symptom is a runtime +"no cache builder registered". diff --git a/extension/llm/cache/cache.h b/extension/llm/cache/cache.h index fb25e43f4a6..986c7a1b6a4 100644 --- a/extension/llm/cache/cache.h +++ b/extension/llm/cache/cache.h @@ -10,45 +10,72 @@ // Neutral, tensor-free, ET-independent KV-cache core shared across backends. A // cache exposes a runner-facing control face and a backend-facing planner face, -// recovered from the owning CacheBase*. Which pair it implements depends on the -// layout: one sequence over per-layer runs, or many sequences over a pool of -// per-token cells. +// recovered from the owning Cache* with as(). Which pair it implements +// depends on the layout: one sequence over per-layer runs, or many sequences +// over a pool of per-token cells. +// +// Here: the face machinery, the runner-facing faces, and the config a caller +// fills in, all usable without picking a layout. Each backend-facing planner +// face lives with its layout, in sequence_cache.h or cell_cache.h. #include +#include #include #include +#include // ET_EXPERIMENTAL + namespace executorch { namespace extension { namespace llm { namespace cache { -class SequenceControl; -class SequencePlanner; -class BatchControl; -class CellStepper; +// A face is named by a non-null string it declares itself, so a backend can +// add one without this header learning about it. Names are stable ABI +// identifiers: each must identify exactly one interface across all binaries. +using FaceId = const char*; + +// Pointer equality covers the common case. The strcmp catches a cache built in +// one shared object and queried from another, where the literals may differ. +ET_EXPERIMENTAL inline bool same_face(FaceId a, FaceId b) { + return a != nullptr && b != nullptr && (a == b || std::strcmp(a, b) == 0); +} + +// Hands back `self` as each face it names, or nullptr for one it does not. +// static_cast applies the pointer adjustment a face at a non-zero offset needs +// and refuses to compile if Self does not derive from it. Each cast is bound to +// its own name in the pack, so a name cannot be paired with the wrong face. +template +ET_EXPERIMENTAL void* expose(Self* self, FaceId id) { + void* out = nullptr; + const bool matched[] = { + (same_face(id, Fs::kFaceName) ? (out = static_cast(self), true) + : false)...}; + (void)matched; + return out; +} -// Registry ownership anchor. A cache returns `this` from the faces it -// implements and leaves the rest null. -class CacheBase { +// Registry ownership anchor. A cache names the faces it implements from +// face(); everything else it is asked for comes back null. +class ET_EXPERIMENTAL Cache { public: - virtual ~CacheBase() = default; - virtual SequenceControl* as_control() { - return nullptr; - } - virtual SequencePlanner* as_planner() { - return nullptr; - } - virtual BatchControl* as_batch_control() { - return nullptr; - } - virtual CellStepper* as_cell_stepper() { - return nullptr; + virtual ~Cache() = default; + + // Naming T::kFaceName means a type that is not a face fails to compile, + // rather than quietly returning null at run time. + template + T* as() { + return static_cast(face(T::kFaceName)); } + + protected: + // Implemented with expose<...>(this, id). Kept behind as() so callers do + // not handle erased pointers or face names directly. + virtual void* face(FaceId id) = 0; }; // Lifecycle and admission, tensor-free. -class CacheControl { +class ET_EXPERIMENTAL CacheControl { public: virtual ~CacheControl() = default; virtual bool can_extend(int n = 1) const = 0; // admission / hard-stop @@ -57,59 +84,21 @@ class CacheControl { }; // Application face of a single-sequence cache: one length to rewind. -class SequenceControl : public CacheControl { +class ET_EXPERIMENTAL SequenceControl : public CacheControl { public: + static constexpr const char* kFaceName = "et.cache.SequenceControl"; + // Truncate to new_len; false = cannot grow, or the target is older than an // evicting layer still retains. virtual bool rewind(int new_len) = 0; }; -// A contiguous span of physical rows in a layer's pool. -struct Run { - int start; - int len; -}; - -// Integer-only handoff to the backend byte layer. Runs are in logical order -// (oldest -> newest); a flat layer uses one, a ring layer two when it wraps. -// read_base_pos is the logical position of read[0].start. -struct SeqStepPlan { - Run write[2]; - int n_write; - Run read[2]; - int n_read; - int read_base_pos; -}; - -// Backend face. plan() is const: it computes a layer's layout without changing -// state, and commit() advances the shared logical length. nullopt = the step -// exceeds capacity, or `layer` is out of range. -class SequencePlanner { - public: - virtual ~SequencePlanner() = default; - virtual std::optional plan(int layer, int position, int T) - const = 0; - // Advance the logical length past this step. Idempotent, so once per step - // suffices. - virtual void commit(const SeqStepPlan& plan) = 0; -}; - -// Per-layer layout: flat keeps all history, ring slides a window. Stateless. -class LayoutPolicy { - public: - virtual ~LayoutPolicy() = default; - // Write/read runs for T cells at logical `position`. Precondition: T fits the - // policy's window. - virtual SeqStepPlan plan(int position, int T) const = 0; - // Oldest logical position still retained at this length: 0 for flat, - // length - window for ring. - virtual int retained_from(int length) const = 0; -}; - // Application face of any multi-sequence cache: the sequence verbs. They run // between forwards, never during one. -class BatchControl : public CacheControl { +class ET_EXPERIMENTAL BatchControl : public CacheControl { public: + static constexpr const char* kFaceName = "et.cache.BatchControl"; + // Which sequence each of the next forward's tokens belongs to, one entry per // token; every id must be one seq_new handed out. Also the admission gate: // false = rejected and nothing changed, and a step that passes has room for @@ -137,7 +126,7 @@ class BatchControl : public CacheControl { }; // Per-layer cache kind and its parameters. -struct LayerPolicy { +struct ET_EXPERIMENTAL LayerPolicy { enum class Kind : int { Flat = 0, Ring = 1 @@ -147,7 +136,7 @@ struct LayerPolicy { }; // Per-layer architecture facts + cache policy. -struct LayerConfig { +struct ET_EXPERIMENTAL LayerConfig { LayerPolicy policy; // default Flat int n_kv_heads; int head_dim; @@ -155,7 +144,7 @@ struct LayerConfig { // Model facts and the policy the byte layer sizes its pools from. `layers` is // per-layer: size 1 applies to every layer, else one entry each. -struct CacheConfig { +struct ET_EXPERIMENTAL CacheConfig { int capacity; // logical cap in cells int n_layers; std::vector layers; @@ -167,7 +156,7 @@ struct CacheConfig { }; // Whether `cfg` satisfies the contract above. -inline bool valid(const CacheConfig& cfg) { +ET_EXPERIMENTAL inline bool valid(const CacheConfig& cfg) { // initial_capacity may be 0 but not negative, and may exceed capacity -- the // byte layer clamps it. return cfg.capacity > 0 && cfg.n_layers > 0 && cfg.initial_capacity >= 0 && diff --git a/extension/llm/cache/cache_et.h b/extension/llm/cache/cache_et.h deleted file mode 100644 index 1c157f8ed60..00000000000 --- a/extension/llm/cache/cache_et.h +++ /dev/null @@ -1,64 +0,0 @@ -/* - * Copyright (c) Meta Platforms, Inc. and affiliates. - * All rights reserved. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#pragma once - -// ExecuTorch adapter for the neutral cache core. The core (cache.h / -// sequence_cache.h) is ET-independent and reports failures as -// bool/std::optional so it is usable outside an ET runner. These thin inline -// adapters map those results to ExecuTorch Error/Result (logging on failure) -// for ET consumers -- the runner and the delegate byte layer. (The registry is -// delegate-specific and already returns Result directly, so it needs no -// adapter.) - -#include - -#include -#include -#include - -namespace executorch { -namespace extension { -namespace llm { -namespace cache { -namespace et { - -using ::executorch::runtime::Error; -using ::executorch::runtime::Result; - -// Plan a layer's step, or OutOfResources if it would exceed capacity (or the -// layer is out of range). -inline Result -plan(const SequencePlanner& planner, int layer, int position, int T) { - std::optional p = planner.plan(layer, position, T); - ET_CHECK_OR_RETURN_ERROR( - p.has_value(), - OutOfResources, - "cache: plan(layer=%d, position=%d, T=%d) exceeds capacity or bad layer", - layer, - position, - T); - return *p; -} - -// Truncate the history, or InvalidArgument if new_len would grow it (or is -// older than an evicting layer retains). -inline Error rewind(SequenceControl& control, int new_len) { - ET_CHECK_OR_RETURN_ERROR( - control.rewind(new_len), - InvalidArgument, - "rewind: cannot grow to %d", - new_len); - return Error::Ok; -} - -} // namespace et -} // namespace cache -} // namespace llm -} // namespace extension -} // namespace executorch diff --git a/extension/llm/cache/cache_registry.cpp b/extension/llm/cache/cache_registry.cpp index d54d05ddc73..51cb9cd711b 100644 --- a/extension/llm/cache/cache_registry.cpp +++ b/extension/llm/cache/cache_registry.cpp @@ -11,6 +11,7 @@ #include #include +#include namespace executorch { namespace extension { @@ -24,12 +25,12 @@ CacheRegistry& CacheRegistry::global() { void CacheRegistry::install( const std::string& key, - std::shared_ptr cache) { + std::shared_ptr cache) { std::lock_guard lock(mu_); caches_[key] = std::move(cache); } -std::shared_ptr CacheRegistry::get(const std::string& key) const { +std::shared_ptr CacheRegistry::get(const std::string& key) const { std::lock_guard lock(mu_); const auto it = caches_.find(key); return it == caches_.end() ? nullptr : it->second; @@ -40,20 +41,25 @@ void CacheRegistry::erase(const std::string& key) { caches_.erase(key); } -CacheBuilderRegistry& CacheBuilderRegistry::global() { - static CacheBuilderRegistry registry; +CacheFactory& CacheFactory::global() { + static CacheFactory registry; return registry; } -void CacheBuilderRegistry::register_builder( +Error CacheFactory::register_builder( const std::string& backend_id, const std::string& kind, CacheBuilder builder) { + if (!builder) { + return Error::InvalidArgument; + } std::lock_guard lock(mu_); - builders_[{backend_id, kind}] = std::move(builder); + const auto inserted = + builders_.emplace(std::make_pair(backend_id, kind), std::move(builder)); + return inserted.second ? Error::Ok : Error::InvalidArgument; } -Result> CacheBuilderRegistry::build( +Result> CacheFactory::build( const std::string& backend_id, const std::string& kind, const CacheConfig& cfg) const { @@ -61,12 +67,27 @@ Result> CacheBuilderRegistry::build( { std::lock_guard lock(mu_); const auto it = builders_.find({backend_id, kind}); - ET_CHECK_OR_RETURN_ERROR( - it != builders_.end(), - NotFound, - "no cache builder registered for %s:%s", - backend_id.c_str(), - kind.c_str()); + if (it == builders_.end()) { + // Name what is registered. A kind is a string, so a typo is otherwise a + // dead end. builders_ is ordered, so these come out sorted. + std::string known; + for (const auto& entry : builders_) { + if (entry.first.first != backend_id) { + continue; + } + if (!known.empty()) { + known += ", "; + } + known += entry.first.second; + } + ET_LOG( + Error, + "no '%s' cache registered for '%s'; registered: %s", + kind.c_str(), + backend_id.c_str(), + known.empty() ? "(none)" : known.c_str()); + return Error::NotFound; + } builder = it->second; } // Checked here rather than in each cache: `layers` is indexed directly, so a @@ -77,13 +98,36 @@ Result> CacheBuilderRegistry::build( "cache: invalid CacheConfig for %s:%s", backend_id.c_str(), kind.c_str()); - return builder(cfg); + // A builder that hands back null would otherwise travel as an ok() Result + // and be dereferenced by the caller. + auto cache = builder(cfg); + ET_CHECK_OR_RETURN_ERROR( + cache != nullptr, + Internal, + "cache: builder for %s:%s returned null", + backend_id.c_str(), + kind.c_str()); + return cache; } -std::string make_unique_key() { +namespace { +// Process-global atomic counter -> "cache-N". Internal: InstallGuard is the +// only thing that publishes, so nothing outside needs to mint a key. +std::string new_cache_key() { static std::atomic counter{0}; return "cache-" + std::to_string(counter.fetch_add(1)); } +} // namespace + +InstallGuard::InstallGuard(std::shared_ptr cache) + : key_(new_cache_key()), cache_(std::move(cache)) { + ET_CHECK_MSG(cache_ != nullptr, "Cannot install a null cache"); + CacheRegistry::global().install(key_, cache_); +} + +InstallGuard::~InstallGuard() { + CacheRegistry::global().erase(key_); +} } // namespace cache } // namespace llm diff --git a/extension/llm/cache/cache_registry.h b/extension/llm/cache/cache_registry.h index ae730344c8a..020a6fc2124 100644 --- a/extension/llm/cache/cache_registry.h +++ b/extension/llm/cache/cache_registry.h @@ -12,8 +12,8 @@ // is opaque to the host, so the runner (which knows the cache kind) creates the // cache and binds it to the delegate through a process-global registry; the two // sides rendezvous on a cache_key passed as a runtime backend-load option. -// Caches are owned as CacheBase* and the faces are recovered through its as_* -// accessors (no RTTI), each null for a face the cache does not implement. +// Caches are owned as Cache*; a face comes from as(), null when the cache +// does not implement it. #include #include @@ -24,8 +24,10 @@ #include #include +#include #include #include +#include namespace executorch { namespace extension { @@ -35,82 +37,100 @@ namespace cache { using ::executorch::runtime::Error; using ::executorch::runtime::Result; -// Process-global map>. Ownership is shared: -// the registry entry, the runner's session guard, and the delegate handle all -// hold the cache, so erasing the entry mid-method is safe. -class CacheRegistry { +// Backend-load option carrying the key of an installed cache. This name is the +// rendezvous contract shared by cache-owning runners and cache-aware backends. +inline constexpr char kCacheKeyOption[] = "llm_cache_registry_key"; + +// Process-global map>. Ownership is shared: +// the registry entry, the runner's guard, and the delegate handle all hold +// the cache, so erasing the entry mid-method is safe. +class ET_EXPERIMENTAL CacheRegistry { public: static CacheRegistry& global(); - void install(const std::string& key, std::shared_ptr cache); - std::shared_ptr get(const std::string& key) const; - void erase(const std::string& key); + // The delegate's half of the rendezvous: resolve a key it was handed as a + // backend option. Null if no cache is published under it. + std::shared_ptr get(const std::string& key) const; private: CacheRegistry() = default; + // Only InstallGuard may publish, so an entry cannot outlive its owner, two + // callers cannot collide on a key, and no erase can go unpaired. + friend class InstallGuard; + void install(const std::string& key, std::shared_ptr cache); + void erase(const std::string& key); + mutable std::mutex mu_; - std::unordered_map> caches_; + std::unordered_map> caches_; }; +// The registered cache kinds. Spelling one inline is a runtime NotFound rather +// than a compile error, so go through these. +namespace kind { +// One sequence over per-layer runs. +inline constexpr const char* kSingle = "single"; +// Many sequences sharing one pool of per-token cells. +inline constexpr const char* kBatchedCell = "batched-cell"; +} // namespace kind + // Cache kind is expressed by which factory you call: backends register a // builder per (backend_id, kind) and the kind survives only as an internal // lookup tag. -using CacheBuilder = - std::function(const CacheConfig&)>; +using CacheBuilder = std::function(const CacheConfig&)>; -class CacheBuilderRegistry { +class ET_EXPERIMENTAL CacheFactory { public: - static CacheBuilderRegistry& global(); + static CacheFactory& global(); + + // Public so a test can hold its own rather than registering builders into + // the process-global one, where they outlive it. + CacheFactory() = default; - void register_builder( + // Registers one builder without replacing an existing entry. Returns + // InvalidArgument if builder is empty or the pair is already registered. + ET_NODISCARD Error register_builder( const std::string& backend_id, const std::string& kind, CacheBuilder builder); - // Returns Error::NotFound if no builder is registered for (backend_id, kind). - Result> build( + // Returns NotFound if no builder is registered for (backend_id, kind), and + // Internal if the registered builder returns null. + Result> build( const std::string& backend_id, const std::string& kind, const CacheConfig& cfg) const; private: - CacheBuilderRegistry() = default; - mutable std::mutex mu_; - std::map, CacheBuilder> - builders_; // keyed by (backend_id, kind) + std::map, CacheBuilder> builders_; }; -// Process-global atomic counter -> "cache-N"; centralizes key generation so -// keys never collide. -std::string make_unique_key(); - -// RAII: installs the cache into the global registry under a unique key on -// construction and erases it on destruction (no leak on any exit path). Holds -// the runner's shared_ptr and exposes the control face for the generation loop. -class CacheSession { +// RAII over one registry entry: installs the cache under a key of its own +// making on construction and erases it on destruction (no leak on any exit +// path). Minting the key here rather than taking one means two live guards +// cannot collide on it. Must outlive the load_method() whose backend init +// resolves the key. +// +// Destruction removes discoverability only. A shared_ptr already returned by +// CacheRegistry::get() remains valid independently. +class ET_EXPERIMENTAL InstallGuard { public: - CacheSession(std::string key, std::shared_ptr cache) - : key_(std::move(key)), cache_(std::move(cache)) { - CacheRegistry::global().install(key_, cache_); - } - ~CacheSession() { - CacheRegistry::global().erase(key_); - } + explicit InstallGuard(std::shared_ptr cache); + ~InstallGuard(); - CacheSession(const CacheSession&) = delete; - CacheSession& operator=(const CacheSession&) = delete; + InstallGuard(const InstallGuard&) = delete; + InstallGuard& operator=(const InstallGuard&) = delete; - SequenceControl* control() const { - return cache_->as_control(); - } - const std::string& key() const { - return key_; + // Adds the complete cache rendezvous option. BackendOptions copies the key + // and value, so the resulting option remains valid independently. + template + Error set_option(::executorch::runtime::BackendOptions& options) const { + return options.set_option(kCacheKeyOption, key_.c_str()); } private: std::string key_; - std::shared_ptr cache_; + std::shared_ptr cache_; }; } // namespace cache diff --git a/extension/llm/cache/cell_cache.h b/extension/llm/cache/cell_cache.h index 2c519d9d28c..a307c72b1cd 100644 --- a/extension/llm/cache/cell_cache.h +++ b/extension/llm/cache/cell_cache.h @@ -29,7 +29,7 @@ namespace cache { // Integer-only handoff to the byte layer, covering the whole forward: a cell // means the same token in every layer's pool. -struct CellStep { +struct ET_EXPERIMENTAL CellStep { int length; int read_len; // the window is cells [0, read_len) std::vector cells; // cell per query token @@ -42,32 +42,26 @@ struct CellStep { // forward. The returned step is owned by the cache and valid until the next // verb. nullptr = no declaration, a token count disagreeing with it, a position // a sequence already holds, a layer out of range, or a layer served twice. -class CellStepper { +class ET_EXPERIMENTAL CellStepper { public: + static constexpr const char* kFaceName = "et.cache.CellStepper"; + virtual ~CellStepper() = default; virtual const CellStep* place_step(int layer, const int32_t* positions, int length) = 0; }; -class CellCache : public CacheBase, public BatchControl, public CellStepper { +class ET_EXPERIMENTAL CellCache : public Cache, + public BatchControl, + public CellStepper { public: // One bit per sequence in the owner bitset. static constexpr int kMaxSeqs = 64; - // Precondition: valid(cfg). CacheBuilderRegistry::build enforces it for + // Precondition: valid(cfg). CacheFactory::build enforces it for // registry-created caches; direct construction must check first. explicit CellCache(const CacheConfig& cfg); - CacheBase* base() { - return this; - } - BatchControl* as_batch_control() override { - return this; - } - CellStepper* as_cell_stepper() override { - return this; - } - // -- CacheControl ------------------------------------------------------ bool can_extend(int n = 1) const override; @@ -92,6 +86,11 @@ class CellCache : public CacheBase, public BatchControl, public CellStepper { const CellStep* place_step(int layer, const int32_t* positions, int length) override; + protected: + void* face(FaceId id) override { + return expose(this, id); + } + private: struct SeqInfo { int count = 0; diff --git a/extension/llm/cache/sequence_cache.h b/extension/llm/cache/sequence_cache.h index 5af756e6e24..75aa02f8987 100644 --- a/extension/llm/cache/sequence_cache.h +++ b/extension/llm/cache/sequence_cache.h @@ -12,6 +12,11 @@ // policies (FlatPolicy / RingPolicy). SequenceCache owns the one logical length // for the whole model and dispatches per-layer layout to a policy, so a mixed // flat/ring model (gemma4) stays coherent. Tensor-free / ET-independent. +// +// The backend-facing planner face (SequencePlanner) and the types it hands over +// live here rather than in cache.h: only this layout implements them, and only +// a byte layer that already includes this header calls them. cell_cache.h holds +// CellStepper for the same reason. #include #include @@ -26,8 +31,52 @@ namespace extension { namespace llm { namespace cache { +// A contiguous span of physical rows in a layer's pool. +struct ET_EXPERIMENTAL Run { + int start; + int len; +}; + +// Integer-only handoff to the backend byte layer. Runs are in logical order +// (oldest -> newest); a flat layer uses one, a ring layer two when it wraps. +// read_base_pos is the logical position of read[0].start. +struct ET_EXPERIMENTAL SeqStepPlan { + Run write[2]; + int n_write; + Run read[2]; + int n_read; + int read_base_pos; +}; + +// Backend face. plan() is const: it computes a layer's layout without changing +// state, and commit() advances the shared logical length. nullopt = the step +// exceeds capacity, or `layer` is out of range. +class ET_EXPERIMENTAL SequencePlanner { + public: + static constexpr const char* kFaceName = "et.cache.SequencePlanner"; + + virtual ~SequencePlanner() = default; + virtual std::optional plan(int layer, int position, int T) + const = 0; + // Advance the logical length past this step. Idempotent, so once per step + // suffices. + virtual void commit(const SeqStepPlan& plan) = 0; +}; + +// Per-layer layout: flat keeps all history, ring slides a window. Stateless. +class ET_EXPERIMENTAL LayoutPolicy { + public: + virtual ~LayoutPolicy() = default; + // Write/read runs for T cells at logical `position`. Precondition: T fits the + // policy's window. + virtual SeqStepPlan plan(int position, int T) const = 0; + // Oldest logical position still retained at this length: 0 for flat, + // length - window for ring. + virtual int retained_from(int length) const = 0; +}; + // Full history [0, length): one contiguous write run, read over all history. -class FlatPolicy final : public LayoutPolicy { +class ET_EXPERIMENTAL FlatPolicy final : public LayoutPolicy { public: int retained_from(int /*length*/) const override { return 0; // keeps all history @@ -48,7 +97,7 @@ class FlatPolicy final : public LayoutPolicy { // slots. The ring is oversized so a step of up to max_write tokens fits without // overwriting cells earlier queries in the same step still attend to; the // backend masks each query to its own window within the read span. -class RingPolicy final : public LayoutPolicy { +class ET_EXPERIMENTAL RingPolicy final : public LayoutPolicy { public: RingPolicy(int window, int max_write) : window_(window), ring_size_(window + max_write - 1) {} @@ -92,9 +141,9 @@ class RingPolicy final : public LayoutPolicy { // rewind; dispatches per-layer layout to a shared LayoutPolicy. Policies are // deduped by (kind, window), so a uniform or two-kind (gemma4) model holds one // or two policy objects. -class SequenceCache : public CacheBase, - public SequenceControl, - public SequencePlanner { +class ET_EXPERIMENTAL SequenceCache : public Cache, + public SequenceControl, + public SequencePlanner { public: explicit SequenceCache(const CacheConfig& cfg) : capacity_(cfg.capacity), max_write_(cfg.max_write) { @@ -108,14 +157,6 @@ class SequenceCache : public CacheBase, } } - // CacheBase: face recovery without RTTI. - SequenceControl* as_control() override { - return this; - } - SequencePlanner* as_planner() override { - return this; - } - // SequenceControl. bool can_extend(int n = 1) const override { return length_ + n <= @@ -173,6 +214,11 @@ class SequenceCache : public CacheBase, length_ = std::max(length_, end); } + protected: + void* face(FaceId id) override { + return expose(this, id); + } + private: int policy_index(const LayerPolicy& lp) { for (std::size_t i = 0; i < specs_.size(); ++i) { diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index e41e91a6089..fda446412a7 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -7,7 +7,6 @@ */ #include -#include #include #include #include @@ -19,31 +18,53 @@ #include #include +#include #include using executorch::extension::llm::cache::BatchControl; -using executorch::extension::llm::cache::CacheBase; -using executorch::extension::llm::cache::CacheBuilderRegistry; +using executorch::extension::llm::cache::Cache; +using executorch::extension::llm::cache::CacheBuilder; using executorch::extension::llm::cache::CacheConfig; +using executorch::extension::llm::cache::CacheFactory; using executorch::extension::llm::cache::CacheRegistry; -using executorch::extension::llm::cache::CacheSession; using executorch::extension::llm::cache::CellCache; using executorch::extension::llm::cache::CellStep; using executorch::extension::llm::cache::CellStepper; +using executorch::extension::llm::cache::InstallGuard; +using executorch::extension::llm::cache::SequenceControl; +using executorch::extension::llm::cache::SequencePlanner; +namespace kind = executorch::extension::llm::cache::kind; using executorch::extension::llm::cache::LayerConfig; using executorch::extension::llm::cache::LayerPolicy; -using executorch::extension::llm::cache::make_unique_key; using executorch::extension::llm::cache::SequenceCache; +using executorch::runtime::BackendOptions; using executorch::runtime::Error; -namespace et = executorch::extension::llm::cache::et; namespace { +std::string installed_key(const InstallGuard& guard) { + BackendOptions<1> options; + if (guard.set_option(options) != Error::Ok) { + return {}; + } + const char* key = nullptr; + if (options.get_option( + executorch::extension::llm::cache::kCacheKeyOption, key) != + Error::Ok) { + return {}; + } + return key; +} + LayerConfig flat_layer() { return LayerConfig{LayerPolicy{LayerPolicy::Kind::Flat, 0}, 2, 8}; } LayerConfig ring_layer(int window) { return LayerConfig{LayerPolicy{LayerPolicy::Kind::Ring, window}, 2, 8}; } + +struct UnsupportedFace { + static constexpr const char* kFaceName = "test.UnsupportedFace"; +}; } // namespace // Initializes the ExecuTorch PAL so the ET adapter's error paths (which ET_LOG) @@ -167,92 +188,169 @@ TEST_F(CacheTest, RewindBoundedByRingWindow) { EXPECT_FALSE(cache.rewind(11)); // cannot grow } -// ---- Faces / registry / session -------------------------------------------- +// ---- Faces / registry / lease ---------------------------------------------- TEST_F(CacheTest, FaceRecoveryReturnsSameObject) { SequenceCache cache(CacheConfig{4, 1, {flat_layer()}}); - CacheBase* base = &cache; - ASSERT_NE(base->as_control(), nullptr); - ASSERT_NE(base->as_planner(), nullptr); - EXPECT_TRUE(base->as_control()->can_extend(4)); - auto plan = base->as_planner()->plan(0, 0, 1); + Cache* base = &cache; + ASSERT_NE(base->as(), nullptr); + ASSERT_NE(base->as(), nullptr); + EXPECT_EQ(base->as(), nullptr); + EXPECT_TRUE(base->as()->can_extend(4)); + auto plan = base->as()->plan(0, 0, 1); ASSERT_TRUE(plan.has_value()); EXPECT_EQ(plan->read[0].len, 1); } -TEST_F(CacheTest, RegistryInstallGetErase) { - auto& reg = CacheRegistry::global(); - const std::string key = make_unique_key(); - EXPECT_EQ(reg.get(key), nullptr); - - std::shared_ptr cache = - std::make_shared(CacheConfig{16, 1, {flat_layer()}}); - reg.install(key, cache); - EXPECT_EQ(reg.get(key), cache); - EXPECT_TRUE(reg.get(key)->as_control()->can_extend(16)); - - reg.erase(key); - EXPECT_EQ(reg.get(key), nullptr); +TEST_F(CacheTest, NullFaceIdsNeverMatch) { + using executorch::extension::llm::cache::same_face; + EXPECT_FALSE(same_face(nullptr, nullptr)); + EXPECT_FALSE(same_face(nullptr, SequenceControl::kFaceName)); + EXPECT_FALSE(same_face(SequenceControl::kFaceName, nullptr)); } -TEST_F(CacheTest, UniqueKeysDoNotCollide) { - EXPECT_NE(make_unique_key(), make_unique_key()); +TEST_F(CacheTest, LiveGuardsDoNotCollideOnKeys) { + auto a = std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + auto b = std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + InstallGuard ga(a); + InstallGuard gb(b); + + const std::string a_key = installed_key(ga); + const std::string b_key = installed_key(gb); + ASSERT_FALSE(a_key.empty()); + ASSERT_FALSE(b_key.empty()); + EXPECT_NE(a_key, b_key); + // Both entries are live: neither guard displaced the other. + EXPECT_EQ(CacheRegistry::global().get(a_key).get(), a.get()); + EXPECT_EQ(CacheRegistry::global().get(b_key).get(), b.get()); } TEST_F(CacheTest, BuilderBuildsRegisteredKindElseError) { - auto& reg = CacheBuilderRegistry::global(); - reg.register_builder("TestBackend", "seq", [](const CacheConfig& cfg) { - return std::static_pointer_cast( - std::make_shared(cfg)); - }); + // Its own factory: a builder registered into the global one would + // outlive this test and be visible to every case after it. + CacheFactory reg; + EXPECT_EQ( + reg.register_builder( + "TestBackend", + kind::kSingle, + [](const CacheConfig& cfg) { + return std::static_pointer_cast( + std::make_shared(cfg)); + }), + Error::Ok); CacheConfig cfg{32, 1, {flat_layer()}}; - auto cache = reg.build("TestBackend", "seq", cfg); + auto cache = reg.build("TestBackend", kind::kSingle, cfg); ASSERT_TRUE(cache.ok()); - EXPECT_EQ(cache.get()->as_control()->capacity(), 32); + EXPECT_EQ(cache.get()->as()->capacity(), 32); EXPECT_EQ(reg.build("TestBackend", "missing", cfg).error(), Error::NotFound); // A layers list that is neither size 1 nor n_layers would be indexed past // the end, so build refuses it before the cache is constructed. EXPECT_EQ( - reg.build("TestBackend", "seq", CacheConfig{32, 3, {}}).error(), + reg.build("TestBackend", kind::kSingle, CacheConfig{32, 3, {}}).error(), Error::InvalidArgument); EXPECT_EQ( reg.build( "TestBackend", - "seq", + kind::kSingle, CacheConfig{32, 3, {flat_layer(), flat_layer()}}) .error(), Error::InvalidArgument); } -TEST_F(CacheTest, SessionInstallsOnCtorErasesOnDtor) { - const std::string key = make_unique_key(); +TEST_F(CacheTest, BuilderRegistrationRejectsInvalidEntries) { + CacheFactory factory; + CacheConfig cfg{32, 1, {flat_layer()}}; + + EXPECT_EQ( + factory.register_builder("TestBackend", "empty", CacheBuilder{}), + Error::InvalidArgument); + EXPECT_EQ( + factory.build("TestBackend", "empty", cfg).error(), Error::NotFound); + + EXPECT_EQ( + factory.register_builder( + "TestBackend", + "duplicate", + [](const CacheConfig& config) { + return std::static_pointer_cast( + std::make_shared(config)); + }), + Error::Ok); + EXPECT_EQ( + factory.register_builder( + "TestBackend", + "duplicate", + [](const CacheConfig&) { return std::shared_ptr{}; }), + Error::InvalidArgument); + auto original = factory.build("TestBackend", "duplicate", cfg); + ASSERT_TRUE(original.ok()); + EXPECT_NE(original.get()->as(), nullptr); +} + +TEST_F(CacheTest, BuilderReturningNullIsAnError) { + CacheFactory factory; + EXPECT_EQ( + factory.register_builder( + "TestBackend", + "null", + [](const CacheConfig&) { return std::shared_ptr{}; }), + Error::Ok); + EXPECT_EQ( + factory.build("TestBackend", "null", CacheConfig{32, 1, {flat_layer()}}) + .error(), + Error::Internal); +} + +TEST_F(CacheTest, NullCacheCannotBeInstalled) { + ET_EXPECT_DEATH( + { InstallGuard guard{std::shared_ptr{}}; }, + "Cannot install a null cache"); +} + +TEST_F(CacheTest, GuardInstallsOnCtorErasesOnDtor) { + auto cache = + std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + std::string key; { - CacheSession session( - key, - std::make_shared(CacheConfig{4, 1, {flat_layer()}})); + InstallGuard guard(cache); + key = installed_key(guard); + ASSERT_FALSE(key.empty()); + // Publishing is the guard's alone, since CacheRegistry::install is + // private: an entry cannot outlive its owner or be clobbered. EXPECT_NE(CacheRegistry::global().get(key), nullptr); - EXPECT_TRUE(session.control()->can_extend(4)); + // The published entry is the same object the caller still holds. + EXPECT_EQ(CacheRegistry::global().get(key).get(), cache.get()); + EXPECT_TRUE(cache->as()->can_extend(4)); } EXPECT_EQ(CacheRegistry::global().get(key), nullptr); + // The guard held only the registry entry; the cache outlives it. + EXPECT_TRUE(cache->as()->can_extend(4)); } -// ---- ET adapter (maps core bool/optional to Error/Result) ------------------ - -TEST_F(CacheTest, EtAdapterMapsResultsAndCodes) { - SequenceCache cache(CacheConfig{2, 1, {flat_layer()}}); - auto ok = et::plan(cache, /*layer=*/0, /*position=*/0, /*T=*/2); - ASSERT_TRUE(ok.ok()); - EXPECT_EQ(ok->read[0].len, 2); - cache.commit(ok.get()); // accept the step so rewind has history to truncate - EXPECT_EQ( - et::plan(cache, 0, 2, 1).error(), Error::OutOfResources); // over capacity - EXPECT_FALSE(et::plan(cache, 5, 0, 1).ok()); // bad layer +TEST_F(CacheTest, AcquiredCacheOutlivesRegistryEntry) { + std::weak_ptr weak; + std::shared_ptr acquired; + std::string key; + { + auto cache = + std::make_shared(CacheConfig{4, 1, {flat_layer()}}); + weak = cache; + InstallGuard guard(cache); + key = installed_key(guard); + ASSERT_FALSE(key.empty()); + acquired = CacheRegistry::global().get(key); + ASSERT_NE(acquired, nullptr); + cache.reset(); + } - EXPECT_EQ(et::rewind(cache, 9), Error::InvalidArgument); // cannot grow - EXPECT_EQ(et::rewind(cache, 1), Error::Ok); + EXPECT_EQ(CacheRegistry::global().get(key), nullptr); + ASSERT_FALSE(weak.expired()); + EXPECT_TRUE(acquired->as()->can_extend(4)); + acquired.reset(); + EXPECT_TRUE(weak.expired()); } // ---- Cell layout ----------------------------------------------------------- @@ -291,8 +389,8 @@ struct Cells { int capacity, std::vector layers = {flat_layer(), flat_layer()}) : cache(CacheConfig{capacity, static_cast(layers.size()), layers}), - ctl(cache.as_batch_control()), - stepper(cache.as_cell_stepper()) {} + ctl(cache.as()), + stepper(cache.as()) {} // Ids come from the cache, so a test names sequences by allocating them. // Unlike the face's, this one unwraps and fails the test if none is free. From 1473d0ee0ae06597a4cfe80bc0b88d564453ef13 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Tue, 8 Sep 2026 17:28:40 -0700 Subject: [PATCH 086/190] Revert "Move pull.yml to linux_job_v3" (#22247) (#22626) ### Summary Reverts #22247 (a519efa64e9f5aea9b449fb21a0abbfd59761396) to restore the previous linux_job_v2 / EC2 configuration in pull.yml. Earlier migrations outside that file are unchanged. The first post-merge trunk run has ten migrated Linux jobs failing on writes to OSDC's read-only /mnt/hf_cache mount. Dataset loaders fail creating lock files, checkpoint conversion fails creating .cache/meta_checkpoints, and the Gemma3 job fails saving Hugging Face login state. All ten corresponding jobs passed in the immediately preceding trunk run. Example failure: https://github.com/pytorch/executorch/actions/runs/34283582096/job/102253931559 Previous trunk run: https://github.com/pytorch/executorch/actions/runs/34282655535 The original PR had the ci-refresh-hf-cache label. Its passing Wikitext job ran with HF_CACHE_REFRESH=1 and writable job-local HF_HOME / HF_DATASETS_CACHE paths, while the post-merge push ran with HF_CACHE_REFRESH=0 and retained the read-only mount. Both used the same test-infra revision. Passing PR job: https://github.com/pytorch/executorch/actions/runs/34256077690/job/102238123181 This is containment rather than a cache bypass. Before relanding, normal non-refresh jobs should be able to consume the shared model cache while keeping dataset locks, converted checkpoints, and authentication state writable. Validation needs to cover the ordinary unlabeled path, not only refresh mode. Authored with assistance from OpenAI Codex. ### Test plan Verified that the resulting pull.yml is byte-for-byte identical to its version immediately before #22247 (blob 2522cca31138ff7bde05049fcbde01d21fac5cb5). Ran git diff --check and confirmed the original migration patch applies cleanly on top of the revert. The local commit hook could not run lintrunner because it is not installed. No end-to-end jobs have been rerun locally; CI validation is pending on this draft. The ci-refresh-hf-cache label is intentionally not applied. --- .github/workflows/pull.yml | 361 ++++++++++++++++--------------------- 1 file changed, 155 insertions(+), 206 deletions(-) diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index d17be95b3aa..2522cca3113 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -13,10 +13,6 @@ concurrency: cancel-in-progress: true jobs: - docker-image: - name: Resolve CI docker image - uses: ./.github/workflows/_docker-image.yml - # Emits the list of changed files for the current PR or push commit. # On PR: PR diff. On push: diff against `github.event.before`. # On events without a diff base (workflow_dispatch, tag creation, @@ -37,9 +33,8 @@ jobs: uses: ./.github/workflows/_ci-run-decision.yml test-qnn-wheel-packages-linux: - needs: docker-image name: test-qnn-wheel-packages-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -48,8 +43,8 @@ jobs: matrix: python-version: [ "3.10", "3.11", "3.12", "3.13" ] with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -82,7 +77,7 @@ jobs: contents: read test-minimal-wheel-linux: - needs: [docker-image, changed-files] + needs: changed-files if: | github.event_name != 'pull_request' || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') || @@ -95,13 +90,13 @@ jobs: contains(needs.changed-files.outputs.changed-files, 'setup.py') || contains(needs.changed-files.outputs.changed-files, 'tools/cmake/') name: test-minimal-wheel-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -112,17 +107,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh test-setup-linux-gcc: - needs: docker-image name: test-setup-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-gcc11 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -138,9 +132,8 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "add" "${BUILD_TOOL}" "portable" test-models-linux-basic: - needs: docker-image name: test-models-linux-basic - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -149,23 +142,23 @@ jobs: model: [mv3, vit] backend: [portable, xnnpack-quantization-delegation] build-tool: [cmake, buck2] - runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] + runner: [linux.2xlarge, linux.arm64.2xlarge] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner - # - Excluding the x86 clang image on the ARM64 runner + # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) + # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) exclude: - - runner: mt-l-x86iavx512-8-64 + - runner: linux.2xlarge docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: mt-l-arm64g4-16-62 + - runner: linux.arm64.2xlarge docker-image: executorch-ubuntu-22.04-clang12 # TODO: Need to figure out why buck2 doesnt work on Graviton instances. - - runner: mt-l-arm64g4-16-62 + - runner: linux.arm64.2xlarge build-tool: buck2 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} + docker-image: ci-image:${{ matrix.docker-image }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -183,9 +176,8 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-models-linux: - needs: docker-image name: test-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -193,33 +185,33 @@ jobs: matrix: model: [linear, add, add_mul, ic3, mv2, resnet18, resnet50, mobilebert, emformer_transcribe] backend: [portable, xnnpack-quantization-delegation] - runner: [mt-l-x86iavx512-8-64] + runner: [linux.2xlarge] include: - model: ic4 backend: portable - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory - model: ic4 backend: xnnpack-quantization-delegation - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory - model: emformer_join backend: portable - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory - model: emformer_join backend: xnnpack-quantization-delegation - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory - model: phi_4_mini backend: portable - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory - model: llama3_2_vision_encoder backend: portable - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory - model: w2l backend: portable - runner: mt-l-x86iavx512-16-128 + runner: linux.4xlarge.memory fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -237,17 +229,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-parakeet-xnnpack-linux: - needs: docker-image name: test-parakeet-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-16-128 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.4xlarge.memory + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -271,17 +262,16 @@ jobs: echo "::endgroup::" test-voxtral-realtime-xnnpack-linux: - needs: docker-image name: test-voxtral-realtime-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-16-128 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.4xlarge.memory + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -308,10 +298,9 @@ jobs: echo "::endgroup::" test-llama-runner-linux: - needs: docker-image # Test Both linux x86 and linux aarch64 name: test-llama-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -319,25 +308,25 @@ jobs: matrix: dtype: [fp32] mode: [xnnpack+custom+qe,xnnpack+custom+quantize_kv,xnnpack+quantize_kv] - runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] + runner: [linux.2xlarge, linux.arm64.2xlarge] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] include: - dtype: bf16 mode: custom - runner: mt-l-x86iavx512-8-64 + runner: linux.2xlarge docker-image: executorch-ubuntu-22.04-clang12 # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner - # - Excluding the x86 clang image on the ARM64 runner + # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) + # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) exclude: - - runner: mt-l-x86iavx512-8-64 + - runner: linux.2xlarge docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: mt-l-arm64g4-16-62 + - runner: linux.arm64.2xlarge docker-image: executorch-ubuntu-22.04-clang12 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} + docker-image: ci-image:${{ matrix.docker-image }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -360,17 +349,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -dtype "${DTYPE}" -mode "${MODE}" -upload "${ARTIFACTS_DIR_NAME}" test-llama-runner-linux-android: - needs: docker-image name: test-llama-runner-linux-android - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12-android submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -386,17 +374,16 @@ jobs: bash .ci/scripts/build_llama_android.sh "${BUILD_TOOL}" test-custom-ops-linux: - needs: docker-image name: test-custom-ops-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -411,17 +398,16 @@ jobs: PYTHON_EXECUTABLE=python bash examples/portable/custom_ops/test_custom_ops.sh "${BUILD_TOOL}" test-selective-build-linux: - needs: docker-image name: test-selective-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -436,10 +422,9 @@ jobs: PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" test-multimodal-linux: - needs: docker-image if: ${{ !github.event.pull_request.head.repo.fork }} name: test-multimodal-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -450,8 +435,8 @@ jobs: model: ["gemma3-4b"] # llava gives segfault so not covering. with: secrets-env: EXECUTORCH_HF_TOKEN - runner: mt-l-x86iavx512-94-768 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge.memory + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -476,17 +461,16 @@ jobs: echo "::endgroup::" test-moshi-linux: - needs: docker-image name: test-moshi-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -509,17 +493,16 @@ jobs: python -m unittest examples.models.moshi.mimi.test_mimi test-quantized-aot-lib-linux: - needs: docker-image name: test-quantized-aot-lib-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -533,17 +516,16 @@ jobs: PYTHON_EXECUTABLE=python bash examples/xnnpack/quantization/test_quantize.sh "${BUILD_TOOL}" mv2 test-binary-size-linux-gcc: - needs: docker-image name: test-binary-size-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc9-nopytorch-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-gcc9-nopytorch submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -577,17 +559,16 @@ jobs: fi test-binary-size-linux: - needs: docker-image name: test-binary-size-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -622,9 +603,8 @@ jobs: fi test-arm-cortex-m-size-test: - needs: docker-image name: test-arm-cortex-m-size-test - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -633,8 +613,8 @@ jobs: os: [bare_metal, zephyr-preset] fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -716,15 +696,14 @@ jobs: fi test-mcu-cortex-m-backend: - needs: docker-image name: test-mcu-cortex-m-backend - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge.memory + docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -784,17 +763,16 @@ jobs: docker-image: ci-image:executorch-ubuntu-22.04-clang12 test-qnn-buck-build-linux: - needs: docker-image name: test-qnn-buck-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -819,9 +797,8 @@ jobs: buck2 build //backends/qualcomm/... test-arm-backend-no-driver: - needs: docker-image name: test-arm-backend-no-driver - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -834,8 +811,8 @@ jobs: - test_arm_backend: test_run_tosa fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -855,15 +832,14 @@ jobs: backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-public-api-backward-compatibility: - needs: docker-image name: test-arm-backend-public-api-backward-compatibility - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge.memory + docker-image: ci-image:executorch-ubuntu-24.04-arm-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -883,9 +859,8 @@ jobs: python backends/arm/test/public_api_bc/run_public_api_bc_scenarios.py test-llama-runner-qnn-linux: - needs: docker-image name: test-llama-runner-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -896,8 +871,8 @@ jobs: mode: [qnn] fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -923,9 +898,8 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -mode "${MODE}" -dtype "${DTYPE}" -pt2e_quantize "${PT2E_QUANTIZE}" test-static-llama-qnn-linux: - needs: docker-image name: test-static-llama-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -934,8 +908,8 @@ jobs: task: [stories_110m, stories_260k_bc] fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -958,9 +932,8 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} test-sqnr-static-llm-qnn-linux: - needs: docker-image name: test-sqnr-static-llm-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -969,8 +942,8 @@ jobs: task: [smollm2_135m] fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -993,9 +966,8 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} sqnr test-qnn-models-linux: - needs: docker-image name: test-qnn-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -1004,8 +976,8 @@ jobs: model: [mv2, mv3, dl3] fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -1019,15 +991,14 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn" test-qnn-direct-build-linux: - needs: docker-image name: test-qnn-direct-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1050,23 +1021,22 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true - # No runner-linux, so this takes _test_backend.yml's default. The suite - # runs one export worker per core, so what matters is memory per core: - # the instance this used to run on gave each worker about 4 GiB and the - # job was killed. The default label is a little under 8 GiB per core. + # No runner-linux, so this takes the memory-optimized default. The suite + # runs one export worker per core, so what matters is memory per core, not + # core count: the previous instance gave each worker about 4 GiB and the + # job was killed. The memory-optimized default gives each worker more. test-qnn-passes-linux: - needs: docker-image name: test-qnn-passes-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1095,21 +1065,16 @@ jobs: pytest -xvs backends/qualcomm/tests/test_import_side_effects.py test-qnn-delegate-linux: - needs: docker-image name: test-qnn-delegate-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - # Intel, unlike the avx512 labels: those are r7a, i.e. AMD EPYC, and the - # QNN backend disables MKLDNN on an AMD host through a non-bracketed - # torch.backends mutation that raises once the tests have frozen the - # flags. Same 8 vCPU / 64Gi, on r7i. - runner: mt-l-x86iamx-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1138,17 +1103,16 @@ jobs: -k "TestQNNFloatingPointOperator or TestQNNQuantizedOperator" test-phi-3-mini-runner-linux: - needs: docker-image name: test-phi-3-mini-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-94-192 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1169,17 +1133,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_phi_3_mini.sh Release test-qnn-python-imports-linux: - needs: docker-image name: test-qnn-python-imports-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 15 @@ -1218,17 +1181,16 @@ jobs: --module-prefix executorch.examples.qualcomm test-eval_llama-wikitext-linux: - needs: docker-image name: test-eval_llama-wikitext-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-94-192 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1248,7 +1210,7 @@ jobs: # TODO(larryliu0820): Fix this issue before reenabling it: https://gist.github.com/larryliu0820/7377ecd0d79dbc06076cec8d9f2b85d2 # test-eval_llama-mmlu-linux: # name: test-eval_llama-mmlu-linux - # uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + # uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main # permissions: # id-token: write # contents: read @@ -1274,17 +1236,16 @@ jobs: # PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh test-llama_runner_eager-linux: - needs: docker-image name: test-llama_runner_eager-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-94-192 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1302,17 +1263,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama_runner_eager.sh test-lora-linux: - needs: docker-image name: test-lora-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-94-192 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1330,17 +1290,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora.sh test-lora-multimethod-linux: - needs: docker-image name: test-lora-multimethod-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-94-192 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1358,17 +1317,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora_multimethod.sh test-mediatek-models-linux: - needs: docker-image name: test-mediatek-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-94-192 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-mediatek-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.24xlarge + docker-image: ci-image:executorch-ubuntu-22.04-mediatek-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1386,17 +1344,16 @@ jobs: # placeholder for mediatek to add more tests test-openvino-linux: - needs: docker-image name: test-openvino-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-gcc11 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1409,17 +1366,16 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_openvino.sh test-build-wasm-linux: - needs: docker-image name: test-build-wasm-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1438,9 +1394,8 @@ jobs: PYTHON_EXECUTABLE=python bash examples/wasm/test_build_wasm.sh unittest-wasm-bindings: - needs: docker-image name: unittest-wasm-bindings - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read @@ -1449,8 +1404,8 @@ jobs: enable-etdump: ['', '--enable-etdump'] fail-fast: false with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1485,14 +1440,13 @@ jobs: pnpm test unittest-nxp-neutron: - needs: docker-image - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 150 @@ -1530,19 +1484,18 @@ jobs: bash backends/nxp/run_unittests.sh test-samsung-quantmodels-linux: - needs: docker-image name: test-samsung-quantmodels-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12-android submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -1569,19 +1522,18 @@ jobs: done test-samsung-models-linux: - needs: docker-image name: test-samsung-models-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12-android submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 360 @@ -1612,15 +1564,14 @@ jobs: python -m unittest discover -s backends/samsung/test/models -p "test_*.py" test-vulkan-models-linux: - needs: docker-image name: test-vulkan-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1654,15 +1605,14 @@ jobs: done test-vulkan-operators-linux: - needs: docker-image name: test-vulkan-operators-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-clang12 submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1747,15 +1697,14 @@ jobs: echo "::endgroup::" nxp-build-test: - needs: docker-image name: nxp-build-test - uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main permissions: id-token: write contents: read with: - runner: mt-l-x86iavx512-8-64 - docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} + runner: linux.2xlarge + docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 From 0a4ec570024ebd1ad0fef62d29b9877506228e7f Mon Sep 17 00:00:00 2001 From: mcremon-meta <134334895+mcremon-meta@users.noreply.github.com> Date: Tue, 8 Sep 2026 19:37:28 -0700 Subject: [PATCH 087/190] Extend permute removal through singleton and sink views (#22549) Differential Revision: D116686835 Pull Request resolved: https://github.com/pytorch/executorch/pull/22549 --- ...ve_permutes_around_elementwise_tosa_ops.py | 115 ++++++++++- .../remove_permutes_around_elementwise_ops.py | 181 +++++++++++++----- .../test/test_permute_optimization_passes.py | 96 +++++++++- 3 files changed, 330 insertions(+), 62 deletions(-) diff --git a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py index f7af7cc41e2..c9ceccdef89 100644 --- a/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py +++ b/backends/arm/test/passes/test_remove_permutes_around_elementwise_tosa_ops.py @@ -152,7 +152,7 @@ def test_remove_permutes_around_rescale_tosa_INT() -> None: assert _count_nodes(result.graph_module, RESCALE_TARGET) == 1 -def test_sink_view_preserves_layout_through_rescale_to_broadcast_tosa_INT() -> None: +def test_terminal_sink_view_broadcast_is_optimized_tosa_INT() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") x.meta["val"] = torch.randn(1, 4, 1, 1) @@ -181,11 +181,120 @@ def test_sink_view_preserves_layout_through_rescale_to_broadcast_tosa_INT() -> N graph_module ) - assert not result.modified - assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 1 + assert result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0 assert sub.args == (direct, rescale) +def test_remove_permutes_around_singleton_view_and_gate_region_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 8, 4) + gate_source = graph.placeholder("gate_source") + gate_source.meta["val"] = torch.randn(1, 1, 1, 4) + skip = graph.placeholder("skip") + skip.meta["val"] = torch.randn(1, 8, 4) + + layout_in = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 1])) + layout_in.meta["val"] = torch.randn(1, 4, 8) + clamp = graph.create_node( + "call_function", exir_ops.edge.aten.clamp.default, args=(layout_in, 0, None) + ) + clamp.meta["val"] = torch.randn(1, 4, 8) + + pool_view = graph.create_node( + "call_function", VIEW_TARGET, args=(clamp, [1, 4, 8, 1]) + ) + pool_view.meta["val"] = torch.randn(1, 4, 8, 1) + pool_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(pool_view, [0, 2, 3, 1]) + ) + pool_layout.meta["val"] = torch.randn(1, 8, 1, 4) + + gate = graph.create_node( + "call_function", VIEW_TARGET, args=(gate_source, [1, 4, 1]) + ) + gate.meta["val"] = torch.randn(1, 4, 1) + gated = graph.create_node("call_function", MUL_TARGET, args=(clamp, gate)) + gated.meta["val"] = torch.randn(1, 4, 8) + + skip_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(skip, [0, 2, 1]) + ) + skip_layout.meta["val"] = torch.randn(1, 4, 8) + residual = graph.create_node("call_function", ADD_TARGET, args=(gated, skip_layout)) + residual.meta["val"] = torch.randn(1, 4, 8) + layout_out = graph.create_node( + "call_function", PERMUTE_TARGET, args=(residual, [0, 2, 1]) + ) + layout_out.meta["val"] = torch.randn(1, 8, 4) + graph.output((pool_layout, layout_out)) + + graph_module = torch.fx.GraphModule({}, graph) + inputs = ( + torch.randn(1, 8, 4), + torch.randn(1, 1, 1, 4), + torch.randn(1, 8, 4), + ) + expected = graph_module(*inputs) + + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0 + assert pool_view.args[1] == [1, 8, 1, 4] + assert gate.args[1] == [1, 1, 4] + actual = result.graph_module(*inputs) + torch.testing.assert_close(actual, expected) + + +def test_remove_permutes_around_amax_reduction_region_tosa_INT() -> None: + graph = torch.fx.Graph() + x = graph.placeholder("x") + x.meta["val"] = torch.randn(1, 8, 4) + + layout_in = graph.create_node("call_function", PERMUTE_TARGET, args=(x, [0, 2, 1])) + layout_in.meta["val"] = torch.randn(1, 4, 8) + clamp = graph.create_node( + "call_function", exir_ops.edge.aten.clamp.default, args=(layout_in, 0, None) + ) + clamp.meta["val"] = torch.randn(1, 4, 8) + + conv_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(clamp, [0, 2, 1]) + ) + conv_layout.meta["val"] = torch.randn(1, 8, 4) + maximum = graph.create_node( + "call_function", exir_ops.edge.aten.amax.default, args=(clamp, 2, True) + ) + maximum.meta["val"] = torch.randn(1, 4, 1) + centered = graph.create_node("call_function", SUB_TARGET, args=(clamp, maximum)) + centered.meta["val"] = torch.randn(1, 4, 8) + stats_layout = graph.create_node( + "call_function", PERMUTE_TARGET, args=(centered, [0, 2, 1]) + ) + stats_layout.meta["val"] = torch.randn(1, 8, 4) + graph.output((conv_layout, stats_layout)) + + graph_module = torch.fx.GraphModule({}, graph) + inputs = (torch.randn(1, 8, 4),) + expected = graph_module(*inputs) + + with TosaLoweringContext(TOSA_INT_SPEC): + result = RemovePermutesAroundElementwiseTosaOps(_fake_exported_program()).call( + graph_module + ) + + assert result.modified + assert _count_nodes(result.graph_module, PERMUTE_TARGET) == 0 + assert maximum.args[1] == 1 + actual = result.graph_module(*inputs) + torch.testing.assert_close(actual, expected) + + def test_remove_permutes_around_gelu_with_folded_scalar_constants_tosa_FP() -> None: graph = torch.fx.Graph() x = graph.placeholder("x") diff --git a/backends/transforms/remove_permutes_around_elementwise_ops.py b/backends/transforms/remove_permutes_around_elementwise_ops.py index 1a3a6ada995..3f698f66562 100644 --- a/backends/transforms/remove_permutes_around_elementwise_ops.py +++ b/backends/transforms/remove_permutes_around_elementwise_ops.py @@ -61,11 +61,22 @@ class Subgraph: interleaves: dict[ torch.fx.Node, tuple[int, int, torch.fx.Node, torch.fx.Node] ] = field(default_factory=dict) + # Views whose target shapes must change when the surrounding layout + # transforms are removed. Boundary views are inside the region; sink + # views feed it from layout-invariant single-non-unit tensors. + view_shape_overrides: dict[torch.fx.Node, list[int]] = field( + default_factory=dict + ) + sink_edges_in: set[tuple[torch.fx.Node, torch.fx.Node]] = field( + default_factory=set + ) def __init__(self, extra_permutable_ops: set | None = None) -> None: super().__init__() self._permutable_ops = { exir_ops.edge.aten.add.Tensor, + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, exir_ops.edge.aten.mul.Tensor, exir_ops.edge.aten.sub.Tensor, exir_ops.edge.aten.hardtanh.default, @@ -186,39 +197,62 @@ def _is_permutation_sink_view(self, node: torch.fx.Node) -> bool: non_unit = [d for d in shape if not (isinstance(d, int) and d == 1)] return len(non_unit) <= 1 - def _sink_users_are_layout_invariant(self, sink: torch.fx.Node) -> bool: - """Return whether dropping layout at ``sink`` is safe for its consumers.""" - frontier = [(user, sink) for user in sink.users] - visited: set[torch.fx.Node] = set() - while frontier: - node, producer = frontier.pop() - if node in visited: - continue - visited.add(node) + def _remapped_sink_shape( + self, sink: torch.fx.Node, start_permute: list[int] + ) -> list[int] | None: + """Return the sink shape in the layout before ``start_permute``. - if node.op == "output": - continue - if node.target == exir_ops.edge.aten.permute_copy.default: - # This explicit transform re-establishes the downstream layout, - # so consumers beyond it do not depend on the sink's layout. - continue - if self._is_permutation_sink_view(node): - continue + A sink input has at most one non-unit dimension, so changing the view's + shape does not change element order. The output rank must match the + region permutation so its broadcast axes can be remapped exactly. + """ + out_shape = self._concrete_shape(sink) + if out_shape is None or len(out_shape) != len(start_permute): + return None + # Reshape can relocate one contiguous non-unit run among singleton axes, + # but it cannot transpose multiple non-unit output axes. + if sum(dim != 1 for dim in out_shape) > 1: + return None + inverse = [start_permute.index(i) for i in range(len(start_permute))] + return [out_shape[index] for index in inverse] - tensor_inputs = [ - input_node - for input_node in node.all_input_nodes - if input_node.meta.get("val") is not None - ] - if any( - input_node is not producer and input_node.meta["val"].numel() != 1 - for input_node in tensor_inputs - ): - return False - if not self.is_node_permutable(node): - return False - frontier.extend((user, node) for user in node.users) - return True + def _singleton_view_boundary_shape( + self, + view: torch.fx.Node, + start_permute: list[int], + end_permute_node: torch.fx.Node, + ) -> list[int] | None: + """Compose a layout pair across a view that only inserts unit dims.""" + shapes = self._view_shapes(view) + end_dims = self.get_permutation(end_permute_node) + end_shape = self._concrete_shape(end_permute_node) + if shapes is None or end_dims is None or end_shape is None: + return None + in_shape, out_shape = shapes + if len(start_permute) != len(in_shape) or len(end_dims) != len(out_shape): + return None + + inserted = self._find_extra_ones(out_shape, in_shape) + if inserted is None: + return None + + # Label each old view-output axis by the corresponding axis before the + # incoming permutation. The outgoing permutation must restore those + # labels to identity order; inserted singleton axes carry no label. + labels: list[int | None] = list(start_permute) + for index in inserted: + labels.insert(index, None) + output_labels = [labels[index] for index in end_dims] + if [label for label in output_labels if label is not None] != list( + range(len(in_shape)) + ): + return None + + inverse = [start_permute.index(i) for i in range(len(start_permute))] + unpermuted_input = [in_shape[index] for index in inverse] + if self._find_extra_ones(end_shape, unpermuted_input) is None: + return None + return end_shape def _inserted_unit_dim(self, node: torch.fx.Node) -> int | None: """Position of the size-1 dim ``node`` inserts, else None. @@ -484,7 +518,19 @@ def visit( # noqa: C901 for user in users_source.users: if user.target in PERMUTE_COPY_TARGETS: user_perm = self.get_permutation(user) - if user_perm == downstream_end: + boundary_shape = None + if ( + triple is None + and self._is_squeeze_unsqueeze_view(node) + and len(node.users) == 1 + ): + boundary_shape = self._singleton_view_boundary_shape( + node, current_start_permute, user + ) + if boundary_shape is not None: + subgraph.view_shape_overrides[node] = boundary_shape + subgraph.edges_out.add((users_source, user)) + elif user_perm == downstream_end: subgraph.edges_out.add((users_source, user)) else: # Non-matching permute: keep it and fold the start permute into it @@ -501,10 +547,11 @@ def visit( # noqa: C901 elif user.op == "output": return False elif self._is_permutation_sink_view(user): - # The tensor's element order is invariant at this reshape, but - # its output shape can still carry broadcast-axis meaning. - if not self._sink_users_are_layout_invariant(user): - return False + # A sink with no other path into this region can terminate it: + # its single non-unit run has layout-invariant element order. + # If a later consumer is reached through another region branch, + # upstream traversal records and remaps the sink via + # ``sink_edges_in`` before any boundary is removed. continue elif not self.visit( user, subgraph, processed_nodes, downstream_end, downstream_start @@ -524,6 +571,15 @@ def visit( # noqa: C901 # stays wired directly. Notably this keeps lifted per-tensor # qparam placeholders as placeholders, which lowering requires. continue + elif inp in subgraph.nodes: + # Already part of the region; it is rewritten as a region node. + continue + elif self._is_permutation_sink_view(inp): + remapped_shape = self._remapped_sink_shape(inp, current_start_permute) + if remapped_shape is None or len(inp.users) != 1: + return False + subgraph.view_shape_overrides[inp] = remapped_shape + subgraph.sink_edges_in.add((inp, node)) elif self._is_constant(inp): const_rank = self._get_node_rank(inp) permute_rank = len(current_end_permute) @@ -603,6 +659,8 @@ def is_node_permutable(self, node: torch.fx.Node) -> bool: return False if node.target in self._permutable_ops: if node.target in ( + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.sum.dim_IntList, ): @@ -693,16 +751,24 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 if node.target == exir_ops.edge.aten.cat.default: self.update_cat(node, node_start_perm) elif node.target in ( + exir_ops.edge.aten.amax.default, + exir_ops.edge.aten.amin.default, exir_ops.edge.aten.mean.dim, exir_ops.edge.aten.sum.dim_IntList, ): - self.update_mean_dim(node, node_start_perm) + self.update_reduction_dim(node, node_start_perm) elif node.target == exir_ops.edge.aten.slice_copy.Tensor: self.update_slice_copy(node, node_start_perm) elif node.target in self._PAD_OPS: self.update_pad(node, node_start_perm) elif node.target in self._VIEW_OPS: - self.update_view_copy(node, node_start_perm) + if node in subgraph.view_shape_overrides: + node.update_arg(1, subgraph.view_shape_overrides[node]) + else: + self.update_view_copy(node, node_start_perm) + + for sink, _ in subgraph.sink_edges_in: + sink.update_arg(1, subgraph.view_shape_overrides[sink]) for head, triple in subgraph.interleaves.items(): self.update_interleave( @@ -769,20 +835,19 @@ def permute_subgraph(self, subgraph: Subgraph) -> bool: # noqa: C901 def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: """Return false if an earlier rewrite invalidated this candidate.""" for inp, out in subgraph.edges_in: - if inp.target not in PERMUTE_COPY_TARGETS or inp not in out.all_input_nodes: - return False - - # edges_out_to_update can rewrite a permute in place, leaving it wired. - if self.get_permutation(inp) != subgraph.node_start_permute.get( - out, subgraph.start_permute + if ( + inp.target not in PERMUTE_COPY_TARGETS + or inp not in out.all_input_nodes + # edges_out_to_update can rewrite a permute in place, leaving it wired. + or self.get_permutation(inp) + != subgraph.node_start_permute.get(out, subgraph.start_permute) ): return False - for inp, out in subgraph.edges_out: - if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: - return False - - for inp, out, _ in subgraph.edges_out_to_update: + outgoing = list(subgraph.edges_out) + [ + (inp, out) for inp, out, _ in subgraph.edges_out_to_update + ] + for inp, out in outgoing: if out.target not in PERMUTE_COPY_TARGETS or out not in inp.users: return False @@ -790,6 +855,10 @@ def _subgraph_edges_are_current(self, subgraph: Subgraph) -> bool: if const_node not in user_node.all_input_nodes: return False + for sink, user_node in subgraph.sink_edges_in: + if sink not in user_node.all_input_nodes or len(sink.users) != 1: + return False + for head, (_, _, expand_node, view_node) in subgraph.interleaves.items(): if ( len(head.users) != 1 @@ -839,9 +908,19 @@ def update_cat(self, node: torch.fx.Node, start_permute: list[int]) -> None: dim = get_arg(node, "dim", int) set_arg(node, "dim", start_permute[dim]) - def update_mean_dim(self, node: torch.fx.Node, start_permute: list[int]) -> None: + def update_reduction_dim( + self, node: torch.fx.Node, start_permute: list[int] + ) -> None: dims = get_arg(node, "dim") - set_arg(node, "dim", [start_permute[d] for d in cast(list[int], dims)]) + rank = len(start_permute) + if isinstance(dims, int): + set_arg(node, "dim", start_permute[dims % rank]) + else: + set_arg( + node, + "dim", + [start_permute[d % rank] for d in cast(list[int], dims)], + ) def update_slice_copy(self, node: torch.fx.Node, start_permute: list[int]) -> None: dim = get_arg(node, "dim", int) diff --git a/backends/transforms/test/test_permute_optimization_passes.py b/backends/transforms/test/test_permute_optimization_passes.py index 1f357deb171..da47fe8186e 100644 --- a/backends/transforms/test/test_permute_optimization_passes.py +++ b/backends/transforms/test/test_permute_optimization_passes.py @@ -1600,7 +1600,7 @@ def test_permutation_sink_view_splitting_the_non_unit_dim(self) -> None: "permutation_sink_view_splitting_the_non_unit_dim", ) - def test_permutation_sink_view_preserves_broadcast_layout(self) -> None: + def test_permutation_sink_view_terminal_broadcast_is_optimized(self) -> None: x_data = torch.randn(1, 4, 1, 1) direct_data = torch.randn(1, 8, 4) builder = GraphBuilder() @@ -1623,15 +1623,87 @@ def test_permutation_sink_view_preserves_broadcast_layout(self) -> None: gm_before = copy.deepcopy(original) result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) self.assertEqual( - count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 ) validate_numerics( gm_before, result.graph_module, [x_data, direct_data], - "permutation_sink_view_preserves_broadcast_layout", + "permutation_sink_view_terminal_broadcast_is_optimized", + ) + + def test_split_sink_view_is_not_remapped_for_broadcast(self) -> None: + x_data = torch.randn(1, 4, 2) + sink_data = torch.randn(1, 1, 1, 8) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + sink_source = builder.placeholder("sink_source", sink_data) + permute_in = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + split_sink = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(sink_source, [1, 2, 4]), + ) + mul = builder.call_operator( + op=exir_ops.edge.aten.mul.Tensor, args=(permute_in, split_sink) + ) + permute_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(mul, [0, 2, 1]) + ) + builder.output([permute_out]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + self.assertFalse(result.modified) + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data, sink_data], + "split_sink_view_is_not_remapped_for_broadcast", + ) + + def test_shared_singleton_view_boundaries_are_not_remapped(self) -> None: + x_data = torch.randn(1, 4, 8) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + permute_in = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, args=(x, [0, 2, 1]) + ) + view = builder.call_operator( + op=exir_ops.edge.aten.view_copy.default, + args=(permute_in, [1, 8, 4, 1]), + ) + first_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(view, [0, 2, 3, 1]), + ) + second_out = builder.call_operator( + op=exir_ops.edge.aten.permute_copy.default, + args=(view, [0, 3, 1, 2]), + ) + builder.output([first_out, second_out]) + original = builder.get_graph_module() + gm_before = copy.deepcopy(original) + + result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) + # The view feeds two differing outgoing permutes, so its shape cannot be + # composed across the boundary. Both outgoing permutes survive; only the + # incoming permute is folded into them. + self.assertEqual( + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 2 + ) + validate_numerics( + gm_before, + result.graph_module, + [x_data], + "shared_singleton_view_boundaries_are_not_remapped", ) def test_permutation_sink_view_preserves_cat_layout(self) -> None: @@ -1657,10 +1729,14 @@ def test_permutation_sink_view_preserves_cat_layout(self) -> None: gm_before = copy.deepcopy(original) result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) self.assertEqual( - count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + (view_after,) = result.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.view_copy.default ) + self.assertEqual(view_after.args[1], [1, 1, 4]) validate_numerics( gm_before, result.graph_module, @@ -1693,10 +1769,14 @@ def test_permutation_sink_view_preserves_keyword_broadcast_layout(self) -> None: gm_before = copy.deepcopy(original) result = cast(PassResult, RemovePermutesAroundElementwiseOps()(original)) - self.assertFalse(result.modified) + self.assertTrue(result.modified) self.assertEqual( - count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 1 + count_node(result.graph_module, exir_ops.edge.aten.permute_copy.default), 0 + ) + (view_after,) = result.graph_module.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.view_copy.default ) + self.assertEqual(view_after.args[1], [1, 1, 4]) validate_numerics( gm_before, result.graph_module, From fd985d516878659e1050bd9b7d7bb00df2ad6ed5 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= Date: Wed, 9 Sep 2026 07:49:37 +0200 Subject: [PATCH 088/190] XNNPACK: Preserve FP16 SiLU through XNNPACK partitioning (#22604) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit PyTorch's default decomposition can lower FP16 aten.silu to a float32 sigmoid and multiply surrounded by dtype copies. For SiLU-heavy models, this fragments XNNPACK delegation and adds portable conversion overhead. Add an XNNPACK partitioner configuration that preserves aten.silu only for FP16 inputs and outputs in to_edge_transform_and_lower(). During delegate preprocessing, rewrite the preserved Edge SiLU as x * sigmoid(x), keeping computation in FP16. FP32 SiLU continues through the default decomposition. Add an XNNPACK transform pass that recognizes this decomposition and rewrites it as x * sigmoid(x) in FP16 before partitioning. Restrict the rewrite to the exact FP16-to-FP32 pattern and support both Edge copy forms. In an Yolo26 BundleIO benchmark, mean latency fell from about 15.6 ms to about 8.8 ms. This is approximately 7 ms faster, a 45% reduction and a 1.8x speedup. Delegated subgraphs fell from 143 to 19 and non-delegated nodes from 425 to 86, while delegated nodes remained at 363. All runs passed BundleIO. cc @GregoryComer @digantdesai @cbilgin @JakeStevens @freddan80 @per @zingo @oscarandersson8218 @Sebastian-Larsson @robell @rascani Signed-off-by: Måns Nilsson --- backends/xnnpack/_passes/__init__.py | 3 + backends/xnnpack/_passes/rewrite_fp16_silu.py | 44 ++++++++ backends/xnnpack/partition/config/__init__.py | 2 + .../partition/config/generic_node_configs.py | 23 ++++ .../test/passes/test_rewrite_fp16_silu.py | 104 ++++++++++++++++++ 5 files changed, 176 insertions(+) create mode 100644 backends/xnnpack/_passes/rewrite_fp16_silu.py create mode 100644 backends/xnnpack/test/passes/test_rewrite_fp16_silu.py diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index 22147fa4215..f55144395af 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -35,6 +36,7 @@ from executorch.backends.xnnpack._passes.remove_redundant_copy_pass import ( RemoveRedundantCopyPass, ) +from executorch.backends.xnnpack._passes.rewrite_fp16_silu import RewriteFp16SiluPass from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.exir.pass_base import ExportPass @@ -69,6 +71,7 @@ def __init__( if not passes: # All the XNNPACK passes self.passes = [ + RewriteFp16SiluPass, XNNPACKRemoveCloneOpsTransform, # TODO - remove this pass once we have a better support for dim_order ops lowering DimOrderOpsRevertPass, diff --git a/backends/xnnpack/_passes/rewrite_fp16_silu.py b/backends/xnnpack/_passes/rewrite_fp16_silu.py new file mode 100644 index 00000000000..775dcf7bc18 --- /dev/null +++ b/backends/xnnpack/_passes/rewrite_fp16_silu.py @@ -0,0 +1,44 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult + + +class RewriteFp16SiluPass(ExportPass): + """Rewrite preserved FP16 SiLU into FP16 sigmoid and multiply.""" + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + graph = graph_module.graph + modified = False + + for node in list(graph.nodes): + if node.target != exir_ops.edge.aten.silu.default: + continue + + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + continue + + with graph.inserting_before(node): + sigmoid = graph.call_function( + exir_ops.edge.aten.sigmoid.default, (input_node,) + ) + mul = graph.call_function( + exir_ops.edge.aten.mul.Tensor, (input_node, sigmoid) + ) + node.replace_all_uses_with(mul) + modified = True + + if not modified: + return PassResult(graph_module, False) + + graph.eliminate_dead_code() + graph.lint() + graph_module.recompile() + graph_module = super().call(graph_module).graph_module + return PassResult(graph_module, True) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index c6c54f083d6..7775a674510 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/backends/xnnpack/partition/config/__init__.py @@ -50,6 +50,7 @@ ReciprocalSquareRootConfig, ReLUConfig, SigmoidConfig, + SiluConfig, SinConfig, SliceCopyConfig, SoftmaxConfig, @@ -115,6 +116,7 @@ TanhConfig, ToDimOrderCopyConfig, SigmoidConfig, + SiluConfig, SinConfig, CosConfig, SliceCopyConfig, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index f2b946e412d..8741962db38 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -409,6 +409,29 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32] +class SiluConfig(GenericNodePartitionerConfig): + target_name = "silu.default" + + def supported_precision_types(self) -> List[ConfigPrecisionType]: + return [ConfigPrecisionType.FP32] + + def get_original_aten(self) -> Optional[torch._ops.OpOverload]: + return torch.ops.aten.silu.default + + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + if not self.check_common_constraints(node, ep): + return False + + input_value = node.args[0].meta.get("val") + output_value = node.meta.get("val") + return ( + isinstance(input_value, torch.Tensor) + and input_value.dtype == torch.float16 + and isinstance(output_value, torch.Tensor) + and output_value.dtype == torch.float16 + ) + + class MulConfig(GenericNodePartitionerConfig): target_name = "mul.Tensor" diff --git a/backends/xnnpack/test/passes/test_rewrite_fp16_silu.py b/backends/xnnpack/test/passes/test_rewrite_fp16_silu.py new file mode 100644 index 00000000000..131be95f846 --- /dev/null +++ b/backends/xnnpack/test/passes/test_rewrite_fp16_silu.py @@ -0,0 +1,104 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch + +from executorch.backends.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner +from executorch.backends.xnnpack.test.tester import Tester +from executorch.backends.xnnpack.utils.configs import get_xnnpack_edge_compile_config +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class TestRewriteFp16SiluPass(unittest.TestCase): + edge_mul = "executorch_exir_dialects_edge__ops_aten_mul_Tensor" + edge_sigmoid = "executorch_exir_dialects_edge__ops_aten_sigmoid_default" + edge_silu = "executorch_exir_dialects_edge__ops_aten_silu_default" + + class Silu(torch.nn.Module): + def forward(self, x): + return torch.nn.functional.silu(x) + + def setUp(self): + torch._dynamo.reset() + + def _export_silu(self, dtype): + return torch.export.export( + self.Silu(), + (torch.randn(2, 3, dtype=dtype),), + strict=True, + ) + + def _get_preserved_edge_fp16_silu(self): + edge_program = to_edge( + self._export_silu(torch.float16), + compile_config=get_xnnpack_edge_compile_config(), + ).exported_program() + graph = edge_program.graph_module.graph + input_node = next(node for node in graph.nodes if node.op == "placeholder") + output_node = next(node for node in graph.nodes if node.op == "output") + silu = output_node.args[0][0] + silu.target = exir_ops.edge.aten.silu.default + silu.args = (input_node,) + silu.kwargs = {} + graph.eliminate_dead_code() + graph.lint() + edge_program.graph_module.recompile() + return edge_program + + def test_ops_to_not_decompose_filters_for_fp16(self): + partitioner = XnnpackPartitioner() + + fp16_program = self._export_silu(torch.float16) + preserved_ops, filter_fn = partitioner.ops_to_not_decompose(fp16_program) + fp16_silu = next( + node + for node in fp16_program.graph.nodes + if node.target == torch.ops.aten.silu.default + ) + self.assertIn(torch.ops.aten.silu.default, preserved_ops) + self.assertIsNotNone(filter_fn) + self.assertTrue(filter_fn(fp16_silu)) + + fp32_program = self._export_silu(torch.float32) + _, filter_fn = partitioner.ops_to_not_decompose(fp32_program) + fp32_silu = next( + node + for node in fp32_program.graph.nodes + if node.target == torch.ops.aten.silu.default + ) + self.assertIsNotNone(filter_fn) + self.assertFalse(filter_fn(fp32_silu)) + + def test_preprocess_rewrites_preserved_fp16_silu(self): + result = XNNPACKPassManager(self._get_preserved_edge_fp16_silu()).transform() + targets = [node.target for node in result.graph.nodes] + + self.assertNotIn(exir_ops.edge.aten.silu.default, targets) + self.assertEqual(targets.count(exir_ops.edge.aten.sigmoid.default), 1) + self.assertEqual(targets.count(exir_ops.edge.aten.mul.Tensor), 1) + + def test_fp32_silu_uses_default_decomposition(self): + ( + Tester(self.Silu(), (torch.randn(2, 3),)) + .export() + .to_edge() + .check_count({self.edge_silu: 0, self.edge_sigmoid: 1, self.edge_mul: 1}) + ) + + def test_to_edge_transform_and_lower_delegates_fp16_silu(self): + ( + Tester(self.Silu(), (torch.randn(2, 3, dtype=torch.float16),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not([self.edge_silu]) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) From 63289c198f811bffb59bf45cd85cec82c228aeed Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Wed, 9 Sep 2026 08:25:49 +0200 Subject: [PATCH 089/190] NXP backend: Add recipes for Neutron backend lowering. (#21516) ### Summary This PR introduces a declarative recipe-based lowering for the NXP Neutron backend. The previous solution was implemented as `executorch_pipeline.py:to_quantized_executorch_program()`. The new solution provides the same functionailty and produces the same results. The benefit of the recipe-based approach is compatibility with other backends (recipe fusing) and adhering to ExecuTorch standards. Once this is merged, the old Neutron lowering pipeline can be removed completely. ### Test plan `pytest backends/nxp/tests/generic_tests/test_recipe_export.py ` cc @robert-kalmar @JakeStevens @digantdesai @rascani --- .../edge_passes/neutron_edge_pass_manager.py | 2 +- backends/nxp/nxp_backend.py | 8 + backends/nxp/recipes/nxp_recipe_provider.py | 332 ++++++++++++++ backends/nxp/recipes/nxp_recipe_types.py | 30 ++ backends/nxp/tests/executorch_pipeline.py | 15 +- .../tests/generic_tests/test_recipe_export.py | 415 ++++++++++++++++++ examples/nxp/aot_neutron_compile.py | 2 + export/recipe.py | 28 ++ export/stages.py | 9 + 9 files changed, 829 insertions(+), 12 deletions(-) create mode 100644 backends/nxp/recipes/nxp_recipe_provider.py create mode 100644 backends/nxp/recipes/nxp_recipe_types.py create mode 100644 backends/nxp/tests/generic_tests/test_recipe_export.py diff --git a/backends/nxp/edge_passes/neutron_edge_pass_manager.py b/backends/nxp/edge_passes/neutron_edge_pass_manager.py index 3a7fc3ffbfa..e95f8d18144 100644 --- a/backends/nxp/edge_passes/neutron_edge_pass_manager.py +++ b/backends/nxp/edge_passes/neutron_edge_pass_manager.py @@ -17,7 +17,7 @@ from executorch.backends.nxp.edge_passes.remove_as_strided_copy_nodes import ( RemoveUselessAsStridedCopyNodes, ) -from torch.fx.passes.infra.pass_manager import PassManager +from executorch.exir.pass_manager import PassManager class NeutronEdgePassManager(PassManager): diff --git a/backends/nxp/nxp_backend.py b/backends/nxp/nxp_backend.py index 16430afada5..0d4324d57c2 100644 --- a/backends/nxp/nxp_backend.py +++ b/backends/nxp/nxp_backend.py @@ -45,6 +45,14 @@ torch.ops.aten.prelu.default, ] +# Aten operators that must be preserved (not decomposed) during lowering to the edge dialect, because the Neutron +# backend can handle them natively. +default_preserve_ops = [ + torch.ops.aten.hardswish.default, + torch.ops.aten.pad.default, + torch.ops.aten.prelu.default, +] + class NeutronCompileSpecBuilder: config: NeutronTargetSpec diff --git a/backends/nxp/recipes/nxp_recipe_provider.py b/backends/nxp/recipes/nxp_recipe_provider.py new file mode 100644 index 00000000000..5617bcdb6f2 --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_provider.py @@ -0,0 +1,332 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging +from copy import deepcopy +from dataclasses import dataclass +from functools import partial +from typing import Any, Callable, cast, Iterable, Optional, Sequence + +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.neutron_target_spec import NeutronTargetSpec +from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass +from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( + NeutronEdgePassManager, +) +from executorch.backends.nxp.edge_passes.remove_additional_quantize_dequantize_nodes_pass import ( + RemoveAdditionalQDQClustersPass, +) +from executorch.backends.nxp.edge_passes.remove_io_quant_ops_pass import ( + RemoveIOQuantOpsPass, +) +from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner +from executorch.backends.nxp.nxp_backend import ( + core_aten_ops_exception_list, + default_preserve_ops, + generate_neutron_compile_spec, +) +from executorch.backends.nxp.recipes.nxp_recipe_types import NXP_BACKEND, NXPRecipeType +from executorch.backends.nxp.tests.executorch_pipeline import ( + get_default_quantizer, + handle_kernel_selection, + ModelInputSpec, + to_model_input_spec, +) +from executorch.exir import ( + EdgeCompileConfig, + EdgeProgramManager, + ExecutorchBackendConfig, + ExportedProgram, +) + +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import Partitioner +from executorch.export import ( + BackendRecipeProvider, + ExportRecipe, + LoweringRecipe, + QuantizationRecipe, + RecipeType, +) +from torchao.quantization.pt2e.quantizer import Quantizer + + +class NeutronEdgePassManagerWrapper: + def __init__(self, passes: list[NeutronEdgePass] | None = None): + self.neutron_edge_pass_manager = NeutronEdgePassManager(passes) + + def __call__( + self, method_name: str, exported_program: ExportedProgram + ) -> NeutronEdgePassManager: + return self.neutron_edge_pass_manager + + +NEUTRON_RECIPE_CONFIG_KEY = "neutron_recipe_config" + + +@dataclass +class NeutronRecipeConfig: + """Configuration shared by all NXP recipe types. + + Parameters that vary the *type* of export (delegate vs no-delegate) + are expressed by choosing a different NXPRecipeType rather than by flags here. + + Attributes: + input_spec: Model input description. Accepts a single shape tuple, a list of + shape tuples (one per input), or a list of ModelInputSpec objects. + target: Neutron hardware target string. Default: "imxrt700". + operators_not_to_delegate: Optional list of op names excluded from NPU delegation. + For example ["aten::convolution"]. + intermediates_dir: Optional directory to dump intermediate compilation artifacts. + get_quantizer_fn: Optional factory that returns a custom Quantizer. When None, + the default NeutronQuantizer is used. + custom_delegation_options: Optional fine-grained control over which ops are + delegated. Default: CustomDelegationOptions(). + remove_quant_io_ops: If True, remove quantize/dequantize ops at the IO boundary + (useful for integer-IO deployments). + use_quant_state_dict: If False, the post-quantization parameter values are not + passed to NeutronPartitioner. + use_neutron_for_format_conversion: Whether Neutron handles data-format conversion. + fetch_constants_to_sram: Place constant tensors in SRAM on the target. + dump_kernel_selection_code: Generate kernel-selection files after compilation. + use_profiling: Enable Neutron execution profiling. IMPORTANT: To also generate an + ETRecord, pass generate_etrecord=True to export() separately. + """ + + input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]] + target: str = "imxrt700" + operators_not_to_delegate: list[str] | None = None + intermediates_dir: str | None = None + get_quantizer_fn: Callable[[], Quantizer] | None = None + custom_delegation_options: CustomDelegationOptions | None = None + remove_quant_io_ops: bool = False + use_quant_state_dict: bool = True + use_neutron_for_format_conversion: bool = True + fetch_constants_to_sram: bool = False + dump_kernel_selection_code: bool = False + use_profiling: bool = False + + +class NXPRecipeProvider(BackendRecipeProvider): + + @property + def backend_name(self) -> str: + return NXP_BACKEND + + def get_supported_recipes(self) -> Sequence[RecipeType]: + return list(NXPRecipeType) + + def create_recipe( + self, recipe_type: RecipeType, **kwargs: Any + ) -> Optional[ExportRecipe]: + if recipe_type not in self.get_supported_recipes(): + logging.warning(f"NXP backend: Recipe `{recipe_type}` is not valid.") + return None + + original_rc = kwargs.get(NEUTRON_RECIPE_CONFIG_KEY) + if original_rc is None: + raise KeyError( + f"NXP backend: create_recipe() requires `{NEUTRON_RECIPE_CONFIG_KEY}=`." + ) + if not isinstance(original_rc, NeutronRecipeConfig): + raise TypeError( + f"NXP backend: `{NEUTRON_RECIPE_CONFIG_KEY}` must be a NeutronRecipeConfig, " + f"got {type(original_rc).__name__}." + ) + + rc = cast(NeutronRecipeConfig, deepcopy(original_rc)) + if rc.custom_delegation_options is None: + rc.custom_delegation_options = CustomDelegationOptions() + + rc.input_spec = to_model_input_spec(rc.input_spec) + + match recipe_type: + case NXPRecipeType.INT8_PTQ_NEUTRON: + return self._build_recipe(recipe_type, rc, is_qat=False, delegate=True) + case NXPRecipeType.INT8_PTQ_NO_DELEGATE: + return self._build_recipe(recipe_type, rc, is_qat=False, delegate=False) + case _: + raise NotImplementedError( + f"NXP backend: Recipe `{recipe_type}` is not supported." + ) + + def _build_recipe( + self, + recipe_type: NXPRecipeType, + rc: NeutronRecipeConfig, + *, + is_qat: bool, + delegate: bool, + ) -> ExportRecipe: + # is_qat=True is reserved for future QAT support; always False for now. + if is_qat: + raise NotImplementedError( + "NXP recipe with QAT (quantization aware training) is not yet supported." + ) + + neutron_target_spec = NeutronTargetSpec(rc.target) + + if rc.get_quantizer_fn is None: + rc.get_quantizer_fn = partial( + get_default_quantizer, neutron_target_spec, is_qat + ) + + quantization_recipe = _build_quantization_recipe(rc) + compile_spec = generate_neutron_compile_spec( + rc.target, + intermediates_dir=rc.intermediates_dir, + operators_not_to_delegate=rc.operators_not_to_delegate, + use_neutron_for_format_conversion=rc.use_neutron_for_format_conversion, + fetch_constants_to_sram=rc.fetch_constants_to_sram, + dump_kernel_selection_code=rc.dump_kernel_selection_code, + use_profiling=rc.use_profiling, + ) + lowering_recipe = _build_lowering_recipe( + compile_spec, neutron_target_spec, rc, delegate=delegate + ) + + return ExportRecipe( + name=recipe_type.value, + quantization_recipe=quantization_recipe, + lowering_recipe=lowering_recipe, + executorch_backend_config=ExecutorchBackendConfig( + extract_delegate_segments=False + ), + ) + + +# --------------------------------------------------------------------------- +# Module-level builder helpers +# --------------------------------------------------------------------------- + + +def _build_quantization_recipe(rc: NeutronRecipeConfig) -> QuantizationRecipe: + """Build the QuantizationRecipe for PTQ. + + PTQ uses the standard QuantizeStage flow (prepare_pt2e -> calibrate -> convert_pt2e). + The example_inputs passed to the export session are used directly for calibration. + Multiple PTQ recipes can be combined with ExportRecipe.combine(). + """ + _quantizer = rc.get_quantizer_fn() + + return QuantizationRecipe( + quantizers=[_quantizer], + ) + + +def _build_lowering_recipe( + compile_spec: list[CompileSpec], + neutron_target_spec: NeutronTargetSpec, + rc: NeutronRecipeConfig, + *, + delegate: bool, +) -> LoweringRecipe: + """Build the LoweringRecipe, optionally including NPU delegation.""" + partitioners = _build_partitioners(compile_spec, neutron_target_spec, rc, delegate) + pre_partitioning_callback = _build_pre_partitioning_callback(rc) + edge_manager_transform_passes = _build_edge_manager_transform_passes(rc) + + # The edge pass manager must be wrapped: EdgeTransformAndLowerStage calls + # edge_transform_passes with (method_name, ep) and expects a PassManager back. + return LoweringRecipe( + partitioners=partitioners, + edge_transform_passes=[NeutronEdgePassManagerWrapper()], + edge_compile_config=EdgeCompileConfig( + _check_ir_validity=False, + _core_aten_ops_exception_list=core_aten_ops_exception_list, + ), + pre_partitioning_callback=pre_partitioning_callback, + edge_manager_transform_passes=edge_manager_transform_passes, + ) + + +def _build_partitioners( + compile_spec: list[CompileSpec], + neutron_target_spec: NeutronTargetSpec, + rc: NeutronRecipeConfig, + delegate: bool, +) -> list: + """Create the NeutronPartitioner list. Empty when delegate=False.""" + if not delegate: + return [] + return [ + NeutronPartitioner( + compile_spec, + neutron_target_spec, + rc.custom_delegation_options, + preserve_ops=default_preserve_ops, + ) + ] + + +def _build_pre_partitioning_callback(rc: NeutronRecipeConfig): + """Return a callback that assigns the post-quantization state_dict to NeutronPartitioner. + + NeutronPartitioner requires static parameter data. Since the partitioner is instantiated + during recipe creation (before model data is available), assignment is deferred to a + callback invoked just before partitioning. + """ + _use_quant_state_dict = rc.use_quant_state_dict + + def _callback( + _partitioners: list[Partitioner] | None, + programs: dict[str, ExportedProgram], + ) -> None: + if not _partitioners: + return + + if _use_quant_state_dict: + post_quant_state_dict: dict | None = {} + for _, program in programs.items(): + post_quant_state_dict.update(program.state_dict) + else: + post_quant_state_dict = None + + for _partitioner in _partitioners: + if isinstance(_partitioner, NeutronPartitioner): + _partitioner.post_quantization_state_dict = post_quant_state_dict + + return _callback + + +def _build_edge_manager_transform_passes(rc: NeutronRecipeConfig) -> list: + """Build edge_manager_transform_passes for the post-partitioning graph cleanup. + + These run in EdgeProgramManagerTransformStage, after to_edge_transform_and_lower: + - RemoveIOQuantOpsPass (optional, when remove_quant_io_ops=True) + - RemoveAdditionalQDQClustersPass (always applied) + - handle_kernel_selection side-effect (optional, when dump_kernel_selection_code=True) + + Each callable receives EdgeProgramManager and returns passes for epm.transform(), + or an empty list when no graph transformation is needed (side-effect only). + """ + passes = [] + + if rc.remove_quant_io_ops: + + def _remove_io_quant_ops(epm: EdgeProgramManager) -> list: + return [RemoveIOQuantOpsPass(edge_program_manager=epm)] + + passes.append(_remove_io_quant_ops) + + def _remove_additional_qdq_clusters( + epm: EdgeProgramManager, + ) -> NeutronEdgePassManager: + return NeutronEdgePassManager([RemoveAdditionalQDQClustersPass()]) + + passes.append(_remove_additional_qdq_clusters) + + if rc.dump_kernel_selection_code: + + def _handle_kernel_selection_side_effect(_epm: EdgeProgramManager) -> list: + # Side-effect only: write kernel-selection files. No graph transform needed. + handle_kernel_selection() + return [] + + passes.append(_handle_kernel_selection_side_effect) + + return passes diff --git a/backends/nxp/recipes/nxp_recipe_types.py b/backends/nxp/recipes/nxp_recipe_types.py new file mode 100644 index 00000000000..a5c655c2fce --- /dev/null +++ b/backends/nxp/recipes/nxp_recipe_types.py @@ -0,0 +1,30 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.export import RecipeType + + +NXP_BACKEND: str = "nxp" + + +class NXPRecipeType(RecipeType): + """NXP-specific recipe types for Neutron NPU export. + + Choose the recipe that matches your intended export configuration: + - INT8_PTQ_NEUTRON: standard post-training quantization, delegates to Neutron NPU. + - INT8_PTQ_NO_DELEGATE: PTQ without NPU delegation (useful for debugging or CPU-only deployment). + """ + + # INT8 static PTQ (weights + activations). Calibration dataset required. + # Applicable operators are delegated to the Neutron NPU. + INT8_PTQ_NEUTRON = "nxp_int8_ptq_neutron" + + # INT8 PTQ without NPU delegation. Produces a quantized graph that runs on CPU. + # Useful for accuracy evaluation or debugging before enabling delegation. + INT8_PTQ_NO_DELEGATE = "nxp_int8_ptq_no_delegate" + + @classmethod + def get_backend_name(cls) -> str: + return NXP_BACKEND diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py index 964b8de159c..896bef9a3d5 100644 --- a/backends/nxp/tests/executorch_pipeline.py +++ b/backends/nxp/tests/executorch_pipeline.py @@ -33,6 +33,7 @@ from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import ( core_aten_ops_exception_list, + default_preserve_ops, generate_neutron_compile_spec, ) from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer @@ -92,7 +93,7 @@ def get_random_calibration_inputs( ] -def _get_default_quantizer(target_spec: NeutronTargetSpec, use_qat: bool) -> Quantizer: +def get_default_quantizer(target_spec: NeutronTargetSpec, use_qat: bool) -> Quantizer: return NeutronQuantizer(target_spec, is_qat=use_qat) @@ -195,9 +196,7 @@ def to_quantized_edge_program( ) -> EdgeProgramManager: _neutron_target_spec = NeutronTargetSpec(target) if get_quantizer_fn is None: - get_quantizer_fn = partial( - _get_default_quantizer, _neutron_target_spec, use_qat - ) + get_quantizer_fn = partial(get_default_quantizer, _neutron_target_spec, use_qat) input_spec = to_model_input_spec(input_spec) calibration_inputs = get_calibration_inputs_fn(input_spec) example_input = _get_example_input(input_spec) @@ -215,12 +214,6 @@ def to_quantized_edge_program( train_fn=train_fn, ) - # List of operators to not decompose during the lowering. - preserve_ops = [ - torch.ops.aten.prelu.default, - torch.ops.aten.pad.default, - torch.ops.aten.hardswish.default, - ] compile_spec = generate_neutron_compile_spec( target, intermediates_dir=intermediates_dir, @@ -240,7 +233,7 @@ def to_quantized_edge_program( _neutron_target_spec, custom_delegation_options, post_quant_state_dict, - preserve_ops=preserve_ops, + preserve_ops=default_preserve_ops, ) ] else: diff --git a/backends/nxp/tests/generic_tests/test_recipe_export.py b/backends/nxp/tests/generic_tests/test_recipe_export.py new file mode 100644 index 00000000000..e820b71d69c --- /dev/null +++ b/backends/nxp/tests/generic_tests/test_recipe_export.py @@ -0,0 +1,415 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch +import torch.nn + +from executorch.backends.nxp.backend.custom_delegation_options import ( + CustomDelegationOptions, +) +from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall +from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( + NeutronEdgePassManager, +) +from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner +from executorch.backends.nxp.recipes.nxp_recipe_provider import ( + NEUTRON_RECIPE_CONFIG_KEY, + NeutronRecipeConfig, + NXPRecipeProvider, +) +from executorch.backends.nxp.recipes.nxp_recipe_types import NXPRecipeType +from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec +from executorch.backends.nxp.tests.executors import ( + graph_contains_any, + graph_contains_any_of_ops, +) +from executorch.export import export +from executorch.export.recipe import ExportRecipe +from torch._inductor.lowering import quantized_decomposed + + +class SimpleCNN(torch.nn.Module): + def __init__(self, channels=3): + super().__init__() + self.conv = torch.nn.Conv2d(channels, channels, kernel_size=3) + + def forward(self, x): + x = self.conv(x) + x = torch.relu(x) + x = x.reshape(1, -1) + x = x + x + return x + + +INPUT_SHAPE = (1, 3, 8, 8) + + +def _run_export( + model, rc, recipe_type=NXPRecipeType.INT8_PTQ_NEUTRON, input_shape=INPUT_SHAPE +): + example_inputs = [(torch.randn(input_shape),)] + recipe = NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + return export(model, example_inputs=example_inputs, export_recipe=recipe) + + +def _get_graph(sess): + return sess.get_edge_program_manager().exported_program().graph + + +def test_ptq_neutron_basic(): + """Baseline PTQ: whole model delegated, IO is quantized.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + sess = _run_export(model, rc) + graph = _get_graph(sess) + + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def is_cnn_op(n): + return any(op in n.name.lower() for op in ["conv", "relu", "view", "add"]) + + assert not graph_contains_any(graph, is_cnn_op) + + nodes = list(graph.nodes) + first_call = next(n for n in nodes if n.op == "call_function" and n.name != "alloc") + last_call = next(n for n in reversed(nodes) if n.op == "call_function") + assert first_call.target == quantized_decomposed.quantize_per_tensor.out + assert last_call.target == quantized_decomposed.dequantize_per_tensor.out + + +class TestInt8PTQNoDelegate: + + def test__basic(self): + """INT8_PTQ_NO_DELEGATE: model is quantized but no delegate call is present in the graph.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + sess = _run_export(model, rc, recipe_type=NXPRecipeType.INT8_PTQ_NO_DELEGATE) + graph = _get_graph(sess) + + assert not graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def is_cnn_op(n): + return any(op in n.name.lower() for op in ["conv", "relu", "view", "add"]) + + # With no delegation, original ops should be visible in the graph. + assert graph_contains_any(graph, is_cnn_op) + + def test__recipe_has_empty_partitioners(self): + """INT8_PTQ_NO_DELEGATE recipe has an empty partitioner list.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NO_DELEGATE, neutron_recipe_config=rc + ) + assert recipe.lowering_recipe.partitioners == [] + + +class TestNeutronRecipeConfigFlags: + def test_operators_not_to_delegate(self): + """Ops listed in operators_not_to_delegate are not lowered to Neutron.""" + model = SimpleCNN() + rc = NeutronRecipeConfig( + INPUT_SHAPE, operators_not_to_delegate=["aten::convolution"] + ) + sess = _run_export(model, rc) + graph = _get_graph(sess) + + assert graph_contains_any_of_ops( + graph, [torch.ops.aten.convolution.out] + ) # Convolution was not delegated. + assert graph_contains_any_of_ops( + graph, [ExecutorchDelegateCall] + ) # Other operators were delegated. + + def _is_relu_add_or_view(n: torch.fx.Node) -> bool: + return any(op in n.name.lower() for op in ["relu", "add", "view"]) + + assert not graph_contains_any(graph, _is_relu_add_or_view) + + def test_remove_quant_io_ops(self): + """remove_quant_io_ops=True: no quantize op at the IO boundary.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, remove_quant_io_ops=True) + sess = _run_export(model, rc) + graph = _get_graph(sess) + nodes = list(graph.nodes) + + real_nodes = [n for n in nodes if n.op not in ("placeholder", "output")] + assert real_nodes[0].target != quantized_decomposed.quantize_per_tensor.out + assert real_nodes[-1].target != quantized_decomposed.dequantize_per_tensor.out + assert real_nodes[-1].meta["val"].dtype == torch.int8 + placeholder_nodes = [n for n in nodes if n.op == "placeholder"] + assert placeholder_nodes[0].name == "x" # Main input + assert placeholder_nodes[0].meta["val"].dtype == torch.int8 + + def test_use_quant_state_dict_false(self, mocker): + """use_quant_state_dict=False: the NeutronPartitioner used during lowering has + post_quantization_state_dict=None, confirmed by intercepting the constructor.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_quant_state_dict=False) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + + _run_export(model, rc) + + assert ( + len(captured) == 1 + ), "Expected exactly one NeutronPartitioner to be created." + assert captured[0].post_quantization_state_dict is None + + def test_custom_delegation_options_explicit(self, mocker): + """Explicitly provided CustomDelegationOptions are forwarded to NeutronPartitioner.""" + model = SimpleCNN() + opts = CustomDelegationOptions() + rc = NeutronRecipeConfig(INPUT_SHAPE, custom_delegation_options=opts) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + _run_export(model, rc) + + assert len(captured) == 1 + assert captured[0].custom_delegation_options == opts + + def test_intermediates_dir(self, tmp_path): + """intermediates_dir: intermediate compilation files are written to the directory.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, intermediates_dir=str(tmp_path)) + _run_export(model, rc) + assert any( + tmp_path.iterdir() + ), "No intermediate files written to intermediates_dir." + + def test_fetch_constants_to_sram_flag(self, mocker): + """fetch_constants_to_sram=True reaches the NeutronPartitioner used during export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, fetch_constants_to_sram=True) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + _run_export(model, rc) + + assert ( + len(captured) == 1 + ), "Expected exactly one NeutronPartitioner to be created." + spec_map = {s.key: s.value.decode() for s in captured[0].delegation_spec[1]} + assert spec_map["fetch_constants_to_sram"] == "True" + + def test_use_profiling_flag(self, mocker): + """use_profiling=True reaches the NeutronPartitioner used during export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_profiling=True) + + captured = [] + original_init = NeutronPartitioner.__init__ + + def capturing_init(self_, *args, **kwargs): + original_init(self_, *args, **kwargs) + captured.append(self_) + + mocker.patch.object(NeutronPartitioner, "__init__", capturing_init) + _run_export(model, rc) + + assert ( + len(captured) == 1 + ), "Expected exactly one NeutronPartitioner to be created." + spec_map = {s.key: s.value.decode() for s in captured[0].delegation_spec[1]} + assert spec_map["use_profiling"] == "True" + + def test_dump_kernel_selection_code(self, tmp_path, monkeypatch): + """dump_kernel_selection_code=True causes a kernel selection C file to be written.""" + monkeypatch.chdir(tmp_path) + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, dump_kernel_selection_code=True) + _run_export(model, rc) + assert ( + tmp_path / "_kernel_selection.c" + ).exists(), "_kernel_selection.c was not created in the working directory." + + def test_custom_quantizer_fn(self): + """get_quantizer_fn overrides the default NeutronQuantizer.""" + from executorch.backends.nxp.backend.neutron_target_spec import ( + NeutronTargetSpec, + ) + from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer + + custom_quantizer_called = [] + + def my_quantizer_fn(): + q = NeutronQuantizer(NeutronTargetSpec("imxrt700")) + custom_quantizer_called.append(True) + return q + + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, get_quantizer_fn=my_quantizer_fn) + sess = _run_export(model, rc) + assert custom_quantizer_called, "Custom quantizer factory was not called." + assert sess.get_edge_program_manager() is not None + + def test_use_neutron_for_format_conversion_false(self): + """use_neutron_for_format_conversion=False still produces a valid export.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, use_neutron_for_format_conversion=False) + sess = _run_export(model, rc) + assert sess.get_edge_program_manager() is not None + + def test_target_explicit(self): + """Specifying fake target to make sure an error is raised.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE, target="FAKE") + with pytest.raises(ValueError, match="`FAKE` is not a valid target"): + _run_export(model, rc) + + +class TestInputSpecForms: + def test__single_tuple(self): + """input_spec as a plain shape tuple works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig((1, 3, 8, 8))) + assert sess.get_edge_program_manager() is not None + + def test__list_of_tuples(self): + """input_spec as list of shape tuples works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig([(1, 3, 8, 8)])) + assert sess.get_edge_program_manager() is not None + + def test__model_input_spec(self): + """input_spec as list of ModelInputSpec objects works.""" + model = SimpleCNN() + sess = _run_export(model, NeutronRecipeConfig([ModelInputSpec((1, 3, 8, 8))])) + assert sess.get_edge_program_manager() is not None + + def test__multi_input(self): + """input_spec with multiple inputs (two tensors) works.""" + + class AddModel(torch.nn.Module): + def forward(self, x, y): + return x + y + + model = AddModel() + rc = NeutronRecipeConfig([(1, 3, 8, 8), (1, 3, 8, 8)]) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + example_inputs = [(torch.randn(1, 3, 8, 8), torch.randn(1, 3, 8, 8))] + sess = export(model, example_inputs=example_inputs, export_recipe=recipe) + assert sess.get_edge_program_manager() is not None + + +class TestErrorHandling: + def test_create_recipe_missing_config_key(self): + """create_recipe without neutron_recipe_config kwarg raises KeyError.""" + with pytest.raises(KeyError, match=NEUTRON_RECIPE_CONFIG_KEY): + NXPRecipeProvider().create_recipe(NXPRecipeType.INT8_PTQ_NEUTRON) + + def test_create_recipe_invalid_recipe_type(self): + """create_recipe with an unsupported recipe type returns None with a warning.""" + from executorch.export.recipe import RecipeType + + class FakeRecipeType(RecipeType): + FAKE = "fake" + + @classmethod + def get_backend_name(cls): + return "fake_backend" + + rc = NeutronRecipeConfig(INPUT_SHAPE) + result = NXPRecipeProvider().create_recipe( + FakeRecipeType.FAKE, neutron_recipe_config=rc + ) + assert result is None + + +class TestRecipeStructureValidation: + + def test_ptq_neutron_recipe_structure(self): + """INT8_PTQ_NEUTRON recipe: correct quantizer and partitioner are set.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.quantization_recipe is not None + assert len(recipe.quantization_recipe.quantizers) == 1 + assert recipe.lowering_recipe.partitioners is not None + assert len(recipe.lowering_recipe.partitioners) == 1 + + def test_ptq_neutron_recipe_name(self): + """INT8_PTQ_NEUTRON recipe has the expected name.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.name == NXPRecipeType.INT8_PTQ_NEUTRON.value + + +class TestRecipeCombination: + def test__chains_pre_partitioning_callbacks(self): + """Combining two NXP recipes chains both pre_partitioning_callbacks.""" + recipe1 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + recipe2 = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, + neutron_recipe_config=NeutronRecipeConfig(INPUT_SHAPE), + ) + combined = ExportRecipe.combine([recipe1, recipe2]) + assert combined.lowering_recipe.pre_partitioning_callback is not None + # Calling the combined callback should not raise. + combined.lowering_recipe.pre_partitioning_callback(None, {}) + + +class TestEdgeManagerTransformPasses: + def test__executed(self): + """edge_manager_transform_passes are called after partitioning.""" + model = SimpleCNN() + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + + transform_called = [] + + def tracking_pass(epm): + transform_called.append(True) + return [] + + recipe.lowering_recipe.edge_manager_transform_passes = [tracking_pass] + + example_inputs = [(torch.randn(INPUT_SHAPE),)] + export(model, example_inputs=example_inputs, export_recipe=recipe) + assert transform_called, "edge_manager_transform_passes were not executed." + + def test__qdq_pass_callable_returns_pass_manager(self, mocker): + """_remove_additional_qdq_clusters returns a bare NeutronEdgePassManager, not a + list containing one. EdgeProgramManagerTransformStage calls epm.transform(passes) + directly, so a list-of-PassManager would be silently mis-applied.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + # remove_quant_io_ops=False (default): first callable is _remove_additional_qdq_clusters. + qdq_callable = recipe.lowering_recipe.edge_manager_transform_passes[0] + result = qdq_callable(mocker.MagicMock()) + assert isinstance(result, NeutronEdgePassManager) diff --git a/examples/nxp/aot_neutron_compile.py b/examples/nxp/aot_neutron_compile.py index 7b7eeff8d19..1efdf03a0d6 100644 --- a/examples/nxp/aot_neutron_compile.py +++ b/examples/nxp/aot_neutron_compile.py @@ -27,6 +27,7 @@ from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import ( core_aten_ops_exception_list, + default_preserve_ops, generate_neutron_compile_spec, ) from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer @@ -384,6 +385,7 @@ def _get_arg_parser(): compile_spec, neutron_target_spec, post_quantization_state_dict=module.state_dict(), + preserve_ops=default_preserve_ops, ) ] if args.delegate diff --git a/export/recipe.py b/export/recipe.py index 8ea2256b7c3..0d244b76556 100644 --- a/export/recipe.py +++ b/export/recipe.py @@ -209,6 +209,8 @@ class LoweringRecipe: edge_manager_transform_passes: Optional list of callables that take EdgeProgramManager as argument and return passes to be applied. Applied sequentially after TO_EDGE stage. edge_compile_config: Optional edge compilation configuration + pre_partitioning_callback: Optional callable invoked just before partitioning with + `(partitioners, programs)` arguments. """ partitioners: Optional[Union[List[Partitioner], Dict[str, List[Partitioner]]]] = ( @@ -223,6 +225,9 @@ class LoweringRecipe: ) = None # pyre-ignore[11]: Type not defined edge_compile_config: Optional[EdgeCompileConfig] = None + pre_partitioning_callback: Optional[ + Callable[[Optional[list[Partitioner]], dict[str, ExportedProgram]], None] + ] = None @dataclass @@ -249,6 +254,7 @@ class _CombineAccumulator: pipeline_stages_values: list = field(default_factory=list) source_transform_in_place_values: list = field(default_factory=list) backend_config: object = None + pre_partitioning_callbacks: list = field(default_factory=list) @experimental( @@ -437,6 +443,7 @@ def _combine_lowering_recipe( all_partitioners_by_method: dict, all_edge_transform_passes: list, all_edge_manager_transform_passes: list, + all_pre_partitioning_callbacks: list, ) -> "Optional[LoweringRecipe]": """ Build the combined LoweringRecipe from per-recipe collected lists. @@ -474,11 +481,28 @@ def _combine_lowering_recipe( ) edge_compile_config = copy.deepcopy(distinct[0][1]) if distinct else None + combined_pre_partitioning_callback = None + if all_pre_partitioning_callbacks: + _cbs = all_pre_partitioning_callbacks + + def _chained_pre_partitioning_callback(partitioners, programs): + for cb in _cbs: + try: + cb(partitioners, programs) + except Exception as e: + name = getattr(cb, "__qualname__", repr(cb)) + raise RuntimeError( + f"Pre-partitioning callback `{name}` failed: {e}" + ) from e + + combined_pre_partitioning_callback = _chained_pre_partitioning_callback + if not ( combined_partitioners or all_edge_transform_passes or all_edge_manager_transform_passes or edge_compile_config + or combined_pre_partitioning_callback ): logging.info( "Combined recipe has no lowering fields; lowering_recipe will be None." @@ -490,6 +514,7 @@ def _combine_lowering_recipe( edge_transform_passes=all_edge_transform_passes or None, edge_manager_transform_passes=all_edge_manager_transform_passes or None, edge_compile_config=edge_compile_config or EdgeCompileConfig(), + pre_partitioning_callback=combined_pre_partitioning_callback, ) @staticmethod @@ -509,6 +534,8 @@ def _collect_lowering_fields( acc.edge_transform_passes.extend(lr.edge_transform_passes) if lr.edge_manager_transform_passes: acc.edge_manager_transform_passes.extend(lr.edge_manager_transform_passes) + if lr.pre_partitioning_callback: + acc.pre_partitioning_callbacks.append(lr.pre_partitioning_callback) @staticmethod def _collect_quantization_fields( @@ -611,6 +638,7 @@ def _combine_recipes( all_partitioners_by_method=acc.partitioners_by_method, all_edge_transform_passes=acc.edge_transform_passes, all_edge_manager_transform_passes=acc.edge_manager_transform_passes, + all_pre_partitioning_callbacks=acc.pre_partitioning_callbacks, ) recipe_name = recipe_name or "_".join( diff --git a/export/stages.py b/export/stages.py index c08a72cdba7..06bcf8a0b74 100644 --- a/export/stages.py +++ b/export/stages.py @@ -293,6 +293,15 @@ def run(self, artifact: PipelineArtifact) -> None: # method the dict does not name, so it would copy to apply nothing. final_passes = pass_manager or _drop_empty(transform_passes) or None + export_recipe = artifact.context.get("export_recipe") + lowering_recipe = getattr(export_recipe, "lowering_recipe", None) + + if ( + lowering_recipe is not None + and lowering_recipe.pre_partitioning_callback is not None + ): + lowering_recipe.pre_partitioning_callback(self._partitioners, artifact.data) + with validation_disabled(): edge_program_manager = to_edge_transform_and_lower( exported_programs, From 1851e2d186f35e52000bb508241b22b0ac99f02b Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Wed, 9 Sep 2026 08:01:06 +0100 Subject: [PATCH 090/190] Arm backend: Add INT test for grid_sampler. (#22602) Added test for grid_sampler cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Elena Zhelezina --- backends/arm/test/ops/test_grid_sampler.py | 15 +++++++++++++++ docs/source/backends/arm-vgf/VGF_op_support.md | 4 ++-- 2 files changed, 17 insertions(+), 2 deletions(-) diff --git a/backends/arm/test/ops/test_grid_sampler.py b/backends/arm/test/ops/test_grid_sampler.py index c5a1f3560bd..4c8ab0de8ce 100644 --- a/backends/arm/test/ops/test_grid_sampler.py +++ b/backends/arm/test/ops/test_grid_sampler.py @@ -60,3 +60,18 @@ def test_grid_sampler_vgf_no_quant(test_data): run_on_vulkan_runtime=False, ) pipeline.run() + + +@common.parametrize("test_data", test_data_suite, xfails=xfails, strict=False) +@common.SkipIfNoModelConverter +def test_grid_sampler_vgf_quant(test_data): + test_data = test_data() + pipeline = VgfPipeline[input_t]( + GridSampler2d(), + test_data, + aten_op, + exir_op, + quantize=True, + run_on_vulkan_runtime=False, + ) + pipeline.run() diff --git a/docs/source/backends/arm-vgf/VGF_op_support.md b/docs/source/backends/arm-vgf/VGF_op_support.md index fc77f589bbc..7ae696a7642 100644 --- a/docs/source/backends/arm-vgf/VGF_op_support.md +++ b/docs/source/backends/arm-vgf/VGF_op_support.md @@ -63,8 +63,8 @@ Total supported PyTorch APIs: **154**. | `torch.full_like` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.gather` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `BOOL` | 8x8 | | `torch.ge` / `>=` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | -| `torch.grid_sampler` | FP | `FP32` | - | -| `torch.grid_sampler_2d` | FP | `FP32` | - | +| `torch.grid_sampler` | FP, INT | `FP32`, `INT8` | 8x8 | +| `torch.grid_sampler_2d` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.group_norm` | FP | `FP32` | - | | `torch.gt` / `>` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.index_put_` | INT | `INT8` | 8x8 | From 1e6b6b80e1bc37f0541639aa8851558cfb6d7766 Mon Sep 17 00:00:00 2001 From: Per Held Date: Mon, 17 Aug 2026 09:30:14 +0200 Subject: [PATCH 091/190] Arm backend: Delegate constant index select on U55 Lower constant index_select operations to U55-supported ops. Contiguous indices such as [1, 2, 3] become one slice: [1:4]. Indices such as [1, 3, 1] become one-element slices followed by concatenation. This supports noncontiguous, reordered, and repeated indices. Preserve CPU fallback for runtime, empty, and symbolic cases. Authored with Codex. Change-Id: I1073680c7844925f905bd3d2de2b47fc3a0d62c1 Signed-off-by: Per Held --- backends/arm/_passes/arm_pass_manager.py | 2 +- .../decompose_index_select_to_gather_pass.py | 55 ++++++ .../arm/operator_support/ethos_u55_support.py | 46 ++++- .../tosa_supported_operators.py | 2 + backends/arm/test/ops/test_index_select.py | 161 ++++++++++++++++++ 5 files changed, 263 insertions(+), 3 deletions(-) diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 041296ec36f..d9a667aff4f 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -593,8 +593,8 @@ def _tosa_pipeline( DecomposeGroupedConvPass(), DecomposeUnfoldToGatherPass(use_slice=self.tosa_spec.is_U55_subset), DecomposeEmbeddingPass(), - DecomposeIndexSelectToGatherPass(), CastInt64BuffersToInt32Pass(exported_program), + DecomposeIndexSelectToGatherPass(exported_program), DecomposeStridedSliceCopyPass(), DecomposeSliceScatterPass(), AccumulateIndexPutPass(), diff --git a/backends/arm/_passes/decompose_index_select_to_gather_pass.py b/backends/arm/_passes/decompose_index_select_to_gather_pass.py index be0d4dbb07c..0a356b6c1d4 100644 --- a/backends/arm/_passes/decompose_index_select_to_gather_pass.py +++ b/backends/arm/_passes/decompose_index_select_to_gather_pass.py @@ -9,12 +9,17 @@ import torch from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.backends.arm._passes.arm_pass_utils import ( + get_param_tensor, + is_param_node, +) from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( ConvertExpandCopyToRepeatPass, ) from executorch.backends.arm._passes.convert_squeezes_to_view import ( ConvertSqueezesToViewPass, ) +from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -71,6 +76,12 @@ class DecomposeIndexSelectToGatherPass(ArmOpTargetedPass): exir_ops.edge.aten.index_select.default, } + def __init__( + self, exported_program: ExportedProgram | None = None, *args, **kwargs + ) -> None: + super().__init__(*args, **kwargs) + self.exported_program = exported_program + def call_operator(self, op, args, kwargs, meta): if op not in self.target_ops: return super().call_operator(op, args, kwargs, meta) @@ -81,6 +92,50 @@ def call_operator(self, op, args, kwargs, meta): x_shape, idx_shape = tuple(x_t.shape), tuple(idx_t.shape) x_rank, idx_rank = len(x_shape), len(idx_shape) + if ( + x_rank >= 1 + and idx_rank == 1 + and self.exported_program is not None + and is_param_node(self.exported_program, index.node) + and all(isinstance(size, int) for size in x_shape) + ): + constant_index = get_param_tensor(self.exported_program, index.node) + if constant_index is not None and constant_index.numel() > 0: + indices = constant_index.tolist() + dim_norm = dim % x_rank + dim_size = x_shape[dim_norm] + if any(index < 0 or index >= dim_size for index in indices): + raise RuntimeError( + f"index_select index out of range for dimension of size {dim_size}" + ) + if indices == list(range(indices[0], indices[0] + len(indices))): + return super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, dim_norm, indices[0], indices[-1] + 1), + {}, + meta, + updated=True, + ) + + slices = [] + for index_value in indices: + slices.append( + super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, dim_norm, index_value, index_value + 1), + {}, + meta, + updated=True, + ) + ) + return super().call_operator( + exir_ops.edge.aten.cat.default, + (slices, dim_norm), + {}, + meta, + updated=True, + ) + assert x_rank >= 1 and idx_rank == 1 and idx_t.dtype == torch.int32, ( f"[{self.__class__.__name__}] unsupported index_select signature: " f"x_rank={x_rank}, index_rank={idx_rank}, index_dtype={idx_t.dtype} " diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py index 8cf8e56560a..62631893636 100644 --- a/backends/arm/operator_support/ethos_u55_support.py +++ b/backends/arm/operator_support/ethos_u55_support.py @@ -14,8 +14,12 @@ import torch import torch.fx as fx -from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor +from executorch.backends.arm._passes.arm_pass_utils import ( + get_first_fake_tensor, + is_param_node, +) from executorch.backends.arm._passes.insert_table_ops import TableOps +from executorch.exir import ExportedProgram from executorch.exir.backend.utils import WhyNoPartitionReporter from executorch.exir.dialects._ops import ops as exir_ops from torch.fx.passes.operator_support import OperatorSupportBase @@ -200,7 +204,6 @@ class EthosU55NotSupported(OperatorSupportBase): exir_ops.edge.aten.gather.default, # GATHER exir_ops.edge.aten.grid_sampler_2d, # GATHER exir_ops.edge.aten.index.Tensor, # GATHER - exir_ops.edge.aten.index_select.default, # GATHER exir_ops.edge.aten.index_put.default, # SCATTER exir_ops.edge.aten.scatter.src, exir_ops.edge.aten.scatter.value, @@ -428,6 +431,45 @@ def is_node_supported( return True +class EthosU55IndexSelectCheck(OperatorSupportBase): + """Accept constant contiguous index_select cases that lower to a slice.""" + + def __init__( + self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter + ): + self.exported_program = exported_program + self.reporter = reporter + + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + del submodules + if node.target != exir_ops.edge.aten.index_select.default: + return True + + input_arg, dim, index_arg = node.args + input_node = typing.cast(fx.Node, input_arg) + index_node = typing.cast(fx.Node, index_arg) + input_shape = get_first_fake_tensor(input_node).shape + index_shape = get_first_fake_tensor(index_node).shape + if ( + not isinstance(dim, int) + or len(input_shape) == 0 + or not is_param_node(self.exported_program, index_node) + or len(index_shape) != 1 + or index_shape[0] == 0 + or any(not isinstance(size, int) for size in input_shape) + ): + self.reporter.report_reject( + node, + "U55 index_select requires static input shape and nonempty " + "constant indices.", + ) + return False + + return True + + class EthosU55CastCheck(OperatorSupportBase): """Reject unsupported casts on U55. diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 9a6149b0a7f..04d1e416d48 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -39,6 +39,7 @@ from executorch.backends.arm.operator_support.ethos_u55_support import ( EthosU55CastCheck, EthosU55DtypeSupport, + EthosU55IndexSelectCheck, EthosU55NotSupported, EthosU55ResizeCheck, EthosU55ReverseCheck, @@ -413,6 +414,7 @@ def _negative_checks( checks.append(EthosU55ResizeCheck(reporter)) checks.append(EthosU55ReverseCheck(reporter)) checks.append(EthosU55UnfoldCopyCheck(reporter)) + checks.append(EthosU55IndexSelectCheck(exported_program, reporter)) checks.append(EthosU55DtypeSupport(reporter)) checks.append(EthosU55CastCheck(reporter)) diff --git a/backends/arm/test/ops/test_index_select.py b/backends/arm/test/ops/test_index_select.py index 5410bc09a4e..0729b9dbd08 100644 --- a/backends/arm/test/ops/test_index_select.py +++ b/backends/arm/test/ops/test_index_select.py @@ -6,10 +6,19 @@ from typing import Tuple +import pytest import torch +from executorch.backends.arm._passes.decompose_index_select_to_gather_pass import ( + DecomposeIndexSelectToGatherPass, +) +from executorch.backends.arm.operator_support.ethos_u55_support import ( + EthosU55IndexSelectCheck, +) from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, EthosU85PipelineINT, OpNotSupportedPipeline, TosaPipelineFP, @@ -17,6 +26,9 @@ VgfPipeline, ) +from executorch.exir.backend.utils import WhyNoPartitionReporter +from executorch.exir.dialects._ops import ops as exir_ops + class IndexSelect(torch.nn.Module): aten_op = "torch.ops.aten.index_select.default" @@ -26,6 +38,17 @@ def forward(self, input_: torch.Tensor, dim: int, index_: torch.Tensor): return torch.index_select(input_, dim=dim, index=index_) +class ConstantIndexSelect(torch.nn.Module): + def __init__(self, dim: int, indices: list[int], dtype: torch.dtype = torch.int32): + super().__init__() + self.dim = dim + self.register_buffer("indices", torch.tensor(indices, dtype=dtype)) + + def forward(self, input_: torch.Tensor): + return torch.index_select(input_, dim=self.dim, index=self.indices) + + +input_t1 = Tuple[torch.Tensor] input_params = Tuple[torch.Tensor, int, torch.Tensor] # ---- FP profile: only float inputs ---- @@ -229,3 +252,141 @@ def test_index_select_vgf_quant(test_data: input_params): quantize=True, ) pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_contiguous(): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(2, [1, 2, 3]), + (torch.rand(1, 2, 5, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_contiguous_negative_dim(): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(-1, [1, 2]), + (torch.rand(1, 2, 4, 5),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.parametrize( + "indices", + { + "noncontiguous": [1, 3], + "descending": [3, 1], + "duplicate": [1, 1], + }, +) +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_slices_concat(indices): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(2, indices), + (torch.rand(1, 2, 5, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +def test_index_select_u55_INT_constant_contiguous_symbolic_dim_not_delegated(): + selected_dim = torch.export.Dim("selected_dim", min=4, max=8) + tester = ArmTester( + ConstantIndexSelect(2, [1, 2, 3]), + (torch.rand(1, 2, 5, 3),), + common.get_u55_compile_spec(), + dynamic_shapes={"input_": {2: selected_dim}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert exir_ops.edge.aten.index_select.default in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets + + +@common.parametrize( + "indices", + { + "negative_index": [-1], + "upper_bound_index": [5], + }, +) +def test_index_select_u55_constant_out_of_bounds_raises(indices): + tester = ArmTester( + ConstantIndexSelect(2, indices), + (torch.rand(1, 2, 5, 3),), + common.get_u55_compile_spec(), + ) + tester.export().to_edge() + exported_program = tester.stages[tester.cur].artifact.exported_program() + + with pytest.raises(RuntimeError, match="index_select index out of range"): + DecomposeIndexSelectToGatherPass(exported_program).call( + exported_program.graph_module + ) + + +def test_index_select_u55_scalar_not_supported(): + tester = ArmTester( + ConstantIndexSelect(0, [0]), + (torch.tensor(1.0),), + common.get_u55_compile_spec(), + ) + tester.export().to_edge() + exported_program = tester.stages[tester.cur].artifact.exported_program() + index_select_node = next( + node + for node in exported_program.graph.nodes + if node.target == exir_ops.edge.aten.index_select.default + ) + + assert not EthosU55IndexSelectCheck( + exported_program, WhyNoPartitionReporter() + ).is_node_supported({}, index_select_node) + + +def test_index_select_u55_INT_constant_int64_delegated(): + tester = ArmTester( + ConstantIndexSelect(2, [1, 2, 3], torch.int64), + (torch.rand(1, 2, 5, 3),), + common.get_u55_compile_spec(), + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert torch.ops.higher_order.executorch_call_delegate in targets + + +@common.XfailIfNoCorstone300 +def test_index_select_u55_INT_constant_contiguous_a16w8(): + pipeline = EthosU55PipelineINT[input_t1]( + ConstantIndexSelect(2, [1, 2, 3]), + (torch.rand(1, 2, 5, 3),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_index_select_u55_INT_constant_empty_not_delegated(): + pipeline = OpNotSupportedPipeline[input_t1]( + ConstantIndexSelect(2, []), + (torch.rand(1, 2, 5, 3),), + {IndexSelect.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() From aa3d4f35f2980ce72ab29600005134ed9a1e8360 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 9 Sep 2026 01:29:29 -0700 Subject: [PATCH 092/190] Stop shipping files nothing can reach in the wheel (#22584) A pip install of ExecuTorch carries about 1150 test modules, 188 shader templates, and the Python sources of vendored third-party checkouts. Nothing in an installed wheel can reach any of it. Test cases are the largest group. pytest loads them from a path in the checkout, never through the installed name. Shared helpers are the opposite: the suites import each other by installed name, and one helper has 267 importers. A file name cannot tell them apart, since test_pipeline.py and test_add.py look alike and only one can go. ### What ships now Vendored trees are found by asking git for its submodules, so a hand-edited list cannot go stale. The paths are read with `git config -z`, which separates a key from its value with a newline rather than a space, so a submodule whose path contains a space is not cut in half. A test module ships only if something can still reach it. Reaching means an import from anywhere in the checkout followed through, including from a directory the wheel does not ship, a relative import, an import written without the `executorch.` prefix, a literal name passed to importlib.import_module, or a name run as `python -m` by a workflow, a script, a README, or a module's own docstring. The unprefixed spelling matters because this repository imports itself both ways, and following only the prefixed one dropped two helpers that eight test files import. Workflow names are listed explicitly, since a source distribution has no .github directory. Each listed directory is paired with the file that drives it, and a test checks the pairing both ways: an entry whose driver is gone fails, and so does a driver with no entry. A yaml file ships unless it carries a shader template key, matched on content rather than a path list. Editable installs are untouched. A rebuild also deletes what an earlier build staged and this one does not want, limited to Python and yaml. Without that, building twice into the same directory keeps the old files and the wheel packages them, with no sign anything went wrong. The removed files are Python and yaml, apart from one marker file, so the same 1646 files leave every platform: Linux CPU 18.2 -> 15.0 MB 3.2 MB smaller, 17.4% Linux CUDA 27.1 -> 24.0 MB 3.1 MB smaller, 11.4% macOS 18.8 -> 15.7 MB 3.1 MB smaller, 16.4% Windows 14.1 -> 11.0 MB 3.1 MB smaller, 22.0% The Linux CPU row is a build of this branch against a build of its base. The other three come from rewriting the released wheel without the removed files, which understates the saving a little, because rewriting a zip does not reproduce the original compressor. ### Test plan Built the wheel on Linux x86_64 for CPU and CUDA, on Linux aarch64 for CUDA, and on macOS, and built the base of this branch the same way for comparison. Between those two builds 1646 files leave and none arrive. Nothing vendored remains, no C++ sources, no build files, no markdown. Of the yaml, 37 survive and every one is operator or model data that is read at run time. Installed each wheel into a clean environment, from a directory with no checkout in it so the installed package cannot be shadowed, and exported and ran a model through each. Ran the suites CI runs against the installed wheel, collected from the checkout. Every failure also fails against a wheel built from the base, so none of them is caused by this change. Added unit tests beside the existing wheel checks. Each was confirmed to go red when the behaviour it covers is switched off, including when a filter is left in place but never called, which an earlier version of these tests did not catch. Co-authored-by: PyTorch Bot --- .ci/scripts/tests/test_wheel_test_modules.py | 867 ++++++++++++++++++ .../tests/test_wheel_vendored_packages.py | 494 ++++++++++ pyproject.toml | 10 +- setup.py | 517 ++++++++++- 4 files changed, 1879 insertions(+), 9 deletions(-) create mode 100644 .ci/scripts/tests/test_wheel_test_modules.py create mode 100644 .ci/scripts/tests/test_wheel_vendored_packages.py diff --git a/.ci/scripts/tests/test_wheel_test_modules.py b/.ci/scripts/tests/test_wheel_test_modules.py new file mode 100644 index 00000000000..39c1d3f250f --- /dev/null +++ b/.ci/scripts/tests/test_wheel_test_modules.py @@ -0,0 +1,867 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the test modules the full wheel drops. + +The wheel used to carry every test file in the repository, about 9.7 MB of Python that nothing +in an installed wheel can reach. A test case is only ever loaded by pytest from a path in the +checkout, never through the installed name, so shipping it buys nothing. + +Shared helpers are the opposite. The suites here import each other by installed name, for +example `from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP`, so a +helper has to ship or collection breaks. That is why the keep set is computed from the import +graph and not from file names: `test_pipeline.py` and `test_add.py` are indistinguishable by +name and only one of them can go. + +setup.py is read rather than imported. It calls setup() at module scope, so importing it under a +test runner hands setup() the runner's own arguments and the session dies on an invalid command +name. +""" + +import ast +import functools +import os +import re +import shutil +import subprocess +import tempfile +import unittest +from pathlib import Path +from typing import Dict, FrozenSet, List, Optional, Set, Tuple + +from setuptools import find_namespace_packages + +SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py" +REPO_ROOT = SETUP_PY.parent + +# Enough of setup.py to exercise the keep set, and nothing that builds anything. +_WANTED = ( + "_WALK_SKIP_DIRS", + "_TEST_DIR_NAMES", + "_CI_ENTRY_POINTS", + "_CI_ENTRY_POINT_DIRS", + "_SHADER_TEMPLATE_MARKERS", + "_is_shader_template", + "_VENDORED_DIR_NAMES", + "_VENDORED_SUBMODULE_FALLBACK", + "_top_level_package_dirs", + "_first_party_module", + "_is_test_module", + "_module_name", + "_import_targets", + "_scan_imports", + "_GENERATED_DIR_NAMES", + "_unshipped_directories", + "_import_graph", + "_reachable_test_modules", + "_vendored_prefixes", + "_is_vendored_path", + "_full_packages", + "_minimal_packages", +) + + +def _setup_py_module() -> ast.Module: + # Name the encoding: these tests are collected on Windows too (pytest-windows.ini line 19), + # where the default is cp1252, and setup.py holds a non-ascii apostrophe that would decode to + # the wrong characters without a word rather than raising. + return ast.parse(SETUP_PY.read_text(encoding="utf-8")) + + +def _load_from_setup_py() -> Dict[str, object]: + """Run only the named definitions from setup.py, not its build logic.""" + selected: List[ast.stmt] = [] + found: Set[str] = set() + for node in _setup_py_module().body: + if isinstance(node, ast.FunctionDef) and node.name in _WANTED: + selected.append(node) + found.add(node.name) + elif isinstance(node, ast.Assign): + names = { + target.id + for target in node.targets + if isinstance(target, ast.Name) and target.id in _WANTED + } + if names: + selected.append(node) + found |= names + + assert found == set( + _WANTED + ), f"setup.py no longer defines {sorted(set(_WANTED) - found)}, so this test checks nothing" + + namespace: Dict[str, object] = { + "__file__": str(SETUP_PY), + "ast": ast, + "os": os, + "Path": Path, + "functools": functools, + "subprocess": subprocess, + "Dict": Dict, + "FrozenSet": FrozenSet, + "List": List, + "Optional": Optional, + "Set": Set, + "Tuple": Tuple, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +_NAMESPACE = _load_from_setup_py() +_is_test_module = _NAMESPACE["_is_test_module"] +_import_graph = _NAMESPACE["_import_graph"] +_reachable_test_modules = _NAMESPACE["_reachable_test_modules"] +_CI_ENTRY_POINTS = _NAMESPACE["_CI_ENTRY_POINTS"] + + +@functools.lru_cache(maxsize=None) +def _graph() -> Tuple[Set[str], Dict[str, Set[str]], Set[str]]: + return _import_graph(REPO_ROOT / "src" / "executorch") + + +def _a_dropped_test_module() -> str: + """A real test module the keep set excludes, so the wiring tests assert on real data. + + Must be a leaf inside a test package, because find_package_modules is only given a chance to + drop something when the package it is asked about is itself under a test directory. + """ + modules, _edges, _dynamic = _graph() + keep = _reachable_test_modules() + dropped = sorted( + name + for name in modules + if name not in keep + and _is_test_module(name.rsplit(".", 1)[0]) + and name.rsplit(".", 1)[1] != "__init__" + ) + assert dropped, "nothing is dropped, so the wiring tests would be vacuous" + return dropped[0] + + +_UNREACHABLE_TEST_MODULE = _a_dropped_test_module() + + +@functools.lru_cache(maxsize=None) +def _reachable_from_imports_only() -> FrozenSet[str]: + """The keep set the import graph produces on its own, with no directory entries applied. + + Used to tell a load-bearing directory entry from a redundant one: a module the graph already + reaches would ship whether or not its directory is listed. + """ + modules, edges, dynamic = _graph() + referenced = set(dynamic) + referenced.update(_CI_ENTRY_POINTS) + for targets in edges.values(): + referenced.update(targets) + return frozenset(name for name in referenced if _is_test_module(name)) & modules + + +@functools.lru_cache(maxsize=None) +def _entry_point_dir_drivers() -> Dict[str, str]: + """Why each `_CI_ENTRY_POINT_DIRS` entry exists, as the file that drives it. + + These directories cannot be re-derived from the source, which is the whole reason they are + listed by hand: each is walked by something that never spells out a module name, so there is + no import to find and no literal to grep for. What CAN be checked is that the thing doing the + walking still exists and still refers to the directory. If a driver is deleted or stops + mentioning its directory, the entry has outlived its reason and this pairing fails. + + Keyed by the dotted prefix, valued by a repository-relative path. + """ + return { + "executorch.backends.mlx.custom_kernel_ops": ".github/workflows/mlx.yml", + "executorch.backends.webgpu.test": "backends/webgpu/scripts/test_webgpu_native_ci.sh", + "executorch.backends.test.suite": "backends/test/suite/runner.py", + "executorch.examples.models.llava.test": "examples/models/llava/README.md", + } + + +@functools.lru_cache(maxsize=None) +def _tracked_shell_scripts() -> Tuple[str, ...]: + """Shell scripts this repository actually owns, as repository-relative paths. + + Asked of git rather than found by walking. CI checks other repositories out INSIDE this one, + for example a `pytorch/` sibling clone, and a walk cannot tell those files from ours. It found + `pytorch/.ci/pytorch/test.sh` and reported a module belonging to a different project, so the + walk failed on CI while passing in every local checkout. + + An archive with no git available yields nothing, which makes this test vacuous rather than + wrong. It is a drift guard, so silence in an environment that cannot check is the safe way to + fail. + """ + try: + listed = subprocess.run( + ["git", "-C", str(REPO_ROOT), "ls-files", "-z", "*.sh"], + capture_output=True, + text=True, + check=False, + ) + except OSError: + # No git on PATH, as in an unpacked source archive. + return () + if listed.returncode: + return () + return tuple(name for name in listed.stdout.split("\0") if name) + + +def _fake_prune(build_lib, source_root): + """A CustomBuildPy whose prune runs, with build_lib and the source tree given directly. + + The real method resolves the source tree from setup.py's own location, so the lifted body is + bound to a stand-in whose __file__ points at the fixture instead. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + assert len(classes) == 1, "setup.py no longer defines CustomBuildPy" + bodies = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name == "_prune_unstaged_files" + ] + assert len(bodies) == 1, "the stale-file prune is gone" + + namespace = { + "os": os, + "Path": Path, + "__file__": str(source_root.parent / "setup.py"), + } + exec(compile(ast.unparse(bodies[0]), "prune", "exec"), namespace) + + class Stub: + editable_mode = False + packages = ["executorch", "executorch.pkg"] + + def __init__(self): + self.build_lib = str(build_lib) + + def find_all_modules(self): + # Deliberately omits stale.py, which is what marks it unwanted. + return [("executorch", "__init__", ""), ("executorch.pkg", "__init__", "")] + + def get_package_dir(self, package): + return str(source_root / Path(*package.split("."))) + + def find_data_files(self, package, src_dir): + return [] + + Stub._prune_unstaged_files = namespace["_prune_unstaged_files"] + return Stub() + + +def _fake_build_py(): + """A CustomBuildPy whose overrides run, without configuring a real distribution. + + The overrides are lifted from setup.py and bound to a stand-in so they can be CALLED. The + point is to exercise the real bodies: a test that only reads their syntax passes on code that + never runs, which is the hole this helper exists to close. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + assert len(classes) == 1, "setup.py no longer defines CustomBuildPy" + wanted = ("find_package_modules", "find_data_files") + overrides = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name in wanted + ] + assert {node.name for node in overrides} == set( + wanted + ), f"CustomBuildPy no longer overrides {sorted(set(wanted) - {n.name for n in overrides})}" + + package = _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[0] + leaf = _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[1] + modules = [(package, "__init__", "x"), (package, leaf, "y")] + + class Stub: + editable_mode = False + + def __init__(self) -> None: + self._data_files_to_return: List[str] = [] + + # Stands in for build_py's own implementations, which need a configured distribution. + def _super_find_package_modules(self, _package, _package_dir): + return list(modules) + + def _super_find_data_files(self, _package, _src_dir): + return list(self._data_files_to_return) + + namespace = dict(_NAMESPACE) + namespace["os"] = os + # `super()` needs a real base, so give the lifted bodies one that returns the fixtures above. + source = "\n".join( + ast.unparse(node) + .replace( + "super().find_package_modules(package, package_dir)", + "self._super_find_package_modules(package, package_dir)", + ) + .replace( + "super().find_data_files(package, src_dir)", + "self._super_find_data_files(package, src_dir)", + ) + for node in overrides + ) + exec(compile(source, "overrides", "exec"), namespace) + for name in wanted: + setattr(Stub, name, namespace[name]) + return Stub(), package, modules + + +class TestDroppedTestModules(unittest.TestCase): + def test_something_is_actually_dropped(self) -> None: + """The rule removes a substantial number of modules. + + Without this, every assertion below is vacuous on a keep set that happens to contain + everything, and the whole change could be reverted with the suite still green. + """ + modules, _, _ = _graph() + tests = {name for name in modules if _is_test_module(name)} + keep = _reachable_test_modules() + self.assertGreater(len(tests), 500, "no test modules discovered at all") + self.assertLess( + len(keep), + len(tests) // 2, + f"keeping {len(keep)} of {len(tests)} test modules, so almost nothing is dropped", + ) + + def test_shared_helpers_are_kept(self) -> None: + """Modules the suites import by installed name still ship. + + These are the ones whose removal breaks collection rather than a single test. Each is + imported from outside its own directory, which is what makes the installed name matter. + """ + keep = _reachable_test_modules() + for helper in ( + "executorch.backends.arm.test.tester.test_pipeline", + "executorch.backends.xnnpack.test.tester.tester", + "executorch.backends.test.harness.stages", + "executorch.backends.test.graph_builder", + "executorch.exir.backend.test.op_partitioner_demo", + ): + self.assertIn(helper, keep) + + def test_leaf_cases_are_dropped(self) -> None: + """A test case nothing imports does not ship. + + Chosen from different suites, because one backend getting this right says nothing about + the others. + """ + keep = _reachable_test_modules() + modules, _, _ = _graph() + for leaf in ( + "executorch.backends.arm.test.ops.test_add", + "executorch.backends.xnnpack.test.ops.test_bilinear2d", + ): + self.assertIn( + leaf, modules, f"{leaf} no longer exists, pick another example" + ) + self.assertNotIn(leaf, keep) + + def test_relative_imports_are_followed(self) -> None: + """A submodule reached only by a relative import is kept. + + backends/test/harness/stages/__init__.py does `from .export import Export`, so treating + a relative import as reaching nothing new drops stages.export and breaks every importer + of that package. This is a regression guard: it failed exactly that way once. + """ + self.assertIn( + "executorch.backends.test.harness.stages.export", _reachable_test_modules() + ) + + def test_dynamic_imports_are_followed(self) -> None: + """A module named only as a string to importlib is kept. + + backends/mlx/test/run_all_tests.py does + `importlib.import_module(".test_ops", package=__package__)`, which an import scan that + only reads import statements cannot see. Note test_ops is also named like a leaf, so a + file name rule would drop it. + """ + self.assertIn( + "executorch.backends.mlx.test.test_ops", _reachable_test_modules() + ) + + def test_ci_entry_points_are_kept(self) -> None: + """The modules only a workflow names are kept.""" + keep = _reachable_test_modules() + for name in _CI_ENTRY_POINTS: + self.assertIn(name, keep) + + def test_ci_entry_points_still_match_the_workflows(self) -> None: + """The hand-written CI list has not drifted from what the workflows actually run. + + The list is explicit rather than scanned at build time, because a source distribution + carries no .github directory and a scan there would silently keep nothing. The cost of + being explicit is drift, so it is checked here instead. + """ + pattern = re.compile(r"executorch(?:\.[A-Za-z0-9_]+)+") + referenced: Set[str] = set() + # The whole tree, not just .github and .ci. A workflow often calls a script that lives + # beside the backend it tests, and those name modules too: + # backends/webgpu/scripts/test_webgpu_native_ci.sh runs six of them by dotted name. + skip = { + ".git", + "pip-out", + "cmake-out", + "third-party", + "third_party", + "__pycache__", + } + for dirpath, dirnames, filenames in os.walk(REPO_ROOT, followlinks=False): + dirnames[:] = [d for d in dirnames if d not in skip] + for filename in filenames: + # Markdown too: a README documenting `python -m executorch.x.test.y` is a + # promise to users, and dropping that module breaks the documented command. + # Python as well, because several modules document their own `python -m` + # invocation in a docstring rather than in a README, and that is the same + # promise written somewhere else. + if not filename.endswith( + (".yml", ".yaml", ".sh", ".ps1", ".md", ".py") + ): + continue + path = Path(dirpath) / filename + if path.resolve() == Path(__file__).resolve(): + # This file names dropped modules as examples of what the rule removes, so + # reading itself would report them as promised and contradict its own tests. + continue + text = path.read_text(encoding="utf-8", errors="replace") + referenced.update(pattern.findall(text)) + + modules, _, _ = _graph() + + # A reference like executorch.a.test.b.SomeClass.some_method is one dotted run to the + # regex, and it is not a module, so trim each match back to its longest real module + # prefix. Without this the class-suffixed entries silently drop out of the comparison + # and the guard protects fewer names than it appears to. + def longest_module(name: str) -> str: + parts = name.split(".") + while parts: + candidate = ".".join(parts) + if candidate in modules: + return candidate + parts.pop() + return name + + expected = { + trimmed + for trimmed in (longest_module(name) for name in referenced) + if _is_test_module(trimmed) and trimmed in modules + } + missing = sorted(expected - set(_CI_ENTRY_POINTS) - _reachable_test_modules()) + self.assertEqual( + missing, + [], + f"a workflow names these test modules but nothing keeps them: {missing}", + ) + + def test_parent_packages_of_kept_modules_are_kept(self) -> None: + """Every kept module's package chain is kept, or the dotted path cannot resolve.""" + keep = _reachable_test_modules() + modules, _, _ = _graph() + for name in keep: + parts = name.split(".") + for end in range(2, len(parts)): + parent = ".".join(parts[:end]) + if _is_test_module(parent) and parent in modules: + self.assertIn(parent, keep, f"{parent} missing but {name} is kept") + + def test_shader_templates_do_not_ship(self) -> None: + """Shader codegen inputs are dropped, and op definitions are not. + + The cmake build expands these into SPIR-V and WGSL headers, so the wheel already carries + the compiled result. Matched on content, so the two examples below are the real + distinction: one is a template, the other is read at run time through + importlib.resources and must survive. + """ + is_template = _NAMESPACE["_is_shader_template"] + self.assertTrue( + is_template("backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml") + ) + self.assertTrue( + is_template("backends/webgpu/runtime/ops/binary_op/binary_op.yaml") + ) + for needed in ( + "exir/dialects/edge/edge.yaml", + "kernels/portable/functions.yaml", + "backends/cadence/aot/functions.yaml", + ): + self.assertFalse(is_template(needed), f"{needed} would stop shipping") + + def test_build_py_is_wired_to_the_custom_class(self) -> None: + """setup() receives CustomBuildPy, not the stock build_py. + + Every other test here exercises the class directly, so all of them stay green when the + cmdclass entry is pointed back at setuptools' own build_py. That single edit disables + the module filter, the data file filter and the prune at once, and the wheel then ships + everything again. + """ + assignments = [ + node + for node in ast.walk(_setup_py_module()) + if isinstance(node, ast.Call) + and isinstance(node.func, ast.Name) + and node.func.id == "setup" + ] + self.assertEqual(len(assignments), 1, "expected exactly one setup() call") + + mapping = [kw.value for kw in assignments[0].keywords if kw.arg == "cmdclass"] + self.assertEqual(len(mapping), 1, "setup() no longer passes cmdclass") + wired = { + key.value: value.id + for key, value in zip(mapping[0].keys, mapping[0].values) + if isinstance(key, ast.Constant) and isinstance(value, ast.Name) + } + self.assertEqual( + wired.get("build_py"), + "CustomBuildPy", + "build_py is not wired to CustomBuildPy, so none of the filters run", + ) + + def test_both_package_lists_are_anchored_on_this_file(self) -> None: + """Neither package list depends on the working directory. + + A cwd-relative `where` returns nothing when the build runs from anywhere but the + repository root, and an empty package list makes the prune treat every staged file as + unwanted. The full list was anchored for this reason; the minimal one has to match. + + Both lists are CALLED from a directory that is not the repository root, because reading + the syntax of the `where=` argument only proves it is not a literal. Swapping the anchor + for `Path.cwd()` leaves the syntax test green and breaks every build started elsewhere. + """ + original = os.getcwd() + os.chdir(tempfile.gettempdir()) + try: + full = _NAMESPACE["_full_packages"]() + minimal = _NAMESPACE["_minimal_packages"]() + finally: + os.chdir(original) + self.assertIn("executorch", full) + self.assertGreater( + len(full), 100, "the full list collapsed when built from another directory" + ) + self.assertIn("executorch", minimal) + self.assertGreater( + len(minimal), + 1, + "the minimal list collapsed when built from another directory", + ) + + def test_stale_staged_files_are_pruned(self) -> None: + """A rebuild removes what an earlier build staged and this one does not want. + + build_py only copies, so without this a second build into the same directory keeps + every file the first one put there and the wheel packages it. The failure is silent: + the build succeeds and the wheel quietly contains the dropped files. + + The prune is CALLED against a real staging directory, because checking that the method + and its call site exist leaves an early `return` inside the body undetected, and the + prune then does nothing while this test stays green. + """ + staging = Path(tempfile.mkdtemp(prefix="prunetest-")) + self.addCleanup(shutil.rmtree, staging, ignore_errors=True) + source = staging / "src" + (source / "executorch" / "pkg").mkdir(parents=True) + for name in ("executorch/__init__.py", "executorch/pkg/__init__.py"): + (source / name).write_text("") + # Exists in the source tree and is NOT in build_py's file list, so the prune wants it + # gone. That is the whole contract. + (source / "executorch" / "pkg" / "stale.py").write_text( + "# left by an earlier build\n" + ) + build_lib = staging / "lib" + shutil.copytree(source, build_lib) + # Generated by a later build command, absent from src/, and must survive. + (build_lib / "executorch" / "pkg" / "generated.py").write_text( + "# from a template\n" + ) + + command = _fake_prune(build_lib, source) + command._prune_unstaged_files() + + remaining = sorted(p.name for p in (build_lib / "executorch" / "pkg").iterdir()) + self.assertNotIn( + "stale.py", + remaining, + "the prune left a file the current build does not want", + ) + self.assertIn( + "generated.py", + remaining, + "the prune deleted a file another command generated", + ) + self.assertIn("__init__.py", remaining, "the prune deleted a wanted module") + + def test_build_py_applies_the_keep_set(self) -> None: + """The drop is actually wired into the build, checked by CALLING the override. + + An earlier version of this test read the override's syntax tree instead. That passes on + code that is present but never runs, so an early `return modules` at the top of the + override left the filter dead with every assertion here still true. Build a real command + and look at what it returns. + """ + command, package, modules = _fake_build_py() + result = command.find_package_modules(package, "unused") + returned = {entry[1] for entry in result} + offered = {entry[1] for entry in modules} + self.assertIn("__init__", returned, "a kept package must still import") + self.assertTrue( + offered - returned, + "find_package_modules returned everything it was offered, so nothing is dropped", + ) + self.assertNotIn( + _UNREACHABLE_TEST_MODULE.rsplit(".", 1)[1], + returned, + "an unreachable test module was not dropped", + ) + + def test_shader_filter_is_wired_into_find_data_files(self) -> None: + """The shader classifier is actually CALLED, not merely correct. + + test_shader_templates_do_not_ship above checks the predicate. That is not the same thing: + deleting the filtering line in find_data_files leaves the predicate perfect and unused, + and every shader template ships again. + """ + command, _package, _modules = _fake_build_py() + template = "backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml" + needed = "kernels/portable/functions.yaml" + root = str(REPO_ROOT) + command._data_files_to_return = [ + os.path.join(root, template), + os.path.join(root, needed), + ] + kept = command.find_data_files("executorch", root) + self.assertNotIn( + os.path.join(root, template), + kept, + "find_data_files does not drop shader templates, so the filter is not wired in", + ) + self.assertIn( + os.path.join(root, needed), + kept, + "find_data_files dropped a yaml the runtime reads", + ) + + def test_ci_entry_point_dirs_are_all_load_bearing(self) -> None: + """Every listed directory still has a driver, and still keeps something. + + Nothing referenced `_CI_ENTRY_POINT_DIRS`, so an entry could be deleted with the whole + suite green: removing the backend suite line silently stopped 86 modules shipping. Two + checks close that. Each entry must be paired with the file that walks it, which fails when + an entry is added or removed without updating the pairing, and each entry must keep modules + the import graph cannot reach on its own, which fails when an entry becomes dead weight. + """ + listed = set(_NAMESPACE["_CI_ENTRY_POINT_DIRS"]) + self.assertTrue(listed, "the list is empty, so nothing is protected") + + drivers = _entry_point_dir_drivers() + self.assertEqual( + listed, + set(drivers), + "_CI_ENTRY_POINT_DIRS and its list of drivers disagree. Add the new entry with the " + "file that walks it, or drop the driver for the entry that went away", + ) + + for prefix, driver in sorted(drivers.items()): + path = REPO_ROOT / driver + self.assertTrue( + path.is_file(), + f"{prefix} is kept for {driver}, which no longer exists, so the entry may be " + "obsolete", + ) + tail = prefix.split(".")[-1] + self.assertIn( + tail, + path.read_text(encoding="utf-8", errors="replace"), + f"{driver} no longer mentions {tail}, so it may have stopped driving {prefix}", + ) + + # And the other direction: an entry that keeps nothing new is dead weight. + reached_anyway = _reachable_from_imports_only() + keep = _reachable_test_modules() + for entry in sorted(listed): + covered = { + name for name in keep if name == entry or name.startswith(f"{entry}.") + } + self.assertTrue( + covered - reached_anyway, + f"{entry} keeps nothing the import graph does not already reach, so the entry " + "is redundant and should be removed", + ) + + def test_ci_entry_points_cover_constructed_module_names(self) -> None: + """A runner that BUILDS a dotted name is covered too. + + The drift test above searches for a literal dotted name, so it cannot see a script that + assembles one, and a directory whose tests are only reached that way would be dropped + with nothing to warn about. + + A script that runs from a checkout by design is exempt, and says so in its own header. + `backends/apple/coreai/run_all_tests.sh` is the current example: it cds to the repository + root, so it always finds the files on disk and never needs them installed. + """ + pattern = re.compile(r"find\s+([A-Za-z0-9_./-]+)\s+-name\s+'?test_\*\.py'?") + keep = _reachable_test_modules() + listed = _NAMESPACE["_CI_ENTRY_POINT_DIRS"] + unprotected = [] + for relative in _tracked_shell_scripts(): + path = REPO_ROOT / relative + text = path.read_text(encoding="utf-8", errors="replace") + walked = pattern.findall(text) + if not walked: + continue + if "not a landing artifact" in text: + continue + for entry in walked: + dotted = "executorch." + entry.strip("./").replace("/", ".") + covered = any( + dotted == prefix or dotted.startswith(f"{prefix}.") + for prefix in listed + ) or any(name.startswith(f"{dotted}.") for name in keep) + if not covered: + unprotected.append(f"{relative} -> {dotted}") + self.assertEqual( + unprotected, + [], + "a script discovers test modules under these paths by building dotted names, and " + "nothing keeps them. Either add the directory to _CI_ENTRY_POINT_DIRS in setup.py, " + "or say in the script's header that it is not a landing artifact if it only ever " + f"runs from a checkout: {unprotected}", + ) + + def test_unprefixed_first_party_imports_count_as_references(self) -> None: + """`from backends.x import y` keeps y, the same as the prefixed spelling. + + This repository imports itself both ways: most code says `executorch.backends.x`, but the + Arm suites say `backends.arm.test...`, which resolves because the repository root is on + sys.path. Both name the same file. Following only the prefixed spelling dropped two shared + helpers with eight importers between them, which is the invariant this change exists to + preserve. + """ + first_party = _NAMESPACE["_first_party_module"] + self.assertEqual( + first_party("backends.arm.test.common"), + "executorch.backends.arm.test.common", + ) + self.assertEqual( + first_party("executorch.exir.tests.common"), "executorch.exir.tests.common" + ) + # A third-party module whose first component is not one of ours stays out. + self.assertIsNone(first_party("torch.nn.functional")) + self.assertIsNone(first_party("numpy")) + + keep = _reachable_test_modules() + for helper in ( + "executorch.backends.arm.test._custom_vgf_test_utils", + "executorch.backends.arm.test.runtime._vgf_runtime_test_utils", + ): + self.assertIn( + helper, + keep, + f"{helper} is imported without the executorch prefix and must still ship", + ) + + def test_importers_outside_the_shipped_tree_are_followed(self) -> None: + """A file the wheel does not carry can still import one that it does. + + `src/executorch` is a subset of the checkout, so an importer in a directory that is never + packaged is invisible to a walk of the shipped tree alone. Its imports still have to keep + their targets: test/end2end/test_end2end.py imports two model helpers out of exir/tests. + """ + keep = _reachable_test_modules() + importer = REPO_ROOT / "test" / "end2end" / "test_end2end.py" + self.assertTrue( + importer.is_file(), "this test needs a different example importer" + ) + for helper in ( + "executorch.exir.tests.dynamic_shape_models", + "executorch.exir.tests.transformer", + ): + self.assertIn( + helper, + keep, + f"{helper} is imported from outside the shipped tree and must still ship", + ) + + def test_vendored_trees_are_not_read_as_import_evidence(self) -> None: + """A vendored submodule's own imports do not keep anything. + + The package list excludes vendored trees, so nothing in one ships. The import scan has to + agree, or the two disagree about the same directory: a submodule checked out under an + ordinary name, rather than under `third-party`, was read as first-party and its imports + kept test modules the wheel never carries. + + Skipping by directory name alone is not enough, which is why this asserts on the scan's + output rather than on the skip list. + """ + modules, _edges, _dynamic = _graph() + is_vendored = _NAMESPACE["_is_vendored_path"] + vendored = sorted( + name for name in modules if is_vendored(name.replace(".", "/")) + ) + self.assertEqual( + vendored, + [], + "the import scan read these vendored modules as first-party, so their imports can " + f"keep test modules nothing shipped reaches: {vendored[:5]}", + ) + + def test_generated_directories_are_not_read_as_import_evidence(self) -> None: + """A build tree or an in-tree virtualenv does not vote on what ships. + + Those hold an INSTALLED copy of this package, so reading one lets the last wheel decide + what the next carries: a file that shipped once keeps itself alive. A clean checkout has + none of them, so the guard is exercised here by creating one. + """ + unshipped = _NAMESPACE["_unshipped_directories"] + root = REPO_ROOT / "src" / "executorch" + planted = REPO_ROOT / ".venv" + created = not planted.exists() + if created: + (planted / "lib").mkdir(parents=True) + self.addCleanup(shutil.rmtree, planted, ignore_errors=True) + walked = {entry.name for entry in unshipped(root)} + self.assertNotIn( + ".venv", + walked, + "a generated directory is read as import evidence, so an installed copy of this " + "package can keep test modules alive across builds", + ) + self.assertIn( + "test", walked, "the guard also dropped a real unshipped directory" + ) + + def test_editable_installs_are_left_alone(self) -> None: + """An editable install still exposes every test module. + + It maps the package root to a directory, so the suites resolve from the source tree + whatever is listed, and dropping modules there would only make the two install modes + disagree for no benefit. + """ + classes = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.ClassDef) and node.name == "CustomBuildPy" + ] + overrides = [ + node + for node in classes[0].body + if isinstance(node, ast.FunctionDef) and node.name == "find_package_modules" + ] + source = ast.unparse(overrides[0]) + self.assertIn("editable_mode", source) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/tests/test_wheel_vendored_packages.py b/.ci/scripts/tests/test_wheel_vendored_packages.py new file mode 100644 index 00000000000..a3d81a940b0 --- /dev/null +++ b/.ci/scripts/tests/test_wheel_vendored_packages.py @@ -0,0 +1,494 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Unit tests for the packages the full wheel publishes. + +The wheel used to carry the Python files and codegen scripts of every vendored third-party +checkout, because the full build passed no `packages` list and setuptools then discovered +everything under src/executorch. Those files exist to build the C++ targets, so nothing in +an installed wheel imports them. + +Asserting on the discovery result rather than on a built wheel, because the behaviour under +test is a pure function of the source tree plus the exclude patterns, and a full build takes +minutes to exercise one filter. `.ci/scripts/test_minimal_wheel.sh` already covers the +built-artifact side for the minimal wheel. + +setup.py is read rather than imported. It calls setup() at module scope, so importing it under +a test runner hands setup() the runner's own arguments and the session dies on an invalid +command name. +""" + +import ast +import functools +import os +import subprocess +import tempfile +import unittest +from pathlib import Path +from typing import Dict, FrozenSet, List, Set, Tuple + +from setuptools import find_namespace_packages +from setuptools.command.build_py import build_py + +SETUP_PY = Path(__file__).resolve().parents[3] / "setup.py" +REPO_ROOT = SETUP_PY.parent +# Discovery is anchored on this file's location, not on the working directory, so the result +# does not depend on where the runner was started. +PACKAGE_ROOT = str(SETUP_PY.parent / "src") + + +# The helpers this test drives, shared by both loaders below. +_HELPERS = ( + # _VENDORED_DIR_NAMES is not used directly below, but _is_vendored_path closes over it. + "_VENDORED_DIR_NAMES", + "_VENDORED_SUBMODULE_FALLBACK", + "_vendored_prefixes", + "_is_vendored_path", + # CustomBuildPy calls this, so the class cannot be exec'd without it. + "_SHADER_TEMPLATE_MARKERS", + "_is_shader_template", + "_full_packages", +) + + +def _setup_py_module() -> ast.Module: + # Name the encoding: these tests are collected on Windows too (pytest-windows.ini line 19), + # where the default is cp1252, and setup.py holds a non-ascii apostrophe that would decode to + # the wrong characters without a word rather than raising. + return ast.parse(SETUP_PY.read_text(encoding="utf-8")) + + +def _load_from_setup_py(root: Path = None) -> Dict[str, object]: + """The vendored-path helpers and the package list builder, from setup.py's source. + + Only those definitions are executed, so none of setup.py's module level build logic runs. + """ + wanted = _HELPERS + + selected: List[ast.stmt] = [] + found = set() + for node in _setup_py_module().body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted: + selected.append(node) + found.add(node.name) + elif isinstance(node, ast.Assign): + names = { + target.id + for target in node.targets + if isinstance(target, ast.Name) and target.id in wanted + } + if names: + selected.append(node) + found |= names + + assert found == set( + wanted + ), f"setup.py no longer defines {sorted(set(wanted) - found)}, so this test checks nothing" + + namespace: Dict[str, object] = { + "__file__": str((root or SETUP_PY.parent) / "setup.py"), + "Path": Path, + "List": List, + "Tuple": Tuple, + "functools": functools, + "subprocess": subprocess, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +@functools.lru_cache(maxsize=None) +def _load_build_py() -> Dict[str, object]: + """CustomBuildPy plus the helpers it calls, so analyze_manifest can be driven directly. + + Only the class body and those helpers run. Its methods reference names from setup.py's own + imports, so the ones analyze_manifest touches are supplied here. + """ + wanted = {"CustomBuildPy"} | set(_HELPERS) + selected: List[ast.stmt] = [] + for node in _setup_py_module().body: + if isinstance(node, (ast.FunctionDef, ast.ClassDef)) and node.name in wanted: + selected.append(node) + elif isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id in wanted + for target in node.targets + ): + selected.append(node) + + namespace: Dict[str, object] = { + "__file__": str(SETUP_PY), + "os": os, + "ast": ast, + "Path": Path, + "functools": functools, + "subprocess": subprocess, + "build_py": build_py, + "Dict": Dict, + "FrozenSet": FrozenSet, + "List": List, + "Set": Set, + "Tuple": Tuple, + "find_namespace_packages": find_namespace_packages, + } + exec( + compile(ast.Module(body=selected, type_ignores=[]), str(SETUP_PY), "exec"), + namespace, + ) + return namespace + + +_NAMESPACE = _load_from_setup_py() +_vendored_prefixes = _NAMESPACE["_vendored_prefixes"] +_is_vendored_path = _NAMESPACE["_is_vendored_path"] +_full_packages = _NAMESPACE["_full_packages"] + + +def _discovered_packages() -> List[str]: + """Everything setuptools finds, before any of this change's filtering.""" + return sorted( + find_namespace_packages( + where=PACKAGE_ROOT, include=["executorch", "executorch.*"] + ) + ) + + +def _vendored(packages: List[str]) -> List[str]: + return [ + package for package in packages if _is_vendored_path(package.replace(".", "/")) + ] + + +class TestFullWheelPackages(unittest.TestCase): + def test_the_tree_has_vendored_packages_to_exclude(self) -> None: + """Fail rather than skip when there is nothing to exclude. + + Every other test here is vacuous on a tree with no vendored checkouts: an empty + package list contains no vendored package, so the exclusion would look correct even + if it had been deleted. Assert the premise instead of quietly passing on it. + """ + discovered = _discovered_packages() + self.assertNotEqual( + discovered, [], f"no packages discovered under {PACKAGE_ROOT}" + ) + self.assertNotEqual( + _vendored(discovered), + [], + "no vendored third-party packages in this tree, so the exclusion below cannot " + "be shown to do anything. Initialize the submodules before running this.", + ) + + def test_no_vendored_package_ships(self) -> None: + """No package from another repository is published. + + Compares against what discovery finds rather than re-filtering the helper's own output. + Filtering the result with the same predicate the helper already applied is a tautology: + it is empty whatever the helper did, so it would pass even with the exclusion removed. + """ + discovered = set(_discovered_packages()) + shipped = set(_full_packages()) + dropped = discovered - shipped + + leaked = sorted(shipped & set(_vendored(discovered))) + # Only the count and a few names, because a regression here leaks hundreds of + # packages and the default diff would bury the message. + self.assertEqual( + len(leaked), + 0, + f"the wheel would publish {len(leaked)} vendored packages, " + f"e.g. {leaked[:3]}", + ) + # And the helper really removed them, rather than discovery never having found them. + self.assertEqual( + dropped, + set(_vendored(discovered)), + "the set the helper drops is not the set of vendored packages on disk", + ) + + def test_the_exclusion_is_load_bearing(self) -> None: + """Discovery without the exclusion finds the packages the exclusion removes.""" + self.assertLess( + len(_full_packages()), + len(_discovered_packages()), + "the exclusion dropped nothing, so it is no longer doing any work", + ) + + def test_setup_passes_the_package_list(self) -> None: + """The helper is actually wired into the full build. + + Without this, every test above still passes when the assignment that hands the list + to setuptools is deleted, which is the whole of the change. The sibling wheel test + asserts its own wiring the same way and for the same reason. + + The search is limited to the else branch of the minimal-build check, because an + unrestricted walk also matches an assignment that can never run: moved into the + minimal branch it is overwritten by the next line, and wrapped in a false condition + it is dead, and both of those leave the full wheel discovering everything. + """ + minimal_checks = [ + node + for node in _setup_py_module().body + if isinstance(node, ast.If) + and isinstance(node.test, ast.Call) + and isinstance(node.test.func, ast.Name) + and node.test.func.id == "_is_minimal_build" + ] + self.assertEqual( + len(minimal_checks), + 1, + "expected exactly one module level `if _is_minimal_build():`", + ) + + assigned = [ + node + for node in minimal_checks[0].orelse + if isinstance(node, ast.Assign) + for target in node.targets + if isinstance(target, ast.Subscript) + and isinstance(target.value, ast.Name) + and target.value.id == "setup_kwargs" + and isinstance(target.slice, ast.Constant) + and target.slice.value == "packages" + and isinstance(node.value, ast.Call) + and isinstance(node.value.func, ast.Name) + and node.value.func.id == "_full_packages" + ] + self.assertEqual( + len(assigned), + 1, + "setup.py does not assign _full_packages() to setup_kwargs['packages'], " + "so the full build falls back to discovering every package", + ) + + def test_first_party_packages_still_ship(self) -> None: + """A named first-party package survives the exclusion. + + Every other test here asks whether unwanted packages left. This one asks whether + wanted ones stayed, which is the failure mode a too-greedy filter produces and the + one nothing else would notice. + """ + packages = _full_packages() + for package in ( + "executorch.exir", + "executorch.backends.xnnpack", + "executorch.extension.pybindings", + "executorch.devtools", + ): + self.assertIn(package, packages) + + def test_only_submodule_sections_are_read(self) -> None: + """A `path` line outside a submodule section is not an exclusion prefix. + + Written against a file with a stray entry rather than by comparing git's output to git's + own output. That comparison holds for any reader on today's clean file, so it would pass + just as well for a line scanner that accepts `path =` from any section, which is the + failure this is meant to catch: one stray line silently removes a real package. + """ + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text( + '[submodule "real"]\n' + "\tpath = extension/llm/tokenizers\n" + "[core]\n" + "\tpath = executorch/exir\n" + ) + self.assertEqual( + _load_from_setup_py(root)["_vendored_prefixes"](), + ("extension/llm/tokenizers",), + "a path line outside a submodule section became an exclusion prefix", + ) + + def test_prefixes_are_normalized(self) -> None: + """A legal but unusual spelling in .gitmodules still matches the real directory. + + git treats a trailing slash, a leading ./ and a doubled separator as the same path, + so storing the raw text would silently disable the exclusion for that entry. Asserted + against a written file rather than against today's values, because today's are already + tidy and would pass either way. + """ + for spelling in ( + "extension/llm/tokenizers/", + "./extension/llm/tokenizers", + "extension//llm/tokenizers", + ): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + # Only the path is read, so the entry needs no url. + (root / ".gitmodules").write_text( + f'[submodule "t"]\n\tpath = {spelling}\n' + ) + # The helper reads .gitmodules beside its own setup.py, so it is loaded + # against the temporary tree rather than the real one. + prefixes = _load_from_setup_py(root)["_vendored_prefixes"]() + self.assertEqual( + prefixes, + ("extension/llm/tokenizers",), + f"{spelling!r} did not normalize", + ) + + def test_the_fallback_matches_gitmodules(self) -> None: + """The hardcoded fallback still lists the same submodules the file does. + + It is only used when .gitmodules cannot be read, which is the case in a source + distribution, so nothing else would notice it drifting out of date. + """ + self.assertEqual( + _vendored_prefixes(), _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"] + ) + + # And it is actually returned when the file is missing, which is the only case it + # exists for. Without this the fallback could be replaced by an empty tuple and the + # comparison above would still hold. + with tempfile.TemporaryDirectory() as tmp: + self.assertEqual( + _load_from_setup_py(Path(tmp))["_vendored_prefixes"](), + _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"], + "with no .gitmodules the submodule exclusion silently does nothing", + ) + + def test_a_submodule_name_with_a_space_is_read(self) -> None: + """A submodule whose NAME contains a space still yields its path. + + git prints " " and permits spaces in the name, so splitting on the first + space truncates the key and leaves a value that matches no directory, turning the + exclusion off for that entry. + """ + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text( + '[submodule "my module"]\n\tpath = extension/llm/tokenizers\n' + ) + self.assertEqual( + _load_from_setup_py(root)["_vendored_prefixes"](), + ("extension/llm/tokenizers",), + ) + + def test_a_broken_gitmodules_falls_back(self) -> None: + """An unreadable .gitmodules reaches the fallback rather than excluding nothing. + + git exits non-zero with empty output on a bad section header or on conflict markers. + Reading that as "this repository has no submodules" would turn the exclusion off with + no warning, which is the one failure the fallback exists to prevent. + """ + for broken in ( + '[submodule "x"\n\tpath = extension/llm/tokenizers\n', + '<<<<<<< HEAD\n[submodule "x"]\n\tpath = a/b\n=======\n', + "", + ): + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + (root / ".gitmodules").write_text(broken) + prefixes = _load_from_setup_py(root)["_vendored_prefixes"]() + self.assertEqual( + prefixes, + _NAMESPACE["_VENDORED_SUBMODULE_FALLBACK"], + f"a broken .gitmodules ({broken[:20]!r}) silently excluded nothing", + ) + + def test_manifest_filter_actually_drops_vendored_data(self) -> None: + """The data-file half of the fix removes files, through the real build code path. + + `packages` only governs Python modules. Non-Python files arrive through the + package_data manifest, and setuptools attributes a file under an unlisted directory to + its nearest listed parent, so vendored data returns unless the manifest is filtered too. + + Drives CustomBuildPy.analyze_manifest itself rather than reimplementing the filter here. + Checking the predicate in isolation is not enough: inverting the editable guard or + short-circuiting the condition leaves the predicate correct and the build unfiltered, + and both of those left an earlier version of this test green. + """ + namespace = _load_build_py() + build_py_class = namespace["CustomBuildPy"] + + vendored = ( + "src/executorch/backends/xnnpack/third-party/generate-cpuinfo-wrappers.py" + ) + # A shader template goes through this same filter, and it needs its own example here: + # deleting the shader line leaves the vendored assertions below green, so the manifest + # half of the shader fix was unprotected. + shader = "src/executorch/backends/vulkan/runtime/graph/ops/glsl/adamw_step.yaml" + ordinary = "setup.py" + # All of them have to exist on disk, because the filter also drops anything that is not a + # file, and a missing path would be removed for that reason instead of this one. + for relative in (vendored, shader, ordinary): + self.assertTrue( + (REPO_ROOT / relative).is_file(), + f"{relative} is gone, so this test needs a different example", + ) + + # analyze_manifest calls up into setuptools first, which needs the full command + # machinery. Only the filtering after that call is under test, so the parent's method + # is replaced with a no-op for the duration and the manifest seeded directly. This runs + # the shipped code path rather than a copy of it, which is the point: a filter that has + # been turned off still reads correctly in the source. + parent = build_py_class.__mro__[1] + original = parent.analyze_manifest + parent.analyze_manifest = lambda self: None + try: + stub = build_py_class.__new__(build_py_class) + stub.editable_mode = False + stub.manifest_files = {"executorch": [vendored, shader, ordinary]} + stub.analyze_manifest() + kept = stub.manifest_files["executorch"] + finally: + parent.analyze_manifest = original + + self.assertNotIn( + vendored, kept, "a vendored data file survived the manifest filter" + ) + self.assertNotIn(shader, kept, "a shader template survived the manifest filter") + self.assertIn(ordinary, kept, "the filter dropped an ordinary file") + + def test_is_vendored_path_matches_whole_components(self) -> None: + """The filter matches a path component, not a substring.""" + self.assertTrue( + _is_vendored_path( + "src/executorch/backends/xnnpack/third-party/XNNPACK/a.py" + ) + ) + self.assertTrue(_is_vendored_path("src/executorch/x/third_party/y.yaml")) + self.assertFalse(_is_vendored_path("src/executorch/exir/program/_program.py")) + # "third-party" as part of a longer name is a different directory. + self.assertFalse(_is_vendored_path("src/executorch/x/third-party-tools/y.py")) + + def test_submodules_outside_a_vendored_dir_are_recognized(self) -> None: + """A submodule checked out under an ordinary name is still another repository. + + These are not matched by the directory name, so they are read from .gitmodules. Their + nested copies also cannot satisfy the imports the code uses: the FACTO helper imports + facto.specdb from the top level, and the tokenizers ship as a declared dependency. + """ + prefixes = _vendored_prefixes() + self.assertIn("backends/cadence/utils/FACTO", prefixes) + self.assertIn("extension/llm/tokenizers", prefixes) + for prefix in ("backends/cadence/utils/FACTO", "extension/llm/tokenizers"): + self.assertTrue(_is_vendored_path(f"executorch/{prefix}")) + self.assertTrue(_is_vendored_path(f"src/executorch/{prefix}/setup.py")) + self.assertFalse( + _is_vendored_path("executorch/extension/llm/custom_ops/op_sdpa.py") + ) + + def test_root_level_submodules_are_not_listed(self) -> None: + """A submodule at the repository root is not a wheel path. + + Those are build tooling, never copied into the package, and listing one would put a + bare single-word name into the matcher. That would then drop any directory sharing the + name, anywhere in the tree, which is a much wider rule than intended. + """ + for prefix in _vendored_prefixes(): + self.assertIn( + "/", + prefix, + f"{prefix!r} is a root-level submodule and must not be listed", + ) + self.assertFalse(_is_vendored_path("executorch/some/nested/shim")) + + +if __name__ == "__main__": + unittest.main() diff --git a/pyproject.toml b/pyproject.toml index 105565862df..f7cf98679fb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,9 +100,6 @@ Changelog = "https://github.com/pytorch/executorch/releases" [project.scripts] flatc = "executorch.data.bin:flatc" -# TODO(dbort): Could use py_modules to restrict the set of modules we -# package, and package_data to restrict the set up non-python files we -# include. See also setuptools/discovery.py for custom finders. [tool.setuptools] license-files = ["LICENSE"] @@ -118,10 +115,9 @@ license-files = ["LICENSE"] "executorch" = "src/executorch" [tool.setuptools.package-data] -# TODO(dbort): Prune /test[s]/ dirs, /third-party/ dirs, yaml files that we -# don't need. -# TODO(RobertKalmar): When test[s] dirs pruned the PROJECT_DIR resolution in backends.nxp.tests_models.config.py can -# avoid exporting and reading env variable. +# TODO(RobertKalmar): the PROJECT_DIR env variable is still needed. Test directories still install, +# since the suites import shared helpers from them, and the artifacts that config.py resolves are +# not in the wheel, so a path derived from __file__ would point at files that are not there. "*" = [ # Some backends like XNNPACK need their .fbs files. "*.fbs", diff --git a/setup.py b/setup.py index 8980d1bd94a..132e9bcd1e3 100644 --- a/setup.py +++ b/setup.py @@ -47,7 +47,9 @@ # other computer software, distribute, and sublicense such enhancements or # derivative works thereof, in binary and source code form. +import ast import contextlib +import functools # Import this before distutils so that setuptools can intercept the distuils # imports. @@ -64,7 +66,7 @@ from distutils import log # type: ignore[import-not-found] from distutils.sysconfig import get_python_lib # type: ignore[import-not-found] from pathlib import Path, PurePosixPath -from typing import List, Optional +from typing import Dict, FrozenSet, List, Optional, Set, Tuple # Clean dynamic import using importlib _install_utils_path = Path(__file__).parent / "install_utils.py" @@ -177,10 +179,117 @@ def _minimal_cmake_flags() -> List[str]: ] +_VENDORED_DIR_NAMES = frozenset({"third-party", "third_party"}) + +# Used only when .gitmodules cannot be read, as in a source distribution. A test keeps it in step. +_VENDORED_SUBMODULE_FALLBACK = ( + "backends/cadence/utils/FACTO", + "extension/llm/tokenizers", +) + + +@functools.lru_cache(maxsize=None) +def _vendored_prefixes() -> Tuple[str, ...]: + """Source-tree prefixes holding code from another repository. + + Two shapes reach the wheel. Most vendored code sits in a directory named third-party, + which the name above covers wherever it appears. The rest are git submodules checked out + under an ordinary name, so they can only be recognized by asking git what they are. + + None of them are importable from where they sit. FACTO is pure Python but its nested copy + cannot satisfy backends/cadence/utils/facto_util.py, which imports the top level facto.specdb, + and the tokenizers ship separately as pytorch-tokenizers in the dependency list. The rest, + XNNPACK and the Vulkan headers among them, are C++ sources that the wheel has no use for once + the libraries are built. + + Read through git rather than by scanning the file, so only real submodule entries count. + A hand-rolled reader accepts a `path` line from any section, and one stray line elsewhere + in the file would drop a first-party package from the wheel with nothing to warn about. + + Submodules at the repository root are skipped. Those are build tooling, never copied into + the package, and carrying a bare single-word name here would make the match below drop any + directory that happened to share it. + """ + root = Path(__file__).parent + if not (root / ".gitmodules").is_file(): + # A source distribution carries no .gitmodules, so nothing can be read there. Fall + # back to the directories the vendored trees occupy, or the exclusion would quietly + # do half its job and those files would ship again. + return _VENDORED_SUBMODULE_FALLBACK + try: + listed = subprocess.run( + [ + "git", + "config", + "-z", + "-f", + ".gitmodules", + "--get-regexp", + r"^submodule\..*\.path$", + ], + cwd=root, + capture_output=True, + text=True, + check=False, + ) + except OSError: + # No git on PATH, so fall back for the same reason as above. + return _VENDORED_SUBMODULE_FALLBACK + + if listed.returncode or not listed.stdout.strip(): + # git ran and told us nothing useful, which happens when the file has a bad section + # header or conflict markers in it. Reading that as "no submodules" would turn the + # exclusion off without a word, so fall back rather than trust an empty answer. + return _VENDORED_SUBMODULE_FALLBACK + + prefixes = [] + # -z separates each record with NUL and its key from its value with a newline, so neither a + # name nor a path containing a space can be misread. Splitting the default space-separated + # output cannot do that: "submodule.a b.path c d/e" is ambiguous either way round. + for record in listed.stdout.split("\0"): + if not record: + continue + _, separator, value = record.partition("\n") + if not separator: + continue + # Normalize, because git accepts a trailing slash, a ./ prefix and doubled + # separators as the same path, and the raw text would stop matching the real + # directory. + parts = Path(value).parts + if len(parts) < 2 or any(part in _VENDORED_DIR_NAMES for part in parts): + continue + prefixes.append("/".join(parts)) + return tuple(sorted(prefixes)) + + +def _is_vendored_path(path: str) -> bool: + """Whether a source-tree path holds code from another repository.""" + parts = Path(path).parts + if any(part in _VENDORED_DIR_NAMES for part in parts): + return True + # A submodule path is relative to the repository root, while a path here may be relative + # to src/executorch or carry a src/executorch prefix, so match on any suffix boundary. + # Whole-component match: the prefix must be the entire path, or sit at its start, end, or + # middle bounded by separators. Substring matching would let a directory whose name merely + # begins with a prefix be dropped. + posix = "/".join(parts) + return any( + posix == prefix + or posix.startswith(f"{prefix}/") + or posix.endswith(f"/{prefix}") + or f"/{prefix}/" in posix + for prefix in _vendored_prefixes() + ) + + def _minimal_packages() -> List[str]: return sorted( find_namespace_packages( - where="src", + # Anchored on this file, not the working directory, so the list does not change + # with where the build was started from. A cwd-relative path returns nothing when + # the build runs from anywhere but the repository root, and a wheel with no packages + # in it ships no Python at all. + where=str(Path(__file__).parent / "src"), include=[ "executorch", "executorch.data", @@ -204,6 +313,334 @@ def _minimal_packages() -> List[str]: ) +_WALK_SKIP_DIRS = frozenset( + {".git", "pip-out", "cmake-out", "third-party", "third_party", "__pycache__"} +) + +_TEST_DIR_NAMES = frozenset({"test", "tests"}) + + +@functools.lru_cache(maxsize=None) +def _top_level_package_dirs() -> FrozenSet[str]: + """The first path component of every package the wheel ships. + + Derived from the tree rather than listed, so a new top-level directory is covered without an + edit here. Used to recognize the unprefixed spelling of a first-party import. + """ + root = Path(__file__).parent / "src" / "executorch" + if not root.is_dir(): + return frozenset() + return frozenset(entry.name for entry in root.iterdir() if entry.is_dir()) + + +# Named only by a workflow or by a documented command, so no import reaches them. Listed here +# rather than scanned from .github, which a source distribution does not carry; a test re-derives +# the list so it cannot drift. +_CI_ENTRY_POINTS = ( + "executorch.backends.mlx.test.run_all_tests", + "executorch.backends.mlx.test.test_sample", + "executorch.backends.mlx.test.test_slot_recycling", + "executorch.backends.samsung.test.utils.run_tests", + "executorch.backends.test.suite.generate_markdown_summary_json", + "executorch.examples.models.muse_glimmer.tests.gen_prompt_golden", + "executorch.examples.models.muse_glimmer.tests.test_mlx_pipeline", + "executorch.examples.models.muse_glimmer.tests.test_prompt_tokens", + "executorch.extension.pybindings.test.test_pybindings", +) + +# Directories whose test modules are reached without any import statement naming them, so no scan +# of the source can find them: mlx.yml runs each file it discovers under custom_kernel_ops, the +# webgpu scripts import one module per operator, runner.py resolves a suite root out of a dict and +# then walks it, and the llava README documents a `python -m` command. Directories rather than file +# names, so a new test is covered when it is added. +_CI_ENTRY_POINT_DIRS = ( + "executorch.backends.mlx.custom_kernel_ops", + "executorch.backends.webgpu.test", + "executorch.backends.test.suite", + "executorch.examples.models.llava.test", +) + + +def _is_test_module(dotted: str) -> bool: + return any(part in _TEST_DIR_NAMES for part in dotted.split(".")) + + +def _module_name(root: Path, path: Path) -> str: + parts = list(path.relative_to(root).parts) + if parts[-1] == "__init__.py": + parts = parts[:-1] + else: + parts[-1] = parts[-1][: -len(".py")] + return ".".join(["executorch"] + parts) + + +def _first_party_module(name: str) -> Optional[str]: + """The `executorch.`-prefixed spelling of an import target, or None if it is not ours. + + This repository imports itself two ways. Most code says `executorch.backends.x`, but some + says `backends.x`, which resolves because pytest puts the repository root on sys.path. Both + name the same file, so both have to count as a reference or a helper reached only by the + second spelling is dropped from the wheel while its importers still expect it. + """ + if name.startswith("executorch."): + return name + if name.split(".", 1)[0] in _top_level_package_dirs(): + return f"executorch.{name}" + return None + + +def _scan_imports(path: Path, package: str, out: Set[str], dynamic: Set[str]) -> None: + """Collect into out the executorch modules one file refers to.""" + try: + tree = ast.parse(path.read_text(encoding="utf-8", errors="replace")) + except SyntaxError: + return + for node in ast.walk(tree): + out.update(_import_targets(node, package, dynamic)) + + +_GENERATED_DIR_NAMES = frozenset( + { + ".venv", + "venv", + "build", + "dist", + "buck-out", + ".cache", + ".hypothesis", + ".mypy_cache", + ".pytest_cache", + ".tox", + "test-build", + "arm_test", + "riscv_test", + } +) + + +def _unshipped_directories(root: Path) -> List[Path]: + """Checkout directories the wheel does not carry, whose imports still have to be followed. + + src/executorch is a subset of the repository, so a file under test/ or tools/ is never + packaged, yet a module it imports still has to ship. + + Generated directories are left out, because a build tree or an in-tree virtualenv holds an + INSTALLED copy of this package, and reading it would let the last wheel vote on what the next + one ships. Listed by name rather than asked of git, because `git check-ignore` needs a working + repository and answers differently for a pattern with a trailing slash depending on whether the + directory exists yet, which made the same build behave differently on two platforms. + """ + repository = Path(__file__).parent + if not repository.is_dir() or not root.is_dir(): + return [] + shipped = {entry.name for entry in root.iterdir()} + return [ + entry + for entry in sorted(repository.iterdir()) + if entry.is_dir() + and entry.name not in shipped + and entry.name not in _WALK_SKIP_DIRS + and entry.name not in _GENERATED_DIR_NAMES + and entry.name not in ("src", ".github") + ] + + +def _import_graph(root: Path) -> Tuple[Set[str], Dict[str, Set[str]], Set[str]]: + """Every module under root, what each imports, and literal importlib targets. + + The walk covers root, but the SEED covers more: a file elsewhere in the checkout can import + a module that ships, so its imports are collected too and attributed to a synthetic name. + Without that, a helper whose only importer lives outside the shipped tree looks unreachable. + """ + modules: Set[str] = set() + edges: Dict[str, Set[str]] = {} + dynamic: Set[str] = set() + + # followlinks, because src/executorch is a tree of symlinks into the repository root. + for dirpath, dirnames, filenames in os.walk(root, followlinks=True): + # Vendored trees are skipped by the same test that excludes them from the package list, + # not only by directory name. A submodule checked out under an ordinary name, FACTO and + # the tokenizers among them, is otherwise read as first-party, and its imports would keep + # test modules the wheel has no reason to carry. + dirnames[:] = [ + d + for d in dirnames + if d not in _WALK_SKIP_DIRS + and not _is_vendored_path(os.path.relpath(os.path.join(dirpath, d), root)) + ] + for filename in filenames: + if not filename.endswith(".py"): + continue + path = Path(dirpath) / filename + me = _module_name(root, path) + modules.add(me) + package = me if filename == "__init__.py" else me.rsplit(".", 1)[0] + _scan_imports(path, package, edges.setdefault(me, set()), dynamic) + + # Directories of the checkout that the wheel does not ship, test/ among them. Their files are + # never packaged, so they are not modules, but what they import must still ship: for example + # test/end2end/test_end2end.py imports two model helpers out of exir/tests. + for entry in _unshipped_directories(root): + for dirpath, dirnames, filenames in os.walk(entry, followlinks=False): + dirnames[:] = [d for d in dirnames if d not in _WALK_SKIP_DIRS] + for filename in filenames: + if not filename.endswith(".py"): + continue + outside = f"{dirpath}/{filename}" + _scan_imports( + Path(dirpath) / filename, + "", + edges.setdefault(outside, set()), + dynamic, + ) + + return modules, edges, dynamic + + +def _import_targets(node: ast.AST, package: str, dynamic: Set[str]) -> Set[str]: + """The executorch modules one AST node refers to.""" + found: Set[str] = set() + if isinstance(node, ast.Import): + found.update( + name for name in (_first_party_module(a.name) for a in node.names) if name + ) + elif isinstance(node, ast.ImportFrom): + if node.level: + if not package: + # A file outside the shipped tree, so a relative import stays inside that tree + # and cannot name anything the wheel carries. + return found + # A relative import names a real module too, and inside a kept package its target + # has to ship: stages/__init__.py does `from .export import Export`, so dropping + # stages.export would break every importer of that package. + parts = package.split(".") + if node.level > 1: + parts = parts[: len(parts) - (node.level - 1)] + base = ".".join(parts + (node.module.split(".") if node.module else [])) + elif node.module and (prefixed := _first_party_module(node.module)): + base = prefixed + else: + return found + if base.startswith("executorch"): + found.add(base) + # `from pkg import name` may name a submodule rather than an attribute, and there + # is no way to tell without importing, so both readings are kept. + found.update(f"{base}.{a.name}" for a in node.names) + elif ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and node.func.attr == "import_module" + and node.args + and isinstance(node.args[0], ast.Constant) + and isinstance(node.args[0].value, str) + ): + target = node.args[0].value + if target.startswith("."): + if not package: + return found + target = package + target + resolved = _first_party_module(target) + if resolved: + dynamic.add(resolved) + return found + + +@functools.lru_cache(maxsize=None) +def _reachable_test_modules() -> FrozenSet[str]: + """Test modules something can still reach once the wheel is installed. + + A test case that nothing imports is dead weight in the wheel: pytest loads it from a path in + the checkout, never through the installed name. A shared helper is the opposite, because the + suites import each other by installed name, so it has to ship or collection breaks. + + Reachable means named by something, anywhere in the checkout, including by a test module + itself. That looks circular and is not: a test collected from the checkout still resolves + `from executorch.x.test import helper` through the INSTALLED package, so the helper must be + in the wheel even though the file importing it is not. + """ + root = Path(__file__).parent / "src" / "executorch" + modules, edges, dynamic = _import_graph(root) + + # Every name anything refers to. No transitive walk is needed: this is already the union of + # every edge target, so following an edge could only rediscover a name that is in here. + referenced = set(dynamic) + referenced.update(_CI_ENTRY_POINTS) + for targets in edges.values(): + referenced.update(targets) + + keep = {name for name in referenced if _is_test_module(name)} & modules + # Everything under a directory whose tests are run one file at a time by a discovery loop. + keep |= { + name + for name in modules + if _is_test_module(name) + and any( + name == prefix or name.startswith(f"{prefix}.") + for prefix in _CI_ENTRY_POINT_DIRS + ) + } + # Parent packages of anything kept, or the dotted path cannot resolve. + for name in list(keep): + parts = name.split(".") + for end in range(2, len(parts)): + parent = ".".join(parts[:end]) + if _is_test_module(parent): + keep.add(parent) + return frozenset(keep) + + +_SHADER_TEMPLATE_MARKERS = ( + "parameter_names_with_default_values", + "shader_variants", + "generate_variant_forall", +) + + +@functools.lru_cache(maxsize=None) +def _is_shader_template(path: str) -> bool: + """Whether a yaml file is a shader codegen input rather than data the wheel needs. + + gen_vulkan_spv.py and gen_wgsl_headers.py expand these into SPIR-V and WGSL headers during + the cmake build, so the wheel already carries the compiled result. Matched on content rather + than on a directory list, because the same shape appears under vulkan and webgpu and a path + list goes stale as soon as a backend adds one. The op and kernel definitions that ARE read + at run time, edge.yaml among them, carry none of these keys. + """ + if not path.endswith(".yaml"): + return False + full = Path(__file__).parent / path + try: + head = full.read_text(encoding="utf-8", errors="replace") + except OSError: + return False + return any(marker in head for marker in _SHADER_TEMPLATE_MARKERS) + + +def _full_packages() -> List[str]: + """Every package the full wheel ships. + + Without an explicit list setuptools discovers all of src/executorch, which pulls in the + Python files and codegen scripts of the vendored third-party checkouts. Those exist to + build the C++ targets, so once the libraries are built no shipped module imports them. + + Test packages deliberately stay. The suites in this repository import each other through + the installed name, for example `from executorch.backends.arm.test import common`, so + dropping them from the wheel stops the suites collecting under a non-editable install. + """ + return sorted( + package + # Anchored on this file rather than the working directory, so the list does not + # change with where the build or a test was started from. + for package in find_namespace_packages( + where=str(Path(__file__).parent / "src"), + include=["executorch", "executorch.*"], + ) + # The include patterns above DO match these, since they are ordinary dotted names, + # which is exactly why they have to be removed here instead. + if not _is_vendored_path(package.replace(".", "/")) + ) + + # The published project names for the CUDA runtime components a CUDA wheel links but # does not bundle, keyed by CUDA major version. Not derivable from a suffix rule: the # CUDA 12 wheels carry a "-cu12" suffix while the CUDA 13 ones are published under @@ -1628,6 +2065,68 @@ class CustomBuildPy(build_py): a file to a different relative location under the output package directory. """ + def _prune_unstaged_files(self) -> None: + """Delete .py and .yaml this command staged in an earlier build and no longer wants. + + Restricted to files that exist in the source tree, because build_py is the first of + build's sub commands and everything a later one stages is still sitting in the build + directory when this runs. build_ext generates executorch/data/bin/__init__.py, the + target of the flatc console script, from a template that lives elsewhere, so a walk + that removed anything absent from build_py's own file list would delete it. + """ + if self.editable_mode: + return + if not self.packages: + # Nothing to compare against, so every staged file would look unwanted. Refuse + # rather than empty the build directory. + return + + wanted = set() + for package, module, _ in self.find_all_modules(): + parts = package.split(".") if package else [] + wanted.add(os.path.join(self.build_lib, *parts, f"{module}.py")) + for package in self.packages or (): + src_dir = self.get_package_dir(package) + build_dir = os.path.join(*([self.build_lib] + package.split("."))) + for filename in self.find_data_files(package, src_dir): + wanted.add(os.path.join(build_dir, os.path.relpath(filename, src_dir))) + + source_root = Path(__file__).parent / "src" + for dirpath, _, filenames in os.walk(self.build_lib): + for filename in filenames: + if not filename.endswith((".py", ".yaml")): + continue + staged = os.path.join(dirpath, filename) + if staged in wanted: + continue + relative = os.path.relpath(staged, self.build_lib) + if not (source_root / relative).is_file(): + # Generated by another command, so build_py must not remove it. + continue + os.remove(staged) + + def find_package_modules(self, package, package_dir): + modules = super().find_package_modules(package, package_dir) + if self.editable_mode or not _is_test_module(package): + # An editable install exposes the whole source tree whatever is listed here, and a + # package outside a test directory has nothing to drop. + return modules + keep = _reachable_test_modules() + return [ + entry + for entry in modules + if entry[1] == "__init__" or f"{package}.{entry[1]}" in keep + ] + + def find_data_files(self, package, src_dir): + files = super().find_data_files(package, src_dir) + if self.editable_mode: + return files + root = os.path.dirname(os.path.abspath(__file__)) + return [ + _f for _f in files if not _is_shader_template(os.path.relpath(_f, root)) + ] + def analyze_manifest(self): super().analyze_manifest() # Recent versions of setuptools may include bare directory symlinks from version @@ -1642,6 +2141,13 @@ def analyze_manifest(self): _f for _f in self.manifest_files[_pkg] if os.path.isfile(os.path.join(_root, _f)) + # A directory left out of `packages` is not simply skipped. setuptools + # walks up to the nearest listed package and records the file as that + # package's data, so a vendored *.yaml still arrives under its parent. + # Filter with the same list so the two agree. + and not _is_vendored_path(_f) + # Shader templates are consumed by the cmake build, not at run time. + and not _is_shader_template(_f) ] def _copy_extra_files(self, src_to_dst, dst_root: str) -> None: @@ -1684,6 +2190,12 @@ def run(self): # defined by the py_module list and package_data patterns. build_py.run(self) + # A rebuild over a staging directory left by an earlier build keeps whatever that build + # put there, because build_py only ever copies and never deletes. So a file this build + # deliberately leaves out is still present from last time, and the wheel packages it. + # Remove what is no longer wanted rather than only skipping the copy. + self._prune_unstaged_files() + # dst_root is the root of the `executorch` module in the output package # directory. build_lib is the platform-independent root of the output # package, and will look like `pip-out/lib`. It can contain multiple @@ -2325,6 +2837,7 @@ def iter_distribution_names(self): setup_kwargs["packages"] = _minimal_packages() setup_kwargs["install_requires"] = _minimal_dependencies() else: + setup_kwargs["packages"] = _full_packages() # A CUDA wheel links the CUDA runtime but does not bundle it, so the wheels that # carry it are declared here. A CPU wheel adds nothing. setup_kwargs["install_requires"] = _base_dependencies() + _cuda_dependencies() From 56a461f7b7c8eeb2b025ff043f8d1d7f330930e9 Mon Sep 17 00:00:00 2001 From: Tom Allsop <72802373+tom-arm@users.noreply.github.com> Date: Wed, 9 Sep 2026 10:13:18 +0100 Subject: [PATCH 093/190] Arm backend: Improve handling of inputting mutable buffers to ref model (#22603) * Use _graph_module_flat_inputs to get all inputs in correct order Signed-off-by: Tom Allsop Change-Id: I1b22ac2f9af53a120f969ef67bd42fc986eea60e cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- backends/arm/test/tester/arm_tester.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/backends/arm/test/tester/arm_tester.py b/backends/arm/test/tester/arm_tester.py index ba84e30c809..490262ff169 100644 --- a/backends/arm/test/tester/arm_tester.py +++ b/backends/arm/test/tester/arm_tester.py @@ -232,18 +232,13 @@ def run( class ToExecutorch(BaseStages.ToExecutorch): def run_artifact(self, inputs): with TosaReferenceModelDispatch(): - # Check if the model has mutable buffers. These are not delegated to the backend - # and are handled by core ExecuTorch as I/O. In other words, the mutable buffer - # is outputted and re-inputted into the model. As we are calling the graph module - # directly, we need to ensure we handle these extra mutable inputs. - if ( - len(self.artifact.exported_program().graph_signature.buffers_to_mutate) - > 0 - ): - buffers = list(self.artifact.exported_program().buffers()) - buffers.extend(inputs) - - return self.artifact.exported_program().graph_module(*buffers) + program = self.artifact.exported_program() + # Mutable inputs and other parameters become inputs to the graph + # so we need to input these in the correct order. + # Also, execute the raw graph to preserve mutation outputs for comparison. + if program.graph_signature.buffers_to_mutate: + flat_inputs = program._graph_module_flat_inputs(inputs, {}) + return program.graph_module(*flat_inputs) else: return super().run_artifact(inputs) From 78c85d7553df47fb38c89ee3923e7123af54ac71 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Wed, 9 Sep 2026 12:57:07 +0200 Subject: [PATCH 094/190] NXP backend: Register Neutron backend recipes. (#22635) ### Summary Register Neutron backend recipes. ### Test plan N/A cc @robert-kalmar @JakeStevens @digantdesai @rascani --- backends/nxp/recipes/__init__.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) create mode 100644 backends/nxp/recipes/__init__.py diff --git a/backends/nxp/recipes/__init__.py b/backends/nxp/recipes/__init__.py new file mode 100644 index 00000000000..f768d4fa88d --- /dev/null +++ b/backends/nxp/recipes/__init__.py @@ -0,0 +1,18 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.export import recipe_registry + +from .nxp_recipe_provider import NeutronRecipeConfig, NXPRecipeProvider +from .nxp_recipe_types import NXPRecipeType + +# Auto-register NXP recipe provider +recipe_registry.register_backend_recipe_provider(NXPRecipeProvider()) + +__all__ = [ + "NeutronRecipeConfig", + "NXPRecipeProvider", + "NXPRecipeType", +] From 5dc1c7274ed298e0a259ba76ba320ad889d21d41 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Wed, 9 Sep 2026 13:40:09 +0100 Subject: [PATCH 095/190] Arm backend: Add automatic SDPA safe-softmax guard removal (#22637) Implement the AUTO policy by replacing safe softmax with regular softmax for unmasked, noncausal, zero-dropout SDPA with a statically nonempty key sequence. Preserve the guard otherwise. Assisted by Codex. Change-Id: I1386d5dbd0463f36916beab8af445a923f133ba1 Signed-off-by: Yufeng Shi --- backends/arm/_passes/__init__.py | 3 + backends/arm/_passes/arm_pass_manager.py | 15 ++- ...ecompose_sdpa_with_regular_softmax_pass.py | 111 +++++++++++++++++ backends/arm/common/pipeline_config.py | 11 +- .../test/misc/test_pass_pipeline_config.py | 21 +++- .../arm/test/models/test_deit_tiny_arm.py | 22 ++++ .../test_remove_safe_softmax_guard_pass.py | 112 ++++++++++++++++-- backends/arm/test/tester/arm_tester.py | 2 - 8 files changed, 276 insertions(+), 21 deletions(-) create mode 100644 backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index fc48569a3da..57fe9ff414f 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -97,6 +97,9 @@ from .decompose_rnn_pass import DecomposeRnnPass # noqa from .decompose_round_pass import DecomposeRoundPass # noqa from .decompose_sdpa_pass import DecomposeScaledDotProductAttentionPass # noqa +from .decompose_sdpa_with_regular_softmax_pass import ( # noqa + DecomposeSDPAWithRegularSoftmaxPass, +) from .decompose_select import DecomposeSelectPass # noqa from .decompose_select_scatter_pass import DecomposeSelectScatterPass # noqa from .decompose_sign_pass import DecomposeSignPass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index d9a667aff4f..12d766c2ffa 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -90,6 +90,7 @@ DecomposeRnnPass, DecomposeRoundPass, DecomposeScaledDotProductAttentionPass, + DecomposeSDPAWithRegularSoftmaxPass, DecomposeSelectPass, DecomposeSelectScatterPass, DecomposeSignPass, @@ -311,10 +312,7 @@ def configure_skip_passes(self) -> tuple[type, ...]: skip_set.add(DecomposeLeakyReLUPass) match config.sdpa_safe_softmax_guard: # type: ignore[attr-defined] - case ( - SDPASafeSoftmaxGuardPolicy.PRESERVE - | SDPASafeSoftmaxGuardPolicy.REMOVE_WHEN_PROVEN - ): + case SDPASafeSoftmaxGuardPolicy.PRESERVE | SDPASafeSoftmaxGuardPolicy.AUTO: skip_set.add(RemoveSafeSoftmaxGuardPass) case SDPASafeSoftmaxGuardPolicy.REMOVE: pass @@ -461,6 +459,15 @@ def transform_for_pre_decomposition_pipeline( self, exported_program: ExportedProgram ) -> ExportedProgram: """Apply Arm passes before default ATen decompositions.""" + config = self.compile_spec._get_pass_pipeline_config() + passes: list[ExportPass] = [] + + if config.sdpa_safe_softmax_guard is SDPASafeSoftmaxGuardPolicy.AUTO: + passes.append(DecomposeSDPAWithRegularSoftmaxPass()) + + if passes: + self.add_passes(passes) + self._transform(exported_program, exported_program.graph_module) return exported_program def _transform_graph_module(self, graph_module: GraphModule): diff --git a/backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py b/backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py new file mode 100644 index 00000000000..9d2e6da7fa9 --- /dev/null +++ b/backends/arm/_passes/decompose_sdpa_with_regular_softmax_pass.py @@ -0,0 +1,111 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Set, Type + +import torch +from executorch.backends.arm._passes.arm_pass import ArmPass +from executorch.backends.transforms import decompose_sdpa +from executorch.exir.pass_base import ExportPass, PassResult + + +class DecomposeSDPAWithRegularSoftmaxPass( + ArmPass, decompose_sdpa.DecomposeScaledDotProductAttention +): + """Decompose eligible SDPA calls using regular softmax. + + Matches unmasked, noncausal, zero-dropout SDPA calls whose key sequence + length is statically known to be nonzero. The matched form is conceptually:: + + scaled_dot_product_attention( + query, + key, + value, + attn_mask=None, + dropout_p=0.0, + is_causal=False, + scale=scale, + enable_gqa=enable_gqa, + ) + + Other SDPA options, such as ``scale`` and ``enable_gqa``, are preserved. + + The generated subgraph is approximately:: + + scores = (query @ key.transpose(-2, -1)) * scale + output = softmax(scores, dim=-1) @ value + + The standard SDPA decomposition initially generates ``_safe_softmax``. + This pass replaces that operator with regular ``softmax`` only in the + newly generated subgraph. The eligibility checks prevent masks or empty + key sequences from producing all-negative-infinity score rows. They assume + such rows are not produced by nonfinite inputs or numerical overflow. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + + def call( + self, graph_module: torch.fx.GraphModule, allow_non_fake_inputs: bool = True + ) -> PassResult: + graph = graph_module.graph + modified = False + for node in list(graph.nodes): + if node.target != torch.ops.aten.scaled_dot_product_attention.default: + continue + if not self._is_auto_guard_removal_candidate(node): + continue + + existing_nodes = set(graph.nodes) + super()._decompose_sdpa_node(graph_module, node, allow_non_fake_inputs) + self._remove_safe_softmax_guard(graph, existing_nodes) + modified = True + + if modified: + graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) + + @classmethod + def _is_auto_guard_removal_candidate(cls, node: torch.fx.Node) -> bool: + """Return true when SDPA meets automatic removal constraints. + + These structural checks exclude fully masked rows. They assume scores + do not become all ``-inf`` through nonfinite inputs or overflow. + + """ + canonical_args, _, _ = cls._canonicalize_sdpa_call(node) + _, key, _, attn_mask, dropout_p, is_causal, _ = canonical_args + + if attn_mask is not None: + return False + if cls._extract_arg_value(is_causal) is not False: + return False + if cls._extract_arg_value(dropout_p) != 0.0: + return False + return cls._has_nonzero_key_sequence_length(key) + + @staticmethod + def _has_nonzero_key_sequence_length(key: object) -> bool: + if not isinstance(key, torch.fx.Node): + return False + + val = key.meta.get("val") + shape = getattr(val, "shape", None) + if shape is None or len(shape) < 2: + return False + + key_sequence_length = shape[-2] + return isinstance(key_sequence_length, int) and key_sequence_length > 0 + + @staticmethod + def _remove_safe_softmax_guard( + graph: torch.fx.Graph, existing_nodes: set[torch.fx.Node] + ) -> None: + for decomposed_node in graph.nodes: + if decomposed_node in existing_nodes: + continue + if decomposed_node.target == torch.ops.aten._safe_softmax.default: + decomposed_node.target = torch.ops.aten.softmax.int diff --git a/backends/arm/common/pipeline_config.py b/backends/arm/common/pipeline_config.py index 3784849556b..4d1563cc328 100644 --- a/backends/arm/common/pipeline_config.py +++ b/backends/arm/common/pipeline_config.py @@ -24,11 +24,18 @@ class LeakyReLULoweringConfig(Enum): class SDPASafeSoftmaxGuardPolicy(Enum): - """Options for preserving or removing SDPA safe-softmax guards.""" + """Options for preserving or removing SDPA safe-softmax guards. + + ``AUTO`` removes guards only for structurally eligible, unmasked, + noncausal, zero-dropout SDPA calls with a nonempty key sequence. ``AUTO`` + and ``REMOVE`` assume attention scores do not become all ``-inf`` through + nonfinite inputs or overflow. + + """ PRESERVE = auto() # Preserve safe-softmax all--inf row guards REMOVE = auto() # Remove exact expanded safe-softmax guards - REMOVE_WHEN_PROVEN = auto() # Preserve unless a proof is available + AUTO = auto() # Remove eligible guards; preserve uncertain cases @dataclass diff --git a/backends/arm/test/misc/test_pass_pipeline_config.py b/backends/arm/test/misc/test_pass_pipeline_config.py index 5b820102d4f..2ac52322125 100644 --- a/backends/arm/test/misc/test_pass_pipeline_config.py +++ b/backends/arm/test/misc/test_pass_pipeline_config.py @@ -118,14 +118,25 @@ def test_sdpa_safe_softmax_guard_config_controls_guard_removal_pass(): assert RemoveSafeSoftmaxGuardPass in manager._skip_pass_types - stable_compile_spec = TosaCompileSpec( + auto_compile_spec = TosaCompileSpec( + TosaSpecification.create_from_string("TOSA-1.00+INT") + ) + auto_config = ArmPassPipelineConfig( + sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO + ) + auto_compile_spec.set_pass_pipeline_config(auto_config) + auto_manager = ArmPassManager(auto_compile_spec) + + assert RemoveSafeSoftmaxGuardPass in auto_manager._skip_pass_types + + remove_compile_spec = TosaCompileSpec( TosaSpecification.create_from_string("TOSA-1.00+INT") ) remove_config = ArmPassPipelineConfig( sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.REMOVE ) - stable_compile_spec.set_pass_pipeline_config(remove_config) - remove_manager = ArmPassManager(stable_compile_spec) + remove_compile_spec.set_pass_pipeline_config(remove_config) + remove_manager = ArmPassManager(remove_compile_spec) assert RemoveSafeSoftmaxGuardPass not in remove_manager._skip_pass_types @@ -260,11 +271,11 @@ def test_leaky_relu_decompose_config_reaches_backend_pipeline(): def test_sdpa_safe_softmax_guard_config_serializes(): config = ArmPassPipelineConfig( - sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.REMOVE + sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO ) roundtripped = ArmPassPipelineConfig.from_dict(config.to_dict()) - assert roundtripped.sdpa_safe_softmax_guard is SDPASafeSoftmaxGuardPolicy.REMOVE + assert roundtripped.sdpa_safe_softmax_guard is SDPASafeSoftmaxGuardPolicy.AUTO def test_sdpa_safe_softmax_guard_preserves_positional_config_arguments(): diff --git a/backends/arm/test/models/test_deit_tiny_arm.py b/backends/arm/test/models/test_deit_tiny_arm.py index bfbef6cb835..35462e6e6da 100644 --- a/backends/arm/test/models/test_deit_tiny_arm.py +++ b/backends/arm/test/models/test_deit_tiny_arm.py @@ -119,6 +119,28 @@ def test_deit_tiny_tosa_FP_remove_sdpa_safe_softmax_guard(deit_tiny): pipeline.run() +def test_deit_tiny_tosa_FP_auto_remove_sdpa_safe_softmax_guard(deit_tiny): + pipeline = TosaPipelineFP[input_t]( + deit_tiny, + model_inputs, + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.tester.compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + pipeline.count_tosa_ops( + { + "EQUAL": 0, + "LOGICAL_NOT": 0, + "REDUCE_ANY": 0, + "SELECT": 0, + } + ) + pipeline.run() + + def test_deit_tiny_tosa_INT(deit_tiny): pipeline = TosaPipelineINT[input_t]( deit_tiny, diff --git a/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py b/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py index 772b5d39eb8..8064607de59 100644 --- a/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py +++ b/backends/arm/test/passes/test_remove_safe_softmax_guard_pass.py @@ -18,20 +18,24 @@ _get_tosa_operator_distribution, ArmTester, ) +from executorch.backends.arm.tosa.partitioner import TOSAPartitioner from executorch.backends.test.harness.stages import StageType -from executorch.exir import to_edge +from executorch.exir import to_edge, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export from torch.fx import GraphModule, Node class SDPA(torch.nn.Module): - def __init__(self, attn_mask: torch.Tensor | None = None) -> None: + def __init__( + self, attn_mask: torch.Tensor | None = None, is_causal: bool = False + ) -> None: super().__init__() if attn_mask is not None: self.register_buffer("attn_mask", attn_mask) else: self.attn_mask = None + self.is_causal = is_causal def forward( self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor @@ -41,6 +45,35 @@ def forward( key, value, attn_mask=self.attn_mask, + is_causal=self.is_causal, + ) + + +class DynamicMaskSDPA(torch.nn.Module): + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attn_mask: torch.Tensor, + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + attn_mask=attn_mask, + ) + + +class DropoutSDPA(torch.nn.Module): + def forward( + self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor + ) -> torch.Tensor: + return torch.nn.functional.scaled_dot_product_attention( + query, + key, + value, + dropout_p=0.1, ) @@ -217,14 +250,35 @@ def test_sdpa_safe_softmax_guard_preserve_keeps_guard_before_tosa_lowering(): assert counts["SELECT"] == 1 -def test_sdpa_safe_softmax_guard_remove_when_proven_keeps_guard(): +def test_auto_removal_runs_through_partitioner_hook(): compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") compile_spec.set_pass_pipeline_config( - ArmPassPipelineConfig( - sdpa_safe_softmax_guard=(SDPASafeSoftmaxGuardPolicy.REMOVE_WHEN_PROVEN) - ) + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) ) - tester = ArmTester(SDPA(), _sdpa_inputs(), compile_spec) + edge_program = to_edge_transform_and_lower( + export(SDPA(), _sdpa_inputs(), strict=True), + partitioner=[TOSAPartitioner(compile_spec)], + ) + graph_module = edge_program.exported_program().graph_module + counts = dict(_get_tosa_operator_distribution(graph_module)) + + assert counts.get("EQUAL", 0) == 0 + assert counts.get("LOGICAL_NOT", 0) == 0 + assert counts.get("REDUCE_ANY", 0) == 0 + assert counts.get("SELECT", 0) == 0 + assert counts["REDUCE_MAX"] == 1 + assert counts["EXP"] == 1 + assert counts["REDUCE_SUM"] == 1 + assert counts["RECIPROCAL"] == 1 + + +def test_sdpa_safe_softmax_guard_auto_keeps_dynamic_mask_guard(): + compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") + compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + attn_mask = torch.zeros(1, 1, 4, 4) + tester = ArmTester(DynamicMaskSDPA(), (*_sdpa_inputs(), attn_mask), compile_spec) tester.export().to_edge_transform_and_lower() graph_module = ( @@ -237,7 +291,49 @@ def test_sdpa_safe_softmax_guard_remove_when_proven_keeps_guard(): assert counts["EQUAL"] == 1 assert counts["LOGICAL_NOT"] == 2 assert counts["REDUCE_ANY"] == 1 - assert counts["SELECT"] == 1 + assert counts["SELECT"] >= 1 + + +def test_sdpa_safe_softmax_guard_auto_keeps_causal_guard(): + compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") + compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + tester = ArmTester(SDPA(is_causal=True), _sdpa_inputs(), compile_spec) + + tester.export().to_edge_transform_and_lower() + graph_module = ( + tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + .exported_program() + .graph_module + ) + counts = dict(_get_tosa_operator_distribution(graph_module)) + + assert counts["EQUAL"] == 1 + assert counts["LOGICAL_NOT"] == 2 + assert counts["REDUCE_ANY"] == 1 + assert counts["SELECT"] >= 1 + + +def test_sdpa_safe_softmax_guard_auto_keeps_dropout_guard(): + compile_spec = common.get_tosa_compile_spec("TOSA-1.0+FP") + compile_spec.set_pass_pipeline_config( + ArmPassPipelineConfig(sdpa_safe_softmax_guard=SDPASafeSoftmaxGuardPolicy.AUTO) + ) + tester = ArmTester(DropoutSDPA(), _sdpa_inputs(), compile_spec) + + tester.export().to_edge_transform_and_lower() + graph_module = ( + tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER) + .exported_program() + .graph_module + ) + counts = dict(_get_tosa_operator_distribution(graph_module)) + + assert counts["EQUAL"] == 1 + assert counts["LOGICAL_NOT"] == 2 + assert counts["REDUCE_ANY"] == 1 + assert counts["SELECT"] >= 1 def test_remove_safe_softmax_guard_pass_does_not_rewrite_regular_softmax(): diff --git a/backends/arm/test/tester/arm_tester.py b/backends/arm/test/tester/arm_tester.py index 490262ff169..c7e09129f27 100644 --- a/backends/arm/test/tester/arm_tester.py +++ b/backends/arm/test/tester/arm_tester.py @@ -192,7 +192,6 @@ def __init__( transform_passes: Optional[ Union[Sequence[PassType], Dict[str, Sequence[PassType]]] ] = None, - compile_spec: Optional[ArmCompileSpec] = None, ): super().__init__( default_partitioner_cls=None, @@ -469,7 +468,6 @@ def to_edge_transform_and_lower( edge_compile_config, constant_methods=self.constant_methods, transform_passes=self.transform_passes, - compile_spec=self.compile_spec, ) else: if partitioners is not None: From 5c4d2c4a0150a0809bc77be9a67093f80f60a60f Mon Sep 17 00:00:00 2001 From: Usamah Date: Wed, 9 Sep 2026 14:32:50 +0100 Subject: [PATCH 096/190] Arm backend: Simplify MobileSAM first-run workflow (#22640) ### Summary Follow-up to #21397 based on first-run feedback. The MobileSAM example required too much model-specific setup, configuration, and runtime code for an educational example. This change gives it one tested 448x448 Ethos-U85-256 configuration and one command for preparation, export, quantization, FVP execution, and mask validation. Start with the example README and `run.sh`, then review the exporter and visualization scripts. The flow reuses the standard Arm executor runner with semihosted tensor input/output, replacing the custom C++ runtime, CMake project, image-header generation, and RLE log protocol. The smoke test invokes the same command as the tutorial. The documentation puts the fixed-prompt contract up front, clarifies source installation and host AOT dependencies, and gives concrete artifact paths. MobileSAM source and weights remain outside the repository. Cortex-M custom semihosting-command support is left to a separate follow-up. ### Test plan The end-to-end flow passed on this commit on 2026-09-08, using macOS, Python 3.12, the Arm GNU 15.2 toolchain, and Corstone-320 through FVPs-on-Mac/Docker: ```bash bash backends/arm/test/test_arm_ootb.sh run_mobilesam_e2e_ethos_u ``` Preparation, export, quantization, lowering, runtime build, FVP inference, and visualization completed successfully. Host FP32/quantized mask IoU was 0.9550; FVP/host-quantized mask IoU was 0.9809, above the required 0.9. The raw target output was finite and the mask non-degenerate. The FVP run used fast simulation mode; these are correctness results, not hardware performance measurements. Rechecked before opening this PR: ```bash python -m pytest -q backends/arm/test/misc/test_mobilesam.py bash .githooks/pre-commit bash backends/arm/scripts/pre-push ``` All three focused tests and both hooks passed, including lint/type checks, license and commit-message checks, generated documentation, and Arm public API validation. Authored with assistance from Codex. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Usamah Zaheer --- README-wheel.md | 10 + README.md | 4 + backends/arm/scripts/corstone_utils.cmake | 11 +- backends/arm/scripts/fvp_utils.sh | 12 +- backends/arm/scripts/run_fvp.sh | 13 + backends/arm/test/misc/test_mobilesam.py | 48 + backends/arm/test/pytest.ini | 1 + backends/arm/test/test_arm_ootb.sh | 135 +-- .../using-executorch-building-from-source.md | 21 + examples/arm/README.md | 24 +- examples/arm/ethos_u_minimal_example.ipynb | 14 +- .../README.md | 122 ++- .../model_export/README.md | 203 +--- .../model_export/export_mobilesam.py | 906 ++---------------- .../model_export/prepare_mobilesam.py | 167 ++-- .../requirements.txt | 6 - .../run.sh | 60 ++ .../runtime/CMakeLists.txt | 225 ----- .../runtime/README.md | 116 +-- .../runtime/image_to_array.py | 123 --- .../runtime/main.cpp | 413 -------- .../runtime/visualize_fvp_output.py | 224 ++--- examples/arm/setup.sh | 10 +- 23 files changed, 542 insertions(+), 2326 deletions(-) create mode 100644 backends/arm/test/misc/test_mobilesam.py delete mode 100644 examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt create mode 100755 examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh delete mode 100644 examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt delete mode 100644 examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py delete mode 100644 examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp diff --git a/README-wheel.md b/README-wheel.md index 8ee58b56f2f..86edce654cb 100644 --- a/README-wheel.md +++ b/README-wheel.md @@ -8,6 +8,16 @@ The `executorch` pip package is in beta. * Supported python versions: 3.10, 3.11, 3.12, 3.13, 3.14 * Compatible systems: Linux x86_64, Linux aarch64, macOS aarch64 +Backend export tools can require optional Python dependencies. For example, +install the dependencies needed for Ethos-U ahead-of-time (AOT) export with: + +```bash +pip install 'executorch[ethos_u]' +``` + +The `ethos_u` extra does not install components needed to build or run a target +runtime. + To build a minimal wheel from source, set `EXECUTORCH_BUILD_MINIMAL=1` when running `pip wheel` or `pip install`. That wheel contains the Python EXIR export path and `flatc` for `.pte` diff --git a/README.md b/README.md index 3e7037640cf..cf35f80e2bc 100644 --- a/README.md +++ b/README.md @@ -64,6 +64,10 @@ Learn more: [How ExecuTorch Works](https://docs.pytorch.org/executorch/main/intr pip install executorch ``` +Backend export tools can require optional dependencies. For example, use +`pip install 'executorch[ethos_u]'` for Ethos-U AOT export. Embedded +toolchains, simulators, and target runtimes are installed separately. + For platform-specific setup (Android, iOS, embedded systems), see the [Quick Start](https://docs.pytorch.org/executorch/main/quick-start-section.html) documentation for additional info. ### Export and Deploy in 3 Steps diff --git a/backends/arm/scripts/corstone_utils.cmake b/backends/arm/scripts/corstone_utils.cmake index 95ef6f8b866..2dba2052555 100644 --- a/backends/arm/scripts/corstone_utils.cmake +++ b/backends/arm/scripts/corstone_utils.cmake @@ -32,14 +32,9 @@ function(fetch_ethos_u_content ETHOS_SDK_PATH ET_DIR_PATH) GIT_REPOSITORY https://git.gitlab.arm.com/artificial-intelligence/ethos-u/ethos-u.git GIT_TAG ${ethos_u_base_tag} - SOURCE_DIR - ${ETHOS_SDK_PATH} - BINARY_DIR - ${ETHOS_SDK_PATH} - SUBBUILD_DIR - ${ETHOS_SDK_PATH}/../ethos_u-subbuild - SOURCE_SUBDIR - none + SOURCE_DIR ${ETHOS_SDK_PATH} BINARY_DIR ${ETHOS_SDK_PATH} + # Keep the generator-specific population project local to this build. + SOURCE_SUBDIR none ) FetchContent_MakeAvailable(ethos_u) # Patch manifest to remove unused projects. diff --git a/backends/arm/scripts/fvp_utils.sh b/backends/arm/scripts/fvp_utils.sh index 73f67112efd..7c66845d908 100644 --- a/backends/arm/scripts/fvp_utils.sh +++ b/backends/arm/scripts/fvp_utils.sh @@ -138,8 +138,12 @@ function setup_path_fvp() { # Fixup for Corstone-320 python dependency append_env_in_setup_path LD_LIBRARY_PATH "${root_dir}/FVP-corstone320/python/lib/" - echo "hash FVP_Corstone_SSE-300_Ethos-U55" >> ${setup_path_script}.sh - echo "hash FVP_Corstone_SSE-300_Ethos-U65" >> ${setup_path_script}.sh - echo "hash FVP_Corstone_SSE-320" >> ${setup_path_script}.sh - echo "hash FVP_Corstone-1000-A320" >> ${setup_path_script}.sh + local fvp_command + for fvp_command in \ + FVP_Corstone_SSE-300_Ethos-U55 \ + FVP_Corstone_SSE-300_Ethos-U65 \ + FVP_Corstone_SSE-320 \ + FVP_Corstone-1000-A320; do + echo "hash ${fvp_command} 2>/dev/null || true" >> "${setup_path_script}.sh" + done } diff --git a/backends/arm/scripts/run_fvp.sh b/backends/arm/scripts/run_fvp.sh index 7289f37484a..e0e2ba2aa7b 100755 --- a/backends/arm/scripts/run_fvp.sh +++ b/backends/arm/scripts/run_fvp.sh @@ -25,6 +25,7 @@ timeout="600" etrecord_file="" trace_file="" semihosting_cwd="" +semihosting_cmd_line="" ethosu_fast=0 help() { @@ -38,6 +39,7 @@ help() { echo " --etrecord= If ETDump is used you can supply a ETRecord file matching the PTE" echo " --trace_file= File to write PMU trace output to" echo " --semihosting-cwd=

Enable target semihosting with this host working directory" + echo " --semihosting-cmd-line= Command line passed to a semihosting runner" echo " --fast Use fast Ethos-U model simulation for Ethos-U targets" exit 0 } @@ -53,6 +55,7 @@ for arg in "$@"; do --etrecord=*) etrecord_file="${arg#*=}";; --trace_file=*) trace_file="${arg#*=}";; --semihosting-cwd=*) semihosting_cwd="${arg#*=}";; + --semihosting-cmd-line=*) semihosting_cmd_line="${arg#*=}";; --fast) ethosu_fast=1;; *) ;; @@ -139,6 +142,16 @@ if [[ -n "${semihosting_cwd}" ]]; then -C "mps4_board.subsystem.cpu0.semihosting-cwd=${semihosting_cwd}" ) fi +if [[ -n "${semihosting_cmd_line}" ]]; then + [[ -n "${semihosting_cwd}" ]] \ + || { echo "--semihosting-cmd-line requires --semihosting-cwd"; exit 1; } + semihosting_args_u55+=( + -C "cpu0.semihosting-cmd_line=${semihosting_cmd_line}" + ) + semihosting_args_u85+=( + -C "mps4_board.subsystem.cpu0.semihosting-cmd_line=${semihosting_cmd_line}" + ) +fi if [[ ${target} == cortex-m* ]]; then [[ -z "${bundle_file}" ]] \ diff --git a/backends/arm/test/misc/test_mobilesam.py b/backends/arm/test/misc/test_mobilesam.py new file mode 100644 index 00000000000..323549c017b --- /dev/null +++ b/backends/arm/test/misc/test_mobilesam.py @@ -0,0 +1,48 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from examples.arm.mobilesam_prompt_segmentation_example_ethos_u.model_export.export_mobilesam import ( + iou, + mask, +) +from examples.arm.mobilesam_prompt_segmentation_example_ethos_u.runtime.visualize_fvp_output import ( + load_mask, + OUTPUT_SIZE, +) + + +def test_mask_and_iou() -> None: + first = mask(torch.tensor([[[[-1.0, 1.0], [1.0, -1.0]]]])) + second = np.array([[0, 1], [0, 1]], dtype=np.uint8) + + assert first.tolist() == [[0, 1], [1, 0]] + assert iou(first, second) == pytest.approx(1 / 3) + + +def test_load_mask(tmp_path: Path) -> None: + output = np.ones((OUTPUT_SIZE, OUTPUT_SIZE), dtype=np.float32) + output[0, 0] = -1 + path = tmp_path / "output.bin" + output.tofile(path) + + loaded = load_mask(path) + + assert loaded.shape == (OUTPUT_SIZE, OUTPUT_SIZE) + assert loaded[0, 0] == 0 + assert loaded[1, 1] == 1 + + +def test_load_mask_rejects_wrong_output_size(tmp_path: Path) -> None: + path = tmp_path / "output.bin" + np.ones(3, dtype=np.float32).tofile(path) + + with pytest.raises(ValueError, match="Expected"): + load_mask(path) diff --git a/backends/arm/test/pytest.ini b/backends/arm/test/pytest.ini index 09b26752421..301b06d4917 100644 --- a/backends/arm/test/pytest.ini +++ b/backends/arm/test/pytest.ini @@ -1,6 +1,7 @@ [pytest] timeout = 1800 addopts = --strict-markers +pythonpath = ../../.. markers = slow: Tests that take long time xlarge: Tests that are known to use a lot of memory diff --git a/backends/arm/test/test_arm_ootb.sh b/backends/arm/test/test_arm_ootb.sh index 7e3d110855b..d61606ba707 100755 --- a/backends/arm/test/test_arm_ootb.sh +++ b/backends/arm/test/test_arm_ootb.sh @@ -173,140 +173,7 @@ run_deit_e2e_ethos_u() { } run_mobilesam_e2e_ethos_u() { - echo "$FUNCNAME: Export, build, and run the MobileSAM e2e test" - - local example_dir="${et_root_dir}/examples/arm/mobilesam_prompt_segmentation_example_ethos_u" - local work_root="${et_root_dir}/arm_test/mobilesam_ootb_smoke" - local export_dir="${work_root}/export" - local artifact_dir="${work_root}/artifacts" - local debug_dir="${work_root}/debug" - local et_build_dir="${work_root}/cmake-out-arm" - local quantized_aot_build_dir="${work_root}/quantized_ops_aot" - local build_dir="${work_root}/runtime" - local mobile_sam_source="${work_root}/mobile_sam/source" - local image_path="${et_root_dir}/examples/models/dinov2/dog.jpg" - local pte_path="${export_dir}/mobilesam_prompt_smoke.pte" - local metadata_path="${export_dir}/mobilesam_prompt_smoke.json" - local fvp_log="${work_root}/fvp.log" - local toolchain_file="${et_root_dir}/examples/arm/ethos-u-setup/arm-none-eabi-gcc.cmake" - local input_size=448 - local fvp_timelimit="${FVP_TIMELIMIT:-300}" - echo "${FUNCNAME}: Work directory: ${work_root}; existing artifacts will be reused if present" - - mkdir -p "${export_dir}" "${artifact_dir}" "${debug_dir}" "${build_dir}" - - setup_path_script=${et_root_dir}/examples/arm/arm-scratch/setup_path.sh - source ${setup_path_script} - - source ${et_root_dir}/backends/arm/scripts/utils.sh - local n_proc="$(get_parallel_jobs)" - - echo "${FUNCNAME}: Building ExecuTorch (if needed)" - cmake --preset arm-baremetal -B "${et_build_dir}" - cmake --build "${et_build_dir}" --target install -j"$n_proc" - - echo "${FUNCNAME}: Building host quantized AOT library" - local python_executable - python_executable="$(python3 -c 'import sys; print(sys.executable)')" - cmake \ - -S "${et_root_dir}" \ - -B "${quantized_aot_build_dir}" \ - -DCMAKE_BUILD_TYPE=Release \ - -DEXECUTORCH_BUILD_KERNELS_QUANTIZED=ON \ - -DEXECUTORCH_BUILD_KERNELS_QUANTIZED_AOT=ON \ - -DEXECUTORCH_BUILD_XNNPACK=OFF \ - -DPYTHON_EXECUTABLE="${python_executable}" - cmake --build "${quantized_aot_build_dir}" --target quantized_ops_aot_lib -j"$n_proc" - - local quantized_ops_library - quantized_ops_library="$( - find "${quantized_aot_build_dir}/kernels/quantized" \ - -name 'libquantized_ops_aot_lib.*' \ - -type f \ - -print \ - -quit - )" - [[ -n "${quantized_ops_library}" ]] || { - echo "${FUNCNAME}: Missing quantized AOT library under ${quantized_aot_build_dir}" - return 1 - } - - echo "${FUNCNAME}: Installing example requirements" - pip install -r "${example_dir}/requirements.txt" - - echo "${FUNCNAME}: Preparing pinned MobileSAM source" - python3 "${example_dir}/model_export/prepare_mobilesam.py" \ - --source-dir "${mobile_sam_source}" - - echo "${FUNCNAME}: Exporting quantized MobileSAM PTE" - env EXECUTORCH_QUANTIZED_OPS_AOT_LIBRARY="${quantized_ops_library}" \ - python3 "${example_dir}/model_export/export_mobilesam.py" \ - --output-path "${pte_path}" \ - --calibration-image "${image_path}" \ - --eval-image "${image_path}" \ - --mobile-sam-source "${mobile_sam_source}" \ - --input-size "${input_size}" \ - --point 219 193 \ - --num-calibration-samples 1 \ - --num-eval-samples 1 \ - --num-debug-samples 1 \ - --minimum-fp32-quantized-iou 0.9 \ - --artifact-dir "${artifact_dir}" \ - --debug-output-dir "${debug_dir}" - - for artifact in \ - "${pte_path}" \ - "${metadata_path}" \ - "${export_dir}/mobilesam_prompt_smoke_delegation.txt" \ - "${export_dir}/mobilesam_prompt_smoke_metrics.json"; do - [[ -f "${artifact}" ]] || { - echo "${FUNCNAME}: Missing export artifact ${artifact}" - return 1 - } - done - - echo "${FUNCNAME}: Configuring the MobileSAM application" - cmake \ - -U "LIB_*" \ - -U executorch_DIR \ - -S "${example_dir}/runtime" \ - -B "${build_dir}" \ - -DCMAKE_TOOLCHAIN_FILE="${toolchain_file}" \ - -DET_PTE_FILE_PATH="${pte_path}" \ - -DMODEL_METADATA_PATH="${metadata_path}" \ - -DIMAGE_PATH="${image_path}" \ - -DMASK_THRESHOLD=0.0 \ - -DET_SEGMENTATION_DUMP_MASK=ON \ - -DPYTHON_EXECUTABLE="${python_executable}" \ - -DET_BUILD_DIR_PATH="${et_build_dir}" - - echo "${FUNCNAME}: Building mobilesam_prompt_segmentation_example" - cmake --build "${build_dir}" -j"$n_proc" --target mobilesam_prompt_segmentation_example - - local elf="${build_dir}/mobilesam_prompt_segmentation_example" - - echo "${FUNCNAME}: Running on FVP" - backends/arm/scripts/run_fvp.sh \ - --elf="${elf}" \ - --target=ethos-u85-256 \ - --timeout="${fvp_timelimit}" \ - --semihosting-cwd="${build_dir}" \ - --fast | tee "${fvp_log}" - - grep -q "Model executed successfully." "${fvp_log}" || { - echo "${FUNCNAME}: FVP run did not report successful execution" - return 1 - } - - python3 "${example_dir}/runtime/visualize_fvp_output.py" \ - --fvp-log "${fvp_log}" \ - --input-image "${image_path}" \ - --metadata "${metadata_path}" \ - --reference-mask "${debug_dir}/dog/quantized_mask.png" \ - --minimum-iou 0.9 \ - --output-dir "${work_root}/fvp_visual" - - echo "${FUNCNAME}: PASS" + "${et_root_dir}/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh" } run_swin2sr_e2e_vgf() { diff --git a/docs/source/using-executorch-building-from-source.md b/docs/source/using-executorch-building-from-source.md index e8c7f0e5a3b..03d73420bb9 100644 --- a/docs/source/using-executorch-building-from-source.md +++ b/docs/source/using-executorch-building-from-source.md @@ -79,8 +79,29 @@ portability details. * `--clean`: Removes build artifacts. * `--editable`: Install the ExecuTorch python package in editable mode (see [Editable Install](#editable-install)). * `--minimal`: Install only the minimal set of dependencies required to run ExecuTorch. Do not install dependencies for examples. + * `--optional-dependency `: Install an optional Python dependency set. + Repeat the flag to select more than one. Supported names are `ethos_u`, + `vgf`, and `openvino`. * `--use-pt-pinned-commit`: Install the pinned PyTorch commit or release version. When not specified, the latest PyTorch nightly build is installed. + For example, install the current checkout with the dependencies needed for + Ethos-U ahead-of-time (AOT) export: + + ```bash + ./install_executorch.sh --optional-dependency ethos_u + ``` + + After the base dependencies have already been installed, the equivalent + editable package command is: + + ```bash + pip install -e '.[ethos_u]' --no-build-isolation + ``` + + The `ethos_u` optional dependencies are host-side Python tools used during + AOT export. Embedded toolchains, simulators, and target runtimes are + configured separately by the backend setup and build instructions. + For Intel-based macOS systems, use `--use-pt-pinned-commit --minimal`. As PyTorch does not provide pre-built binaries for Intel Mac, installation requires building PyTorch from source. Instructions can be found in [PyTorch Installation](https://github.com/pytorch/pytorch#installation). Note that only the XNNPACK and CoreML backends are built by default. Additional backends can be enabled or disabled by setting the corresponding CMake flags: diff --git a/examples/arm/README.md b/examples/arm/README.md index 1a0923f1ab3..d830356717c 100644 --- a/examples/arm/README.md +++ b/examples/arm/README.md @@ -11,13 +11,29 @@ This directory contains documentation and scripts to help you setup and run a PyTorch model on the Arm backend via ExecuTorch. +## Python package setup + +For Ethos-U examples, install the current checkout and the dependencies needed +for ahead-of-time (AOT) export in a clean Python environment: + +```bash +python3.12 -m venv .venv +source .venv/bin/activate +./install_executorch.sh --optional-dependency ethos_u +``` + +After the base dependencies are installed, the equivalent editable package +command is `pip install -e '.[ethos_u]' --no-build-isolation`. The `ethos_u` +extra provides host-side export dependencies; it does not install the Arm +toolchain, FVPs, or target runtime. Run `setup.sh` below to install the cross +compiler, FVPs, and backend tools used by these examples. + ## setup.sh `setup.sh` downloads the Arm cross-compilation toolchain and Corstone FVP -simulators, installs the Python dependencies for TOSA, Ethos-U Vela, and -Cortex-M/CMSIS-NN, and generates `setup_path.sh` scripts for adding those tools -to your environment. Optional flags also install VGF/MLSDK and Vulkan -dependencies. +simulators, installs the backend dependencies, and generates `setup_path.sh` +scripts for adding those tools to your environment. Optional flags also install +VGF/MLSDK and Vulkan dependencies. Example to install the default Arm backend dependencies and add them to your current shell: diff --git a/examples/arm/ethos_u_minimal_example.ipynb b/examples/arm/ethos_u_minimal_example.ipynb index 11f24019d23..d9298237b95 100644 --- a/examples/arm/ethos_u_minimal_example.ipynb +++ b/examples/arm/ethos_u_minimal_example.ipynb @@ -21,11 +21,17 @@ "This guide demonstrates the full flow for running a module on Arm Ethos-U55 using ExecuTorch.\n", "Tested on Linux x86_64 and macOS aarch64. If something is not working for you, please raise a GitHub issue and tag Arm.\n", "\n", - "Before you begin:\n", - "1. (In a clean virtual environment with a compatible Python version) Install executorch using `./install_executorch.sh`\n", - "2. Install Arm cross-compilation toolchain and simulators using `./examples/arm/setup.sh --i-agree-to-the-contained-eula`\n", + "Before you begin, run these commands from the base `executorch` folder:\n", "\n", - "With all commands executed from the base `executorch` folder.\n", + "```bash\n", + "python3.12 -m venv .venv\n", + "source .venv/bin/activate\n", + "./install_executorch.sh --optional-dependency ethos_u\n", + "./examples/arm/setup.sh --i-agree-to-the-contained-eula\n", + "source examples/arm/arm-scratch/setup_path.sh\n", + "```\n", + "\n", + "`--optional-dependency ethos_u` installs the Python tools needed to export models for Ethos-U, including Vela. The Arm setup script separately installs the cross compiler and FVPs used later in the notebook.\n", "\n", "\n", "\n", diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md index 58eae2b0201..ee2359c171d 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/README.md @@ -1,43 +1,79 @@ -# MobileSAM Prompt Segmentation Example Application - -This end-to-end example shows how to use the Arm Ethos-U backend in -ExecuTorch for transformer-based prompt segmentation. MobileSAM predicts a -binary mask for fixed positive point prompts rather than semantic class IDs. -The host debug flow validates quantization by comparing FP32 and quantized -masks, with an optional binary reference mask when one is available. - -It covers: - -- Loading the MobileSAM `vit_t` checkpoint. -- Freezing one or more positive point prompts into the exported graph. -- Applying post-training quantization with the Ethos-U quantizer. -- Lowering the quantized model to an Ethos-U85-256 ExecuTorch program. -- Producing validation and debugging artifacts such as masks, overlays, - mismatch heatmaps, metrics, and delegation summaries. -- Building a bare-metal Corstone-320 runtime app and running it on FVP. - -The default export uses a reduced `448x448` image input and returns one -low-resolution `[1, 1, 112, 112]` mask-logit tensor. The example prepares the -official MobileSAM GitHub source at a pinned revision in an external checkout -and applies a small configurable-input patch there. Neither the MobileSAM -source nor checkpoint is redistributed in ExecuTorch. - -The export uses int8 activations and int8 weights globally, and A16W8 -quantization for TinyViT attention modules. This keeps the transformer -attention numerically stable while still producing one Ethos-U delegate. - -The exported graph intentionally uses `multimask_output=False` and leaves -mask thresholding outside the model. SAM-style candidate-mask selection can be -numerically sensitive after export and quantization, so this example keeps the -target graph focused on the fixed-prompt image encoder and mask decoder. - -## Layout - -- `model_export/prepare_mobilesam.py` - Prepares the pinned external MobileSAM - checkout and applies the configurable-input patch. -- `model_export/README.md` - Model loading, quantization, lowering, - validation, and debug artifact generation. -- `runtime/README.md` - Bare-metal runtime build, image header generation, and - Corstone-320 FVP execution. -- `runtime/visualize_fvp_output.py` - Decodes the target mask dump, creates an - overlay, and compares FVP output with the host quantized mask. +# MobileSAM Prompt Segmentation on Ethos-U + +This example turns a point on an image into an object mask. It shows the full +ExecuTorch flow: export MobileSAM, quantize it, delegate it to Ethos-U85, run it +on the Corstone-320 FVP, and compare the target result with the host result. + +There is one tested configuration: MobileSAM `vit_t`, a `448x448` input, and +Ethos-U85-256. The image can change at runtime, but the point prompt is embedded +in the exported model. Changing the prompt requires re-exporting the model. + +## Run It + +From the ExecuTorch repository root: + +```bash +python3.12 -m venv .venv +source .venv/bin/activate +./install_executorch.sh --optional-dependency ethos_u +./examples/arm/setup.sh --i-agree-to-the-contained-eula +./examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh +``` + +The final command performs the complete flow and prints +`MobileSAM example: PASS`. Its main result is: + +`arm_test/mobilesam/result/fvp_comparison.png` + +## What It Does + +1. Fetches the pinned official MobileSAM source and checkpoint outside the + repository. +2. Runs `torch.export`, PT2E quantization, and `EthosUPartitioner` to create a + `.pte` containing one Ethos-U delegate. +3. Builds the standard Arm ExecuTorch runner and runs one inference on FVP. +4. Compares the FVP mask with the host quantized mask and requires `0.9` IoU. + +Successful completion creates: + +- Program: `arm_test/mobilesam/export/mobilesam.pte` +- Host masks: `arm_test/mobilesam/export/fp32_mask.png` and + `arm_test/mobilesam/export/quantized_mask.png` +- FVP log: `arm_test/mobilesam/fvp.log` +- Comparison: `arm_test/mobilesam/result/fvp_comparison.png` +- FVP validation: `arm_test/mobilesam/result/metrics.json` +- TOSA and Vela artifacts: `arm_test/mobilesam/export/artifacts` + +The Python installer uses this source checkout and installs the dependencies +needed for ahead-of-time Ethos-U export. The Arm setup script installs the +cross compiler and FVP. Do not install a separate PyPI `executorch` wheel for +this source example. + +On macOS, Docker must be running and the +[FVPs-on-Mac](https://github.com/Arm-Examples/FVPs-on-Mac) wrapper must be on +`PATH`. + +## Code Map + +- [`prepare_mobilesam.py`](model_export/prepare_mobilesam.py) fetches and + verifies the external model. +- [`export_mobilesam.py`](model_export/export_mobilesam.py) contains the model, + quantization, validation, and lowering flow. +- [`run.sh`](run.sh) uses ExecuTorch's standard Arm runner for target execution. +- [`visualize_fvp_output.py`](runtime/visualize_fvp_output.py) checks and plots + the raw output tensor. + +There is no MobileSAM-specific C++ runtime or CMake project. + +## Limitations + +- The exported model accepts one image tensor and uses one fixed positive point. +- It returns a low-resolution mask. Upsampling and thresholding are host-side + post-processing. +- The demo image is also the calibration image. Product use requires a + representative calibration set. +- The default fast FVP mode validates correctness. Its counters are not a + performance benchmark or a measurement of real-device latency. + +See [model export](model_export/README.md) and +[runtime](runtime/README.md) for details of each stage. diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md index 8069c0d8cc4..15c6043f3d0 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/README.md @@ -1,192 +1,29 @@ -# MobileSAM Export, Quantization, and Debugging +# MobileSAM Export -This directory exports MobileSAM `vit_t` for the ExecuTorch Ethos-U backend. -The exporter freezes one or more positive point prompts into the graph, so the -runtime app has one tensor input: the preprocessed image. The model returns one -mask-logit tensor for those prompts. The default export uses a `448x448` input, -which produces a `112x112` mask-logit tensor. +The exporter keeps one tested model configuration so the ExecuTorch steps are +visible without a layer of command-line configuration. -Production SAM applications usually pass prompts dynamically from a UI or -tracking pipeline. This example freezes the prompt embeddings so the FVP -runtime stays small and demonstrates the Ethos-U flow with the same image -encoder and mask decoder used by MobileSAM. - -## Requirements - -- Python 3.10+ with `executorch`. -- Dependencies from - `examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt`. -- Git and internet access to prepare the pinned external MobileSAM checkout and - download the official checkpoint, unless both are already cached. -- Ethos-U dependencies from `examples/arm/setup.sh`. - -MobileSAM's A16W8 attention requires `ethos-u-vela>=5.1.0`. Vela 5.0 produces -incorrect Ethos-U85 INT16 reductions, so the exporter rejects that version -before generating a `.pte`. - -## Export - -Run from the ExecuTorch repo root: - -```bash -MOBILE_SAM_SOURCE="$HOME/.cache/executorch/mobilesam/f706ad9c4eb7f219c00d9050e46328518ffb65d2/source" -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py \ - --source-dir "$MOBILE_SAM_SOURCE" - -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py \ - --output-path ./mobilesam_point_ethos_u85_448.pte \ - --mobile-sam-source "$MOBILE_SAM_SOURCE" \ - --calibration-image examples/models/dinov2/dog.jpg \ - --eval-image examples/models/dinov2/dog.jpg \ - --point 219 193 \ - --artifact-dir ./mobilesam_point_artifacts \ - --debug-output-dir ./mobilesam_point_debug -``` - -The default configuration is: - -- Model source: `https://github.com/ChaoningZhang/MobileSAM` -- MobileSAM source revision: `f706ad9c4eb7f219c00d9050e46328518ffb65d2` -- External source patch: `0001-Make-TinyViT-image-size-configurable.patch` -- Checkpoint URL: - `https://github.com/ChaoningZhang/MobileSAM/raw/f706ad9c4eb7f219c00d9050e46328518ffb65d2/weights/mobile_sam.pt` -- Checkpoint SHA256: - `6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f` -- MobileSAM source-code license: Apache-2.0 -- Static input shape: `[1, 3, 448, 448]` -- Output mask logits shape: `[1, 1, 112, 112]` -- Positive point prompt: `(224, 224)` in the padded `448x448` model input - frame -- Target: `ethos-u85-256` -- System config: `Ethos_U85_SYS_DRAM_Mid` -- Memory mode: `Dedicated_Sram_384KB` -- Calibration samples: `4` -- Validation samples: `4` - -The MobileSAM source and checkpoint are not redistributed by this example. The -preparation script clones the pinned official source into a managed external -cache and applies the configurable-input patch. The exporter downloads the -pinned official checkpoint and verifies its SHA256 unless `--checkpoint-path` -is provided. The source repository and checkpoint may have separate terms; the -export metadata records only the source-code license and does not assign a -license to the checkpoint. - -For a quick offline smoke test after caching the official checkpoint and source -checkout: - -```bash -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py \ - --source-dir "$MOBILE_SAM_SOURCE" \ - --local-files-only - -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py \ - --output-path /tmp/mobilesam_smoke.pte \ - --local-files-only \ - --mobile-sam-source "$MOBILE_SAM_SOURCE" \ - --calibration-image examples/models/dinov2/dog.jpg \ - --eval-image examples/models/dinov2/dog.jpg \ - --point 219 193 \ - --num-calibration-samples 1 \ - --num-eval-samples 1 \ - --minimum-fp32-quantized-iou 0.9 -``` - -Repeat `--point X Y` to freeze a multi-point prompt into the graph. Validate -every prompt set because adding positive points can substantially change the -predicted object after quantization: +From the repository root, run: ```bash -python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py \ - --output-path /tmp/mobilesam_multipoint_smoke.pte \ - --local-files-only \ - --mobile-sam-source "$MOBILE_SAM_SOURCE" \ - --calibration-image examples/models/dinov2/dog.jpg \ - --eval-image examples/models/dinov2/dog.jpg \ - --point 166 158 \ - --point 219 193 \ - --point 289 184 \ - --num-calibration-samples 1 \ - --num-eval-samples 1 \ - --debug-output-dir /tmp/mobilesam_multipoint_debug +python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py +python examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py ``` -Pass `--input-size 1024` to reproduce the original MobileSAM resolution. Smaller -inputs export and run faster, but very small inputs such as `224` or `256` -usually produce lower-quality masks because the checkpoint was trained for -`1024x1024` images. - -To validate against a known binary mask, pass one `--eval-mask` per -`--eval-image`. Non-zero mask pixels are treated as foreground. When no -reference mask is provided, validation reports FP32/quantized mask agreement. -Use `--minimum-fp32-quantized-iou` in automated runs to reject an inaccurate -quantized model before lowering. - -The export flow: - -1. Loads the MobileSAM `vit_t` checkpoint through the patched external API. -2. Builds a fixed-prompt wrapper containing the image encoder and mask decoder. -3. Calibrates PT2E quantization with `EthosUQuantizer`. -4. Uses stable softmax decomposition for transformer attention blocks. -5. Lowers the quantized graph with `EthosUPartitioner`. -6. Writes an ExecuTorch `.pte` program. - -The quantization recipe uses int8 activations and int8 weights globally, with -A16W8 quantization for the TinyViT attention modules. Static int8 activation -quantization collapses MobileSAM attention features, while the selective A16W8 -attention path preserves mask quality and still lowers as one Ethos-U delegate. - -## Outputs - -For `--output-path ./mobilesam_point_ethos_u85_448.pte`, the script writes: - -- `mobilesam_point_ethos_u85_448.pte` - Ethos-U-ready ExecuTorch program. -- `mobilesam_point_ethos_u85_448.json` - Export metadata. -- `mobilesam_point_ethos_u85_448_metrics.json` - FP32/quantized mask - agreement and optional reference-mask IoU. -- `mobilesam_point_ethos_u85_448_delegation.txt` - Operator delegation - summary. -- `mobilesam_point_artifacts/` - Optional TOSA/Vela intermediate artifacts. -- `mobilesam_point_debug/` - Optional per-sample masks, overlays, mismatch - heatmaps, and mask summaries. - -## Interpreting the Debug Artifacts - -Each debug sample contains: - -- `input.png` - The resized RGB input used by the exported model. -- `reference_mask.png` - Optional binary reference mask resized to the model - output-mask size. -- `fp32_mask.png` - Host-side FP32 model prediction. -- `fp32_overlay.png` - Colored FP32 prediction blended over the input image. -- `quantized_mask.png` - Host-side PT2E quantized prediction before lowering. -- `quantized_overlay.png` - Colored quantized prediction blended over the input - image. -- `mismatch_heatmap.png` - Green for FP32/quantized agreement and red for - mismatch. -- `mask_summary.json` - Foreground/background pixel counts and FP32/quantized - IoU. +The first script prepares the pinned official MobileSAM source and verified +checkpoint in `~/.cache/executorch/mobilesam`. It applies the included patch +there to support the smaller input size. Neither source nor checkpoint is +copied into the ExecuTorch repository. -The runtime app thresholds mask logits on target, logs a mask hash and -foreground/background counts, and can dump the thresholded mask as RLE chunks. -Host-side debug masks are intentionally generated before lowering so users can -inspect quantization quality without needing target-side image output. +The second script performs the model flow directly: -## Limitations +1. Wrap MobileSAM with the fixed point prompt. +2. Export with `torch.export`. +3. Calibrate and convert with PT2E, using A8W8 generally and A16W8 attention + activations to preserve mask quality. +4. Check the FP32 and quantized masks have at least `0.9` IoU. +5. Lower with `EthosUPartitioner` and require one delegated subgraph. +6. Save `arm_test/mobilesam/export/mobilesam.pte`. -- The example freezes positive point prompts into the exported graph to keep - the target app to a single image input. Changing the prompt requires - re-exporting the `.pte`. -- The export uses `multimask_output=False` and does not include candidate-mask - argmax selection, upsampling, or thresholding in the graph. Keep those steps - in host-side or target-side post-processing when comparing mask quality. -- The runtime app logs a mask hash and foreground/background counts; it does - not render a color image on target. -- The first supported runtime target is Corstone-320/Ethos-U85-256. -- Reduced input sizes require the patch applied by `prepare_mobilesam.py`; the - MobileSAM source remains outside the ExecuTorch checkout. -- The default `448x448` input is the smallest size in the local sweep that - retained at least `0.95` host quantized/FP32 mask IoU on the demo image. - `512x512` retained about `0.975` IoU, while smaller sizes were inconsistent. -- MobileSAM PTQ is sensitive to the calibration set and the mask-logit - threshold. Inspect `*_metrics.json` and the debug overlays before treating - the quantized mask as an accuracy result. +The export directory also contains the input tensor, host masks, validation +metrics, delegation summary, and TOSA/Vela intermediate artifacts. diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py index e7c22518f00..c2fc992bfe7 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/export_mobilesam.py @@ -3,25 +3,16 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from __future__ import annotations - -import argparse -import hashlib -import importlib -import inspect import json -import os import sys -import urllib.error -import urllib.request -from dataclasses import dataclass -from importlib.metadata import PackageNotFoundError, version from pathlib import Path from typing import Any, cast +import executorch.kernels.quantized # noqa: F401 + import numpy as np import torch -import tqdm # type: ignore[import] +import torch.nn.functional as F from executorch.backends.arm.common.pipeline_config import ( ArmPassPipelineConfig, SoftmaxDecompositionConfig, @@ -39,856 +30,165 @@ to_edge_transform_and_lower, ) from executorch.extension.export_util.utils import save_pte_program -from packaging.version import Version from PIL import Image -from torchao.quantization.pt2e.quantize_pt2e import ( # type: ignore[import] - convert_pt2e, - prepare_pt2e, +from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + +ROOT = Path(__file__).resolve().parents[4] +WORK_DIR = ROOT / "arm_test" / "mobilesam" +SOURCE_DIR = ( + Path.home() + / ".cache" + / "executorch" + / "mobilesam" + / "f706ad9c4eb7f219c00d9050e46328518ffb65d2" + / "source" ) - -MOBILE_SAM_SOURCE_URL = "https://github.com/ChaoningZhang/MobileSAM" -MOBILE_SAM_SOURCE_REVISION = "f706ad9c4eb7f219c00d9050e46328518ffb65d2" -MOBILE_SAM_PATCH = "0001-Make-TinyViT-image-size-configurable.patch" -DEFAULT_CHECKPOINT_FILENAME = "mobile_sam.pt" -DEFAULT_CHECKPOINT_URL = ( - f"{MOBILE_SAM_SOURCE_URL}/raw/{MOBILE_SAM_SOURCE_REVISION}/weights/" - f"{DEFAULT_CHECKPOINT_FILENAME}" -) -DEFAULT_CHECKPOINT_SHA256 = ( - "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f" -) -MOBILE_SAM_SOURCE_LICENSE = "Apache-2.0" -MINIMUM_VELA_VERSION = Version("5.1.0") -DEFAULT_INPUT_SIZE = 448 -MOBILE_SAM_INPUT_ALIGNMENT = 16 - - -@dataclass -class PreparedSample: - name: str - image: Image.Image - pixel_values: torch.Tensor - labels: torch.Tensor | None - - -def load_mobile_sam( - checkpoint_path: str, - mobile_sam_source: str | None, - input_size: int, -) -> torch.nn.Module: - mobile_sam_module = import_mobile_sam_module(mobile_sam_source) - builder = mobile_sam_module.sam_model_registry["vit_t"] - if "image_size" not in inspect.signature(builder).parameters: - raise RuntimeError( - "The MobileSAM checkout does not provide configurable image sizes. " - "Run model_export/prepare_mobilesam.py and pass the prepared checkout " - "with --mobile-sam-source." - ) - return builder(checkpoint=checkpoint_path, image_size=input_size).eval() +CHECKPOINT = SOURCE_DIR.parent / "mobile_sam.pt" +IMAGE = ROOT / "examples" / "models" / "dinov2" / "dog.jpg" +POINT = (219.0, 193.0) +INPUT_SIZE = 448 +MINIMUM_IOU = 0.9 class MobileSAMFixedPrompt(torch.nn.Module): - image_encoder: Any - mask_decoder: Any - - def __init__( - self, - sam: torch.nn.Module, - point_prompts: list[tuple[float, float]], - ) -> None: + def __init__(self, sam: torch.nn.Module) -> None: super().__init__() sam = cast(Any, sam) - self.image_encoder = sam.image_encoder self.mask_decoder = sam.mask_decoder with torch.no_grad(): points = ( - torch.tensor([point_prompts], dtype=torch.float32), - torch.ones((1, len(point_prompts)), dtype=torch.int64), - ) - sparse_embeddings, dense_embeddings = sam.prompt_encoder( - points=points, - boxes=None, - masks=None, + torch.tensor([[POINT]], dtype=torch.float32), + torch.ones((1, 1), dtype=torch.int64), ) + sparse, dense = sam.prompt_encoder(points=points, boxes=None, masks=None) image_pe = sam.prompt_encoder.get_dense_pe() - - self.register_buffer("sparse_prompt_embeddings", sparse_embeddings) - self.register_buffer("dense_prompt_embeddings", dense_embeddings) + self.register_buffer("sparse_prompt", sparse) + self.register_buffer("dense_prompt", dense) self.register_buffer("image_pe", image_pe) - def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: - image_embeddings = self.image_encoder(pixel_values) - low_res_masks, _ = self.mask_decoder( - image_embeddings=image_embeddings, + def forward(self, image: torch.Tensor) -> torch.Tensor: + masks, _ = self.mask_decoder( + image_embeddings=self.image_encoder(image), image_pe=self.image_pe, - sparse_prompt_embeddings=self.sparse_prompt_embeddings, - dense_prompt_embeddings=self.dense_prompt_embeddings, + sparse_prompt_embeddings=self.sparse_prompt, + dense_prompt_embeddings=self.dense_prompt, multimask_output=False, ) - return low_res_masks - - -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Export fixed-prompt MobileSAM segmentation for Ethos-U." - ) - parser.add_argument( - "--checkpoint-path", - default=None, - help="Optional local MobileSAM checkpoint path.", - ) - parser.add_argument( - "--mobile-sam-source", - default=None, - help=( - "Optional MobileSAM checkout prepared by prepare_mobilesam.py. " - "Used before importing the mobile_sam package." - ), - ) - parser.add_argument( - "--local-files-only", - action="store_true", - help="Load the checkpoint from the local cache only; do not download.", - ) - parser.add_argument( - "--calibration-image", - action="append", - default=[], - help="Local RGB image used for PTQ calibration. Can be repeated.", - ) - parser.add_argument( - "--eval-image", - action="append", - default=[], - help="Local RGB image used for validation/debugging. Can be repeated.", - ) - parser.add_argument( - "--eval-mask", - action="append", - default=[], - help=( - "Optional binary reference mask used for validation/debugging. " - "Can be repeated and must match --eval-image count when provided." - ), - ) - parser.add_argument( - "--point", - type=float, - nargs=2, - action="append", - default=[], - metavar=("X", "Y"), - help=( - "Positive point prompt in the resized square input frame. " - "Can be repeated for multi-point prompts." - ), - ) - parser.add_argument( - "--mask-threshold", - type=float, - default=0.0, - help="Mask-logit threshold used for metrics and debug masks.", - ) - parser.add_argument( - "--output-path", - type=str, - required=True, - help="Path to save the exported ExecuTorch program.", - ) - parser.add_argument( - "--input-size", - type=int, - default=DEFAULT_INPUT_SIZE, - help="Square MobileSAM input size. Must be divisible by 16.", - ) - parser.add_argument( - "--num-calibration-samples", - type=int, - default=4, - help="Number of local samples used for PTQ calibration.", - ) - parser.add_argument( - "--num-eval-samples", - type=int, - default=4, - help="Number of local samples used for host-side validation.", - ) - parser.add_argument( - "--num-debug-samples", - type=int, - default=4, - help="Number of validation samples written as visual debug artifacts.", - ) - parser.add_argument( - "--minimum-fp32-quantized-iou", - type=float, - default=None, - help="Fail before lowering when host quantized/FP32 mask IoU is lower.", - ) - parser.add_argument( - "--target", - default="ethos-u85-256", - help="Ethos-U target passed to Vela.", - ) - parser.add_argument( - "--system-config", - default="Ethos_U85_SYS_DRAM_Mid", - help="Vela system configuration.", - ) - parser.add_argument( - "--memory-mode", - default="Dedicated_Sram_384KB", - help="Vela memory mode.", - ) - parser.add_argument( - "--extra-vela-flag", - action="append", - default=[], - help="Additional Vela flag. Can be provided multiple times.", - ) - parser.add_argument( - "--artifact-dir", - default=None, - help="Optional directory for intermediate TOSA/Vela artifacts.", - ) - parser.add_argument( - "--debug-output-dir", - default=None, - help="Optional directory for masks, overlays, and validation summaries.", - ) - return parser.parse_args() - - -def write_json(path: Path, payload: dict[str, Any]) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n") - - -def validate_vela_version() -> str: - try: - installed_version = version("ethos-u-vela") - except PackageNotFoundError as error: - raise RuntimeError( - "MobileSAM export requires ethos-u-vela 5.1.0 or newer. " - "Run examples/arm/setup.sh and retry." - ) from error - - if Version(installed_version) < MINIMUM_VELA_VERSION: - raise RuntimeError( - "MobileSAM A16W8 attention requires ethos-u-vela 5.1.0 or newer; " - f"found {installed_version}. Run examples/arm/setup.sh and retry." - ) - return installed_version + return masks -def import_mobile_sam_module(mobile_sam_source: str | None) -> Any: - if mobile_sam_source is not None: - sys.path.insert(0, str(Path(mobile_sam_source).expanduser().resolve())) - try: - return importlib.import_module("mobile_sam") - except ImportError as error: - raise ImportError( - "Could not import the patched mobile_sam package. Run " - "model_export/prepare_mobilesam.py and pass its checkout with " - "--mobile-sam-source." - ) from error +def load_model() -> torch.nn.Module: + if not SOURCE_DIR.exists() or not CHECKPOINT.exists(): + raise RuntimeError("Run prepare_mobilesam.py first.") + sys.path.insert(0, str(SOURCE_DIR)) + from mobile_sam import sam_model_registry # type: ignore[import-not-found] + return sam_model_registry["vit_t"]( + checkpoint=str(CHECKPOINT), image_size=INPUT_SIZE + ).eval() -def find_module_type(module: torch.nn.Module, class_name: str) -> type[torch.nn.Module]: - for child in module.modules(): - if child.__class__.__name__ == class_name: - return child.__class__ - raise ValueError(f"Could not find module type {class_name} in {module.__class__}.") +def prepare_image(sam: torch.nn.Module) -> tuple[Image.Image, torch.Tensor]: + image = Image.open(IMAGE).convert("RGB") + scale = INPUT_SIZE / max(image.size) + resized_size = (round(image.width * scale), round(image.height * scale)) + resized = image.resize(resized_size, Image.Resampling.BILINEAR) + padded = Image.new("RGB", (INPUT_SIZE, INPUT_SIZE)) + padded.paste(resized) -def file_sha256(path: Path) -> str: - digest = hashlib.sha256() - with path.open("rb") as file: - for chunk in iter(lambda: file.read(1024 * 1024), b""): - digest.update(chunk) - return digest.hexdigest() - - -def default_checkpoint_cache_dir() -> Path: - return ( - Path.home() / ".cache" / "executorch" / "mobilesam" / MOBILE_SAM_SOURCE_REVISION + sam = cast(Any, sam) + mean = sam.pixel_mean.detach().cpu().reshape(3).numpy() + std = sam.pixel_std.detach().cpu().reshape(3).numpy() + tensor = torch.from_numpy((np.asarray(resized, dtype=np.float32) - mean) / std) + tensor = tensor.permute(2, 0, 1).unsqueeze(0) + tensor = F.pad( + tensor, (0, INPUT_SIZE - resized.width, 0, INPUT_SIZE - resized.height) ) + return padded, tensor.contiguous() -def verify_checkpoint(path: Path, expected_sha256: str) -> None: - actual_sha256 = file_sha256(path) - if actual_sha256 != expected_sha256: - raise RuntimeError( - f"Checkpoint SHA256 mismatch for {path}: expected {expected_sha256}, " - f"got {actual_sha256}." - ) - - -def download_checkpoint(url: str, path: Path) -> None: - path.parent.mkdir(parents=True, exist_ok=True) - temp_path = path.with_suffix(path.suffix + ".tmp") - try: - with ( - urllib.request.urlopen(url, timeout=60) as response, # nosec B310 - temp_path.open("wb") as file, - ): - for chunk in iter(lambda: response.read(1024 * 1024), b""): - file.write(chunk) - temp_path.replace(path) - except (OSError, urllib.error.URLError) as error: - temp_path.unlink(missing_ok=True) - raise RuntimeError( - f"Failed to download MobileSAM checkpoint from {url}." - ) from error - - -def resolve_checkpoint( - args: argparse.Namespace, -) -> tuple[str, str | None, str | None]: - if args.checkpoint_path is not None: - checkpoint_path = Path(args.checkpoint_path).expanduser().resolve() - if not checkpoint_path.is_file(): - raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}") - return str(checkpoint_path), None, None - - checkpoint_path = default_checkpoint_cache_dir() / DEFAULT_CHECKPOINT_FILENAME - if not checkpoint_path.is_file(): - if args.local_files_only: - raise FileNotFoundError( - f"Checkpoint not found in local cache: {checkpoint_path}" - ) - download_checkpoint(DEFAULT_CHECKPOINT_URL, checkpoint_path) - - verify_checkpoint(checkpoint_path, DEFAULT_CHECKPOINT_SHA256) - return ( - str(checkpoint_path.resolve()), - DEFAULT_CHECKPOINT_URL, - DEFAULT_CHECKPOINT_SHA256, - ) +def mask(logits: torch.Tensor) -> np.ndarray: + return (logits.detach().cpu().squeeze().numpy() > 0).astype(np.uint8) -def preprocess_image( - image: Image.Image, - segmentation_map: Image.Image | None, - *, - name: str, - input_size: int, - pixel_mean: np.ndarray, - pixel_std: np.ndarray, - output_mask_size: tuple[int, int] | None, -) -> PreparedSample: - rgb_image = image.convert("RGB") - width, height = rgb_image.size - scale = input_size / max(height, width) - resized_size = (round(width * scale), round(height * scale)) - resized_image = rgb_image.resize(resized_size, Image.Resampling.BILINEAR) - - padded_image = Image.new("RGB", (input_size, input_size)) - padded_image.paste(resized_image, (0, 0)) - - image_np = np.asarray(resized_image, dtype=np.float32) - image_np = (image_np - pixel_mean) / pixel_std - pixel_values = torch.from_numpy(image_np).permute(2, 0, 1).unsqueeze(0) - pixel_values = torch.nn.functional.pad( - pixel_values, - ( - 0, - input_size - resized_size[0], - 0, - input_size - resized_size[1], - ), - ) - pixel_values = pixel_values.contiguous() - - labels = None - if segmentation_map is not None: - if output_mask_size is None: - raise ValueError("An output mask size is required for reference masks.") - resized_mask = segmentation_map.convert("L").resize( - resized_size, - Image.Resampling.NEAREST, - ) - padded_mask = Image.new("L", (input_size, input_size)) - padded_mask.paste(resized_mask, (0, 0)) - resized_mask = padded_mask.resize(output_mask_size, Image.Resampling.NEAREST) - mask_np = np.asarray(resized_mask, dtype=np.uint8) - labels = torch.from_numpy((mask_np > 0).astype(np.uint8)).to(torch.long) - - return PreparedSample( - name=name, - image=padded_image, - pixel_values=pixel_values, - labels=labels, - ) +def iou(first: np.ndarray, second: np.ndarray) -> float: + intersection = np.logical_and(first, second).sum() + union = np.logical_or(first, second).sum() + return 1.0 if union == 0 else float(intersection / union) -def load_local_samples( - image_paths: list[str], - mask_paths: list[str], - limit: int, - *, - include_labels: bool, - input_size: int, - pixel_mean: np.ndarray, - pixel_std: np.ndarray, - output_mask_size: tuple[int, int] | None = None, -) -> list[PreparedSample]: - if include_labels and len(mask_paths) not in (0, len(image_paths)): - raise ValueError("--eval-mask must be omitted or match --eval-image count.") - - samples: list[PreparedSample] = [] - for index, image_path in enumerate(image_paths): - if len(samples) >= limit: - break - image = Image.open(image_path) - mask = None - if include_labels and len(mask_paths) > 0: - mask = Image.open(mask_paths[index]) - samples.append( - preprocess_image( - image, - mask, - name=Path(image_path).stem or f"local_sample_{index:04d}", - input_size=input_size, - pixel_mean=pixel_mean, - pixel_std=pixel_std, - output_mask_size=output_mask_size, - ) - ) - if len(samples) == 0: - raise ValueError("No local samples were loaded.") - return samples - - -def run_mask_logits(model: torch.nn.Module, input_tensor: torch.Tensor) -> torch.Tensor: - output = model(input_tensor) - if isinstance(output, (tuple, list)): - output = output[0] - if not isinstance(output, torch.Tensor): - raise TypeError(f"Expected tensor mask logits, got {type(output)}") - return output - - -def predict_mask(logits: torch.Tensor, threshold: float) -> np.ndarray: - mask = logits.detach().cpu().squeeze(0).squeeze(0).numpy() > threshold - return mask.astype(np.uint8) - - -def binary_iou(mask_a: np.ndarray, mask_b: np.ndarray) -> float: - intersection = np.logical_and(mask_a, mask_b).sum() - union = np.logical_or(mask_a, mask_b).sum() - if union == 0: - return 1.0 - return float(intersection / union) - - -def save_binary_mask(path: Path, mask: np.ndarray) -> Image.Image: - image = Image.fromarray((mask.astype(np.uint8) * 255), mode="L") - image.save(path) - return image.convert("RGB") - - -def save_mask_overlay( - path: Path, - image: Image.Image, - mask: np.ndarray, - color: tuple[int, int, int], -) -> None: - rgb_image = image.convert("RGB") - resized_mask = Image.fromarray(mask.astype(np.uint8) * 255, mode="L").resize( - rgb_image.size, - Image.Resampling.NEAREST, - ) - mask_np = np.asarray(resized_mask, dtype=np.uint8) > 0 - overlay_np = np.asarray(rgb_image, dtype=np.float32) - color_np = np.asarray(color, dtype=np.float32) - overlay_np[mask_np] = overlay_np[mask_np] * 0.55 + color_np * 0.45 - Image.fromarray(np.clip(overlay_np, 0, 255).astype(np.uint8), mode="RGB").save(path) - - -def write_debug_artifacts( - debug_dir: Path, - sample: PreparedSample, - fp32_mask: np.ndarray, - quantized_mask: np.ndarray, -) -> None: - sample_dir = debug_dir / sample.name - sample_dir.mkdir(parents=True, exist_ok=True) - - sample.image.save(sample_dir / "input.png") - save_binary_mask(sample_dir / "fp32_mask.png", fp32_mask) - save_binary_mask(sample_dir / "quantized_mask.png", quantized_mask) - save_mask_overlay( - sample_dir / "fp32_overlay.png", - sample.image, - fp32_mask, - (0, 220, 120), - ) - save_mask_overlay( - sample_dir / "quantized_overlay.png", - sample.image, - quantized_mask, - (0, 170, 255), - ) - if sample.labels is not None: - save_binary_mask( - sample_dir / "reference_mask.png", - sample.labels.detach().cpu().numpy().astype(np.uint8), - ) - - mismatch = fp32_mask != quantized_mask - heatmap = np.zeros((*fp32_mask.shape, 3), dtype=np.uint8) - heatmap[~mismatch] = [0, 128, 0] - heatmap[mismatch] = [255, 0, 0] - Image.fromarray(heatmap, mode="RGB").save(sample_dir / "mismatch_heatmap.png") - - write_json( - sample_dir / "mask_summary.json", - { - "foreground_pixels": int(quantized_mask.sum()), - "background_pixels": int(quantized_mask.size - quantized_mask.sum()), - "fp32_quantized_iou": binary_iou(fp32_mask, quantized_mask), - }, - ) - - -def evaluate_and_debug( - fp32_model: torch.nn.Module, - quantized_model: torch.nn.Module, - eval_samples: list[PreparedSample], - debug_dir: Path | None, - num_debug_samples: int, - threshold: float, -) -> dict[str, float]: - fp32_quantized_ious: list[float] = [] - fp32_quantized_pixel_agreements: list[float] = [] - reference_ious: list[float] = [] - - if debug_dir is not None: - debug_dir.mkdir(parents=True, exist_ok=True) - - print("\nEvaluating quantized MobileSAM on validation samples...") - for index, sample in enumerate(tqdm.tqdm(eval_samples)): - fp32_logits = run_mask_logits(fp32_model, sample.pixel_values) - quantized_logits = run_mask_logits(quantized_model, sample.pixel_values) - fp32_mask = predict_mask(fp32_logits, threshold) - quantized_mask = predict_mask(quantized_logits, threshold) - - fp32_quantized_ious.append(binary_iou(fp32_mask, quantized_mask)) - fp32_quantized_pixel_agreements.append( - float(np.mean(fp32_mask == quantized_mask)) - ) - if sample.labels is not None: - labels = sample.labels.detach().cpu().numpy().astype(np.uint8) - reference_ious.append(binary_iou(quantized_mask, labels)) - - if debug_dir is not None and index < num_debug_samples: - write_debug_artifacts(debug_dir, sample, fp32_mask, quantized_mask) - - metrics = { - "num_samples": float(len(eval_samples)), - "fp32_quantized_mean_iou": float(np.mean(fp32_quantized_ious)), - "fp32_quantized_pixel_agreement": float( - np.mean(fp32_quantized_pixel_agreements) - ), - } - if len(reference_ious) > 0: - metrics["reference_mean_iou"] = float(np.mean(reference_ious)) - return metrics - - -def quantize_model( - model: torch.nn.Module, - quantizer: EthosUQuantizer, - example_inputs: tuple[torch.Tensor], - calibration_samples: list[PreparedSample], +def quantize( + model: torch.nn.Module, image: torch.Tensor, quantizer: EthosUQuantizer ) -> torch.export.ExportedProgram: - exported = torch.export.export(model, example_inputs) + exported = torch.export.export(model, (image,)) prepared = prepare_pt2e(exported.module(), quantizer) - - print("\nCalibrating MobileSAM...") - for sample in tqdm.tqdm(calibration_samples): - prepared(sample.pixel_values) - - quantized = convert_pt2e(prepared) - return torch.export.export(quantized, example_inputs) - - -def has_quantized_out_variants() -> bool: - try: - _ = torch.ops.quantized_decomposed.quantize_per_tensor.out - _ = torch.ops.quantized_decomposed.dequantize_per_tensor.out - return True - except AttributeError: - return False - - -def load_quantized_ops_library(library_path: Path) -> Path: - if not library_path.is_file(): - raise FileNotFoundError(f"Quantized ops library not found: {library_path}") - torch.ops.load_library(str(library_path)) - if has_quantized_out_variants(): - return library_path - raise RuntimeError( - f"Quantized ops library did not register required out variants: {library_path}" - ) - - -def ensure_quantized_ops_loaded() -> Path | None: - if has_quantized_out_variants(): - return None - - quantized_ops_library = os.environ.get("EXECUTORCH_QUANTIZED_OPS_AOT_LIBRARY") - if quantized_ops_library: - return load_quantized_ops_library( - Path(quantized_ops_library).expanduser().resolve() - ) - - try: - import executorch.kernels.quantized # noqa: F401 - except ImportError: - pass - else: - if has_quantized_out_variants(): - return None - - repo_root = Path(__file__).resolve().parents[4] - search_patterns = ( - "cmake-out/kernels/quantized/libquantized_ops_aot_lib.*", - "arm_test/*/kernels/quantized/libquantized_ops_aot_lib.*", - "arm_test/**/kernels/quantized/libquantized_ops_aot_lib.*", - ) - for pattern in search_patterns: - for candidate in sorted(repo_root.glob(pattern)): - if not candidate.is_file(): - continue - return load_quantized_ops_library(candidate) - - raise RuntimeError( - "MobileSAM int8 export requires the quantized ops out-variant library. " - "Build or install ExecuTorch quantized kernels so that " - "`quantized_decomposed::quantize_per_tensor.out` and " - "`quantized_decomposed::dequantize_per_tensor.out` are available." - ) - - -def write_delegation_report(edge_program_manager: Any, report_path: Path) -> None: - delegation_info = get_delegation_info( - edge_program_manager.exported_program().graph_module - ) - report_path.write_text(delegation_info.get_summary() + "\n") - - -def resolve_point_prompts(args: argparse.Namespace) -> list[tuple[float, float]]: - if len(args.point) > 0: - point_prompts = [(float(x), float(y)) for x, y in args.point] - else: - point_prompts = [(args.input_size / 2, args.input_size / 2)] - - for point_x, point_y in point_prompts: - if not (0 <= point_x <= args.input_size and 0 <= point_y <= args.input_size): - raise ValueError("Point prompts must be inside the square input.") - return point_prompts - - -def validate_export_args(args: argparse.Namespace) -> None: - if args.input_size < 224 or args.input_size % MOBILE_SAM_INPUT_ALIGNMENT != 0: - raise ValueError("--input-size must be at least 224 and divisible by 16.") - if args.num_calibration_samples <= 0: - raise ValueError("--num-calibration-samples must be positive.") - if args.num_eval_samples <= 0: - raise ValueError("--num-eval-samples must be positive.") - if args.minimum_fp32_quantized_iou is not None and not ( - 0.0 <= args.minimum_fp32_quantized_iou <= 1.0 - ): - raise ValueError("--minimum-fp32-quantized-iou must be between 0 and 1.") - if len(args.calibration_image) == 0: - raise ValueError("At least one --calibration-image is required.") - if len(args.eval_image) == 0: - args.eval_image = list(args.calibration_image) - if len(args.eval_mask) not in (0, len(args.eval_image)): - raise ValueError("--eval-mask must be omitted or match --eval-image count.") + prepared(image) + return torch.export.export(convert_pt2e(prepared), (image,)) def main() -> None: - args = parse_args() - validate_export_args(args) - vela_version = validate_vela_version() - point_prompts = resolve_point_prompts(args) - quantized_ops_library = ensure_quantized_ops_loaded() - if quantized_ops_library is not None: - print(f"Loaded quantized ops library from {quantized_ops_library}") - - output_path = Path(args.output_path).resolve() - output_path.parent.mkdir(parents=True, exist_ok=True) - metadata_path = output_path.with_suffix(".json") - metrics_path = output_path.with_name(f"{output_path.stem}_metrics.json") - delegation_path = output_path.with_name(f"{output_path.stem}_delegation.txt") - debug_dir = Path(args.debug_output_dir).resolve() if args.debug_output_dir else None - - checkpoint_path, checkpoint_url, checkpoint_sha256 = resolve_checkpoint(args) - mobile_sam = load_mobile_sam( - checkpoint_path, - args.mobile_sam_source, - args.input_size, - ) - pixel_mean = cast(Any, mobile_sam).pixel_mean.detach().cpu().reshape(-1).numpy() - pixel_std = cast(Any, mobile_sam).pixel_std.detach().cpu().reshape(-1).numpy() - if pixel_mean.shape != (3,) or pixel_std.shape != (3,): - raise ValueError("MobileSAM preprocessing must provide three RGB values.") - wrapped_model = MobileSAMFixedPrompt(mobile_sam, point_prompts).eval() - - calibration_samples = load_local_samples( - args.calibration_image, - [], - args.num_calibration_samples, - include_labels=False, - input_size=args.input_size, - pixel_mean=pixel_mean, - pixel_std=pixel_std, - ) - example_inputs = (calibration_samples[0].pixel_values,) - with torch.no_grad(): - output_shape = list(run_mask_logits(wrapped_model, example_inputs[0]).shape) - if len(output_shape) != 4 or output_shape[:2] != [1, 1]: - raise ValueError( - f"Expected MobileSAM output shape [1, 1, height, width], got {output_shape}." - ) - output_mask_size = (output_shape[3], output_shape[2]) - - eval_samples = load_local_samples( - args.eval_image, - args.eval_mask, - args.num_eval_samples, - include_labels=True, - input_size=args.input_size, - pixel_mean=pixel_mean, - pixel_std=pixel_std, - output_mask_size=output_mask_size, - ) + export_dir = WORK_DIR / "export" + export_dir.mkdir(parents=True, exist_ok=True) + + sam = load_model() + input_image, example_input = prepare_image(sam) + model = MobileSAMFixedPrompt(sam).eval() compile_spec = EthosUCompileSpec( - target=args.target, - system_config=args.system_config, - memory_mode=args.memory_mode, - extra_flags=args.extra_vela_flag, + "ethos-u85-256", memory_mode="Dedicated_Sram_384KB" ) compile_spec.set_pass_pipeline_config( ArmPassPipelineConfig(softmax=SoftmaxDecompositionConfig.STABLE) ) - if args.artifact_dir is not None: - artifact_dir = Path(args.artifact_dir).resolve() - artifact_dir.mkdir(parents=True, exist_ok=True) - compile_spec.dump_intermediate_artifacts_to(str(artifact_dir)) + compile_spec.dump_intermediate_artifacts_to(str(export_dir / "artifacts")) quantizer = EthosUQuantizer(compile_spec) quantizer.set_global(get_symmetric_quantization_config()) - attention_module_type = find_module_type(wrapped_model.image_encoder, "Attention") - quantizer.set_module_type( - attention_module_type, - get_symmetric_a16w8_quantization_config(), + attention_type = next( + type(module) + for module in model.image_encoder.modules() + if type(module).__name__ == "Attention" ) + # Int16 attention activations preserve the segmentation mask quality. + quantizer.set_module_type(attention_type, get_symmetric_a16w8_quantization_config()) with torch.no_grad(): - quantized_program = quantize_model( - wrapped_model, - quantizer, - example_inputs, - calibration_samples, - ) - quantized_module = quantized_program.module() - metrics = evaluate_and_debug( - wrapped_model, - quantized_module, - eval_samples, - debug_dir, - args.num_debug_samples, - args.mask_threshold, - ) - write_json(metrics_path, metrics) - print( - "Validation metrics: " - f"fp32_quantized_mean_iou={metrics['fp32_quantized_mean_iou']:.4f} " - "fp32_quantized_pixel_agreement=" - f"{metrics['fp32_quantized_pixel_agreement']:.4f}" - ) - if ( - args.minimum_fp32_quantized_iou is not None - and metrics["fp32_quantized_mean_iou"] < args.minimum_fp32_quantized_iou - ): - raise RuntimeError( - "Host quantized/FP32 mask IoU " - f"{metrics['fp32_quantized_mean_iou']:.4f} is below " - f"{args.minimum_fp32_quantized_iou:.4f}." - ) + fp32_mask = mask(model(example_input)) + quantized = quantize(model, example_input, quantizer) + quantized_mask = mask(quantized.module()(example_input)) + + host_iou = iou(fp32_mask, quantized_mask) + if host_iou < MINIMUM_IOU: + raise RuntimeError(f"FP32/quantized mask IoU is too low: {host_iou:.4f}") - edge_program_manager = to_edge_transform_and_lower( - programs=quantized_program, + edge = to_edge_transform_and_lower( + quantized, partitioner=[EthosUPartitioner(compile_spec)], compile_config=EdgeCompileConfig(_check_ir_validity=False), ) - write_delegation_report(edge_program_manager, delegation_path) + delegation = get_delegation_info(edge.exported_program().graph_module) + if delegation.num_delegated_subgraphs != 1: + raise RuntimeError("Expected one Ethos-U delegate.") - executorch_program_manager = edge_program_manager.to_executorch( + program = edge.to_executorch( config=ExecutorchBackendConfig(extract_delegate_segments=False) ) - save_pte_program( - executorch_program_manager, - str(output_path), - output_dir=str(output_path.parent), - ) + save_pte_program(program, "mobilesam", output_dir=str(export_dir)) - write_json( - metadata_path, - { - "model_name": "MobileSAM vit_t", - "checkpoint_filename": DEFAULT_CHECKPOINT_FILENAME, - "checkpoint_path": checkpoint_path, - "checkpoint_url": checkpoint_url, - "checkpoint_sha256": checkpoint_sha256, - "mobile_sam_source_license": MOBILE_SAM_SOURCE_LICENSE, - "mobile_sam_source_url": MOBILE_SAM_SOURCE_URL, - "mobile_sam_source_revision": MOBILE_SAM_SOURCE_REVISION, - "mobile_sam_patch": MOBILE_SAM_PATCH, - "input_shape": list(example_inputs[0].shape), - "output_shape": output_shape, - "input_size": args.input_size, - "preprocessing": { - "pixel_mean": pixel_mean.tolist(), - "pixel_std": pixel_std.tolist(), - "resize": "longest_side_then_zero_pad", - }, - "point_prompts_xy": point_prompts, - "mask_threshold": args.mask_threshold, - "target": args.target, - "vela_version": vela_version, - "system_config": args.system_config, - "memory_mode": args.memory_mode, - "extra_vela_flags": args.extra_vela_flag, - "quantization": { - "global": "int8 activations and int8 weights", - "tinyvit_attention": "int16 activations and int8 weights", - }, - "num_calibration_samples": len(calibration_samples), - "num_eval_samples": len(eval_samples), - "calibration_images": args.calibration_image, - "eval_images": args.eval_image, - "eval_masks": args.eval_mask, - "output_path": str(output_path), - "metrics_path": str(metrics_path), - "delegation_path": str(delegation_path), - "debug_output_dir": str(debug_dir) if debug_dir is not None else None, - }, + input_image.save(export_dir / "input.png") + Image.fromarray(fp32_mask * 255).save(export_dir / "fp32_mask.png") + Image.fromarray(quantized_mask * 255).save(export_dir / "quantized_mask.png") + example_input.numpy().astype(np.float32).tofile(export_dir / "input.bin") + (export_dir / "delegation.txt").write_text(delegation.get_summary() + "\n") + (export_dir / "metrics.json").write_text( + json.dumps({"fp32_quantized_iou": host_iou}, indent=2) + "\n" ) - print(f"\nExported model saved to {output_path}") - print(f"Metadata saved to {metadata_path}") - print(f"Metrics saved to {metrics_path}") - print(f"Delegation summary saved to {delegation_path}") - if debug_dir is not None: - print(f"Debug artifacts saved to {debug_dir}") + print(f"FP32/quantized mask IoU: {host_iou:.4f}") + print(f"Saved {export_dir / 'mobilesam.pte'}") if __name__ == "__main__": diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py index 73dd6119c39..bcb922ca052 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/model_export/prepare_mobilesam.py @@ -3,119 +3,86 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -from __future__ import annotations - -import argparse +import hashlib import subprocess # nosec B404 +import urllib.request from pathlib import Path -MOBILE_SAM_SOURCE_URL = "https://github.com/ChaoningZhang/MobileSAM.git" -MOBILE_SAM_SOURCE_REVISION = "f706ad9c4eb7f219c00d9050e46328518ffb65d2" -PATCH_DIR = Path(__file__).resolve().parent / "patches" / "mobile_sam" - - -def run(command: list[str], *, cwd: Path | None = None) -> None: +REVISION = "f706ad9c4eb7f219c00d9050e46328518ffb65d2" +SOURCE_URL = "https://github.com/ChaoningZhang/MobileSAM.git" +CHECKPOINT_URL = ( + f"https://github.com/ChaoningZhang/MobileSAM/raw/{REVISION}/weights/mobile_sam.pt" +) +CHECKPOINT_SHA256 = "6dbb90523a35330fedd7f1d3dfc66f995213d81b29a5ca8108dbcdd4e37d6c2f" +CACHE_DIR = Path.home() / ".cache" / "executorch" / "mobilesam" / REVISION +SOURCE_DIR = CACHE_DIR / "source" +CHECKPOINT = CACHE_DIR / "mobile_sam.pt" +PATCH = ( + Path(__file__).parent + / "patches" + / "mobile_sam" + / ("0001-Make-TinyViT-image-size-configurable.patch") +) + + +def run(*command: str, cwd: Path | None = None) -> None: subprocess.run(command, cwd=cwd, check=True) # nosec B603 -def default_source_dir() -> Path: - return ( - Path.home() - / ".cache" - / "executorch" - / "mobilesam" - / MOBILE_SAM_SOURCE_REVISION - / "source" - ) - +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as file: + for chunk in iter(lambda: file.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() -def prepare_source(source_dir: Path, *, local_files_only: bool) -> None: - source_dir = source_dir.expanduser().resolve() - marker = source_dir.parent / f".{source_dir.name}.executorch-managed" - if source_dir.exists() and not marker.is_file(): - raise RuntimeError( - f"Refusing to modify unmanaged MobileSAM directory: {source_dir}" - ) +def prepare_source() -> None: + marker = SOURCE_DIR.parent / ".source.executorch-managed" + if SOURCE_DIR.exists() and not marker.exists(): + raise RuntimeError(f"Refusing to modify unmanaged directory: {SOURCE_DIR}") - if not source_dir.exists(): - if local_files_only: - raise FileNotFoundError( - f"Managed MobileSAM checkout not found: {source_dir}" - ) - source_dir.parent.mkdir(parents=True, exist_ok=True) - run( - [ - "git", - "clone", - "--filter=blob:none", - "--no-checkout", - MOBILE_SAM_SOURCE_URL, - str(source_dir), - ] - ) - marker.write_text(MOBILE_SAM_SOURCE_REVISION + "\n") + if not SOURCE_DIR.exists(): + SOURCE_DIR.parent.mkdir(parents=True, exist_ok=True) run( - [ - "git", - "sparse-checkout", - "set", - "mobile_sam", - ], - cwd=source_dir, + "git", + "clone", + "--filter=blob:none", + "--no-checkout", + SOURCE_URL, + str(SOURCE_DIR), ) - - if not local_files_only: - run( - ["git", "fetch", "--quiet", "origin", MOBILE_SAM_SOURCE_REVISION], - cwd=source_dir, - ) - - try: - run( - ["git", "cat-file", "-e", f"{MOBILE_SAM_SOURCE_REVISION}^{{commit}}"], - cwd=source_dir, - ) - except subprocess.CalledProcessError as error: + marker.write_text(REVISION + "\n") + run("git", "sparse-checkout", "set", "mobile_sam", cwd=SOURCE_DIR) + + run("git", "fetch", "--quiet", "origin", REVISION, cwd=SOURCE_DIR) + run("git", "checkout", "--detach", "--force", REVISION, cwd=SOURCE_DIR) + run("git", "reset", "--hard", REVISION, cwd=SOURCE_DIR) + run("git", "apply", str(PATCH), cwd=SOURCE_DIR) + + +def prepare_checkpoint() -> None: + if not CHECKPOINT.exists(): + CHECKPOINT.parent.mkdir(parents=True, exist_ok=True) + with ( + urllib.request.urlopen( + CHECKPOINT_URL, timeout=60 + ) as response, # nosec B310 + CHECKPOINT.open("wb") as file, + ): + while chunk := response.read(1024 * 1024): + file.write(chunk) + + actual_sha256 = sha256(CHECKPOINT) + if actual_sha256 != CHECKPOINT_SHA256: raise RuntimeError( - f"MobileSAM revision {MOBILE_SAM_SOURCE_REVISION} is unavailable locally." - ) from error - - run( - ["git", "checkout", "--detach", "--force", MOBILE_SAM_SOURCE_REVISION], - cwd=source_dir, - ) - run(["git", "reset", "--hard", MOBILE_SAM_SOURCE_REVISION], cwd=source_dir) - - patches = sorted(PATCH_DIR.glob("*.patch")) - if not patches: - raise FileNotFoundError(f"No MobileSAM patches found in {PATCH_DIR}") - for patch in patches: - run(["git", "apply", "--check", str(patch)], cwd=source_dir) - run(["git", "apply", str(patch)], cwd=source_dir) - - print(f"Prepared patched MobileSAM source at {source_dir}") - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Prepare the pinned MobileSAM source with ExecuTorch patches." - ) - parser.add_argument( - "--source-dir", - type=Path, - default=default_source_dir(), - help="Managed checkout destination.", - ) - parser.add_argument( - "--local-files-only", - action="store_true", - help="Reuse an existing managed checkout without network access.", - ) - args = parser.parse_args() - prepare_source(args.source_dir, local_files_only=args.local_files_only) + f"Checkpoint SHA256 mismatch: expected {CHECKPOINT_SHA256}, " + f"got {actual_sha256}" + ) if __name__ == "__main__": - main() + prepare_source() + prepare_checkpoint() + print(f"MobileSAM ready in {CACHE_DIR}") diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt deleted file mode 100644 index 6eb309f00de..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ -# Copyright 2026 Arm Limited and/or its affiliates. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -tqdm == 4.67.1 diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh new file mode 100755 index 00000000000..511284d2ba7 --- /dev/null +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/run.sh @@ -0,0 +1,60 @@ +#!/usr/bin/env bash +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +set -euo pipefail + +if (($#)); then + echo "This example has one supported configuration; run it without arguments." + exit 2 +fi + +example_dir=$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd) +repo_root=$(cd "${example_dir}/../../.." && pwd) +work_dir="${repo_root}/arm_test/mobilesam" +export_dir="${work_dir}/export" +io_dir="${work_dir}/io" +runner_dir="${work_dir}/runner" + +cd "${repo_root}" +[[ -f examples/arm/arm-scratch/setup_path.sh ]] || { + echo "Arm tools are missing. Run ./examples/arm/setup.sh first." + exit 1 +} +source examples/arm/arm-scratch/setup_path.sh +mkdir -p "${io_dir}" + +if [[ "$(uname -s)" == "Darwin" ]]; then + export FVP_MOUNT_DIR="${FVP_MOUNT_DIR:-${repo_root}}" + export FVP_WORKDIR="${FVP_WORKDIR:-${repo_root}}" +fi + +echo "[1/4] Prepare MobileSAM" +python3 "${example_dir}/model_export/prepare_mobilesam.py" + +echo "[2/4] Export, quantize, and lower to Ethos-U" +python3 "${example_dir}/model_export/export_mobilesam.py" + +echo "[3/4] Build the standard Arm executor runner" +backends/arm/scripts/build_executor_runner.sh \ + --pte="${export_dir}/mobilesam.pte" \ + --target=ethos-u85-256 \ + --output="${runner_dir}" \ + '--extra_build_flags=-DSEMIHOSTING=ON -DET_COMPILED_PTE=ON' + +cp "${export_dir}/input.bin" "${io_dir}/input.bin" +rm -f "${io_dir}/output-0.bin" + +echo "[4/4] Run on FVP and validate the output" +backends/arm/scripts/run_fvp.sh \ + --elf="${runner_dir}/arm_executor_runner" \ + --target=ethos-u85-256 \ + --timeout=300 \ + --semihosting-cwd="${io_dir}" \ + '--semihosting-cmd-line=executor_runner -i input.bin -o output' \ + --fast | tee "${work_dir}/fvp.log" +python3 "${example_dir}/runtime/visualize_fvp_output.py" + +echo "MobileSAM example: PASS" diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt deleted file mode 100644 index ce6786f6a58..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/CMakeLists.txt +++ /dev/null @@ -1,225 +0,0 @@ -# Copyright 2026 Arm Limited and/or its affiliates. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. - -cmake_minimum_required(VERSION 3.20) - -project(mobilesam_prompt_segmentation_ethos_u_application) - -set(ET_DIR_PATH - "${CMAKE_CURRENT_SOURCE_DIR}/../../../.." - CACHE PATH "Path to ExecuTorch dir" -) -set(ET_BUILD_DIR_PATH - "${ET_DIR_PATH}/cmake-out-arm" - CACHE PATH "Path to ExecuTorch build/install dir" -) -set(ET_INCLUDE_PATH - "${ET_DIR_PATH}/.." - CACHE PATH "Path to ExecuTorch headers" -) -set(ET_PTE_FILE_PATH - "" - CACHE PATH "Path to ExecuTorch model pte" -) -set(MODEL_METADATA_PATH - "" - CACHE PATH "Path to MobileSAM exporter metadata" -) -set(IMAGE_PATH - "" - CACHE PATH "Path to an RGB image to use for the application" -) -set(MASK_THRESHOLD - "0.0" - CACHE STRING "MobileSAM mask logit threshold" -) -set(SYSTEM_CONFIG - "Ethos_U85_SYS_DRAM_Mid" - CACHE STRING "Vela system configuration" -) -set(MEMORY_MODE - "Dedicated_Sram_384KB" - CACHE STRING "Vela memory mode" -) -set(ETHOS_SDK_PATH - "${ET_DIR_PATH}/examples/arm/arm-scratch/ethos-u" - CACHE PATH "Path to Ethos-U bare metal driver/env" -) -set(PYTHON_EXECUTABLE - "python" - CACHE PATH "Define to override python executable used" -) -option(ET_SEGMENTATION_DUMP_MASK - "Dump the predicted segmentation mask as run-length encoded log chunks" - OFF -) -option(ET_SEGMENTATION_SEMIHOSTING_OUTPUT - "Emit validation logs through Arm semihosting for FVP smoke runs" ON -) - -if(NOT EXISTS "${IMAGE_PATH}") - message( - FATAL_ERROR - "Image not provided. Please provide -DIMAGE_PATH= and retry." - ) -endif() -if(NOT EXISTS "${ET_PTE_FILE_PATH}") - message( - FATAL_ERROR - "PTE file not provided. Please provide -DET_PTE_FILE_PATH= and retry." - ) -endif() -if(NOT MODEL_METADATA_PATH) - get_filename_component(model_directory "${ET_PTE_FILE_PATH}" DIRECTORY) - get_filename_component(model_name "${ET_PTE_FILE_PATH}" NAME_WE) - set(MODEL_METADATA_PATH "${model_directory}/${model_name}.json") -endif() -if(NOT EXISTS "${MODEL_METADATA_PATH}") - message( - FATAL_ERROR - "Model metadata not found. Provide -DMODEL_METADATA_PATH=." - ) -endif() -if(NOT SYSTEM_CONFIG MATCHES "Ethos_U85") - message(FATAL_ERROR "This example currently supports Corstone-320/Ethos-U85.") -endif() - -include(${ET_DIR_PATH}/backends/arm/scripts/corstone_utils.cmake) -fetch_ethos_u_content(${ETHOS_SDK_PATH} ${ET_DIR_PATH}) - -if(NOT EXISTS "${ETHOS_SDK_PATH}") - message( - FATAL_ERROR - "The ${ETHOS_SDK_PATH} directory does not exist. Please run examples/arm/setup.sh and retry." - ) -endif() - -find_package( - executorch REQUIRED HINTS "${ET_BUILD_DIR_PATH}/lib/cmake/ExecuTorch" -) - -add_corstone_subdirectory( - "${SYSTEM_CONFIG}" "${ETHOS_SDK_PATH}" "${MEMORY_MODE}" -) -configure_timing_adapters("${SYSTEM_CONFIG}" "${MEMORY_MODE}") - -add_executable(mobilesam_prompt_segmentation_example main.cpp) -target_sources( - mobilesam_prompt_segmentation_example - PRIVATE main.cpp ${ET_DIR_PATH}/examples/arm/common/arm_memory_allocator.cpp -) -target_link_libraries( - mobilesam_prompt_segmentation_example - PUBLIC executorch - ethosu_target_init - extension_runner_util - quantized_ops_lib - portable_kernels - cortex_m_kernels - cortex_m_ops_lib -) - -include(${ET_DIR_PATH}/tools/cmake/Utils.cmake) -executorch_target_link_options_shared_lib(executorch_delegate_ethos_u) -target_link_libraries( - mobilesam_prompt_segmentation_example PUBLIC executorch_delegate_ethos_u -) - -if(MEMORY_MODE MATCHES "^Dedicated_Sram($|_)") - set(ETHOSU_ARENA "1") - if(NOT DEFINED ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - set(ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x1000000) - endif() - if(NOT DEFINED ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - set_ethosu_dedicated_sram_fast_scratch_size( - ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE "${MEMORY_MODE}" - ) - endif() -else() - set(ETHOSU_ARENA "0") - if(NOT DEFINED ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - set(ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE 0x400000) - endif() -endif() -if(NOT DEFINED ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE) - set(ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE 0x4000000) -endif() - -target_compile_definitions( - mobilesam_prompt_segmentation_example - PRIVATE - ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE=${ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE} - ET_SEGMENTATION_MASK_THRESHOLD=${MASK_THRESHOLD} - ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE=${ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE} -) -if(DEFINED ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) - target_compile_definitions( - mobilesam_prompt_segmentation_example - PRIVATE - ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE=${ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE} - ) -endif() -if(ET_SEGMENTATION_DUMP_MASK) - target_compile_definitions( - mobilesam_prompt_segmentation_example PRIVATE ET_SEGMENTATION_DUMP_MASK - ) -endif() -if(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) - target_compile_definitions( - mobilesam_prompt_segmentation_example - PRIVATE ET_SEGMENTATION_SEMIHOSTING_OUTPUT - ) -endif() - -if(CMAKE_CXX_COMPILER_ID STREQUAL "GNU") - set(LINK_FILE_EXT ld) - set(COMPILER_PREPROCESSOR_OPTIONS -E -x c -P) -endif() - -set(LINK_FILE_OUT_BASE "platform_linker_script") -set(LINK_FILE_IN - "${ET_DIR_PATH}/backends/arm/cmake/linker_scripts/Corstone-320.ld" -) -set(LINK_FILE_OUT - ${CMAKE_CURRENT_BINARY_DIR}/${LINK_FILE_OUT_BASE}.${LINK_FILE_EXT} -) -execute_process( - COMMAND ${CMAKE_C_COMPILER} ${COMPILER_PREPROCESSOR_OPTIONS} -DETHOSU_MODEL=1 - -DETHOSU_ARENA=${ETHOSU_ARENA} -o ${LINK_FILE_OUT} ${LINK_FILE_IN} -) -target_link_options( - mobilesam_prompt_segmentation_example PRIVATE "-T" "${LINK_FILE_OUT}" -) - -set(MODEL_PTE_HEADER "${CMAKE_CURRENT_BINARY_DIR}/model_pte.h") -add_custom_command( - OUTPUT "${MODEL_PTE_HEADER}" - COMMAND - ${PYTHON_EXECUTABLE} ${ET_DIR_PATH}/examples/arm/common/pte_to_header.py - --pte ${ET_PTE_FILE_PATH} --outdir ${CMAKE_CURRENT_BINARY_DIR} - DEPENDS ${ET_PTE_FILE_PATH} - ${ET_DIR_PATH}/examples/arm/common/pte_to_header.py - VERBATIM -) -set(IMAGE_HEADER "${CMAKE_CURRENT_BINARY_DIR}/image.h") -add_custom_command( - OUTPUT "${IMAGE_HEADER}" - COMMAND - ${PYTHON_EXECUTABLE} ${CMAKE_SOURCE_DIR}/image_to_array.py --image - ${IMAGE_PATH} --metadata ${MODEL_METADATA_PATH} --output ${IMAGE_HEADER} - DEPENDS ${IMAGE_PATH} ${MODEL_METADATA_PATH} - ${CMAKE_SOURCE_DIR}/image_to_array.py - VERBATIM -) -target_sources( - mobilesam_prompt_segmentation_example PRIVATE ${MODEL_PTE_HEADER} - ${IMAGE_HEADER} -) - -target_include_directories( - mobilesam_prompt_segmentation_example - PRIVATE ${ET_INCLUDE_PATH} ${ET_DIR_PATH}/runtime/core/portable_type/c10 - ${ET_DIR_PATH}/examples/arm/common ${CMAKE_CURRENT_BINARY_DIR} -) diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md index f422edc6ad9..5473faa16a2 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/README.md @@ -1,110 +1,16 @@ -# MobileSAM Runtime Example +# MobileSAM Runtime -This directory builds a bare-metal Corstone-320 application for a MobileSAM -fixed-prompt `.pte` generated by `model_export/export_mobilesam.py`. +MobileSAM uses the standard ExecuTorch Arm executor runner. The example has no +model-specific C++ runtime or CMake project. -## Build ExecuTorch for Arm +The top-level [`run.sh`](../run.sh) embeds the exported `.pte` in the runner, +builds it for Ethos-U85-256, and launches it on Corstone-320. Semihosting passes +the raw input and output tensors between the host and FVP. -Run the Arm setup first if it has not already been run: +After inference, [`visualize_fvp_output.py`](visualize_fvp_output.py) thresholds +the raw output tensor, checks it against the host quantized mask, and writes: -```bash -./examples/arm/setup.sh --i-agree-to-the-contained-eula -source examples/arm/arm-scratch/setup_path.sh -``` +`arm_test/mobilesam/result/fvp_comparison.png` -Build and install the Arm bare-metal ExecuTorch libraries from -`examples/arm`: - -```bash -cmake --preset arm-baremetal \ - -DCMAKE_BUILD_TYPE=Release \ - -B../../cmake-out-arm ../.. -cmake --build ../../cmake-out-arm --target install -j$(nproc) -``` - -The Arm bare-metal preset installs into the build directory. If you use a -different `-B` path, pass the same path as `-DET_BUILD_DIR_PATH` when -configuring the runtime app. - -## Configure the Runtime App - -Use the `.pte` from the export step and any RGB image: - -```bash -cmake \ - -DCMAKE_TOOLCHAIN_FILE=$(pwd)/ethos-u-setup/arm-none-eabi-gcc.cmake \ - -DET_BUILD_DIR_PATH=../../cmake-out-arm \ - -DET_PTE_FILE_PATH= \ - -DMODEL_METADATA_PATH= \ - -DIMAGE_PATH= \ - -DMASK_THRESHOLD=0.0 \ - -DSYSTEM_CONFIG=Ethos_U85_SYS_DRAM_Mid \ - -DMEMORY_MODE=Dedicated_Sram_384KB \ - -Bmobilesam_point_runtime \ - mobilesam_prompt_segmentation_example_ethos_u/runtime -``` - -The generated image header reads the input shape and normalization values from -the exporter metadata, resizes the longest side, and pads the shorter side with -zeros. `MODEL_METADATA_PATH` defaults to the `.json` file beside the `.pte`. -The generated `448x448` float input is placed in the Corstone-320 DDR input -section so it does not consume BRAM. - -The `MEMORY_MODE` value must match the value used during export. The default -`Dedicated_Sram_384KB` places the Ethos-U tensor arena in DDR and the fast -scratch buffer in SRAM, matching the Corstone-320 linker script. The default -method allocator pool is sized for MobileSAM's reduced `448x448` input tensor. - -The runtime emits validation logs through Arm semihosting by default so the FVP -smoke run below prints the segmentation summary. Configure with -`-DET_SEGMENTATION_SEMIHOSTING_OUTPUT=OFF` when targeting an environment -without semihosting support. - -The default target-side mask threshold is `0.0`, matching MobileSAM's mask-logit -threshold. The runtime logs foreground counts for several thresholds so users -can inspect and retune post-processing without re-exporting the graph. - -## Compile - -```bash -cmake --build mobilesam_point_runtime -j$(nproc) -- mobilesam_prompt_segmentation_example -``` - -## Run on Corstone-320 FVP - -```bash -../../backends/arm/scripts/run_fvp.sh \ - --elf=mobilesam_point_runtime/mobilesam_prompt_segmentation_example \ - --target=ethos-u85-256 \ - --timeout=300 \ - --semihosting-cwd=$(pwd)/mobilesam_point_runtime \ - --fast 2>&1 | tee mobilesam_fvp.log -``` - -Expected logs include: - -- `MobileSAM Ethos-U example started`. -- Input and output tensor shapes. -- `Mask threshold`. -- `Segmentation mask hash`. -- Foreground/background pixel counts. -- `Model executed successfully.` - -To dump the predicted binary mask as run-length encoded log chunks, configure -with `-DET_SEGMENTATION_DUMP_MASK=ON`. This is useful for debugging but makes -the UART output larger. - -Capture that FVP output and reconstruct the target mask and overlay with: - -```bash -python mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py \ - --fvp-log=mobilesam_fvp.log \ - --input-image= \ - --metadata= \ - --reference-mask= \ - --minimum-iou=0.9 \ - --output-dir=mobilesam_fvp_visual -``` - -The tool writes the decoded target mask, an input/prompt image, a colored FVP -overlay, a side-by-side comparison, and target/reference agreement metrics. +The runtime stage passes when the runner reports successful execution and the +FVP/reference mask IoU is at least `0.9`. diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py deleted file mode 100644 index a3d169907ec..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/image_to_array.py +++ /dev/null @@ -1,123 +0,0 @@ -# Copyright 2026 Arm Limited and/or its affiliates. -# -# This source code is licensed under the BSD-style license found in the -# LICENSE file in the root directory of this source tree. -# -# Copyright (c) Meta Platforms, Inc. and affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -from __future__ import annotations - -import json -import os -from argparse import ArgumentParser -from pathlib import Path - -import numpy as np -from PIL import Image - - -def convert_image_to_c_array( - image_path: str, - output_path: str, - image_size: tuple[int, int], - pixel_mean: tuple[float, float, float], - pixel_std: tuple[float, float, float], - array_name: str = "image_data", -) -> None: - image = Image.open(image_path).convert("RGB") - width, height = image.size - target_width, target_height = image_size - if target_width != target_height: - raise ValueError("MobileSAM runtime preprocessing expects a square input.") - - scale = target_width / max(height, width) - resized_size = (round(width * scale), round(height * scale)) - image = image.resize(resized_size, resample=Image.Resampling.BILINEAR) - - data = np.asarray(image, dtype=np.float32) - data = (data - np.asarray(pixel_mean, dtype=np.float32)) / np.asarray( - pixel_std, dtype=np.float32 - ) - padded_data = np.zeros((target_height, target_width, 3), dtype=np.float32) - padded_data[: resized_size[1], : resized_size[0], :] = data - data = np.transpose(padded_data, (2, 0, 1)).flatten() - - array_lines = [] - for i in range(0, len(data), 12): - line = ", ".join(f"{value:.8f}" for value in data[i : i + 12]) - array_lines.append(" " + line + ",") - - c_array = f"""#include -#include - -const size_t image_width = {image_size[0]}; -const size_t image_height = {image_size[1]}; -const size_t image_channels = 3; -__attribute__((section("input_data_sec"), aligned(16))) float {array_name}[{len(data)}] = {{ -{os.linesep.join(array_lines)} -}}; -""" - with open(output_path, "w") as output_file: - output_file.write(c_array) - print(f"Converted '{image_path}' to '{output_path}' ({len(data)} floats)") - - -def load_model_metadata( - metadata_path: str, -) -> tuple[tuple[int, int], tuple[float, float, float], tuple[float, float, float]]: - metadata = json.loads(Path(metadata_path).read_text()) - input_shape = metadata.get("input_shape") - if not isinstance(input_shape, list) or len(input_shape) != 4: - raise ValueError("Model metadata must contain a four-dimensional input_shape.") - if input_shape[0] != 1 or input_shape[1] != 3: - raise ValueError("MobileSAM runtime expects input shape [1, 3, H, W].") - - preprocessing = metadata.get("preprocessing") - if not isinstance(preprocessing, dict): - raise ValueError("Model metadata does not contain preprocessing values.") - pixel_mean_values = preprocessing.get("pixel_mean") - pixel_std_values = preprocessing.get("pixel_std") - if not isinstance(pixel_mean_values, list) or len(pixel_mean_values) != 3: - raise ValueError("MobileSAM pixel_mean must contain three RGB values.") - if not isinstance(pixel_std_values, list) or len(pixel_std_values) != 3: - raise ValueError("MobileSAM pixel_std must contain three RGB values.") - - image_size = (int(input_shape[3]), int(input_shape[2])) - pixel_mean = tuple(float(value) for value in pixel_mean_values) - pixel_std = tuple(float(value) for value in pixel_std_values) - return ( - image_size, - (pixel_mean[0], pixel_mean[1], pixel_mean[2]), - (pixel_std[0], pixel_std[1], pixel_std[2]), - ) - - -def main() -> None: - parser = ArgumentParser() - parser.add_argument("--image", required=True, help="Path to an RGB image.") - parser.add_argument( - "--output", required=True, help="Output path for the generated C array." - ) - parser.add_argument( - "--metadata", - required=True, - help="Exporter metadata containing input shape and preprocessing values.", - ) - args = parser.parse_args() - - image_size, pixel_mean, pixel_std = load_model_metadata(args.metadata) - convert_image_to_c_array( - args.image, - args.output, - image_size, - pixel_mean, - pixel_std, - ) - - -if __name__ == "__main__": - main() diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp deleted file mode 100644 index 9acf74ec860..00000000000 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/main.cpp +++ /dev/null @@ -1,413 +0,0 @@ -/* - * Copyright 2026 Arm Limited and/or its affiliates. - * - * This source code is licensed under the BSD-style license found in the - * LICENSE file in the root directory of this source tree. - */ - -#include -#include -#include -#include -#include -#include -#include - -#include -#include -#include - -#include -#include -#include -#include -#include -#include -#include - -#include "arm_memory_allocator.h" -#include "image.h" -#include "model_pte.h" - -using executorch::aten::ScalarType; -using executorch::aten::Tensor; -using executorch::extension::BufferDataLoader; -using executorch::runtime::Error; -using executorch::runtime::EValue; -using executorch::runtime::HierarchicalAllocator; -using executorch::runtime::MemoryAllocator; -using executorch::runtime::MemoryManager; -using executorch::runtime::Method; -using executorch::runtime::MethodMeta; -using executorch::runtime::Program; -using executorch::runtime::Result; -using executorch::runtime::Span; - -const size_t method_allocation_pool_size = - ET_SEGMENTATION_METHOD_ALLOCATOR_POOL_SIZE; -unsigned char __attribute__(( - section("method_allocator_sec"), - aligned(16))) method_allocation_pool[method_allocation_pool_size]; - -const size_t temp_allocation_pool_size = - ET_SEGMENTATION_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE; -unsigned char __attribute__(( - section(".bss.tensor_arena"), - aligned(16))) temp_allocation_pool[temp_allocation_pool_size]; - -#if defined(ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE) -extern "C" { -size_t ethosu_fast_scratch_size = - ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE; -unsigned char __attribute__((section(".bss.ethosu_scratch"), aligned(16))) -dedicated_sram[ET_SEGMENTATION_FAST_SCRATCH_TEMP_ALLOCATOR_POOL_SIZE]; -unsigned char* ethosu_fast_scratch = dedicated_sram; -} -#endif - -namespace { - -#if defined(ET_SEGMENTATION_MASK_THRESHOLD) -constexpr float kMaskThreshold = ET_SEGMENTATION_MASK_THRESHOLD; -#else -constexpr float kMaskThreshold = 0.0f; -#endif - -#if defined(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) && \ - (defined(__arm__) || defined(__thumb__)) -constexpr uint32_t kSemihostingSysWrite0 = 0x04; -constexpr uint32_t kSemihostingSysExitExtended = 0x20; -constexpr uint32_t kAdpStoppedApplicationExit = 0x20026; - -uint32_t semihosting_call(uint32_t operation_code, const void* argument) { - uint32_t result; - asm volatile( - "mov r0, %[operation]\n" - "mov r1, %[argument]\n" - "bkpt 0xab\n" - "mov %[result], r0\n" - : [result] "=r"(result) - : [operation] "r"(operation_code), [argument] "r"(argument) - : "r0", "r1", "memory"); - return result; -} - -void semihosting_write0(const char* message) { - semihosting_call(kSemihostingSysWrite0, message); -} -#endif - -void write_runtime_line(const char* message) { -#if defined(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) && \ - (defined(__arm__) || defined(__thumb__)) - semihosting_write0(message); - semihosting_write0("\n"); -#else - (void)message; -#endif -} - -void write_runtime_format(const char* format, ...) { - char line[768]; - va_list args; - va_start(args, format); - vsnprintf(line, sizeof(line), format, args); - va_end(args); - write_runtime_line(line); -} - -void request_runtime_exit(int code) { -#if defined(ET_SEGMENTATION_SEMIHOSTING_OUTPUT) && \ - (defined(__arm__) || defined(__thumb__)) - const uint32_t exit_block[2] = { - kAdpStoppedApplicationExit, - static_cast(code), - }; - semihosting_call(kSemihostingSysExitExtended, exit_block); -#else - (void)code; -#endif -} - -uint32_t update_hash(uint32_t hash, uint8_t value) { - hash ^= value; - hash *= 16777619u; - return hash; -} - -#if defined(ET_SEGMENTATION_DUMP_MASK) -void dump_mask_rle(const std::vector& mask) { - ET_LOG(Info, "Segmentation mask RLE begin"); - write_runtime_line("Segmentation mask RLE begin"); - char line[512]; - size_t line_len = 0; - line[0] = '\0'; - - size_t index = 0; - while (index < mask.size()) { - const uint8_t class_id = mask[index]; - size_t run_len = 1; - while (index + run_len < mask.size() && mask[index + run_len] == class_id) { - ++run_len; - } - - char entry[32]; - const int entry_len = snprintf( - entry, - sizeof(entry), - "%u:%zu,", - static_cast(class_id), - run_len); - if (entry_len <= 0) { - break; - } - if (line_len + static_cast(entry_len) >= sizeof(line)) { - ET_LOG(Info, "Segmentation mask RLE chunk %s", line); - write_runtime_format("Segmentation mask RLE chunk %s", line); - line_len = 0; - line[0] = '\0'; - } - line_len += static_cast( - snprintf(line + line_len, sizeof(line) - line_len, "%s", entry)); - index += run_len; - } - if (line_len > 0) { - ET_LOG(Info, "Segmentation mask RLE chunk %s", line); - write_runtime_format("Segmentation mask RLE chunk %s", line); - } - ET_LOG(Info, "Segmentation mask RLE end"); - write_runtime_line("Segmentation mask RLE end"); -} -#endif - -void summarize_segmentation_output(const Tensor& out) { - ET_CHECK_MSG( - out.dim() == 4, - "Expected mask logits with shape [1, 1, height, width], got rank %zd", - out.dim()); - ET_CHECK_MSG( - out.scalar_type() == ScalarType::Float, - "Expected float mask logits, got dtype %d", - out.scalar_type()); - ET_CHECK_MSG(out.size(0) == 1, "Only batch size 1 is supported."); - ET_CHECK_MSG( - out.size(1) == 1, - "MobileSAM fixed-prompt export expects one mask channel, got %zd", - out.size(1)); - - const size_t height = static_cast(out.size(2)); - const size_t width = static_cast(out.size(3)); - const auto strides = out.strides(); - const float* data = out.const_data_ptr(); - - size_t foreground_pixels = 0; -#if defined(ET_SEGMENTATION_DUMP_MASK) - std::vector mask(height * width, 0); -#endif - uint32_t mask_hash = 2166136261u; - float min_score = data[0]; - float max_score = data[0]; - double score_sum = 0.0; - const float threshold_sweep[] = { - 0.0f, - -2.0f, - -4.0f, - -5.0f, - -6.0f, - -7.0f, - -8.0f, - -10.0f, - -12.0f, - -14.0f, - }; - size_t threshold_sweep_counts - [sizeof(threshold_sweep) / sizeof(threshold_sweep[0])] = {}; - - for (size_t y = 0; y < height; ++y) { - for (size_t x = 0; x < width; ++x) { - const float score = data[y * strides[2] + x * strides[3]]; - min_score = std::min(min_score, score); - max_score = std::max(max_score, score); - score_sum += score; - for (size_t i = 0; - i < sizeof(threshold_sweep) / sizeof(threshold_sweep[0]); - ++i) { - threshold_sweep_counts[i] += score > threshold_sweep[i] ? 1 : 0; - } - const uint8_t mask_value = score > kMaskThreshold ? 1 : 0; - foreground_pixels += mask_value; -#if defined(ET_SEGMENTATION_DUMP_MASK) - mask[y * width + x] = mask_value; -#endif - mask_hash = update_hash(mask_hash, mask_value); - } - } - - ET_LOG(Info, "Output mask logits shape = [1, 1, %zu, %zu]", height, width); - write_runtime_format( - "Output mask logits shape = [1, 1, %zu, %zu]", height, width); - ET_LOG( - Info, - "Segmentation input image = %zu x %zu x %zu", - image_width, - image_height, - image_channels); - write_runtime_format( - "Segmentation input image = %zu x %zu x %zu", - image_width, - image_height, - image_channels); - ET_LOG(Info, "Mask threshold = %.4f", static_cast(kMaskThreshold)); - write_runtime_format( - "Mask threshold = %.4f", static_cast(kMaskThreshold)); - ET_LOG( - Info, - "Mask logits min/max/mean = %.6f / %.6f / %.6f", - static_cast(min_score), - static_cast(max_score), - score_sum / static_cast(height * width)); - write_runtime_format( - "Mask logits min/max/mean = %.6f / %.6f / %.6f", - static_cast(min_score), - static_cast(max_score), - score_sum / static_cast(height * width)); - for (size_t i = 0; i < sizeof(threshold_sweep) / sizeof(threshold_sweep[0]); - ++i) { - write_runtime_format( - "Threshold %.1f foreground pixels = %zu", - static_cast(threshold_sweep[i]), - threshold_sweep_counts[i]); - } - ET_LOG(Info, "Segmentation mask hash = 0x%08" PRIx32, mask_hash); - write_runtime_format("Segmentation mask hash = 0x%08" PRIx32, mask_hash); - ET_LOG(Info, "Mask foreground pixels = %zu", foreground_pixels); - write_runtime_format("Mask foreground pixels = %zu", foreground_pixels); - ET_LOG( - Info, "Mask background pixels = %zu", height * width - foreground_pixels); - write_runtime_format( - "Mask background pixels = %zu", height * width - foreground_pixels); - -#if defined(ET_SEGMENTATION_DUMP_MASK) - dump_mask_rle(mask); -#endif -} - -} // namespace - -int main() { - executorch::runtime::runtime_init(); - ET_LOG(Info, "Runtime initialized"); - write_runtime_line("MobileSAM Ethos-U example started"); - BufferDataLoader loader(model_pte, sizeof(model_pte)); - ET_LOG(Info, "Size of the model = %zu", sizeof(model_pte)); - write_runtime_format("Model size = %zu bytes", sizeof(model_pte)); - write_runtime_line("Loading ExecuTorch program"); - Result program = Program::load(&loader); - ET_CHECK_MSG(program.ok(), "Program::load failed: 0x%x", program.error()); - write_runtime_line("Program loaded"); - - const auto method_name_result = program->get_method_name(0); - ET_CHECK_MSG(method_name_result.ok(), "Program has no methods"); - const char* method_name = *method_name_result; - ET_LOG(Info, "Running method %s", method_name); - write_runtime_format("Running method %s", method_name); - - Result method_meta_result = program->method_meta(method_name); - ET_CHECK_MSG( - method_meta_result.ok(), - "method_meta lookup failed: 0x%x", - method_meta_result.error()); - - ArmMemoryAllocator method_allocator( - method_allocation_pool_size, method_allocation_pool); - ArmMemoryAllocator temp_allocator( - temp_allocation_pool_size, temp_allocation_pool); - - std::vector planned_buffers; - std::vector> planned_spans; - const size_t num_memory_planned_buffers = - method_meta_result->num_memory_planned_buffers(); - ET_LOG(Info, "num_memory_planned_buffers = %zu", num_memory_planned_buffers); - for (size_t id = 0; id < num_memory_planned_buffers; ++id) { - const size_t buffer_size = - method_meta_result->memory_planned_buffer_size(id).get(); - ET_LOG(Info, "Planned memory buffer_size %zu %zu bytes", id, buffer_size); - - uint8_t* buffer = reinterpret_cast( - method_allocator.allocate(buffer_size, 16UL)); - ET_CHECK_MSG( - buffer != nullptr, - "Could not allocate memory for memory planned buffer size %zu", - buffer_size); - planned_buffers.push_back(buffer); - planned_spans.push_back({planned_buffers.back(), buffer_size}); - } - HierarchicalAllocator planned_memory( - {planned_spans.data(), planned_spans.size()}); - - MemoryManager memory_manager( - &method_allocator, &planned_memory, &temp_allocator); - write_runtime_line("Loading method"); - Result method = program->load_method(method_name, &memory_manager); - ET_CHECK_MSG(method.ok(), "load_method failed: 0x%x", method.error()); - write_runtime_line("Method loaded"); - - const size_t num_inputs = method->inputs_size(); - ET_LOG(Info, "Number of input tensors = %zu", num_inputs); - ET_CHECK_MSG( - num_inputs == 1, - "The segmentation model has a single input tensor, but the provided model has %zu input tensors", - num_inputs); - - EValue* input_evalues = method_allocator.allocateList(num_inputs); - Error err = method->get_inputs(input_evalues, num_inputs); - ET_CHECK_MSG(err == Error::Ok, "get_inputs failed"); - Tensor& input_tensor = input_evalues[0].toTensor(); - const size_t expected_elems = input_tensor.numel(); - const size_t image_elements = sizeof(image_data) / sizeof(image_data[0]); - ET_CHECK_MSG( - expected_elems == image_elements, - "Input tensor expects %zu elements, but image_data has %zu elements", - expected_elems, - image_elements); - ET_CHECK_MSG( - input_tensor.scalar_type() == ScalarType::Float, - "Expected float input tensor, got dtype %d", - input_tensor.scalar_type()); - - float* input_data = input_tensor.mutable_data_ptr(); - write_runtime_format("Copying %zu input elements", expected_elems); - for (size_t i = 0; i < expected_elems; ++i) { - input_data[i] = image_data[i]; - } - - write_runtime_line("Running model execution"); - Error status_inference = method->execute(); - ET_CHECK_MSG( - status_inference == Error::Ok, - "Inference failed 0x%" PRIx32, - status_inference); - write_runtime_line("Inference finished"); - - const size_t num_outputs = method->outputs_size(); - std::vector outputs(num_outputs); - Error status_outputs = method->get_outputs(outputs.data(), outputs.size()); - ET_CHECK_MSG( - status_outputs == Error::Ok, - "get_outputs failed 0x%" PRIx32, - status_outputs); - - for (size_t i = 0; i < outputs.size(); ++i) { - if (outputs[i].isTensor()) { - summarize_segmentation_output(outputs[i].toTensor()); - ET_LOG(Info, "Model executed successfully."); - write_runtime_line("Model executed successfully."); - request_runtime_exit(0); - return 0; - } - } - - ET_CHECK_MSG(false, "No tensor output found."); - request_runtime_exit(1); - return 1; -} diff --git a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py index bb176c3adf4..df8f73a0a1f 100644 --- a/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py +++ b/examples/arm/mobilesam_prompt_segmentation_example_ethos_u/runtime/visualize_fvp_output.py @@ -3,192 +3,80 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import argparse import json -import re from pathlib import Path -from typing import Any +import numpy as np from PIL import Image, ImageDraw -RLE_PREFIX = "Segmentation mask RLE chunk " +ROOT = Path(__file__).resolve().parents[4] +WORK_DIR = ROOT / "arm_test" / "mobilesam" +EXPORT_DIR = WORK_DIR / "export" +RESULT_DIR = WORK_DIR / "result" +OUTPUT_SIZE = 112 +POINT = (219, 193) -def parse_args() -> argparse.Namespace: - parser = argparse.ArgumentParser( - description="Visualize and validate a MobileSAM mask dumped by the FVP." - ) - parser.add_argument("--fvp-log", required=True, type=Path) - parser.add_argument("--input-image", required=True, type=Path) - parser.add_argument("--metadata", required=True, type=Path) - parser.add_argument("--output-dir", required=True, type=Path) - parser.add_argument("--reference-mask", type=Path) - parser.add_argument("--minimum-iou", type=float) - return parser.parse_args() - - -def parse_rle_mask(log_path: Path, expected_pixels: int) -> list[int]: - mask: list[int] = [] - in_dump = False - for line in log_path.read_text().splitlines(): - if "executorch:main.cpp:" in line: - continue - if "Segmentation mask RLE begin" in line: - in_dump = True - continue - if "Segmentation mask RLE end" in line: - break - if not in_dump or RLE_PREFIX not in line: - continue - payload = line.split(RLE_PREFIX, maxsplit=1)[1] - for value, count in re.findall(r"([01]):([0-9]+),", payload): - mask.extend([int(value)] * int(count)) - - if len(mask) != expected_pixels: - raise ValueError( - f"FVP RLE contains {len(mask)} pixels; expected {expected_pixels}." - ) - return mask - - -def prepare_input_image(image_path: Path, input_size: int) -> Image.Image: - image = Image.open(image_path).convert("RGB") - width, height = image.size - scale = input_size / max(width, height) - resized = image.resize( - (round(width * scale), round(height * scale)), - Image.Resampling.BILINEAR, - ) - padded = Image.new("RGB", (input_size, input_size)) - padded.paste(resized, (0, 0)) - return padded +def load_mask(path: Path) -> np.ndarray: + logits = np.fromfile(path, dtype=np.float32) + if logits.size != OUTPUT_SIZE * OUTPUT_SIZE: + raise ValueError(f"Expected {OUTPUT_SIZE**2} logits, got {logits.size}.") + return (logits.reshape(OUTPUT_SIZE, OUTPUT_SIZE) > 0).astype(np.uint8) + +def iou(first: np.ndarray, second: np.ndarray) -> float: + intersection = np.logical_and(first, second).sum() + union = np.logical_or(first, second).sum() + return 1.0 if union == 0 else float(intersection / union) -def create_overlay( - image: Image.Image, mask: Image.Image, color: tuple[int, int, int] + +def overlay( + image: Image.Image, mask: np.ndarray, color: tuple[int, int, int] ) -> Image.Image: - resized_mask = mask.resize(image.size, Image.Resampling.NEAREST) - color_layer = Image.new("RGB", image.size, color) - blended = Image.blend(image, color_layer, 0.45) - overlay = image.copy() - overlay.paste(blended, mask=resized_mask) - return overlay - - -def draw_prompts(image: Image.Image, metadata: dict[str, Any]) -> Image.Image: - result = image.copy() - draw = ImageDraw.Draw(result) - for point_x, point_y in metadata["point_prompts_xy"]: - radius = max(4, metadata["input_size"] // 80) - draw.ellipse( - ( - point_x - radius, - point_y - radius, - point_x + radius, - point_y + radius, - ), - fill=(255, 48, 48), - outline=(255, 255, 255), - width=2, - ) - return result - - -def add_title(image: Image.Image, title: str) -> Image.Image: - title_height = 30 - panel = Image.new("RGB", (image.width, image.height + title_height), "white") - panel.paste(image, (0, title_height)) - ImageDraw.Draw(panel).text((10, 8), title, fill="black") - return panel - - -def binary_metrics(mask: list[int], reference: list[int]) -> tuple[float, float]: - intersection = sum(a == 1 and b == 1 for a, b in zip(mask, reference)) - union = sum(a == 1 or b == 1 for a, b in zip(mask, reference)) - iou = intersection / union if union else 1.0 - agreement = sum(a == b for a, b in zip(mask, reference)) / len(mask) - return iou, agreement + resized = Image.fromarray(mask * 255).resize(image.size, Image.Resampling.NEAREST) + pixels = np.asarray(image, dtype=np.float32).copy() + selected = np.asarray(resized) > 0 + pixels[selected] = pixels[selected] * 0.55 + np.asarray(color) * 0.45 + return Image.fromarray(pixels.astype(np.uint8)) def main() -> None: - args = parse_args() - metadata = json.loads(args.metadata.read_text()) - _, _, mask_height, mask_width = metadata["output_shape"] - mask = parse_rle_mask(args.fvp_log, mask_width * mask_height) - mask_image = Image.new("L", (mask_width, mask_height)) - mask_image.putdata([value * 255 for value in mask]) - - args.output_dir.mkdir(parents=True, exist_ok=True) - mask_image.save(args.output_dir / "fvp_mask.png") - input_image = prepare_input_image(args.input_image, metadata["input_size"]) - prompted_input = draw_prompts(input_image, metadata) - prompted_input.save(args.output_dir / "input_with_prompts.png") - fvp_overlay = draw_prompts( - create_overlay(input_image, mask_image, (0, 220, 120)), metadata + RESULT_DIR.mkdir(parents=True, exist_ok=True) + image = Image.open(EXPORT_DIR / "input.png").convert("RGB") + reference = (np.asarray(Image.open(EXPORT_DIR / "quantized_mask.png")) > 0).astype( + np.uint8 ) - fvp_overlay.save(args.output_dir / "fvp_overlay.png") - - panels = [add_title(prompted_input, "Input and positive prompt")] - metrics: dict[str, Any] = { - "background_pixels": mask.count(0), - "foreground_pixels": mask.count(1), - "output_mask_size": [mask_width, mask_height], - } - if metrics["foreground_pixels"] in (0, len(mask)): - raise RuntimeError( - "FVP produced a degenerate mask with " - f"{metrics['foreground_pixels']} foreground pixels." - ) - - below_minimum_iou = False - if args.reference_mask is not None: - reference_image = ( - Image.open(args.reference_mask) - .convert("L") - .resize((mask_width, mask_height), Image.Resampling.NEAREST) - ) - reference = [int(value > 0) for value in reference_image.tobytes()] - iou, agreement = binary_metrics(mask, reference) - metrics["fvp_reference_iou"] = iou - metrics["fvp_reference_pixel_agreement"] = agreement - reference_overlay = draw_prompts( - create_overlay(input_image, reference_image, (0, 170, 255)), metadata - ) - panels.append(add_title(reference_overlay, "Host quantized mask")) - below_minimum_iou = args.minimum_iou is not None and iou < args.minimum_iou - elif args.minimum_iou is not None: - raise ValueError("--minimum-iou requires --reference-mask.") - - panels.append(add_title(fvp_overlay, "FVP mask")) - comparison = Image.new( - "RGB", - (sum(panel.width for panel in panels), max(panel.height for panel in panels)), - "white", + fvp_mask = load_mask(WORK_DIR / "io" / "output-0.bin") + score = iou(fvp_mask, reference) + if score < 0.9: + raise RuntimeError(f"FVP/reference mask IoU is too low: {score:.4f}") + + prompted = image.copy() + ImageDraw.Draw(prompted).ellipse( + (POINT[0] - 6, POINT[1] - 6, POINT[0] + 6, POINT[1] + 6), fill="red" ) - offset = 0 - for panel in panels: - comparison.paste(panel, (offset, 0)) - offset += panel.width - comparison.save(args.output_dir / "fvp_comparison.png") - (args.output_dir / "metrics.json").write_text( - json.dumps(metrics, indent=2, sort_keys=True) + "\n" + reference_overlay = overlay(image, reference, (0, 170, 255)) + fvp_overlay = overlay(image, fvp_mask, (0, 220, 120)) + panels = ( + ("Input and prompt", prompted), + ("Host quantized mask", reference_overlay), + ("FVP output", fvp_overlay), ) - - print( - f"FVP mask: {metrics['foreground_pixels']} foreground pixels, " - f"artifacts saved to {args.output_dir}" + comparison = Image.new("RGB", (image.width * 3, image.height + 32), "white") + drawing = ImageDraw.Draw(comparison) + for index, (label, panel) in enumerate(panels): + x = index * image.width + drawing.text((x + 10, 10), label, fill="black") + comparison.paste(panel, (x, 32)) + + Image.fromarray(fvp_mask * 255).save(RESULT_DIR / "fvp_mask.png") + comparison.save(RESULT_DIR / "fvp_comparison.png") + (RESULT_DIR / "metrics.json").write_text( + json.dumps({"fvp_reference_iou": score}, indent=2) + "\n" ) - if "fvp_reference_iou" in metrics: - print( - f"FVP/reference IoU={metrics['fvp_reference_iou']:.4f} " - f"agreement={metrics['fvp_reference_pixel_agreement']:.4f}" - ) - if below_minimum_iou: - raise RuntimeError( - f"FVP/reference IoU {metrics['fvp_reference_iou']:.4f} is below " - f"{args.minimum_iou:.4f}." - ) + print(f"FVP/reference mask IoU: {score:.4f}") + print(f"Saved {RESULT_DIR / 'fvp_comparison.png'}") if __name__ == "__main__": diff --git a/examples/arm/setup.sh b/examples/arm/setup.sh index 33bae7c13e1..55e7a873cbf 100755 --- a/examples/arm/setup.sh +++ b/examples/arm/setup.sh @@ -328,6 +328,7 @@ function create_setup_path(){ if [[ $is_script_sourced -eq 0 ]]; then set -e + ARM_SETUP_CURL_PROGRESS_ARGS=(--progress-bar) if [[ -n "$("${et_dir}/.ci/scripts/detect_ci.sh" --and-not-debug)" ]]; then ARM_SETUP_CURL_PROGRESS_ARGS=(--no-progress-meter) export PIP_PROGRESS_BAR=off @@ -376,9 +377,12 @@ if [[ $is_script_sourced -eq 0 ]]; then # Setup FVP if [[ "${enable_fvps}" -eq 1 ]]; then log_step "fvp" "Setting up Arm Fixed Virtual Platforms" - check_fvp_eula - setup_fvp - install_fvp + if [[ "${OS}" == "Linux" ]]; then + check_fvp_eula + install_fvp + else + setup_fvp + fi fi warn_if_mlsdk_python_is_untested From 01cc07564183cf0a48e34c1f36c3ec6cd0db599a Mon Sep 17 00:00:00 2001 From: Sebastian Larsson <38941629+Sebastian-Larsson@users.noreply.github.com> Date: Wed, 9 Sep 2026 15:55:36 +0200 Subject: [PATCH 097/190] Arm backend: Show real pass names in Arm crash reports (#22211) When an Arm pass failed, crash reports named _ExportedProgramGraphPassAdapter instead of the pass that actually failed. The list of passes that ran successfully had the same problem. This made the report useless for finding the broken pass. Copy the wrapped pass name onto each adapter and make the pass manager use that name in error messages. Add a regression test for the failed pass and the passes that ran before it. Signed-off-by: Sebastian Larsson --- backends/arm/_passes/arm_pass_manager.py | 1 + .../passes/test_arm_pass_manager_errors.py | 37 +++++++++++++++++++ exir/pass_manager.py | 6 ++- 3 files changed, 42 insertions(+), 2 deletions(-) create mode 100644 backends/arm/test/passes/test_arm_pass_manager_errors.py diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 12d766c2ffa..3873a92a442 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -224,6 +224,7 @@ def _graph_pass_name(graph_pass: Callable[[GraphModule], PassResult | None]) -> class _ExportedProgramGraphPassAdapter(ExportedProgramPassBase): def __init__(self, graph_pass: Callable[[GraphModule], PassResult | None]) -> None: self.graph_pass = graph_pass + self.__name__ = _graph_pass_name(graph_pass) def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: graph_pass = cast(Any, self.graph_pass) diff --git a/backends/arm/test/passes/test_arm_pass_manager_errors.py b/backends/arm/test/passes/test_arm_pass_manager_errors.py new file mode 100644 index 00000000000..22818af3b7d --- /dev/null +++ b/backends/arm/test/passes/test_arm_pass_manager_errors.py @@ -0,0 +1,37 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest + +import torch + +from executorch.backends.arm._passes.arm_pass_manager import ( + _ExportedProgramGraphPassAdapter, +) +from executorch.exir.pass_manager import ExportedProgramPassManager, PassType +from torch.fx import GraphModule +from torch.fx.passes.infra.pass_base import PassResult + + +def test_exported_program_adapter_preserves_pass_names_in_errors() -> None: + """Report wrapped pass names instead of the adapter class name.""" + + def successful_pass(graph_module: GraphModule) -> PassResult: + return PassResult(graph_module, False) + + def failing_pass(graph_module: GraphModule) -> PassResult: + raise RuntimeError("test failure") + + exported_program = torch.export.export(torch.nn.ReLU(), (torch.randn(2, 3),)) + passes: list[PassType] = [ + _ExportedProgramGraphPassAdapter(successful_pass), + _ExportedProgramGraphPassAdapter(failing_pass), + ] + + with pytest.raises(Exception) as error: + ExportedProgramPassManager(passes)(exported_program) + + assert "running the 'failing_pass' pass" in str(error.value) + assert "following passes: ['successful_pass']" in str(error.value) diff --git a/exir/pass_manager.py b/exir/pass_manager.py index 829486fa0ce..97c27901d73 100644 --- a/exir/pass_manager.py +++ b/exir/pass_manager.py @@ -1,12 +1,12 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. # pyre-strict import copy -import inspect import logging import operator from typing import Callable, List, Optional, Type, TypeAlias, Union @@ -36,7 +36,9 @@ def _get_pass_name(fn: PassType) -> str: """Returns a human-readable name for a pass.""" - return fn.__name__ if inspect.isfunction(fn) else type(fn).__name__ + if hasattr(fn, "__name__"): + return fn.__name__ + return type(fn).__name__ def _can_eliminate_common_getitems(gm: torch.fx.GraphModule) -> bool: From cba82c3445297a1b0dcf9a4623a0ba51427405a7 Mon Sep 17 00:00:00 2001 From: Michiel Olieslagers <44864547+Michiel-Olieslagers@users.noreply.github.com> Date: Wed, 9 Sep 2026 16:58:34 +0100 Subject: [PATCH 098/190] Arm backend: Add dynamic shape support to ArmQuantize (#22643) Pass optional dynamic shape specifications to torch.export during quantization. This allows models with dynamic input dimensions to be calibrated through the Arm test harness. Change-Id: I4d9f0e3a8c7b6e5d4c3b2a190817263544556677 cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Michiel Olieslagers --- backends/arm/test/tester/quantize.py | 6 ++- backends/arm/test/tester/test_quantize.py | 64 +++++++++++++++++++++++ 2 files changed, 69 insertions(+), 1 deletion(-) create mode 100644 backends/arm/test/tester/test_quantize.py diff --git a/backends/arm/test/tester/quantize.py b/backends/arm/test/tester/quantize.py index ae3a216b528..f1b57a8e237 100644 --- a/backends/arm/test/tester/quantize.py +++ b/backends/arm/test/tester/quantize.py @@ -27,6 +27,7 @@ def __init__( is_qat: Optional[bool] = False, set_global: bool = True, fold_quantize: bool = True, + dynamic_shapes: Optional[Tuple[Any, ...]] = None, ): super().__init__( quantizer, @@ -37,6 +38,7 @@ def __init__( set_global, ) self.fold_quantize = fold_quantize + self.dynamic_shapes = dynamic_shapes def run( self, artifact: torch.nn.Module, inputs: Optional[Tuple[torch.Tensor]] @@ -44,7 +46,9 @@ def run( assert inputs is not None if self.is_qat: artifact.train() - captured_graph = export(artifact, inputs, strict=True).module() + captured_graph = export( + artifact, inputs, dynamic_shapes=self.dynamic_shapes, strict=True + ).module() if not isinstance(self.quantizer, TOSAQuantizer): raise ValueError("ArmQuantizer can only run with TOSAQuantizer.") diff --git a/backends/arm/test/tester/test_quantize.py b/backends/arm/test/tester/test_quantize.py new file mode 100644 index 00000000000..62f84db25fa --- /dev/null +++ b/backends/arm/test/tester/test_quantize.py @@ -0,0 +1,64 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch + +from executorch.backends.arm.quantizer import ( + get_symmetric_quantization_config, + TOSAQuantizer, +) +from executorch.backends.arm.test.tester.quantize import ArmQuantize +from executorch.backends.arm.tosa import TosaSpecification + + +class Add(torch.nn.Module): + def forward(self, x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + return x + y + + +def _quantize( + module: torch.nn.Module, + inputs: tuple[torch.Tensor, ...], + dynamic_shapes: tuple[dict[int, torch.export.Dim], ...], +) -> torch.fx.GraphModule: + quantization_config = get_symmetric_quantization_config() + quantizer = TOSAQuantizer(TosaSpecification.create_from_string("TOSA-1.0+INT")) + stage = ArmQuantize( + quantizer, + quantization_config, + dynamic_shapes=dynamic_shapes, + ) + + stage.run(module, inputs) # type: ignore[arg-type] + + return stage.artifact + + +def test_arm_quantize_preserves_dynamic_input_shape() -> None: + inputs = (torch.randn(2, 4),) + batch = torch.export.Dim("batch", min=1, max=4) + + graph_module = _quantize(torch.nn.ReLU(), inputs, ({0: batch},)) + + placeholder = next( + node for node in graph_module.graph.nodes if node.op == "placeholder" + ) + assert isinstance(placeholder.meta["val"].shape[0], torch.SymInt) + assert graph_module(torch.randn(3, 4)).shape == (3, 4) + + +def test_arm_quantize_preserves_shared_dynamic_input_shape() -> None: + inputs = (torch.randn(2, 4), torch.randn(2, 4)) + batch = torch.export.Dim("batch", min=1, max=4) + + graph_module = _quantize(Add(), inputs, ({0: batch}, {0: batch})) + + placeholders = [ + node for node in graph_module.graph.nodes if node.op == "placeholder" + ] + batch_sizes = [node.meta["val"].shape[0] for node in placeholders] + assert all(isinstance(size, torch.SymInt) for size in batch_sizes) + assert batch_sizes[0] == batch_sizes[1] + assert graph_module(torch.randn(3, 4), torch.randn(3, 4)).shape == (3, 4) From 2ecc182fd519454fa0156d83f7f88bbdd1ce53bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Thu, 3 Sep 2026 08:50:06 +0200 Subject: [PATCH 099/190] [ET-VK] Name the int32 eq shader what the dispatcher asks for add_binary_op_node builds its kernel name as "binary_" + op + storage + dtype, so an int32 aten.eq.Tensor looks for binary_eq_buffer_int32. The yaml instead declares that variant as binary_eq_int32_buffer, which generates shaders nothing ever references and leaves the name the dispatcher wants missing. The op is registered as supported, so the partitioner claims it and the model then aborts at dispatch with "Could not find ShaderInfo with name binary_eq_buffer_int32". Declare it as a second binary_eq_* variant restricted to int32 instead. The generated names do not collide with the half and float ones, and the exact comparison stays separate from the float path's epsilon compare. Found by lowering kokoro's synthesizer, whose mask comparison runs on int32. --- .../runtime/graph/ops/glsl/binary_op_buffer.yaml | 10 ++++++++-- .../runtime/graph/ops/glsl/binary_op_texture.yaml | 8 ++++++-- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml index 1f217acb127..d1f3600cfc9 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml @@ -30,9 +30,15 @@ binary_op_buffer: OPERATOR: floor(X / Y) - NAME: binary_minimum_buffer OPERATOR: min(X, Y) - - NAME: binary_eq_int32_buffer + # Named to match the dispatcher, which builds "binary_" + op + storage + + # dtype. A variant called binary_eq_int32_buffer generates shaders nothing + # ever asks for, and aten.eq on an int32 tensor aborts at dispatch with + # "Could not find ShaderInfo with name binary_eq_buffer_int32". + - NAME: binary_eq_buffer OPERATOR: X == Y - DTYPE: int32 + generate_variant_forall: + DTYPE: + - VALUE: int32 - NAME: binary_eq_buffer OPERATOR: abs(X - Y) < 1e-5 generate_variant_forall: diff --git a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml index 289466e7845..07e3905af02 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml @@ -32,9 +32,13 @@ binary_op_texture: MASK_PADDING: 1 - NAME: binary_minimum_texture3d OPERATOR: min(X, Y) - - NAME: binary_eq_int32_texture3d + # See binary_op_buffer.yaml: the name has to match what the dispatcher + # builds, otherwise the generated shader is unreachable. + - NAME: binary_eq_texture3d OPERATOR: equal(X, Y) - DTYPE: int32 + generate_variant_forall: + DTYPE: + - VALUE: int32 - NAME: binary_eq_texture3d OPERATOR: lessThan(abs(X - Y), VEC4_T(1e-5)) generate_variant_forall: From eb714844a5f06777c87c46ccf06e9adff0fb7eaf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Fri, 4 Sep 2026 09:48:50 +0200 Subject: [PATCH 100/190] [ET-VK] Test that binary op shader names match the dispatcher Guards the naming contract the previous commit fixed, without touching the skipped aten.eq.Tensor correctness case. Two assertions. The first names binary_eq_buffer_int32 and binary_eq_texture3d_int32 directly, so the exact regression is covered. The second is the general form: every variant the two binary op templates generate has to end in a storage suffix followed by a dtype suffix, because that is what add_binary_op_node builds. Both suffix sets are parsed out of ShaderNameUtils.cpp rather than restated, so adding a dtype there does not fail the test. Runs at codegen level with no built runtime and no GPU: the generator is loaded by file path and only its variant names are inspected. Reverting the yaml fix fails both assertions. --- .../vulkan/test/test_vulkan_shader_names.py | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 backends/vulkan/test/test_vulkan_shader_names.py diff --git a/backends/vulkan/test/test_vulkan_shader_names.py b/backends/vulkan/test/test_vulkan_shader_names.py new file mode 100644 index 00000000000..1535e6c9653 --- /dev/null +++ b/backends/vulkan/test/test_vulkan_shader_names.py @@ -0,0 +1,148 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Checks that generated shader names match the names the dispatcher asks for. + +A kernel name is built at runtime rather than looked up: `add_binary_op_node` +concatenates "binary_", the op, a storage suffix and a dtype suffix, then hands +the result to `VK_KERNEL_FROM_STR`. The yaml, meanwhile, is free to name a +variant anything at all. Nothing connects the two, so a variant whose name does +not follow that shape compiles a shader nothing references and leaves the name +the dispatcher wants missing. + +That failure is invisible until dispatch. The op is still registered as +supported, so the partitioner claims it, the export succeeds, and the model +aborts on device with "Could not find ShaderInfo with name ...". These tests +close the gap at codegen time instead. + +The generator is loaded by file path, and the suffixes are read out of the C++ +that produces them, so this needs neither a built runtime nor a GPU and does not +drift when a dtype is added. +""" + +import importlib.util +import re +import unittest +from pathlib import Path + +_REPO_ROOT = Path(__file__).resolve().parents[3] +_VULKAN_ROOT = _REPO_ROOT / "backends" / "vulkan" +_GLSL_DIR = _VULKAN_ROOT / "runtime" / "graph" / "ops" / "glsl" +_SHADER_NAME_UTILS = ( + _VULKAN_ROOT / "runtime" / "graph" / "ops" / "utils" / "ShaderNameUtils.cpp" +) + +_spec = importlib.util.spec_from_file_location( + "gen_vulkan_spv", _VULKAN_ROOT / "runtime" / "gen_vulkan_spv.py" +) +gen_vulkan_spv = importlib.util.module_from_spec(_spec) +_spec.loader.exec_module(gen_vulkan_spv) + +# The yaml templates whose variants `add_binary_op_node` dispatches into. +_BINARY_TEMPLATES = ("binary_op_buffer", "binary_op_texture") + + +def _function_body(text: str, name: str) -> str: + """Source of a top-level C++ function, from its signature to its closing brace.""" + start = next( + ( + i + for i, line in enumerate(text.splitlines()) + if re.match(rf"^\w[\w:<>&* ]*\b{re.escape(name)}\(", line) + ), + None, + ) + if start is None: + raise AssertionError(f"{name} not found in {_SHADER_NAME_UTILS}") + lines = text.splitlines()[start:] + end = next(i for i, line in enumerate(lines) if line == "}") + return "\n".join(lines[: end + 1]) + + +def _suffixes(function_name: str) -> tuple: + """Every literal suffix one ShaderNameUtils.cpp helper can append. + + Read from the C++ rather than restated here: a dtype added to + `add_dtype_suffix` becomes legal in a shader name the moment it is added, + and a test carrying its own copy of the list would reject it. + """ + body = _function_body(_SHADER_NAME_UTILS.read_text(), function_name) + found = re.findall(r'kernel_name \+= "(_[a-z0-9]+)";', body) + if not found: + raise AssertionError(f"no suffixes parsed out of {function_name}") + return tuple(dict.fromkeys(found)) + + +STORAGE_SUFFIXES = _suffixes("add_storage_type_suffix") +DTYPE_SUFFIXES = _suffixes("add_dtype_suffix") + + +def _generated_names() -> dict: + """Variant names the shader codegen produces, keyed by yaml template.""" + env = dict(gen_vulkan_spv.DEFAULT_ENV) + env.update(gen_vulkan_spv.TYPE_MAPPINGS) + env.update(gen_vulkan_spv.UTILITY_FNS) + generator = gen_vulkan_spv.SPVGenerator([str(_GLSL_DIR)], env, glslc_path=None) + return { + template: [variant["NAME"] for variant in variants] + for template, variants in generator.shader_template_params.items() + } + + +class TestShaderNames(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.names = _generated_names() + + def test_suffixes_are_parsed_from_the_cpp(self) -> None: + # Guards the two tests below: a parse that silently returned something + # empty or wrong would make them pass by vacuously accepting any name. + self.assertIn("_buffer", STORAGE_SUFFIXES) + self.assertIn("_texture3d", STORAGE_SUFFIXES) + self.assertIn("_float", DTYPE_SUFFIXES) + self.assertIn("_int32", DTYPE_SUFFIXES) + + def test_int32_eq_shaders_are_named_for_the_dispatcher(self) -> None: + """An int32 `aten.eq.Tensor` must find a shader on both storage types. + + Declared as `binary_eq_int32_{buffer,texture3d}` for a while, which put + the dtype in the middle and so generated a name no dispatch could ever + build. Kokoro's synthesizer aborted on it. + """ + for template, expected in ( + ("binary_op_buffer", "binary_eq_buffer_int32"), + ("binary_op_texture", "binary_eq_texture3d_int32"), + ): + with self.subTest(template=template): + self.assertIn(expected, self.names[template]) + + def test_binary_variants_end_in_a_storage_and_dtype_suffix(self) -> None: + """The general form of the same bug, for every binary op at once. + + `add_binary_op_node` appends the storage suffix and then the dtype + suffix, in that order, to every name it builds. A generated variant that + does not end that way cannot be reached from the dispatcher, whatever + else is true about it. + """ + legal = tuple( + storage + dtype for storage in STORAGE_SUFFIXES for dtype in DTYPE_SUFFIXES + ) + unreachable = [ + name + for template in _BINARY_TEMPLATES + for name in self.names[template] + if not name.endswith(legal) + ] + self.assertEqual( + unreachable, + [], + "these shaders are generated but no dispatch can name them; " + "see add_binary_op_node in BinaryOp.cpp", + ) + + +if __name__ == "__main__": + unittest.main() From 5f7ad34d4910c4ad2a2da0d78c47ee6e24a54afe Mon Sep 17 00:00:00 2001 From: Yu-Yuan Chen Date: Wed, 9 Sep 2026 12:36:25 -0400 Subject: [PATCH 101/190] Scope -Wno-missing-prototypes to non-GCC compilers (#22347) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Summary: D100626340 keyed the `-Wno-missing-prototypes` carve-out on `ovr_config//os:zephyr`, but GCC's C++ frontend rejects the flag regardless of OS, so every non-Zephyr GCC config still fails — including the arm32 FVP config in T285774316: ``` cc1plus: error: command-line option '-Wno-missing-prototypes' is valid for C/ObjC but not for C++ [-Werror] ``` Re-key on `ovr_config//compiler:gcc`. The OS axis cannot express this: the failing GCC config and the working Clang config are both `os[embedded]`. Clang keeps the flag — it is load-bearing there (portable kernels define non-static ops whose prototypes live in generated headers the TU does not include). Where `os:windows` is also a key, the gcc check is nested under `DEFAULT` instead of made a sibling, since the two axes do not refine each other. Windows resolution unchanged. Same 12 files as D100626340. Differential Revision: D118128860 --- kernels/aten/cpu/util/targets.bzl | 10 +++-- kernels/portable/cpu/pattern/targets.bzl | 4 +- kernels/portable/cpu/util/targets.bzl | 40 +++++++++---------- shim_et/xplat/executorch/codegen/codegen.bzl | 15 +++---- .../optimized/op_registration_util.bzl | 26 ++++++------ .../kernels/portable/op_registration_util.bzl | 13 +++--- 6 files changed, 52 insertions(+), 56 deletions(-) diff --git a/kernels/aten/cpu/util/targets.bzl b/kernels/aten/cpu/util/targets.bzl index 983391d3613..411203c2b9a 100644 --- a/kernels/aten/cpu/util/targets.bzl +++ b/kernels/aten/cpu/util/targets.bzl @@ -15,11 +15,13 @@ def define_common_targets(): "copy_ops_util.h", ], compiler_flags = select({ - "DEFAULT": ["-Wno-missing-prototypes"], + "DEFAULT": select({ + "DEFAULT": ["-Wno-missing-prototypes"], + # GCC's C++ frontend rejects this C-only flag under -Werror. + # Nested under DEFAULT so windows and gcc can't both match. + "ovr_config//compiler:gcc": [], + }), "ovr_config//os:windows": [], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses - # this branch via runtime.is_oss. - "ovr_config//os:zephyr": [], }) if not runtime.is_oss else select({ "DEFAULT": ["-Wno-missing-prototypes"], "ovr_config//os:windows": [], diff --git a/kernels/portable/cpu/pattern/targets.bzl b/kernels/portable/cpu/pattern/targets.bzl index 10159c7b540..47cd80ddf1b 100644 --- a/kernels/portable/cpu/pattern/targets.bzl +++ b/kernels/portable/cpu/pattern/targets.bzl @@ -56,8 +56,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], exported_deps = [ "//executorch/kernels/portable/cpu/util:broadcast_util", diff --git a/kernels/portable/cpu/util/targets.bzl b/kernels/portable/cpu/util/targets.bzl index 99f1f3adce3..fe30188e3a8 100644 --- a/kernels/portable/cpu/util/targets.bzl +++ b/kernels/portable/cpu/util/targets.bzl @@ -49,8 +49,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/core/exec_aten/util:tensor_shape_to_c_string", @@ -103,8 +103,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -119,8 +119,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], exported_deps = [ ":broadcast_indexes_range", @@ -155,8 +155,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ ":broadcast_util", @@ -174,8 +174,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], exported_deps = [ ":broadcast_util", @@ -194,8 +194,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -214,8 +214,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -231,8 +231,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ ":broadcast_util", @@ -252,8 +252,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", @@ -269,8 +269,8 @@ def define_common_targets(): ], compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - # ovr_config//os:zephyr is fbsource-internal; OSS bypasses this select via runtime.is_oss. - "ovr_config//os:zephyr": [], + # GCC's C++ frontend rejects this C-only flag under -Werror. + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"], deps = [ "//executorch/runtime/kernel:kernel_includes", diff --git a/shim_et/xplat/executorch/codegen/codegen.bzl b/shim_et/xplat/executorch/codegen/codegen.bzl index 318996784a1..9eba82e09e0 100644 --- a/shim_et/xplat/executorch/codegen/codegen.bzl +++ b/shim_et/xplat/executorch/codegen/codegen.bzl @@ -633,13 +633,10 @@ def build_portable_lib( # Currently fbcode links all dependent libraries through shared # library, and it blocks users like unit tests to use kernel # implementation directly. So we enable this for xplat only. - # -Wno-missing-prototypes is Clang-only for C++; GCC (used by Zephyr ARM - # cross-compilation) rejects it with -Werror, so exclude it for Zephyr. - # OSS bypasses the select since ovr_config//os:zephyr is not in the OSS - # buck2 prelude. + # GCC's C++ frontend rejects this C-only flag under -Werror. compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes"], - "ovr_config//os:zephyr": [], + "ovr_config//compiler:gcc": [], }) if not runtime.is_oss else ["-Wno-missing-prototypes"] if not expose_operator_symbols and is_xplat(): # Removing '-fvisibility=hidden' exposes operator symbols. @@ -686,13 +683,11 @@ def build_optimized_lib(name, oplist_header_name, portable_header_lib, feature = # Currently fbcode links all dependent libraries through shared # library, and it blocks users like unit tests to use kernel # implementation directly. So we enable this for xplat only. - # -Wno-missing-prototypes and -Wno-global-constructors are Clang-only for - # C++; GCC (used by Zephyr ARM cross-compilation) rejects them with - # -Werror, so exclude them for Zephyr. OSS bypasses the select since - # ovr_config//os:zephyr is not in the OSS buck2 prelude. + # Drop the Clang-only flags for GCC: its C++ frontend rejects + # -Wno-missing-prototypes under -Werror. compiler_flags = select({ "DEFAULT": ["-Wno-missing-prototypes", "-Wno-pass-failed", "-Wno-global-constructors", "-Wno-shadow"], - "ovr_config//os:zephyr": ["-Wno-pass-failed", "-Wno-shadow"], + "ovr_config//compiler:gcc": ["-Wno-pass-failed", "-Wno-shadow"], }) if not runtime.is_oss else ["-Wno-missing-prototypes", "-Wno-pass-failed", "-Wno-global-constructors", "-Wno-shadow"] if not expose_operator_symbols and is_xplat(): # Removing '-fvisibility=hidden' exposes operator symbols. diff --git a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl index 7e32f5b7473..fba89adde64 100644 --- a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl @@ -96,19 +96,19 @@ def define_op_library(name, compiler_flags, deps): compiler_flags = (select({ # kernels often have helpers with no prototypes just disabling the warning here as the headers # are codegend and linked in later - # -Wno-missing-prototypes is Clang-only for C++; GCC (used by - # Zephyr ARM cross-compilation) rejects it with -Werror, so - # exclude it for Zephyr. OSS bypasses the select since - # ovr_config//os:zephyr is not in the OSS buck2 prelude. - "DEFAULT": [ - "-Wno-missing-prototypes", - # pragma unroll fails with -Os, don't need to warn us and - # fail Werror builds; see https://godbolt.org/z/zvf85vTsr - "-Wno-pass-failed", - ], - "ovr_config//os:zephyr": [ - "-Wno-pass-failed", - ], + # GCC's C++ frontend rejects this C-only flag under -Werror. Nested + # under DEFAULT so the windows (OS) and gcc keys can't both match. + "DEFAULT": select({ + "DEFAULT": [ + "-Wno-missing-prototypes", + # pragma unroll fails with -Os, don't need to warn us and + # fail Werror builds; see https://godbolt.org/z/zvf85vTsr + "-Wno-pass-failed", + ], + "ovr_config//compiler:gcc": [ + "-Wno-pass-failed", + ], + }), # The vendored ATen vec headers trip several -Werror warnings on # the Windows (clang) host, so disable warnings-as-errors there. "ovr_config//os:windows": select({ diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index f1a46616295..0038df45dc1 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -119,21 +119,20 @@ def define_op_library(name, deps, android_deps, aten_target, _allow_third_party_ visibility = ["PUBLIC"], # kernels often have helpers with no prototypes just disabling the warning here as the headers # are codegend and linked in later - # -Wno-missing-prototypes is Clang-only for C++; GCC (used by Zephyr - # ARM cross-compilation) rejects it with -Werror, so exclude it for - # Zephyr and Windows builds. OSS bypasses the zephyr branch via - # runtime.is_oss since ovr_config//os:zephyr is not in the OSS - # buck2 prelude. + # GCC's C++ frontend rejects this C-only flag under -Werror. Nested under + # DEFAULT so the windows (OS) and gcc (compiler) keys can't both match. # The vendored ATen vec headers pulled in on the Windows host trip # several -Werror warnings (e.g. -Wundef on __GNUC__), so disable # warnings-as-errors for the Windows (clang) kernel compiles. compiler_flags = (select({ - "DEFAULT": ["-Wno-missing-prototypes"], + "DEFAULT": select({ + "DEFAULT": ["-Wno-missing-prototypes"], + "ovr_config//compiler:gcc": [], + }), "ovr_config//os:windows": select({ "DEFAULT": ["-Wno-error"], "ovr_config//compiler:msvc": [], }), - "ovr_config//os:zephyr": [], }) if not runtime.is_oss else select({ "DEFAULT": ["-Wno-missing-prototypes"], # OSS buck2 has no compiler constraint (ovr_config//compiler:msvc From 9007da4d6a6450f260a5775bc828417f9f42075d Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Tue, 8 Sep 2026 20:25:35 -0700 Subject: [PATCH 102/190] [ET-VK][testing] Reduce host tensor retention Pull Request resolved: https://github.com/pytorch/executorch/pull/22624 Store `ValueSpec` tensor and reference payloads in shared copy-on-write storage. Defer constant materialization, share payloads among equivalent variants, and release completed test cases so host memory follows active work instead of accumulating across the suite. Add focused tests for lazy generation, sharing, and copy-on-write behavior. Differential Revision: [D119237809](https://our.internmc.facebook.com/intern/diff/D119237809/) ghstack-source-id: 426683623 --- backends/vulkan/test/custom_ops/targets.bzl | 15 +- backends/vulkan/test/custom_ops/utils.cpp | 131 ++++++++----- backends/vulkan/test/custom_ops/utils.h | 144 +++++++++------ .../vulkan/test/custom_ops/utils_test.cpp | 174 ++++++++++++++++++ 4 files changed, 365 insertions(+), 99 deletions(-) create mode 100644 backends/vulkan/test/custom_ops/utils_test.cpp diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index 5d1045173a2..ed463b79a2f 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -1,4 +1,4 @@ -load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID") +load("@fbsource//tools/build_defs:platform_defs.bzl", "ANDROID", "CXX") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") load( "@fbsource//xplat/executorch/backends/vulkan:targets.bzl", @@ -85,6 +85,19 @@ def define_common_targets(is_fbcode = False): link_whole = True, ) + runtime.cxx_test( + name = "utils_test", + srcs = [ + "utils_test.cpp", + ], + contacts = ["oncall+ai_infra_mobile_platform@xmail.facebook.com"], + platforms = [CXX], + deps = [ + ":prototyping_utils", + "//third-party/googletest:gtest_main", + ], + ) + define_custom_op_test_binary("test_add") define_custom_op_test_binary("test_q8csw_linear") define_custom_op_test_binary("test_q8csw_conv2d") diff --git a/backends/vulkan/test/custom_ops/utils.cpp b/backends/vulkan/test/custom_ops/utils.cpp index 60c3ddee30c..fb7121d6d53 100644 --- a/backends/vulkan/test/custom_ops/utils.cpp +++ b/backends/vulkan/test/custom_ops/utils.cpp @@ -172,11 +172,30 @@ void set_debugging(bool enable_debugging) { } // ValueSpec implementation -void ValueSpec::generate_tensor_data(int seed) { +void ValueSpec::ensure_unique_data() const { + if (data_.use_count() != 1) { + data_ = std::make_shared(*data_); + } +} + +void ValueSpec::ensure_unique_reference_data() const { + if (reference_data_.use_count() != 1) { + reference_data_ = std::make_shared(*reference_data_); + } +} + +void ValueSpec::generate_tensor_data(int seed) const { if (spec_type != SpecType::Tensor) { return; } + ensure_unique_data(); + auto& float_data = data_->float_data; + auto& int32_data = data_->int32_data; + auto& half_data = data_->half_data; + auto& int8_data = data_->int8_data; + auto& uint8_data = data_->uint8_data; + int64_t num_elements = numel(); switch (dtype) { @@ -498,81 +517,95 @@ std::string ValueSpec::to_string() const { // Additional ValueSpec methods void ValueSpec::resize_data(size_t new_size) { + // Generate first so a deferred tensor keeps its data-gen pattern (resized, + // not pinned to zeros by the data_generated_ flag set below). + ensure_data_generated(); + ensure_unique_data(); switch (dtype) { case vkapi::kFloat: - float_data.resize(new_size); + data_->float_data.resize(new_size); break; case vkapi::kHalf: - half_data.resize(new_size); + data_->half_data.resize(new_size); break; case vkapi::kInt: - int32_data.resize(new_size); + data_->int32_data.resize(new_size); break; case vkapi::kChar: - int8_data.resize(new_size); + data_->int8_data.resize(new_size); break; case vkapi::kByte: - uint8_data.resize(new_size); + data_->uint8_data.resize(new_size); break; default: - float_data.resize(new_size); + data_->float_data.resize(new_size); break; } + data_generated_ = true; } void* ValueSpec::get_mutable_data_ptr() { + ensure_data_generated(); + ensure_unique_data(); switch (dtype) { case vkapi::kFloat: - return float_data.data(); + return data_->float_data.data(); case vkapi::kHalf: - return half_data.data(); + return data_->half_data.data(); case vkapi::kInt: - return int32_data.data(); + return data_->int32_data.data(); case vkapi::kChar: - return int8_data.data(); + return data_->int8_data.data(); case vkapi::kByte: - return uint8_data.data(); + return data_->uint8_data.data(); default: - return float_data.data(); + return data_->float_data.data(); } } float ValueSpec::get_element(size_t index) const { + ensure_data_generated(); if (index >= static_cast(numel())) { return 0.0f; } switch (dtype) { case vkapi::kFloat: - return index < float_data.size() ? float_data[index] : 0.0f; + return index < data_->float_data.size() ? data_->float_data[index] : 0.0f; case vkapi::kHalf: - return index < half_data.size() ? half_to_float(half_data[index]) : 0.0f; + return index < data_->half_data.size() + ? half_to_float(data_->half_data[index]) + : 0.0f; case vkapi::kInt: - return index < int32_data.size() ? static_cast(int32_data[index]) - : 0.0f; + return index < data_->int32_data.size() + ? static_cast(data_->int32_data[index]) + : 0.0f; case vkapi::kChar: - return index < int8_data.size() ? static_cast(int8_data[index]) - : 0.0f; + return index < data_->int8_data.size() + ? static_cast(data_->int8_data[index]) + : 0.0f; case vkapi::kByte: - return index < uint8_data.size() ? static_cast(uint8_data[index]) - : 0.0f; + return index < data_->uint8_data.size() + ? static_cast(data_->uint8_data[index]) + : 0.0f; default: return 0.0f; } } const void* ValueSpec::get_data_ptr() const { + ensure_data_generated(); switch (dtype) { case vkapi::kFloat: - return float_data.data(); + return data_->float_data.data(); case vkapi::kHalf: - return half_data.data(); + return data_->half_data.data(); case vkapi::kInt: - return int32_data.data(); + return data_->int32_data.data(); case vkapi::kChar: - return int8_data.data(); + return data_->int8_data.data(); case vkapi::kByte: - return uint8_data.data(); + return data_->uint8_data.data(); default: throw std::runtime_error("Unsupported data type for get_data_ptr"); } @@ -801,7 +834,7 @@ bool ValueSpec::validate_against_reference( } // Ensure data is generated for this ValueSpec -void ValueSpec::ensure_data_generated(int seed) { +void ValueSpec::ensure_data_generated(int seed) const { if (data_generated_) { return; } @@ -809,18 +842,22 @@ void ValueSpec::ensure_data_generated(int seed) { data_generated_ = true; } -// Copy input data from another ValueSpec -void ValueSpec::copy_data_from(const ValueSpec& other) { +void ValueSpec::share_data_from(const ValueSpec& other) { + if (!is_tensor() || !other.is_tensor()) { + return; + } + // Materialize the source first: sharing an ungenerated payload would let a + // later access materialize each spec independently under different seeds. + other.ensure_data_generated(); + data_ = other.data_; + data_generated_ = true; +} + +void ValueSpec::share_reference_from(const ValueSpec& other) { if (!is_tensor() || !other.is_tensor()) { return; } - // Copy raw data based on dtype - float_data = other.float_data; - int32_data = other.int32_data; - half_data = other.half_data; - int8_data = other.int8_data; - uint8_data = other.uint8_data; - data_generated_ = other.data_generated_; + reference_data_ = other.reference_data_; } // ReferenceKey implementation @@ -1743,17 +1780,11 @@ TestResult execute_test_cases( // Compute reference once for prototype bool ref_computed = false; - std::vector> ref_data; if (reference_compute_func) { try { reference_compute_func(prototype); ref_computed = true; - - // Cache the reference output for this group - for (const auto& output : prototype.outputs()) { - ref_data.push_back(output.get_ref_float_data()); - } - } catch (const std::invalid_argument& _) { + } catch (const std::invalid_argument&) { // Reference computation skipped for this group } } @@ -1771,15 +1802,21 @@ TestResult execute_test_cases( const auto& src = prototype.inputs()[j]; if (dest.is_tensor() && src.is_tensor() && dest.sizes == src.sizes && dest.dtype == src.dtype) { - dest.copy_data_from(src); + dest.share_data_from(src); } } // Copy reference output data if available if (ref_computed) { - for (size_t j = 0; j < tc.outputs().size() && j < ref_data.size(); + for (size_t j = 0; + j < tc.outputs().size() && j < prototype.outputs().size(); ++j) { - tc.outputs()[j].get_ref_float_data() = ref_data[j]; + const auto& src = prototype.outputs()[j]; + auto& dest = tc.outputs()[j]; + if (dest.is_tensor() && src.is_tensor() && dest.sizes == src.sizes && + dest.dtype == src.dtype) { + dest.share_reference_from(src); + } } } } @@ -1924,6 +1961,8 @@ TestResult execute_test_cases( // Add result to collection results.add_result(std::move(result)); + + test_case.clear(); } } diff --git a/backends/vulkan/test/custom_ops/utils.h b/backends/vulkan/test/custom_ops/utils.h index 2174ceb5618..802d149d8eb 100644 --- a/backends/vulkan/test/custom_ops/utils.h +++ b/backends/vulkan/test/custom_ops/utils.h @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -220,21 +221,8 @@ struct ValueSpec { bool is_constant_tensor; bool is_none_flag; bool is_int4_tensor; - bool data_generated_ = false; - - std::vector float_data; - std::vector int32_data; - std::vector half_data; // Using uint16_t as substitute for half - std::vector int8_data; // For kChar (signed 8-bit) - std::vector uint8_data; // For kByte (unsigned 8-bit) std::string string_data; - std::vector ref_float_data; - std::vector ref_int32_data; - std::vector ref_half_data; - std::vector ref_int8_data; - std::vector ref_uint8_data; - ValueSpec( const std::vector& sizes, vkapi::ScalarType dtype, @@ -250,7 +238,8 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(false) { - // Data generation is deferred until ensure_data_generated() is called + // Data generation is deferred until first access (any data getter or + // ensure_data_generated() triggers it). } // Constructor for tensor with custom data generation type @@ -270,7 +259,8 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(false) { - // Data generation is deferred until ensure_data_generated() is called + // Data generation is deferred until first access (any data getter or + // ensure_data_generated() triggers it). } // Constructor for single int @@ -285,7 +275,7 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(true) { - int32_data.push_back(value); + data_->int32_data.push_back(value); } // Constructor for single float @@ -300,7 +290,7 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(true) { - float_data.push_back(value); + data_->float_data.push_back(value); } // Constructor for single bool @@ -315,7 +305,7 @@ struct ValueSpec { is_none_flag(false), is_int4_tensor(false), data_generated_(true) { - int32_data.push_back(value ? 1 : 0); + data_->int32_data.push_back(value ? 1 : 0); } // Constructor for int list @@ -329,8 +319,9 @@ struct ValueSpec { is_constant_tensor(false), is_none_flag(false), is_int4_tensor(false), - data_generated_(true), - int32_data(values) {} + data_generated_(true) { + data_->int32_data = values; + } // Factory method for string (avoids ambiguity with vector constructor) static ValueSpec make_string(const std::string& value) { @@ -385,98 +376,136 @@ struct ValueSpec { } int32_t get_int_value() const { - return int32_data.empty() ? 0 : int32_data[0]; + ensure_data_generated(); + return data_->int32_data.empty() ? 0 : data_->int32_data[0]; } float get_float_value() const { - return float_data.empty() ? 0.0f : float_data[0]; + ensure_data_generated(); + return data_->float_data.empty() ? 0.0f : data_->float_data[0]; } bool get_bool_value() const { - return int32_data.empty() ? false : (int32_data[0] != 0); + ensure_data_generated(); + return data_->int32_data.empty() ? false : (data_->int32_data[0] != 0); } const std::string& get_string_value() const { return string_data; } const std::vector& get_int_list() const { - return int32_data; + ensure_data_generated(); + return data_->int32_data; } const std::vector& get_tensor_sizes() const { return sizes; } + // References and pointers into tensor data must not be held across any other + // access to the same spec: a mutable access may detach the shared payload, + // leaving a previously returned reference bound to the old payload. Consume + // immediately. const std::vector& get_float_data() const { - return float_data; + ensure_data_generated(); + return data_->float_data; } const std::vector& get_int32_data() const { - return int32_data; + ensure_data_generated(); + return data_->int32_data; } const std::vector& get_half_data() const { - return half_data; + ensure_data_generated(); + return data_->half_data; } const std::vector& get_int8_data() const { - return int8_data; + ensure_data_generated(); + return data_->int8_data; } const std::vector& get_uint8_data() const { - return uint8_data; + ensure_data_generated(); + return data_->uint8_data; } std::vector& get_float_data() { - return float_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->float_data; } std::vector& get_int32_data() { - return int32_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->int32_data; } std::vector& get_half_data() { - return half_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->half_data; } std::vector& get_int8_data() { - return int8_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->int8_data; } std::vector& get_uint8_data() { - return uint8_data; + ensure_data_generated(); + ensure_unique_data(); + return data_->uint8_data; } const std::vector& get_ref_float_data() const { - return ref_float_data; + return reference_data_->float_data; } const std::vector& get_ref_int32_data() const { - return ref_int32_data; + return reference_data_->int32_data; } const std::vector& get_ref_half_data() const { - return ref_half_data; + return reference_data_->half_data; } const std::vector& get_ref_int8_data() const { - return ref_int8_data; + return reference_data_->int8_data; } const std::vector& get_ref_uint8_data() const { - return ref_uint8_data; + return reference_data_->uint8_data; } std::vector& get_ref_float_data() { - return ref_float_data; + ensure_unique_reference_data(); + return reference_data_->float_data; } std::vector& get_ref_int32_data() { - return ref_int32_data; + ensure_unique_reference_data(); + return reference_data_->int32_data; } std::vector& get_ref_half_data() { - return ref_half_data; + ensure_unique_reference_data(); + return reference_data_->half_data; } std::vector& get_ref_int8_data() { - return ref_int8_data; + ensure_unique_reference_data(); + return reference_data_->int8_data; } std::vector& get_ref_uint8_data() { - return ref_uint8_data; + ensure_unique_reference_data(); + return reference_data_->uint8_data; } void resize_data(size_t new_size); void* get_mutable_data_ptr(); float get_element(size_t index) const; - // Data generation methods for deferred generation and caching + // Data generation methods for deferred generation and caching. + // + // ValueSpec is not thread-safe: lazy materialization and copy-on-write + // detach mutate shared state from const methods. Test cases are built and + // executed on a single thread. + // + // Implicit materialization (any data getter, resize_data) consumes the + // global seed counter. Callers needing deterministic data must call + // ensure_data_generated(explicit_seed) before any other access; a later + // seeded call is a no-op once data is generated. bool is_data_generated() const { return data_generated_; } - void ensure_data_generated(int seed = -1); - void copy_data_from(const ValueSpec& other); + void ensure_data_generated(int seed = -1) const; + void share_data_from(const ValueSpec& other); + void share_reference_from(const ValueSpec& other); // Set/get constant flag bool is_constant() const { @@ -484,10 +513,6 @@ struct ValueSpec { } void set_constant(bool is_constant) { is_constant_tensor = is_constant; - // Constant tensors need data immediately for test case setup - if (is_constant && is_tensor()) { - ensure_data_generated(); - } } // Set/get none flag @@ -517,7 +542,22 @@ struct ValueSpec { float rel_tolerance = 1e-3f) const; private: - void generate_tensor_data(int seed = -1); + struct TensorData { + std::vector float_data; + std::vector int32_data; + std::vector half_data; + std::vector int8_data; + std::vector uint8_data; + }; + + void ensure_unique_data() const; + void ensure_unique_reference_data() const; + void generate_tensor_data(int seed = -1) const; + + mutable bool data_generated_ = false; + mutable std::shared_ptr data_ = std::make_shared(); + mutable std::shared_ptr reference_data_ = + std::make_shared(); }; // diff --git a/backends/vulkan/test/custom_ops/utils_test.cpp b/backends/vulkan/test/custom_ops/utils_test.cpp new file mode 100644 index 00000000000..1d7f7e1ba13 --- /dev/null +++ b/backends/vulkan/test/custom_ops/utils_test.cpp @@ -0,0 +1,174 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include "utils.h" + +namespace executorch::vulkan::prototyping { +namespace { + +TEST(ValueSpecTest, SetConstant_DoesNotMaterializeTensorData) { + ValueSpec value( + {16}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + + value.set_constant(true); + + EXPECT_FALSE(value.is_data_generated()); +} + +TEST(ValueSpecTest, ShareDataFrom_SharesImmutableTensorData) { + ValueSpec source( + {16}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + source.ensure_data_generated(); + + ValueSpec copy( + {16}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kTexture3D, + vkcompute::utils::kChannelsPacked, + DataGenType::ONES); + copy.share_data_from(source); + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_EQ( + const_source.get_float_data().data(), const_copy.get_float_data().data()); +} + +TEST(ValueSpecTest, MutableDataAccess_DetachesSharedTensorData) { + ValueSpec source( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + source.ensure_data_generated(); + + ValueSpec copy = source; + copy.get_float_data()[0] = 7.0f; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_FLOAT_EQ(const_source.get_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_copy.get_float_data()[0], 7.0f); + EXPECT_NE( + const_source.get_float_data().data(), const_copy.get_float_data().data()); +} + +TEST(ValueSpecTest, CopyConstruction_SharesImmutableReferenceData) { + ValueSpec source( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ZEROS); + source.get_ref_float_data() = {1.0f, 2.0f, 3.0f, 4.0f}; + + ValueSpec copy = source; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_EQ( + const_source.get_ref_float_data().data(), + const_copy.get_ref_float_data().data()); +} + +TEST(ValueSpecTest, MutableReferenceAccess_DetachesSharedReferenceData) { + ValueSpec source( + {2}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ZEROS); + source.get_ref_float_data() = {1.0f, 2.0f}; + + ValueSpec copy = source; + copy.get_ref_float_data()[0] = 9.0f; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_FLOAT_EQ(const_source.get_ref_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_copy.get_ref_float_data()[0], 9.0f); + EXPECT_NE( + const_source.get_ref_float_data().data(), + const_copy.get_ref_float_data().data()); +} + +TEST(ValueSpecTest, ConstGetter_MaterializesDeferredData) { + ValueSpec value( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + const ValueSpec& const_value = value; + EXPECT_FALSE(const_value.is_data_generated()); + EXPECT_FLOAT_EQ(const_value.get_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_value.get_float_value(), 1.0f); + EXPECT_TRUE(const_value.is_data_generated()); +} + +TEST(ValueSpecTest, ResizeData_PreservesGeneratedPattern) { + ValueSpec value( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + value.resize_data(8); + const ValueSpec& const_value = value; + const std::vector expected( + {1.0f, 1.0f, 1.0f, 1.0f, 0.0f, 0.0f, 0.0f, 0.0f}); + EXPECT_EQ(const_value.get_float_data(), expected); +} + +TEST(ValueSpecTest, MutableDataPtr_DetachesSharedTensorData) { + ValueSpec source( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ONES); + source.ensure_data_generated(); + + ValueSpec copy = source; + auto* mutable_ptr = static_cast(copy.get_mutable_data_ptr()); + ASSERT_NE(mutable_ptr, nullptr); + mutable_ptr[0] = 7.0f; + + const ValueSpec& const_source = source; + const ValueSpec& const_copy = copy; + EXPECT_FLOAT_EQ(const_source.get_float_data()[0], 1.0f); + EXPECT_FLOAT_EQ(const_copy.get_float_data()[0], 7.0f); +} + +TEST(ValueSpecTest, ShareReferenceFrom_IgnoresNonTensorSpecs) { + ValueSpec scalar(3); + ValueSpec tensor( + {4}, + vkcompute::vkapi::kFloat, + vkcompute::utils::kBuffer, + vkcompute::utils::kWidthPacked, + DataGenType::ZEROS); + tensor.get_ref_float_data() = {1.0f, 2.0f, 3.0f, 4.0f}; + const void* before = tensor.get_ref_float_data().data(); + tensor.share_reference_from(scalar); + EXPECT_EQ(tensor.get_ref_float_data().data(), before); + const std::vector expected({1.0f, 2.0f, 3.0f, 4.0f}); + EXPECT_EQ(tensor.get_ref_float_data(), expected); +} + +} // namespace +} // namespace executorch::vulkan::prototyping From dc8c5789273153791486bff0a1b6a413e5d3a41c Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Tue, 8 Sep 2026 20:25:35 -0700 Subject: [PATCH 103/190] [ET-VK][testing] Reuse allocations in chained benchmarks Pull Request resolved: https://github.com/pytorch/executorch/pull/22625 Build each benchmark operator once and replay its `execute_nodes()` in a benchmark-owned reusable command buffer. This prevents chained timing from multiplying persistent prepack and intermediate GPU allocations while leaving production `ComputeGraph::execute()` behavior unchanged. Differential Revision: [D119237811](https://our.internmc.facebook.com/intern/diff/D119237811/) ghstack-source-id: 426683626 --- backends/vulkan/test/custom_ops/utils.cpp | 116 ++++++++++++++---- backends/vulkan/test/custom_ops/utils.h | 84 ++++++++++--- .../vulkan/test/custom_ops/utils_test.cpp | 53 ++++++++ 3 files changed, 208 insertions(+), 45 deletions(-) diff --git a/backends/vulkan/test/custom_ops/utils.cpp b/backends/vulkan/test/custom_ops/utils.cpp index fb7121d6d53..b00e6b175da 100644 --- a/backends/vulkan/test/custom_ops/utils.cpp +++ b/backends/vulkan/test/custom_ops/utils.cpp @@ -31,6 +31,56 @@ constexpr float kMinProbeTimeUs = 1.0f; } // namespace +// Benchmark graphs are immutable after input upload, so the operator nodes can +// be recorded repeatedly. Staging uploads/downloads are encoded once: +// repeating them would redo host-visible copies on every repetition and bias +// per-invocation timings. +RepeatedGraphExecutor::RepeatedGraphExecutor( + ComputeGraph& graph, + int repetitions, + OpNodeRange op_nodes) + : graph_(graph) { + VK_CHECK_COND(repetitions > 0); + VK_CHECK_COND(op_nodes.begin <= op_nodes.end); + VK_CHECK_COND(op_nodes.end <= graph_.execute_nodes().size()); + + api::Context* const context = graph_.context(); + context->flush(); + context->set_cmd(/*reusable=*/true); + context->cmd_reset_querypool(); + + for (size_t i = 0; i < op_nodes.begin; ++i) { + graph_.execute_nodes()[i]->encode(&graph_); + } + for (int i = 0; i < repetitions; ++i) { + for (size_t j = op_nodes.begin; j < op_nodes.end; ++j) { + graph_.execute_nodes()[j]->encode(&graph_); + } + } + for (size_t i = op_nodes.end; i < graph_.execute_nodes().size(); ++i) { + graph_.execute_nodes()[i]->encode(&graph_); + } + + command_ = + std::make_unique(std::move(context->extract_cmd())); +} + +void RepeatedGraphExecutor::execute() { + api::Context* const context = graph_.context(); + command_->end(); + + // Intentionally bypasses Context::submit_cmd_to_gpu(): its submit-count + // bookkeeping and threshold-split path only serve submit_compute_job + // recording, which benchmarks don't use. + vkapi::VulkanFence fence = context->fences().get_fence(); + context->adapter_ptr()->submit_cmd( + context->queue(), + command_->get_submit_handle(/*final_use=*/false), + fence.get_submit_handle()); + fence.wait(); + context->fences().return_fence(fence); +} + int get_seed() { static int seed = 42; return seed++; @@ -1433,17 +1483,23 @@ int64_t default_flop_calculator(const TestCase& test_case) { return total_elements; } -ComputeGraph setup_compute_graph( +BenchmarkGraph setup_compute_graph( TestCase& test_case, std::string op_name, int op_invocations_per_execute) { GraphConfig config; config.enable_querypool = true; + // Pool sizing takes max(execute, prepack) * factor, so scaling the factor + // also over-reserves the prepack side (encoded once). Accepted: precise + // execute-only sizing would need runtime changes. + config.descriptor_pool_safety_factor *= + std::max(1, op_invocations_per_execute); // Default-on (opt-out via TestCase::set_force_resize(false)): force every - // DynamicDispatchNode to run its resize function on each execute(), - // exercising the op's resize formula even when input shapes are unchanged. + // DynamicDispatchNode to run its resize function when execute_test_case + // runs propagate_resize() after prepack, exercising the op's resize formula + // even when input shapes are unchanged. config.force_resize = test_case.get_force_resize(); - ComputeGraph graph(config); + auto graph = std::make_unique(config); std::vector input_values; @@ -1452,17 +1508,17 @@ ComputeGraph setup_compute_graph( const ValueSpec& input_spec = test_case.inputs()[i]; if (input_spec.is_none()) { - input_values.push_back(graph.add_none()); + input_values.push_back(graph->add_none()); } else if (input_spec.is_float()) { ValueRef input_value = - graph.add_scalar(static_cast(input_spec.get_float_value())); + graph->add_scalar(static_cast(input_spec.get_float_value())); input_values.push_back(input_value); } else if (input_spec.is_int()) { ValueRef input_value = - graph.add_scalar(static_cast(input_spec.get_int_value())); + graph->add_scalar(static_cast(input_spec.get_int_value())); input_values.push_back(input_value); } else if (input_spec.is_bool()) { - ValueRef input_value = graph.add_scalar(input_spec.get_bool_value()); + ValueRef input_value = graph->add_scalar(input_spec.get_bool_value()); input_values.push_back(input_value); } else if (input_spec.is_int_list()) { // Convert int32_t list to int64_t list for ComputeGraph @@ -1472,20 +1528,20 @@ ComputeGraph setup_compute_graph( for (int32_t val : int32_list) { int64_list.push_back(static_cast(val)); } - ValueRef input_value = graph.add_scalar_list(std::move(int64_list)); + ValueRef input_value = graph->add_scalar_list(std::move(int64_list)); input_values.push_back(input_value); } else if (input_spec.is_string()) { std::string str_copy = input_spec.get_string_value(); - ValueRef input_value = graph.add_string(std::move(str_copy)); + ValueRef input_value = graph->add_string(std::move(str_copy)); input_values.push_back(input_value); } else if (input_spec.is_constant()) { - ValueRef input_value = graph.add_tensorref( + ValueRef input_value = graph->add_tensorref( input_spec.get_tensor_sizes(), input_spec.dtype, input_spec.get_data_ptr()); input_values.push_back(input_value); } else { - IOValueRef input_io = graph.add_input_tensor( + IOValueRef input_io = graph->add_input_tensor( input_spec.get_tensor_sizes(), input_spec.dtype, input_spec.storage_type, @@ -1505,7 +1561,7 @@ ComputeGraph setup_compute_graph( } // Create output tensor - ValueRef output_value = graph.add_tensor( + ValueRef output_value = graph->add_tensor( output_spec.get_tensor_sizes(), output_spec.dtype, output_spec.storage_type, @@ -1521,17 +1577,16 @@ ComputeGraph setup_compute_graph( std::vector op_args = input_values; op_args.insert(op_args.end(), output_values.begin(), output_values.end()); - // Invoke the op op_invocations_per_execute times to stack dispatches per - // graph.execute(). The output set_output_value() calls below still happen - // exactly once. - for (int i = 0; i < op_invocations_per_execute; ++i) { - opFn(graph, op_args); - } + // Nodes added before the op are staging uploads; nodes added after are + // staging downloads. Only the op's own nodes are repeated by benchmarks. + const size_t op_begin = graph->execute_nodes().size(); + opFn(*graph, op_args); + const size_t op_end = graph->execute_nodes().size(); for (size_t i = 0; i < output_values.size(); ++i) { - graph.set_output_value(output_values[i]); + graph->set_output_value(output_values[i]); } - return graph; + return {std::move(graph), {op_begin, op_end}}; } // Test execution utilities @@ -1549,16 +1604,20 @@ BenchmarkResult execute_test_case( api::context()->initialize_querypool(); } - // Build the measurement graph with the requested chained_dispatches factor. - // The caller (typically execute_test_cases) decides what it should be — - // this function is a pure "run at the given chained_dispatches" primitive. - ComputeGraph graph = setup_compute_graph( + // Build the operator once. Benchmark repetition is encoded separately so + // persistent graph allocations are not duplicated. + BenchmarkGraph benchmark = setup_compute_graph( test_case, test_case.operator_name(), chained_dispatches); + ComputeGraph& graph = *benchmark.graph; // Prepare the graph graph.prepare(); graph.prepack(); + // Run resize functions once so force_resize exercises resize formulas even + // though the record/replay path below never calls propagate_resize(). + graph.propagate_resize(); + // Copy input data into the graph's staging buffers size_t graph_input_idx = 0; for (size_t i = 0; i < test_case.num_inputs(); ++i) { @@ -1605,9 +1664,12 @@ BenchmarkResult execute_test_case( ++graph_input_idx; } + RepeatedGraphExecutor graph_executor( + graph, chained_dispatches, benchmark.op_nodes); + // Warmup runs for (int run = 0; run < warmup_runs; ++run) { - graph.execute(); + graph_executor.execute(); } // Benchmark runs - collect individual iteration timings @@ -1619,7 +1681,7 @@ BenchmarkResult execute_test_case( for (int run = 0; run < benchmark_runs; ++run) { // Measure CPU time for each execute() call auto cpu_start = std::chrono::high_resolution_clock::now(); - graph.execute(); + graph_executor.execute(); auto cpu_end = std::chrono::high_resolution_clock::now(); auto cpu_duration = std::chrono::duration_cast( diff --git a/backends/vulkan/test/custom_ops/utils.h b/backends/vulkan/test/custom_ops/utils.h index 802d149d8eb..c7d537e6a2b 100644 --- a/backends/vulkan/test/custom_ops/utils.h +++ b/backends/vulkan/test/custom_ops/utils.h @@ -621,9 +621,9 @@ class TestCase { return shader_filter_; } - // Manual override for the number of times the op is dispatched per - // graph.execute() (a.k.a. chained_dispatches). If > 0, the framework uses - // this directly and skips the probe phase. 0 (the default) means adaptive + // Manual override for the number of chained dispatches per measurement + // iteration (a.k.a. chained_dispatches). If > 0, the framework uses this + // directly and skips the probe phase. 0 (the default) means adaptive // (probe-then-scale). void set_op_invocations_per_execute(int n) { op_invocations_per_execute_ = n; @@ -645,13 +645,15 @@ class TestCase { // When true, the ComputeGraph built for this test case sets // GraphConfig::force_resize, so every DynamicDispatchNode runs its resize - // function on each execute() even when no input shape changed. Because the - // output is already allocated at the swept shape, the resize must recompute - // the same shape from the current input — a wrong resize formula resizes the - // output to a mismatched shape and surfaces as a test failure. Default true - // (opt-out): every custom_ops test exercises its resize formulas across the - // swept shapes. Call set_force_resize(false) for the rare op whose resize fn - // is intentionally not shape-preserving under a fixed output allocation. + // function once during measurement setup (execute_test_case runs + // propagate_resize() after prepack) even when no input shape changed. + // Because the output is already allocated at the swept shape, the resize + // must recompute the same shape from the current input — a wrong resize + // formula resizes the output to a mismatched shape and surfaces as a test + // failure. Default true (opt-out): every custom_ops test exercises its + // resize formulas across the swept shapes. Call set_force_resize(false) for + // the rare op whose resize fn is intentionally not shape-preserving under a + // fixed output allocation. void set_force_resize(bool force_resize) { force_resize_ = force_resize; } @@ -940,9 +942,55 @@ int64_t default_flop_calculator(const TestCase& test_case); using ReferenceComputeFunc = std::function; -// Runs a measurement at the given chained_dispatches factor (how many times -// the op is stacked inside one graph.execute()). This is a primitive; the -// probe-then-scale orchestration lives in execute_test_cases(). +// Half-open index range of the operator's own dispatch nodes within a +// benchmark graph's execute_nodes(). Staging upload nodes precede it, staging +// download nodes follow it. +struct OpNodeRange { + size_t begin = 0; + size_t end = 0; +}; + +// A benchmark graph plus the location of its repeatable operator nodes. The +// graph is heap-held: ComputeGraph owns its Context and must never be moved +// (a moved-from graph's destructor dereferences a null context). +struct BenchmarkGraph { + std::unique_ptr graph; + OpNodeRange op_nodes; +}; + +// Benchmark-only executor that records a graph's execute nodes into a single +// reusable command buffer, then replays it on every execute(). Staging uploads +// are encoded once, operator nodes N times, staging downloads once, so each +// iteration performs the same work as the old stacked-nodes layout (1 upload + +// N ops + 1 download) and the per-invocation divisor is unchanged. Production +// ComputeGraph::execute() behavior is unchanged. +// +// Notes for interpreting benchmark numbers: +// - Submit granularity differs from stacking N distinct nodes: all encodings +// live in one command buffer with one submit per iteration (the old path +// could split across command buffers at the node-count threshold), so +// per-dispatch times may shift systematically against older data. +// - Resize functions run once via propagate_resize() in execute_test_case +// before recording; replay itself never re-triggers resize. +// - Repeated encodings share one node/dispatch id, so per-repetition +// querypool attribution is unavailable (aggregation keys on kernel name). +class RepeatedGraphExecutor final { + public: + RepeatedGraphExecutor( + ComputeGraph& graph, + int repetitions, + OpNodeRange op_nodes); + void execute(); + + private: + ComputeGraph& graph_; + std::unique_ptr command_; +}; + +// Runs a measurement at the given chained_dispatches factor. The operator is +// built once, then its execute nodes are encoded that many times into a +// benchmark-only reusable command buffer. The probe-then-scale orchestration +// lives in execute_test_cases(). // // write_outputs controls whether the graph's staging output buffers are copied // back into test_case.outputs() at the end of the run. The probe path needs @@ -1020,11 +1068,11 @@ void compute_weight_sums_4bit_grouped( uint16_t float_to_half(float value); float half_to_float(uint16_t half_val); -// Setup compute graph based on TestCase and operation name. The op function -// is invoked op_invocations_per_execute times so that one graph.execute() -// dispatches the op that many times (Google Benchmark-style stacking). The -// output set_output_value() calls still happen once at the end. -ComputeGraph setup_compute_graph( +// Setup compute graph based on TestCase and operation name. The op function is +// invoked once. op_invocations_per_execute is used only to reserve enough +// descriptor capacity for benchmark-only repeated command encoding. Returns +// the graph plus the range of the operator's own nodes for repeated encoding. +BenchmarkGraph setup_compute_graph( TestCase& test_case, std::string op_name, int op_invocations_per_execute = 1); diff --git a/backends/vulkan/test/custom_ops/utils_test.cpp b/backends/vulkan/test/custom_ops/utils_test.cpp index 1d7f7e1ba13..6d2de13d71e 100644 --- a/backends/vulkan/test/custom_ops/utils_test.cpp +++ b/backends/vulkan/test/custom_ops/utils_test.cpp @@ -11,6 +11,32 @@ namespace executorch::vulkan::prototyping { namespace { +int operator_build_count = 0; +int encode_count = 0; + +class CountingEncodeNode final : public ExecuteNode { + public: + explicit CountingEncodeNode(int& count) : count_(count) {} + void encode(ComputeGraph* graph) override { + (void)graph; + ++count_; + } + + private: + int& count_; +}; + +void counting_operator(ComputeGraph& graph, const std::vector& args) { + (void)args; + ++operator_build_count; + graph.execute_nodes().emplace_back( + std::make_unique(encode_count)); +} + +REGISTER_OPERATORS { + VK_REGISTER_OP(test_etvk.counting_operator.default, counting_operator); +} + TEST(ValueSpecTest, SetConstant_DoesNotMaterializeTensorData) { ValueSpec value( {16}, @@ -170,5 +196,32 @@ TEST(ValueSpecTest, ShareReferenceFrom_IgnoresNonTensorSpecs) { EXPECT_EQ(tensor.get_ref_float_data(), expected); } +TEST(BenchmarkGraphTest, ChainedDispatches_BuildOperatorOnce) { + if (!vkcompute::api::available()) { + return; + } + + TestCase test_case; + operator_build_count = 0; + encode_count = 0; + + BenchmarkGraph benchmark = setup_compute_graph( + test_case, + "test_etvk.counting_operator.default", + /*op_invocations_per_execute=*/8); + ComputeGraph& graph = *benchmark.graph; + + EXPECT_EQ(operator_build_count, 1); + ASSERT_EQ(graph.execute_nodes().size(), 1u); + EXPECT_EQ(benchmark.op_nodes.begin, 0u); + EXPECT_EQ(benchmark.op_nodes.end, 1u); + + // The replay path must encode the operator node once per repetition. + RepeatedGraphExecutor graph_executor( + graph, /*repetitions=*/8, benchmark.op_nodes); + EXPECT_EQ(encode_count, 8); + graph_executor.execute(); +} + } // namespace } // namespace executorch::vulkan::prototyping From e1e551bf7ff6f0bfe3c4878e8253633c21ea8752 Mon Sep 17 00:00:00 2001 From: Sangwon Ha <146179778+FabulousSuperDude@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:38:12 +0100 Subject: [PATCH 104/190] Arm backend: Relax DeepSeek VGF BF16 abs tolerance (#22644) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The xlarge DeepSeek VGF BF16 test is numerically stable in its mean error but exceeds the existing 0.1 absolute tolerance across platforms and random seeds. | Platform | Seed | Result | Max |Δ| | Mean |Δ| | |-----------|------|--------|----------|----------| | Ubuntu x86| 0 | Fail | 0.140625 | 0.026189 | | Ubuntu x86| 1 | Fail | 0.150391 | 0.026114 | | Ubuntu x86| 2 | Fail | 0.134766 | 0.025801 | | macOS | 0 | Fail | 0.156250 | 0.025820 | | macOS | 1 | Fail | 0.153320 | 0.025708 | | macOS | 2 | Pass | — | — | Raise atol from 0.1 to 0.17, providing modest headroom above the largest observed difference while retaining rtol=0.1. The exact test passes on Ubuntu x86 with seeds 0, 1, and 2 after this change. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Sangwon Ha --- .../test_deepseek_r1_distill_qwen_model.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py index 34f87c1ca8b..792cc689c39 100644 --- a/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py +++ b/backends/arm/test/models/DeepSeek_R1_Distill_Qwen/test_deepseek_r1_distill_qwen_model.py @@ -197,7 +197,7 @@ class DeepSeekR1DistillQwenModelTestCase: "base_model": DeepSeekR1DistillQwenModelTestCase( model_cls=BaseModelWrapper, config_factory=_make_deepseek_r1_distill_qwen_1_5b_model_config, - atol=0.1, + atol=0.17, rtol=0.1, tosa_spec="TOSA-1.0+FP+bf16", ), From 3dad019b1901a244303f3c4d6d67b43bdc1a95e6 Mon Sep 17 00:00:00 2001 From: Sangwon Ha <146179778+FabulousSuperDude@users.noreply.github.com> Date: Wed, 9 Sep 2026 18:38:25 +0100 Subject: [PATCH 105/190] Arm backend: Add tests for Stable Diffusion 3.5 Large model (#22606) - Add SD3.5 Large test configs for CLIP text encoders, T5 text encoder, MMDiT transformer, and VAE decoder, with optional upstream-sync validation. - Add Arm backend lowering tests for CLIP-L, CLIP-bigG, T5, transformer, and VAE decoder components across TOSA and VGF pipelines. - Add wrapper and model-loader unit tests covering wrapper output signatures, component subfolder loading, selective component loading, and dummy input shape generation. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --------- Signed-off-by: Sangwon Ha --- ..._CLIPTextModelWithProjection_sd35_large.py | 218 +++++++++++ .../test_SD3Transformer2DModel_sd35_large.py | 172 ++++++++ .../test_T5EncoderModel_sd35_large.py | 175 +++++++++ .../test_configs_sd35_large.py | 311 +++++++++++++++ .../test_model_sd35_large.py | 369 ++++++++++++++++++ .../test_vae_AutoencoderKL_sd35_large.py | 152 ++++++++ 6 files changed, 1397 insertions(+) create mode 100644 backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py create mode 100644 backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py create mode 100644 backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py create mode 100644 backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py create mode 100644 backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py create mode 100644 backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py new file mode 100644 index 00000000000..34d61805e4d --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_CLIPTextModelWithProjection_sd35_large.py @@ -0,0 +1,218 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_int64_to_int32_passes, + get_tiny_sd35_large_text_encoder_2_config, + get_tiny_sd35_large_text_encoder_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3CLIPTextEncoderWrapper, +) +from transformers import CLIPTextModelWithProjection + +input_t = Tuple[torch.Tensor] + + +class TestCLIPTextModelWithProjection: + """Test helper for SD3.5 Large CLIPTextModelWithProjection configs.""" + + ops_after_partitioner_FP = { + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_slice_copy_Tensor": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 2, + } + + ops_after_partitioner_vgf_no_quantize = { + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 2, + } + ops_after_partitioner_vgf_quantize = { + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + def create_dummy_inputs( + self, + config, + batch_size: int = 1, + seq_length: int = 2, + dtype: torch.dtype = torch.long, + ) -> tuple[torch.Tensor]: + """Create dummy inputs for the CLIPTextModelWithProjection tests.""" + # SD3.5 Large uses (batch_size, seq_length) = (1, 77) for both CLIP-L + # and CLIP-bigG. Keep this unit-test default smaller for TOSA runtime. + return ( + torch.randint( + low=0, + high=config.vocab_size, + size=(batch_size, seq_length), + dtype=dtype, + ), + ) + + def create_model( + self, + config, + ) -> SD3CLIPTextEncoderWrapper: + """Instantiate wrapped CLIPTextModelWithProjection for tests.""" + return SD3CLIPTextEncoderWrapper( + CLIPTextModelWithProjection(config).to(dtype=config.dtype) # type: ignore[call-arg] + ).eval() + + @staticmethod + def ops_after_partitioner_INT(config) -> dict[str, int]: + if config.num_hidden_layers == 2: + return { + "executorch_exir_dialects_edge__ops_aten_add_Tensor": 2, + "executorch_exir_dialects_edge__ops_aten_argmax_default": 1, + "executorch_exir_dialects_edge__ops_aten_where_self": 2, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 12, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 18, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 12, + "torch.ops.higher_order.executorch_call_delegate": 7, + } + + raise ValueError( + f"Unexpected CLIP config: hidden_act={config.hidden_act}, " + f"num_hidden_layers={config.num_hidden_layers}" + ) + + +@pytest.mark.parametrize( + ("config_factory", "atol"), + ( + (get_tiny_sd35_large_text_encoder_config, 1e-2), # FP atol + (get_tiny_sd35_large_text_encoder_2_config, 1.5e-2), # FP atol + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_tosa_FP(config_factory, atol): + """Run the CLIPTextModelWithProjection TOSA FP test for a given config.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=atol, + transform_passes=get_int64_to_int32_passes(), + ) + pipeline.change_args( + "check_count.exir", TestCLIPTextModelWithProjection.ops_after_partitioner_FP + ) + pipeline.run() + + +@pytest.mark.parametrize( + ("config_factory", "atol"), + ( + (get_tiny_sd35_large_text_encoder_config, 5.5e-2), # INT atol + (get_tiny_sd35_large_text_encoder_2_config, 6e-2), # INT atol + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_tosa_INT(config_factory, atol): + """Run the CLIPTextModelWithProjection TOSA INT test for a given config.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=atol, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", + TestCLIPTextModelWithProjection.ops_after_partitioner_INT(config), + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.parametrize( + ("config_factory",), + ( + (get_tiny_sd35_large_text_encoder_config,), + (get_tiny_sd35_large_text_encoder_2_config,), + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_vgf_no_quant(config_factory): + """Run the CLIPTextModelWithProjection VGF no-quant test.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=5e-3, + transform_passes=get_int64_to_int32_passes(), + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestCLIPTextModelWithProjection.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +@pytest.mark.parametrize( + ("config_factory", "atol"), + ( + (get_tiny_sd35_large_text_encoder_config, 5.5e-2), + (get_tiny_sd35_large_text_encoder_2_config, 6e-2), + ), + ids=("text_encoder", "text_encoder_2"), +) +def test_clip_text_model_with_projection_vgf_quant(config_factory, atol): + """Run the CLIPTextModelWithProjection VGF quant test.""" + test_helper = TestCLIPTextModelWithProjection() + config = config_factory() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=atol, + quantize=True, + ) + pipeline.change_args( + "check_count.exir", + TestCLIPTextModelWithProjection.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py new file mode 100644 index 00000000000..842c40ffb31 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_SD3Transformer2DModel_sd35_large.py @@ -0,0 +1,172 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_tiny_sd35_large_transformer_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3TransformerWrapper, +) + +input_t4 = Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] + + +class TestSD3Transformer2DModel: + """Test helper for SD3.5 Large SD3Transformer2DModel config.""" + + ops_after_partitioner_FP = { + "executorch_exir_dialects_edge__ops_aten_unsqueeze_copy_default": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 3, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_vgf_quantize = { + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_FP + + def create_config(self): + """Create a tiny SD3.5 Large-like MMDiT config for tests.""" + return get_tiny_sd35_large_transformer_config() + + def create_dummy_inputs( + self, + batch_size: int = 2, + latent_channels: int = 4, + latent_size: int = 32, + seq_length: int = 77, + joint_attention_dim: int = 16, + pooled_projection_dim: int = 32, + max_timestep: int = 1000, + dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Create dummy inputs for the SD3Transformer2DModel tests.""" + # SD3.5 Large uses latent channels=16, latent size=128, T5 seq length=256, + # joint_attention_dim=4096, and pooled_projection_dim=2048. Keep this + # unit-test default smaller for TOSA runtime and VGF memory limits. + return ( + torch.randn( + batch_size, + latent_channels, + latent_size, + latent_size, + dtype=dtype, + ), + torch.randint(low=0, high=max_timestep, size=(batch_size,)), + torch.randn( + batch_size, + seq_length, + joint_attention_dim, + dtype=dtype, + ), + torch.randn(batch_size, pooled_projection_dim, dtype=dtype), + ) + + def create_model(self) -> SD3TransformerWrapper: + """Instantiate wrapped SD3Transformer2DModel for tests.""" + SD3Transformer2DModel = pytest.importorskip( + "diffusers.models.transformers" + ).SD3Transformer2DModel + return SD3TransformerWrapper( + SD3Transformer2DModel(**self.create_config()) + ).eval() + + +def test_sd3_transformer_tosa_FP(): + """Run the SD3Transformer2DModel TOSA FP test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.change_args( + "check_count.exir", TestSD3Transformer2DModel.ops_after_partitioner_FP + ) + pipeline.run() + + +def test_sd3_transformer_tosa_INT(): + """Run the SD3Transformer2DModel TOSA INT test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", TestSD3Transformer2DModel.ops_after_partitioner_INT + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sd3_transformer_vgf_no_quant(): + """Run the SD3Transformer2DModel VGF no-quant test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestSD3Transformer2DModel.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_sd3_transformer_vgf_quant(): + """Run the SD3Transformer2DModel VGF quant test.""" + test_helper = TestSD3Transformer2DModel() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t4]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=True, + ) + pipeline.change_args( + "check_count.exir", + TestSD3Transformer2DModel.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py new file mode 100644 index 00000000000..2d850676d0f --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py @@ -0,0 +1,175 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_int64_to_int32_passes, + get_tiny_sd35_large_t5_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3T5TextEncoderWrapper, +) +from transformers import T5EncoderModel + +input_t = Tuple[torch.Tensor] + + +class TestT5EncoderModel: + """Test helper for SD3.5 Large T5EncoderModel config.""" + + ops_after_partitioner_FP = { + "executorch_exir_dialects_edge__ops_aten_clamp_Tensor": 4, + "executorch_exir_dialects_edge__ops_aten_isinf_default": 4, + "executorch_exir_dialects_edge__ops_aten_where_self": 1, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, + "torch.ops.higher_order.executorch_call_delegate": 10, + } + + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_aten_isinf_default": 4, + "executorch_exir_dialects_edge__ops_aten_mul_Tensor": 5, + "executorch_exir_dialects_edge__ops_aten_where_self": 5, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 21, + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 30, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 24, + "aten.scalar_tensor.default": 9, + "torch.ops.higher_order.executorch_call_delegate": 24, + } + + ops_after_partitioner_vgf_quantize = { + "executorch_exir_dialects_edge__ops_aten_clamp_Tensor": 4, + "executorch_exir_dialects_edge__ops_aten_isinf_default": 4, + "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 9, + } + + ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_vgf_quantize + + def create_dummy_inputs( + self, + config, + batch_size: int = 1, + seq_length: int = 2, + dtype: torch.dtype = torch.long, + ) -> tuple[torch.Tensor]: + """Create dummy inputs for the T5EncoderModel tests.""" + # SD3.5 Large uses (batch_size, seq_length) = (1, 256) for T5. + # Keep this unit-test default smaller for TOSA runtime. + return ( + torch.randint( + low=0, + high=config.vocab_size, + size=(batch_size, seq_length), + dtype=dtype, + ), + ) + + def create_config(self): + """Create a tiny SD3.5 Large-like T5 config for tests.""" + return get_tiny_sd35_large_t5_config() + + def create_model(self, config) -> SD3T5TextEncoderWrapper: + """Instantiate wrapped T5EncoderModel for tests.""" + return SD3T5TextEncoderWrapper( + T5EncoderModel(config).to(dtype=config.dtype) # type: ignore[call-arg] + ).eval() + + +def test_t5_encoder_tosa_FP(): + """Run the T5EncoderModel TOSA FP test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=2e-2, + transform_passes=get_int64_to_int32_passes(), + ) + pipeline.change_args( + "check_count.exir", TestT5EncoderModel.ops_after_partitioner_FP + ) + pipeline.run() + + +def test_t5_encoder_tosa_INT(): + """Run the T5EncoderModel TOSA INT test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=3e-2, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", TestT5EncoderModel.ops_after_partitioner_INT + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_t5_encoder_vgf_no_quant(): + """Run the T5EncoderModel VGF no-quant test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=8e-3, + transform_passes=get_int64_to_int32_passes(), + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestT5EncoderModel.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_t5_encoder_vgf_quant(): + """Run the T5EncoderModel VGF quant test.""" + test_helper = TestT5EncoderModel() + config = test_helper.create_config() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(config), + test_helper.create_dummy_inputs(config), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=1.2e-2, + quantize=True, + ) + pipeline.change_args( + "check_count.exir", + TestT5EncoderModel.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py new file mode 100644 index 00000000000..d4de15e5d74 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_configs_sd35_large.py @@ -0,0 +1,311 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import hashlib +import json +import os +import warnings +from pathlib import Path +from typing import Any + +from executorch.backends.arm._passes import ( + ArmPass, + ConvertInt64ConstOpsToInt32Pass, + ConvertInt64OutputOpsToInt32Pass, + InsertInt32CastsAfterInt64PlaceholdersPass, +) +from transformers import CLIPTextConfig, T5Config + + +_EXECUTORCH_SD35_UPSTREAM_SYNC_ENV_VAR = "EXECUTORCH_SD35_UPSTREAM_SYNC" +_sd35_large_upstream_checked: set[str] = set() + +_SD35_LARGE_UPSTREAM_FINGERPRINTS: dict[str, tuple[tuple[str, ...], str]] = { + "text_encoder": ( + ( + "hidden_act", + "hidden_size", + "intermediate_size", + "max_position_embeddings", + "num_attention_heads", + "num_hidden_layers", + "projection_dim", + "vocab_size", + ), + "0cb164627ea428c39166d8ecf70462f7ddba34a57071d383a1116603e114bf50", + ), + "text_encoder_2": ( + ( + "hidden_act", + "hidden_size", + "intermediate_size", + "max_position_embeddings", + "num_attention_heads", + "num_hidden_layers", + "projection_dim", + "vocab_size", + ), + "4c082c890625573879240ef700b1dbe18e2a263ad8e76d222cba759553176e1e", + ), + "text_encoder_3": ( + ( + "d_ff", + "d_kv", + "d_model", + "dense_act_fn", + "feed_forward_proj", + "num_heads", + "num_layers", + "relative_attention_num_buckets", + "vocab_size", + ), + "da069a817a2fe4eb347fc5aefd7690bf97f8661d33dbe689058977809651023d", + ), + "transformer": ( + ( + "sample_size", + "patch_size", + "in_channels", + "num_layers", + "attention_head_dim", + "num_attention_heads", + "caption_projection_dim", + "joint_attention_dim", + "pooled_projection_dim", + "out_channels", + "pos_embed_max_size", + "qk_norm", + ), + "bc87c7eefb80f7bc6e80479f5bd2af929fd14c2331a19ddb4e27a989546ebfb0", + ), + "vae": ( + ( + "sample_size", + "in_channels", + "out_channels", + "down_block_types", + "up_block_types", + "block_out_channels", + "layers_per_block", + "latent_channels", + "norm_num_groups", + "act_fn", + "mid_block_add_attention", + "force_upcast", + "use_quant_conv", + "use_post_quant_conv", + "scaling_factor", + "shift_factor", + ), + "8019be48e6681895b9b7c54c0f2f48c3ddc9975d7b6c7ef2061b7743f8d55b71", + ), +} + + +def _load_upstream_sd35_large_config(subfolder: str) -> dict[str, Any]: + from executorch.examples.models.stable_diffusion_3_5_large.model import MODEL_ID + from huggingface_hub import hf_hub_download + + config_path = hf_hub_download( # nosec B615 + repo_id=MODEL_ID, + filename="config.json", + subfolder=subfolder, + token=os.environ.get("HF_TOKEN"), + etag_timeout=1.0, + ) + return json.loads(Path(config_path).read_text()) + + +def _fingerprint_sd35_large_config( + config: dict[str, Any], fields: tuple[str, ...] +) -> str: + selected_config = {field: config[field] for field in fields} + serialized_config = json.dumps( + selected_config, + sort_keys=True, + separators=(",", ":"), + ) + return hashlib.sha256(serialized_config.encode()).hexdigest() + + +def _warn_if_sd35_large_config_differs_from_upstream( + subfolder: str, warning_name: str +) -> None: + if os.environ.get(_EXECUTORCH_SD35_UPSTREAM_SYNC_ENV_VAR, "0") != "1": + return + if subfolder in _sd35_large_upstream_checked: + return + + _sd35_large_upstream_checked.add(subfolder) + try: + upstream_config = _load_upstream_sd35_large_config(subfolder) + fields, expected_fingerprint = _SD35_LARGE_UPSTREAM_FINGERPRINTS[subfolder] + upstream_fingerprint = _fingerprint_sd35_large_config(upstream_config, fields) + except Exception as exc: + warnings.warn( + f"Unable to validate {warning_name} against upstream metadata: {exc}", + RuntimeWarning, + stacklevel=2, + ) + return + + if upstream_fingerprint != expected_fingerprint: + warnings.warn( + f"Upstream {warning_name} architecture changed; review the tiny test config", + RuntimeWarning, + stacklevel=2, + ) + + +def get_int64_to_int32_passes() -> list[ArmPass]: + return [ + ConvertInt64ConstOpsToInt32Pass(), + ConvertInt64OutputOpsToInt32Pass(), + InsertInt32CastsAfterInt64PlaceholdersPass(), + ] + + +def get_tiny_sd35_large_text_encoder_config() -> CLIPTextConfig: + """Create a tiny SD3.5 Large-like CLIP-L text encoder config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream( + "text_encoder", "SD3.5 Large CLIP text encoder" + ) + return CLIPTextConfig( # type: ignore[call-arg] + architectures=["CLIPTextModelWithProjection"], + attention_dropout=0.0, + bos_token_id=0, + dropout=0.0, + eos_token_id=2, + hidden_act="quick_gelu", + hidden_size=32, + initializer_factor=1.0, + initializer_range=0.02, + intermediate_size=128, + layer_norm_eps=1e-5, + max_position_embeddings=16, + num_attention_heads=4, + num_hidden_layers=2, + pad_token_id=1, + projection_dim=32, + dtype="float16", + vocab_size=256, + ) + + +def get_tiny_sd35_large_text_encoder_2_config() -> CLIPTextConfig: + """Create a tiny SD3.5 Large-like CLIP-bigG text encoder config for + tests. + """ + _warn_if_sd35_large_config_differs_from_upstream( + "text_encoder_2", "SD3.5 Large CLIP text encoder 2" + ) + return CLIPTextConfig( # type: ignore[call-arg] + architectures=["CLIPTextModelWithProjection"], + attention_dropout=0.0, + bos_token_id=0, + dropout=0.0, + eos_token_id=2, + hidden_act="gelu", + hidden_size=48, + initializer_factor=1.0, + initializer_range=0.02, + intermediate_size=192, + layer_norm_eps=1e-5, + max_position_embeddings=16, + num_attention_heads=6, + num_hidden_layers=2, + pad_token_id=1, + projection_dim=48, + dtype="float16", + vocab_size=256, + ) + + +def get_tiny_sd35_large_t5_config() -> T5Config: + """Create a tiny SD3.5 Large-like T5 config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream( + "text_encoder_3", "SD3.5 Large T5 text encoder" + ) + return T5Config( # type: ignore[call-arg] + architectures=["T5EncoderModel"], + classifier_dropout=0.0, + d_ff=64, + d_kv=8, + d_model=32, + decoder_start_token_id=0, + dense_act_fn="gelu_new", + dropout_rate=0.1, + eos_token_id=1, + feed_forward_proj="gated-gelu", + initializer_factor=1.0, + is_encoder_decoder=True, + is_gated_act=True, + layer_norm_epsilon=1e-6, + num_decoder_layers=2, + num_heads=4, + num_layers=2, + output_past=True, + pad_token_id=0, + relative_attention_max_distance=128, + relative_attention_num_buckets=8, + tie_word_embeddings=False, + dtype="float16", + vocab_size=256, + use_cache=True, + ) + + +def get_tiny_sd35_large_transformer_config() -> dict[str, Any]: + """Create a tiny SD3.5 Large-like MMDiT config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream( + "transformer", "SD3.5 Large transformer" + ) + return { + "sample_size": 32, + "patch_size": 2, + "in_channels": 4, + "num_layers": 2, + "attention_head_dim": 8, + "num_attention_heads": 2, + "caption_projection_dim": 16, + "joint_attention_dim": 16, + "pooled_projection_dim": 32, + "out_channels": 4, + "pos_embed_max_size": 32, + "qk_norm": "rms_norm", + } + + +def get_tiny_sd35_large_vae_config() -> dict[str, Any]: + """Create a tiny SD3.5 Large-like VAE config for tests.""" + _warn_if_sd35_large_config_differs_from_upstream("vae", "SD3.5 Large VAE") + return { + "sample_size": 32, + "in_channels": 3, + "out_channels": 3, + "down_block_types": ( + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + "DownEncoderBlock2D", + ), + "up_block_types": ( + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + "UpDecoderBlock2D", + ), + "block_out_channels": (4, 8, 8, 8), + "layers_per_block": 1, + "latent_channels": 16, + "norm_num_groups": 1, + "act_fn": "silu", + "mid_block_add_attention": True, + "force_upcast": False, + "use_quant_conv": False, + "use_post_quant_conv": False, + "scaling_factor": 1.5305, + "shift_factor": 0.0609, + } diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py new file mode 100644 index 00000000000..336e1f8e1b6 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_model_sd35_large.py @@ -0,0 +1,369 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from types import SimpleNamespace + +import pytest +import torch +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_tiny_sd35_large_t5_config, + get_tiny_sd35_large_text_encoder_2_config, + get_tiny_sd35_large_text_encoder_config, + get_tiny_sd35_large_transformer_config, + get_tiny_sd35_large_vae_config, +) +from executorch.examples.models.stable_diffusion_3_5_large import ( + model as sd35_large_model, +) +from transformers import CLIPTextModelWithProjection, T5EncoderModel + + +@pytest.mark.parametrize( + ("clip_skip", "hidden_state_index"), + ( + pytest.param(None, -2, id="default_clip_skip"), + pytest.param(1, -3, id="clip_skip_1"), + ), +) +def test_clip_text_encoder_wrapper_returns_selected_hidden_state_and_pooled_projection( + clip_skip, hidden_state_index +): + """Verify CLIP wrapper outputs.""" + config = get_tiny_sd35_large_text_encoder_config() + config.num_hidden_layers = 3 # Set to 3 layers to test clip_skip=1 behavior + text_encoder = CLIPTextModelWithProjection(config).to(dtype=config.dtype) + text_encoder.eval() + wrapper = sd35_large_model.SD3CLIPTextEncoderWrapper( + text_encoder, clip_skip=clip_skip + ) + input_ids = torch.randint(0, config.vocab_size, (2, 7)) + + with torch.no_grad(): + hidden_states, pooled_projection = wrapper(input_ids) + expected = text_encoder(input_ids, output_hidden_states=True, return_dict=True) + + torch.testing.assert_close( + hidden_states, expected.hidden_states[hidden_state_index] + ) + torch.testing.assert_close(pooled_projection, expected[0]) + + +def test_t5_text_encoder_wrapper_returns_last_hidden_state(): + """Verify T5 text encoder wrapper returns last hidden state.""" + config = get_tiny_sd35_large_t5_config() + text_encoder = T5EncoderModel(config) + text_encoder.eval() + wrapper = sd35_large_model.SD3T5TextEncoderWrapper(text_encoder) + input_ids = torch.randint(0, config.vocab_size, (2, 7)) + + with torch.no_grad(): + hidden_states = wrapper(input_ids) + expected = text_encoder(input_ids, return_dict=True) + + torch.testing.assert_close(hidden_states, expected.last_hidden_state) + + +def test_transformer_wrapper_returns_sample_tensor(): + """Verify transformer wrapper returns the sample tensor.""" + SD3Transformer2DModel = pytest.importorskip( + "diffusers.models.transformers" + ).SD3Transformer2DModel + transformer = SD3Transformer2DModel(**get_tiny_sd35_large_transformer_config()) + transformer.eval() + wrapper = sd35_large_model.SD3TransformerWrapper(transformer) + batch_size = 2 + latents = torch.randn(batch_size, 4, 32, 32) + timestep = torch.randint(0, 1000, (batch_size,)) + encoder_hidden_states = torch.randn(batch_size, 154, 16) + pooled_projections = torch.randn(batch_size, 32) + + with torch.no_grad(): + sample = wrapper( + latents, + timestep, + encoder_hidden_states, + pooled_projections, + ) + expected = transformer( + hidden_states=latents, + timestep=timestep, + encoder_hidden_states=encoder_hidden_states, + pooled_projections=pooled_projections, + return_dict=True, + ) + + torch.testing.assert_close(sample, expected.sample) + + +def test_vae_decoder_wrapper_scales_shifts_decodes_and_clamps(): + """Verify VAE decoder wrapper scales, shifts, decodes, and clamps.""" + AutoencoderKL = pytest.importorskip("diffusers.models.autoencoders").AutoencoderKL + vae_config = get_tiny_sd35_large_vae_config() + vae = AutoencoderKL(**vae_config) + vae.eval() + wrapper = sd35_large_model.SD3VAEDecoderWrapper(vae) + latents = torch.randn(1, vae_config["latent_channels"], 4, 4) + + with torch.no_grad(): + image = wrapper(latents) + expected_latents = latents / vae.config.scaling_factor + vae.config.shift_factor + expected = vae.decode(expected_latents, return_dict=True).sample + # Normalize decoder output from [-1, 1] to image range [0, 1]. + expected = (expected / 2 + 0.5).clamp(0, 1) + + torch.testing.assert_close(image, expected) + assert torch.all(image >= 0) + assert torch.all(image <= 1) + + +@pytest.mark.parametrize( + "getter_name", + ( + "get_text_encoder_wrapper", + "get_text_encoder_2_wrapper", + "get_text_encoder_3_wrapper", + "get_transformer_wrapper", + "get_vae_decoder_wrapper", + ), +) +def test_model_loader_getters_require_loaded_components(getter_name): + """Verify model loader getters require loaded components.""" + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + + with pytest.raises(ValueError, match="Models not loaded"): + getattr(loader, getter_name)() + + +def test_model_loader_text_encoder_getters_wrap_loaded_components(): + """Verify model loader text encoder getters wrap loaded components.""" + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.text_encoder = CLIPTextModelWithProjection( + get_tiny_sd35_large_text_encoder_config() + ) + loader.text_encoder_2 = CLIPTextModelWithProjection( + get_tiny_sd35_large_text_encoder_2_config() + ) + loader.text_encoder_3 = T5EncoderModel(get_tiny_sd35_large_t5_config()) + + assert loader.get_text_encoder_wrapper().text_encoder is loader.text_encoder + assert loader.get_text_encoder_2_wrapper().text_encoder is loader.text_encoder_2 + assert loader.get_text_encoder_3_wrapper().text_encoder is loader.text_encoder_3 + + +def test_model_loader_transformer_getter_wraps_loaded_component(): + """Verify model loader transformer getter wraps loaded component.""" + SD3Transformer2DModel = pytest.importorskip( + "diffusers.models.transformers" + ).SD3Transformer2DModel + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.transformer = SD3Transformer2DModel( + **get_tiny_sd35_large_transformer_config() + ) + + assert loader.get_transformer_wrapper().transformer is loader.transformer + + +def test_model_loader_vae_getter_wraps_loaded_component(): + """Verify model loader VAE getter wraps loaded component.""" + AutoencoderKL = pytest.importorskip("diffusers.models.autoencoders").AutoencoderKL + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.vae = AutoencoderKL(**get_tiny_sd35_large_vae_config()) + + assert loader.get_vae_decoder_wrapper().vae is loader.vae + + +def _patch_model_loaders(monkeypatch): + """Patch model loaders and return call records.""" + + class FakeModel: + def __init__(self): + self.eval_called = False + + def to(self, dtype): + return self + + def eval(self): + self.eval_called = True + return self + + calls = SimpleNamespace( + tokenizer=[], + text_encoder=[], + t5=[], + transformer=[], + vae=[], + ) + + class FakeTokenizer: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.tokenizer.append((model_id, kwargs)) + return SimpleNamespace(model_max_length=77) + + class FakeTextEncoder: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.text_encoder.append((model_id, kwargs)) + return FakeModel() + + class FakeT5: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.t5.append((model_id, kwargs)) + return FakeModel() + + class FakeTransformer: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.transformer.append((model_id, kwargs)) + return FakeModel() + + class FakeVAE: + @staticmethod + def from_pretrained(model_id, **kwargs): + calls.vae.append((model_id, kwargs)) + return FakeModel() + + monkeypatch.setattr(sd35_large_model, "CLIPTokenizer", FakeTokenizer) + monkeypatch.setattr( + sd35_large_model, "CLIPTextModelWithProjection", FakeTextEncoder + ) + monkeypatch.setattr(sd35_large_model, "T5EncoderModel", FakeT5) + monkeypatch.setattr(sd35_large_model, "SD3Transformer2DModel", FakeTransformer) + monkeypatch.setattr(sd35_large_model, "AutoencoderKL", FakeVAE) + + return calls + + +def test_load_models_uses_component_subfolders(monkeypatch): + """Verify model loading uses the expected component subfolders.""" + calls = _patch_model_loaders(monkeypatch) + + loader = sd35_large_model.StableDiffusion3ModelLoader( + model_id="test/sd3", + dtype=torch.float32, + ) + + assert loader.load_models() + assert calls.tokenizer == [ + ("test/sd3", {"subfolder": "tokenizer"}), + ("test/sd3", {"subfolder": "tokenizer_2"}), + ] + assert calls.text_encoder == [ + ("test/sd3", {"subfolder": "text_encoder", "torch_dtype": torch.float32}), + ("test/sd3", {"subfolder": "text_encoder_2", "torch_dtype": torch.float32}), + ] + assert calls.t5 == [ + ("test/sd3", {"subfolder": "text_encoder_3", "torch_dtype": torch.float32}) + ] + assert calls.transformer == [ + ("test/sd3", {"subfolder": "transformer", "torch_dtype": torch.float32}) + ] + assert calls.vae == [ + ("test/sd3", {"subfolder": "vae", "torch_dtype": torch.float32}) + ] + assert loader.text_encoder.eval_called + assert loader.text_encoder_2.eval_called + assert loader.text_encoder_3.eval_called + assert loader.transformer.eval_called + assert loader.vae.eval_called + + +def test_load_models_loads_only_requested_component(monkeypatch): + """Verify model loading can load only requested components.""" + calls = _patch_model_loaders(monkeypatch) + + loader = sd35_large_model.StableDiffusion3ModelLoader( + model_id="test/sd3", + dtype=torch.float32, + ) + + assert loader.load_models( + [sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_3] + ) + assert calls.tokenizer == [] + assert calls.text_encoder == [] + assert calls.t5 == [ + ("test/sd3", {"subfolder": "text_encoder_3", "torch_dtype": torch.float32}) + ] + assert calls.transformer == [] + assert calls.vae == [] + assert loader.text_encoder is None + assert loader.text_encoder_2 is None + assert loader.text_encoder_3 is not None + assert loader.transformer is None + assert loader.vae is None + + +@pytest.mark.parametrize( + ("latent_size", "expected_latent_size"), + ( + pytest.param(None, 32, id="default_latent_size"), + pytest.param(16, 16, id="override_latent_size"), + ), +) +def test_get_dummy_inputs_builds_expected_component_inputs( + latent_size, expected_latent_size +): + """Verify dummy inputs have expected component shapes.""" + loader = sd35_large_model.StableDiffusion3ModelLoader(dtype=torch.float32) + loader.tokenizer = SimpleNamespace(model_max_length=77) + loader.text_encoder = object() + loader.text_encoder_2 = object() + loader.text_encoder_3 = object() + loader.transformer = SimpleNamespace( + config=SimpleNamespace( + in_channels=4, + sample_size=32, + joint_attention_dim=16, + pooled_projection_dim=32, + ) + ) + loader.vae = SimpleNamespace(config=SimpleNamespace(latent_channels=16)) + + dummy_inputs = loader.get_dummy_inputs( + max_sequence_length=256, + latent_size=latent_size, + ) + + assert set(dummy_inputs) == { + sd35_large_model.StableDiffusionComponent.TEXT_ENCODER, + sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_2, + sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_3, + sd35_large_model.StableDiffusionComponent.TRANSFORMER, + sd35_large_model.StableDiffusionComponent.VAE_DECODER, + } + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER][ + 0 + ].shape == (1, 77) + assert ( + dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER][0].dtype + == torch.long + ) + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_2][ + 0 + ].shape == (1, 77) + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.TEXT_ENCODER_3][ + 0 + ].shape == (1, 256) + + transformer_inputs = dummy_inputs[ + sd35_large_model.StableDiffusionComponent.TRANSFORMER + ] + assert transformer_inputs[0].shape == ( + 1, + 4, + expected_latent_size, + expected_latent_size, + ) + assert transformer_inputs[0].dtype == torch.float32 + assert transformer_inputs[1].shape == (1,) + assert transformer_inputs[1].dtype == torch.float32 + assert transformer_inputs[2].shape == (1, 333, 16) + assert transformer_inputs[3].shape == (1, 32) + + assert dummy_inputs[sd35_large_model.StableDiffusionComponent.VAE_DECODER][ + 0 + ].shape == (1, 16, expected_latent_size, expected_latent_size) diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py new file mode 100644 index 00000000000..30c6df3e693 --- /dev/null +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_vae_AutoencoderKL_sd35_large.py @@ -0,0 +1,152 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import pytest +import torch +from executorch.backends.arm.test import common +from executorch.backends.arm.test.models.stable_diffusion_3_5_large.test_configs_sd35_large import ( + get_tiny_sd35_large_vae_config, +) +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) +from executorch.examples.models.stable_diffusion_3_5_large.model import ( + SD3VAEDecoderWrapper, +) + +input_t = Tuple[torch.Tensor] + + +class TestAutoencoderKL: + """Test helper for SD3.5 Large AutoencoderKL config.""" + + ops_after_partitioner_FP = { + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_INT = { + "executorch_exir_dialects_edge__ops_quantized_decomposed_dequantize_per_tensor_default": 1, + "executorch_exir_dialects_edge__ops_quantized_decomposed_quantize_per_tensor_default": 1, + "torch.ops.higher_order.executorch_call_delegate": 1, + } + + ops_after_partitioner_vgf_quantize = ops_after_partitioner_FP + ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_FP + + def create_config(self): + """Create a tiny SD3.5 Large-like AutoencoderKL config for tests.""" + return get_tiny_sd35_large_vae_config() + + def create_dummy_inputs( + self, + batch_size: int = 1, + latent_channels: int = 16, + latent_size: int = 4, + dtype: torch.dtype = torch.float32, + ) -> tuple[torch.Tensor]: + """Create dummy inputs for the SD3 VAE decoder tests.""" + # SD3.5 Large uses VAE decoder latent channels=16 and latent size=128. + # Keep this unit-test default spatial size smaller for TOSA runtime. + return ( + torch.randn( + batch_size, + latent_channels, + latent_size, + latent_size, + dtype=dtype, + ), + ) + + def create_model(self) -> SD3VAEDecoderWrapper: + """Instantiate wrapped AutoencoderKL decoder for tests.""" + diffusers_autoencoders = pytest.importorskip("diffusers.models.autoencoders") + AutoencoderKL = diffusers_autoencoders.AutoencoderKL + return SD3VAEDecoderWrapper(AutoencoderKL(**self.create_config())).eval() + + +def test_vae_tosa_FP(): + """Run the AutoencoderKL TOSA FP test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = TosaPipelineFP[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + ) + pipeline.change_args( + "check_count.exir", TestAutoencoderKL.ops_after_partitioner_FP + ) + pipeline.run() + + +def test_vae_tosa_INT(): + """Run the AutoencoderKL TOSA INT test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = TosaPipelineINT[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + atol=9e-2, + frobenius_threshold=None, + cosine_threshold=None, + ) + pipeline.change_args( + "check_count.exir", TestAutoencoderKL.ops_after_partitioner_INT + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_vae_vgf_no_quant(): + """Run the AutoencoderKL VGF no-quant test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=False, + ) + pipeline.change_args( + "check_count.exir", + TestAutoencoderKL.ops_after_partitioner_vgf_no_quantize, + ) + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_vae_vgf_quant(): + """Run the AutoencoderKL VGF quant test.""" + test_helper = TestAutoencoderKL() + + with torch.no_grad(): + pipeline = VgfPipeline[input_t]( + test_helper.create_model(), + test_helper.create_dummy_inputs(), + aten_op=[], + exir_op=[], + use_to_edge_transform_and_lower=True, + quantize=True, + qtol=2, + ) + pipeline.change_args( + "check_count.exir", + TestAutoencoderKL.ops_after_partitioner_vgf_quantize, + ) + pipeline.run() From d9b3e457ccd5db188fe659655609ae859554d519 Mon Sep 17 00:00:00 2001 From: Digant Desai Date: Wed, 9 Sep 2026 12:41:19 -0500 Subject: [PATCH 106/190] CMake support for quantized MoE optimized build (#22547) For MoE custom op --- examples/models/llama/README.md | 9 +- examples/models/llama/export_llama_lib.py | 4 +- extension/llm/custom_ops/CMakeLists.txt | 88 ++++++++++----- extension/llm/custom_ops/op_moe.cpp | 130 ++++++++++++++++------ extension/llm/custom_ops/targets.bzl | 12 +- extension/llm/custom_ops/test_op_moe.cpp | 10 +- 6 files changed, 168 insertions(+), 85 deletions(-) diff --git a/examples/models/llama/README.md b/examples/models/llama/README.md index 3ee568120a3..6cbdd8559fd 100644 --- a/examples/models/llama/README.md +++ b/examples/models/llama/README.md @@ -561,15 +561,14 @@ registered in `executorch.extension.llm.custom_ops.custom_ops`. The runtime kernel ships in `extension/llm/custom_ops/op_moe.cpp`. It always compiles with a portable reference fallback (unpack + dequant + -`cpublas::gemm`) that works on any platform. `ENABLE_QUANTIZED_MOE_FFN` -is an **optimization gate**, not a correctness requirement — when -defined, the kernel uses torchao's fused `linear_operator` (NEON -i8mm/dotprod on aarch64) instead of the reference path. +`cpublas::gemm`) that works on any platform. The optimized build option +uses torchao's fused `linear_operator` (NEON i8mm/dotprod on aarch64) +instead of the reference path. In CMake, opt in to the optimized path with: ```cmake --DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON +-DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON ``` In Buck, `_get_quantized_moe_deps()` in `targets.bzl` wires: diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index d982eff7468..07240b11d8c 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -594,8 +594,8 @@ def build_args_parser() -> argparse.ArgumentParser: "Replace eager MoE feed-forward modules with the " "`llama::quantized_moe_ffn` portable-runtime custom op (INT4 " "weights, INT8 dyn-quant activations via torchao). On aarch64 " - "with ENABLE_QUANTIZED_MOE_FFN the optimized torchao kernel is " - "used; otherwise a portable reference fallback runs." + "an optimized runtime build uses the torchao kernel; otherwise " + "a portable reference fallback runs." ), ) diff --git a/extension/llm/custom_ops/CMakeLists.txt b/extension/llm/custom_ops/CMakeLists.txt index c640cad6706..63ff2f23e69 100644 --- a/extension/llm/custom_ops/CMakeLists.txt +++ b/extension/llm/custom_ops/CMakeLists.txt @@ -20,7 +20,7 @@ set(_common_compile_options $<$:/wd4996> $<$>:-Wno-deprecated-declarations -fPIC> ) -if(CMAKE_SYSTEM_PROCESSOR MATCHES "arm64|aarch64") +if(CMAKE_SYSTEM_PROCESSOR MATCHES "^(arm64|aarch64)$") list(APPEND _common_compile_options "$<$>:-march=armv8.2-a+dotprod>" ) @@ -86,28 +86,72 @@ target_link_libraries(custom_ops PUBLIC ${custom_ops_libs} executorch_core) # The MoE kernel always compiles with a reference fallback (unpack + dequant + # cpublas::gemm) using the torchao weight_packing headers from third-party/ao # (already on the include path). Pass -# -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON to additionally link the -# optimized torchao linear_operator (fused INT8-dyn-act GEMM, aarch64 NEON +# -DEXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON to additionally link +# the optimized torchao linear_operator (fused INT8-dyn-act GEMM, aarch64 NEON # dotprod). -option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE +option(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED "Link the optimized torchao linear kernel for llama::quantized_moe_ffn" OFF ) -if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE) - if(NOT TARGET torchao_ops_linear_8bit_act_xbit_weight_executorch) + +function(_target_enable_quantized_moe_torchao target) + target_compile_definitions( + ${target} + PRIVATE TORCHAO_BUILD_CPU_AARCH64=1 TORCHAO_ENABLE_ARM_NEON_DOT=1 + TORCHAO_PARALLEL_EXECUTORCH=1 + TORCHAO_SHARED_KERNELS_BUILD_EXECUTORCH=1 + ) +endfunction() + +if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED) + if(NOT CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") message( FATAL_ERROR - "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON requires target " - "torchao_ops_linear_8bit_act_xbit_weight_executorch, which is not " - "defined. Build the torchao ops or set this option OFF." + "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED=ON is supported only on " + "aarch64 and arm64." ) endif() - # Compile definition and link must be gated on the same condition, else the - # ENABLE_QUANTIZED_MOE_FFN path compiles without the library that defines it. - target_compile_definitions(custom_ops PUBLIC ENABLE_QUANTIZED_MOE_FFN=1) - target_link_libraries( - custom_ops PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch + + # The selected target is also reused by custom_ops_aot_lib below. + if(EXECUTORCH_BUILD_KERNELS_TORCHAO) + if(NOT TARGET torchao_ops_executorch) + message(FATAL_ERROR "EXECUTORCH_BUILD_KERNELS_TORCHAO=ON but target " + "torchao_ops_executorch is not defined." + ) + endif() + set(quantized_moe_torchao_target torchao_ops_executorch) + else() + set(quantized_moe_torchao_target torchao_moe_linear) + add_library( + torchao_moe_linear STATIC + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight/linear_8bit_act_xbit_weight.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/quantization/quantize.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/reduction/compute_sum.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/reduction/find_min_and_max.cpp + ${EXECUTORCH_ROOT}/third-party/ao/torchao/csrc/cpu/torch_free_kernels/aarch64/valpacking/interleave.cpp + ) + target_include_directories( + torchao_moe_linear PRIVATE ${EXECUTORCH_ROOT}/third-party/ao + ) + _target_enable_quantized_moe_torchao(torchao_moe_linear) + target_link_libraries( + torchao_moe_linear PRIVATE cpuinfo executorch_core extension_threadpool + ) + target_compile_options( + torchao_moe_linear PRIVATE ${_common_compile_options} + ) + install( + TARGETS torchao_moe_linear + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + ) + endif() + + target_compile_definitions( + custom_ops PRIVATE EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO=1 ) + _target_enable_quantized_moe_torchao(custom_ops) + target_link_libraries(custom_ops PUBLIC ${quantized_moe_torchao_target}) endif() target_compile_options(custom_ops PUBLIC ${_common_compile_options}) @@ -190,21 +234,13 @@ if(EXECUTORCH_BUILD_KERNELS_LLM_AOT) custom_ops_aot_lib PUBLIC cpublas torch extension_tensor extension_threadpool ) - if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE) - if(NOT TARGET torchao_ops_linear_8bit_act_xbit_weight_executorch) - message( - FATAL_ERROR - "EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE=ON requires target " - "torchao_ops_linear_8bit_act_xbit_weight_executorch, which is not " - "defined. Build the torchao ops or set this option OFF." - ) - endif() + if(EXECUTORCH_BUILD_KERNELS_LLM_QUANTIZED_MOE_OPTIMIZED) target_compile_definitions( - custom_ops_aot_lib PUBLIC ENABLE_QUANTIZED_MOE_FFN=1 + custom_ops_aot_lib PRIVATE EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO=1 ) + _target_enable_quantized_moe_torchao(custom_ops_aot_lib) target_link_libraries( - custom_ops_aot_lib - PUBLIC torchao_ops_linear_8bit_act_xbit_weight_executorch + custom_ops_aot_lib PUBLIC ${quantized_moe_torchao_target} ) endif() if(WIN32) diff --git a/extension/llm/custom_ops/op_moe.cpp b/extension/llm/custom_ops/op_moe.cpp index bcb0e88df10..e218d58b35a 100644 --- a/extension/llm/custom_ops/op_moe.cpp +++ b/extension/llm/custom_ops/op_moe.cpp @@ -9,19 +9,28 @@ #include #include +#if defined(EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) && \ + !defined(TORCHAO_PARALLEL_EXECUTORCH) #include +#endif #include #include +#if defined(EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) && \ + !defined(TORCHAO_PARALLEL_EXECUTORCH) #include +#endif #include #include -#ifdef ENABLE_QUANTIZED_MOE_FFN -#include +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO +#include +#include #include +#include +#include #include // std::nullopt, used only by the optimized aarch64 path -#endif // ENABLE_QUANTIZED_MOE_FFN +#endif // EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO #include #include @@ -37,6 +46,49 @@ namespace { using ::executorch::aten::string_view; +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO +template +const torchao::ops::linear_8bit_act_xbit_weight::UKernelConfig& +universal_ukernel_config() { + using torchao::ops::linear_8bit_act_xbit_weight::UKernelConfig; + namespace kernel = torchao::kernels::cpu::aarch64::linear:: + channelwise_8bit_activation_groupwise_lowbit_weight; + + static const auto config = [] { + ET_CHECK_MSG( + cpuinfo_initialize() && cpuinfo_has_arm_neon_dot(), + "quantized_moe_ffn optimized path requires Arm NEON dot product"); + auto result = UKernelConfig::make( + /*preferred_alignment=*/16, + /*n_step=*/8, + /*nr=*/8, + /*kr=*/16, + /*sr=*/2, + kWeightNbit, + /*has_weight_zeros=*/false, + /*has_bias=*/false, + &torchao::weight_packing::packed_weights_size, + &torchao::weight_packing::packed_weights_offset, + &torchao::weight_packing::pack_weights, + {}); + result.linear_configs[0] = UKernelConfig::linear_config_type({ + /*m_step=*/1, + /*mr=*/1, + &kernel::packed_activations_size, + &kernel::packed_activations_offset, + &kernel::pack_activations<1, 16, 2>, + &kernel::kernel_1x8x16_f32_neondot< + kWeightNbit, + /*has_weight_zeros=*/false, + /*has_lut=*/false>, + }); + result.validate(); + return result; + }(); + return config; +} +#endif + // Numerically-stable sigmoid. Branching on sign keeps exp()'s argument // non-positive on both sides, so it can never overflow. inline float stable_sigmoid(float v) { @@ -189,7 +241,7 @@ inline void reference_linear( // Dispatch a single per-expert grouped GEMM through torchao's // linear_operator (optimized, aarch64) or reference unpack+dequant+gemm. -#ifdef ENABLE_QUANTIZED_MOE_FFN +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO template inline void torchao_linear( const uint8_t* packed_w_blob, @@ -205,21 +257,22 @@ inline void torchao_linear( static_cast(torchao::ops::PackedWeightsHeader::size()), "torchao packed blob too small to contain header"); auto header = torchao::ops::PackedWeightsHeader::read(packed_w_blob); - // Select the ukernel from the format declared in the header. This resolves - // the universal or kleidi packing automatically; a format whose kernels are - // not compiled into this build (e.g. kleidi when TORCHAO_ENABLE_KLEIDI is - // unset) throws here instead of being silently mis-read. - // TODO: enable KleidiAI here — build this op on xplat arm64 with - // -DTORCHAO_ENABLE_KLEIDI (+ -DTORCHAO_ENABLE_ARM_I8MM) and link the kleidi - // kernel target so a kleidi header actually resolves to a kleidi ukernel. - // Must be coordinated with the AoT packer emitting kleidi headers (see - // targets.bzl). - auto uk = torchao::ops::linear_8bit_act_xbit_weight::select_ukernel_config< - kWeightNbit>(header); - - // Validate the blob against the *selected* format's layout. nr/kr/sr and the - // size formula differ between universal and kleidi, so derive them from the - // chosen config rather than assuming a fixed layout. + ET_CHECK_MSG( + header.type == + torchao::ops::PackedWeightsType:: + linear_8bit_act_xbit_weight_universal, + "quantized_moe_ffn requires universal torchao packed weights"); + const auto format = torchao::ops::linear_8bit_act_xbit_weight:: + PackedWeightsFormat::from_packed_weights_header(header); + ET_CHECK_MSG( + format.weight_nbit == kWeightNbit && !format.has_weight_zeros && + !format.has_bias && format.nr == 8 && format.kr == 16 && + format.sr == 2, + "quantized_moe_ffn received an unsupported universal weight format"); + const auto& uk = universal_ukernel_config(); + + // Validate the blob against the universal config's layout without + // duplicating its packed-weight size formula. const int64_t required_bytes = static_cast(torchao::ops::PackedWeightsHeader::size()) + static_cast(uk.packed_weights_size( @@ -258,7 +311,7 @@ inline void torchao_linear( /*clamp_min=*/0.0f, /*clamp_max=*/0.0f); } -#endif // ENABLE_QUANTIZED_MOE_FFN +#endif // EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO inline void expert_linear_dispatch( int64_t weight_nbit, @@ -270,12 +323,21 @@ inline void expert_linear_dispatch( int64_t k, int64_t group_size, float* out) { -#ifndef ENABLE_QUANTIZED_MOE_FFN +#ifndef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO // Reference path only: it unpacks the universal layout, so validate the blob // holds the header plus the universal packed weight-data bytes for the // claimed dims before any path dereferences it. The torchao path validates - // against its own selected format (universal or kleidi) inside - // torchao_linear. + // the same required universal format inside torchao_linear. + ET_CHECK_MSG( + packed_blob_bytes >= + static_cast(torchao::ops::PackedWeightsHeader::size()), + "torchao packed blob too small to contain header"); + const auto header = torchao::ops::PackedWeightsHeader::read(packed_w_blob); + ET_CHECK_MSG( + header.type == + torchao::ops::PackedWeightsType:: + linear_8bit_act_xbit_weight_universal, + "quantized_moe_ffn requires universal torchao packed weights"); constexpr int kNr = 8, kKr = 16, kSr = 2; const int64_t required_bytes = static_cast(torchao::ops::PackedWeightsHeader::size()) + @@ -299,10 +361,10 @@ inline void expert_linear_dispatch( static_cast(k), static_cast(group_size), static_cast(weight_nbit)); -#endif // !ENABLE_QUANTIZED_MOE_FFN +#endif // !EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO switch (weight_nbit) { case 4: -#ifdef ENABLE_QUANTIZED_MOE_FFN +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO torchao_linear<4>( packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); #else @@ -311,7 +373,7 @@ inline void expert_linear_dispatch( #endif return; case 8: -#ifdef ENABLE_QUANTIZED_MOE_FFN +#ifdef EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO torchao_linear<8>( packed_w_blob, packed_blob_bytes, x, m, n, k, group_size, out); #else @@ -634,26 +696,20 @@ Tensor& quantized_moe_ffn_out( } }; -#ifdef ENABLE_QUANTIZED_MOE_FFN - // torchao linear path (perf-sensitive). The kernel threads internally on the - // shared pool in the common config, or runs single-threaded when only the - // thread-pool-free variant is linked. Distribute experts across the pool - // ourselves only when the kernel won't and the pool has more than one thread - // -- running both would nest on one pthreadpool and deadlock. +#if defined(EXECUTORCH_QUANTIZED_MOE_USE_TORCHAO) && \ + !defined(TORCHAO_PARALLEL_EXECUTORCH) const bool parallelize_experts = torchao::ops::linear_8bit_act_xbit_weight:: linear_operator_num_threads() == 1 && ::executorch::extension::threadpool::get_threadpool() ->get_thread_count() > 1; -#else - // Portable reference path: prefer simplicity over speed and run experts - // serially. - const bool parallelize_experts = false; -#endif if (parallelize_experts) { torch::executor::parallel_for(0, E, /*grain_size=*/1, run_experts); } else { run_experts(0, E); } +#else + run_experts(0, E); +#endif // ----- 7. Weighted scatter-add unpermute (cross-expert reduction) ----- // Each token sums the contributions of its top-k experts; run serially to diff --git a/extension/llm/custom_ops/targets.bzl b/extension/llm/custom_ops/targets.bzl index d6f251b38bd..4c2efb9adaa 100644 --- a/extension/llm/custom_ops/targets.bzl +++ b/extension/llm/custom_ops/targets.bzl @@ -53,17 +53,13 @@ def _get_quantized_moe_preproc_flags(): if runtime.is_oss: return [] if is_xplat(): - # TODO: enable KleidiAI for the runtime here by adding - # -DTORCHAO_ENABLE_KLEIDI (+ -DTORCHAO_ENABLE_ARM_I8MM=1) on arm64 and - # linking the kleidi kernel target in _get_quantized_moe_deps(). The - # runtime (op_moe.cpp) already selects the ukernel from the header, so - # kleidi headers resolve automatically once the kernels are compiled. - # Must be paired with the AoT packer emitting kleidi headers (see - # _get_quantized_moe_aot_packer_deps()). + # TODO: enable KleidiAI by adding its runtime config to op_moe.cpp, + # compiling and linking its kernels here, and pairing it with an AoT + # packer that emits Kleidi headers. return select({ "DEFAULT": [], "ovr_config//cpu:arm64": [ - "-DENABLE_QUANTIZED_MOE_FFN", + "-DEXECUTORCH_QUANTIZED_MOE_USE_TORCHAO", "-DTORCHAO_BUILD_CPU_AARCH64=1", "-DTORCHAO_ENABLE_ARM_NEON_DOT=1", ], diff --git a/extension/llm/custom_ops/test_op_moe.cpp b/extension/llm/custom_ops/test_op_moe.cpp index 30e05734bf0..81fb15185f4 100644 --- a/extension/llm/custom_ops/test_op_moe.cpp +++ b/extension/llm/custom_ops/test_op_moe.cpp @@ -43,10 +43,8 @@ TEST(OpQuantizedMoeFfnTest, RegistrationSmokeTest) { Tensor gate = tff.zeros({E, D}); Tensor expert_bias = tff.zeros({0}); - // Use empty packed buffers; the kernel will fail loudly if it tries to - // dereference them. With ENABLE_QUANTIZED_MOE_FFN unset (CI x86 build - // without torchao linkage) the kernel ET_CHECK_MSGs out before doing - // any real work, which is what we want this test to verify. + // Empty packed buffers document the expected schema; this test does not pass + // them to the kernel. Tensor packed_w1 = tfb.zeros({E, 1}); Tensor packed_w3 = tfb.zeros({E, 1}); Tensor packed_w2 = tfb.zeros({E, 1}); @@ -54,9 +52,7 @@ TEST(OpQuantizedMoeFfnTest, RegistrationSmokeTest) { Tensor out = tff.zeros({T, D}); executorch::runtime::KernelRuntimeContext ctx{}; - // We don't actually call the kernel here in the registration smoke test - // because the empty packed buffers would not be valid torchao blobs. - // Just verify the op symbol resolves at link time. + // Verify the op symbol resolves at link time. auto fn = &torch::executor::native::quantized_moe_ffn_out; EXPECT_NE(fn, nullptr); // Silence unused-variable warnings on the input tensors above; they From 5c5c9ee57f4623768b2bb00ea9a2f283e6e75476 Mon Sep 17 00:00:00 2001 From: jathu Date: Wed, 9 Sep 2026 10:53:53 -0700 Subject: [PATCH 107/190] Stop uploading Voxtral checkpoints and partial exports when CI exports fail (#22622) ### Summary Internal context: https://docs.google.com/document/d/1CnGfiP5SU0PG0IdJWn9Ca43LS8bGlEMi3l2gjgC7k48/edit When a job fails, the output dir is uploaded to GHA storage. This causes a significant increase in storage for Voxtral models: up to 0.7 TB/day. Like [Qwen3.5](https://github.com/pytorch/executorch/blob/9b3f660b4a74afa87726fbbb4480196d02fcb9e9/.ci/scripts/export_model_artifact.sh#L452-L455) and [Muse Glimmer](https://github.com/pytorch/executorch/blob/9b3f660b4a74afa87726fbbb4480196d02fcb9e9/.ci/scripts/export_model_artifact.sh#L524-L527), let's upload the Voxtral models to a temporary directory instead. In this PR we generalize the download location and cleanup for the models. ### Test plan CI --- .ci/scripts/export_model_artifact.sh | 41 +++++++++++++++++++--------- 1 file changed, 28 insertions(+), 13 deletions(-) diff --git a/.ci/scripts/export_model_artifact.sh b/.ci/scripts/export_model_artifact.sh index dfa4d7dc38b..223dd82951c 100755 --- a/.ci/scripts/export_model_artifact.sh +++ b/.ci/scripts/export_model_artifact.sh @@ -294,6 +294,26 @@ if [ "$MODEL_NAME" = "muse_glimmer" ]; then fi fi +# Downloads and compiler caches go in scratch dirs outside OUTPUT_DIR because the CI job +# templates upload OUTPUT_DIR even when the job fails. A failed export also empties +# OUTPUT_DIR, but only if it started out empty, so a local run with output_dir=. cannot +# delete the checkout. Scratch goes under RUNNER_TEMP, which the runner wipes between +# jobs, with a fallback for containers where RUNNER_TEMP is not writable. +LOCAL_MODEL_DIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/model_XXXXXX" 2>/dev/null || mktemp -d) +SCRATCH_DIRS=("$LOCAL_MODEL_DIR") +OUTPUT_DIR_WAS_EMPTY=0 +[ -n "$(ls -A -- "$OUTPUT_DIR" 2>/dev/null)" ] || OUTPUT_DIR_WAS_EMPTY=1 +cleanup() { + local rc=$? + set +e + rm -rf "${SCRATCH_DIRS[@]}" + if [ "$rc" -ne 0 ] && [ "$OUTPUT_DIR_WAS_EMPTY" = 1 ] && [ -d "$OUTPUT_DIR" ]; then + echo "Export failed with exit code $rc; removing partial output from ${OUTPUT_DIR}" + (cd -- "$OUTPUT_DIR" && find . -mindepth 1 -delete) + fi +} +trap cleanup EXIT + echo "::group::Export $MODEL_NAME" if [ -n "$EXTRA_PIP" ]; then @@ -386,8 +406,7 @@ fi if [ "$MODEL_NAME" = "voxtral_realtime" ]; then pip install safetensors huggingface_hub - # Download model weights from HuggingFace (requires HF_TOKEN for gated model) - LOCAL_MODEL_DIR="${OUTPUT_DIR}/model_weights" + # Download model weights outside OUTPUT_DIR to avoid uploading on failure (requires HF_TOKEN for gated model) python -c "from huggingface_hub import snapshot_download; snapshot_download('${HF_MODEL}', local_dir='${LOCAL_MODEL_DIR}')" # Per-component quantization flags @@ -437,7 +456,6 @@ if [ "$MODEL_NAME" = "voxtral_realtime" ]; then fi # Copy tokenizer from downloaded model weights cp "$LOCAL_MODEL_DIR/tekken.json" "${OUTPUT_DIR}/tekken.json" - rm -rf "$LOCAL_MODEL_DIR" ls -al "${OUTPUT_DIR}" echo "::endgroup::" exit 0 @@ -448,12 +466,11 @@ if [ "$MODEL_NAME" = "qwen3_5_moe" ]; then pip install safetensors huggingface_hub pip install -r examples/models/qwen3_5_moe/requirements.txt - # Download prequantized model outside OUTPUT_DIR to avoid uploading on failure - LOCAL_MODEL_DIR=$(mktemp -d) INDUCTOR_CACHE=$(mktemp -d "${RUNNER_TEMP:-/tmp}/inductor_cache_XXXXXX") INDUCTOR_TMPDIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/tmpdir_XXXXXX") - trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR"' EXIT + SCRATCH_DIRS+=("$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR") + # Download prequantized model outside OUTPUT_DIR to avoid uploading on failure python -c "from huggingface_hub import snapshot_download; snapshot_download('${HF_MODEL}', local_dir='${LOCAL_MODEL_DIR}')" # Sanity check: run inference on the prequantized model @@ -521,10 +538,9 @@ fi if [ "$MODEL_NAME" = "muse_glimmer" ]; then pip install safetensors huggingface_hub gguf - LOCAL_MODEL_DIR=$(mktemp -d) INDUCTOR_CACHE=$(mktemp -d "${RUNNER_TEMP:-/tmp}/inductor_cache_XXXXXX") INDUCTOR_TMPDIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/tmpdir_XXXXXX") - trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR"' EXIT + SCRATCH_DIRS+=("$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR") case "$QUANT_NAME" in kquant-17gb) @@ -588,14 +604,13 @@ fi if [ "$MODEL_NAME" = "gemma4_31b" ]; then pip install safetensors huggingface_hub gguf - # Download GGUF + tokenizer outside OUTPUT_DIR to avoid uploading on failure. - # The unsloth GGUF repo ships the .gguf but no tokenizer.json, so the tokenizer - # is fetched from the (non-GGUF) unsloth/gemma-4-31B-it repo. - LOCAL_MODEL_DIR=$(mktemp -d) INDUCTOR_CACHE=$(mktemp -d "${RUNNER_TEMP:-/tmp}/inductor_cache_XXXXXX") INDUCTOR_TMPDIR=$(mktemp -d "${RUNNER_TEMP:-/tmp}/tmpdir_XXXXXX") - trap 'rm -rf "$LOCAL_MODEL_DIR" "$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR"' EXIT + SCRATCH_DIRS+=("$INDUCTOR_CACHE" "$INDUCTOR_TMPDIR") + # Download GGUF + tokenizer outside OUTPUT_DIR to avoid uploading on failure. + # The unsloth GGUF repo ships the .gguf but no tokenizer.json, so the tokenizer + # is fetched from the (non-GGUF) unsloth/gemma-4-31B-it repo. GGUF_FILE="gemma-4-31B-it-Q4_K_M.gguf" python -c "from huggingface_hub import hf_hub_download; hf_hub_download('unsloth/gemma-4-31B-it-GGUF', '${GGUF_FILE}', local_dir='${LOCAL_MODEL_DIR}')" python -c "from huggingface_hub import hf_hub_download; hf_hub_download('unsloth/gemma-4-31B-it', 'tokenizer.json', local_dir='${LOCAL_MODEL_DIR}')" From afc795776d050356fb6cce26f16652905ec1b5bf Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Thu, 3 Sep 2026 09:19:50 +0200 Subject: [PATCH 108/190] Vulkan: do not partition constant_pad_nd with a symbolic pad A pad amount derived from a dynamic dimension is serialized as a VALUELIST of Int/SymInt, and Pad.cpp reads the list with get_int_list(), which requires a literal IntList. The partitioner claims the node regardless, so the model lowers and then aborts at prepack: Exception raised from toIntList at .../graph/containers/Value.h:272: (isIntList()) is false! Expected value to have type IntList, got VALUELIST This is reachable from ordinary code. Any model that pads a dynamic sequence up to the static length its LSTM wants hits it, which is how kokoro's synthesizer fails: `F.pad(x, (0, 0, 0, target - seq_len))`. Supporting it properly is more than swapping in extract_int_or_symint_list(), the way Split.cpp, View.cpp and Expand.cpp read their symbolic lists. add_constant_pad_nd_node() folds the amounts into a per-dim offset and bakes that into a params buffer at build time, so the dispatch would use stale offsets even if the list were read symbolically; the buffer has to be refreshed on resize first. Decline the node until then, so it falls back rather than aborting. The existing test_vulkan_backend_constant_pad_nd pads by a literal (1, 2, 3, 4, 5, 6), so no test covered a symbolic pad. The new case aborts on main and falls back cleanly with this change. --- backends/vulkan/op_registry.py | 22 +++++++++++++++++ backends/vulkan/test/test_vulkan_delegate.py | 26 ++++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index 91da9134850..eef73ed04a4 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1507,12 +1507,34 @@ def register_arange(): # ============================================================================= +def _check_pad_is_static(node: torch.fx.Node) -> bool: + """Only support constant_pad_nd when the pad amounts are static. + + A symbolic pad list is serialized as a VALUELIST rather than an INTLIST, and + Pad.cpp reads it with get_int_list(), which throws "Expected value to have + type IntList, got VALUELIST instead". + + Supporting it properly is more than swapping in + extract_int_or_symint_list(), the way Split.cpp, View.cpp and Expand.cpp + read their symbolic lists: add_constant_pad_nd_node() folds the amounts into + a per-dim offset and bakes that into a params buffer at BUILD time, so the + dispatch would still use stale offsets even if the list were read + symbolically. The buffer has to be refreshed on resize first. Decline the + node until then. + """ + pad = node.args[1] + if not isinstance(pad, (list, tuple)): + return False + return all(isinstance(p, int) for p in pad) + + @update_features(exir_ops.edge.aten.constant_pad_nd.default) def register_constant_pad_nd(): return OpFeatures( inputs_storage=utils.ANY_STORAGE, inputs_dtypes=utils.FP_INT_BOOL_T, supports_resize=True, + are_node_inputs_supported_fn=_check_pad_is_static, ) diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 34a4f12f62c..50ac763679a 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -1479,6 +1479,32 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_constant_pad_nd_symbolic_pad(self): + """A pad amount derived from a dynamic dim, as LSTM padding produces. + + Without the guard this partitions and then aborts at prepack with + "Expected value to have type IntList, got VALUELIST instead", because + the pad list is serialized as a VALUELIST of Int/SymInt. + """ + + class TestModule(torch.nn.Module): + def forward(self, x): + # Pad up to a static length, the shape every unrolled LSTM + # wants its input in. + return torch.nn.functional.pad(x, (0, 0, 0, 16 - x.shape[1])) + + sample_inputs = (torch.randn(size=(1, 12, 8), dtype=torch.float32),) + seq = Dim("seq", min=2, max=16) + self.lower_module_and_test_output( + TestModule(), + sample_inputs, + dynamic_shapes={"x": {1: seq}}, + test_inputs=[ + (torch.randn(size=(1, 4, 8), dtype=torch.float32),), + (torch.randn(size=(1, 16, 8), dtype=torch.float32),), + ], + ) + def test_vulkan_backend_repeat(self): class TestModule(torch.nn.Module): def __init__(self): From f1adb8b35c05dae9bb82e0c65df7da158997845e Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak Date: Wed, 2 Sep 2026 09:03:53 +0200 Subject: [PATCH 109/190] [ET-VK] Run weight-only int8 linear on the kernel that still exists FuseQuantizedOpsTransform rewrites a weight-only int8 linear to aten._weight_int8pack_mm, op_registry lists it as supported, so the partitioner takes it. At runtime it dispatched to linear_qcs8w_tiled / linear_qcs8w_coop, and those shaders are not in the repo: #14041 replaced them with linear_q8csw_tiled and a different suffix scheme. Every weight-only int8 model therefore aborts on its first execute with Could not find ShaderInfo with name linear_qcs8w_tiled_texture3d_texture3d_texture2d_texture2d_half_o4x3 aten._weight_int8pack_mm carries the same operands as et_vk.linear_q8csw minus the bias, so register it on that implementation and drop the 8-bit half of QuantizedLinearQCSNW.cpp, whose shaders are gone. The 4-bit et_vk.linear_qcs4w path is left untouched; it references missing shaders too, but the group-wise kernel is not a drop-in replacement for it. Tested on a Galaxy S26 Ultra (Adreno 840) with all-mpnet-base-v2 and multi-qa-mpnet-base-dot-v1 lowered through VulkanQuantizer weight-only int8, at the 382-token bound and resized down to 128. Cosine against the fp32 eager reference is 0.998280 / 0.995343 (fp16 for comparison: 0.999972). The int8 file is 133.2 MB against 217.9 MB for fp16, and per execution it runs 59 ms against 79 ms at 382 tokens, 36 vs 49 at 254, 26 vs 29 at 128 (median of 3 rounds, arm order reversed each round, warm-up discarded). --- .../graph/ops/impl/QuantizedLinear.cpp | 35 +++++ .../graph/ops/impl/QuantizedLinearQCSNW.cpp | 121 +----------------- 2 files changed, 40 insertions(+), 116 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp index dc0a0a0837b..1c394deca08 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinear.cpp @@ -944,6 +944,40 @@ void linear_q8csw(ComputeGraph& graph, const std::vector& args) { output); } +// aten._weight_int8pack_mm is what the AOT weight-only int8 fusion +// (FuseQuantizedOpsTransform) emits. It carries the same operands as +// et_vk.linear_q8csw minus the bias, so it runs through the same +// implementation. +void weight_int8pack_mm( + ComputeGraph& graph, + const std::vector& args) { + int32_t idx = 0; + const ValueRef fp_input = args.at(idx++); + const ValueRef weight_data = args.at(idx++); + const ValueRef weight_scales_data = args.at(idx++); + const ValueRef output = args.at(idx++); + + const int64_t K = graph.size_at(-1, fp_input); + + QuantizationConfig input_quant_config(32, kNoQuantization, {}); + QuantizationConfig weight_quant_config(8, kPerChannel, {K}); + + quantized_linear_impl( + graph, + input_quant_config, + weight_quant_config, + fp_input, + kDummyValueRef, // input scale + kDummyValueRef, // input zp + weight_data, + kDummyValueRef, // weight sums + weight_scales_data, + kDummyValueRef, // weight zeros + kDummyValueRef, // group size + kDummyValueRef, // bias + output); +} + void linear_dq8ca_q4gsw( ComputeGraph& graph, const std::vector& args) { @@ -982,6 +1016,7 @@ void linear_dq8ca_q4gsw( REGISTER_OPERATORS { VK_REGISTER_OP(et_vk.linear_q8ta_q8csw.default, linear_q8ta_q8csw); VK_REGISTER_OP(et_vk.linear_q8csw.default, linear_q8csw); + VK_REGISTER_OP(aten._weight_int8pack_mm.default, weight_int8pack_mm); VK_REGISTER_OP(et_vk.linear_dq8ca_q4gsw.default, linear_dq8ca_q4gsw); } diff --git a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp index 5d9311e9761..c2e45b96418 100644 --- a/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/QuantizedLinearQCSNW.cpp @@ -16,34 +16,6 @@ namespace vkcompute { -// Custom global workgroup size function for linear_qcs8w -GlobalWorkGrid linear_qcs8w_gwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; - const ValueRef out = args.at(0).refs.at(0); - return graph->create_linear_gwg( - utils::safe_downcast(graph->numel_of(out))); -} - -// Custom local workgroup size function for linear_qcs8w -LocalWorkGroup linear_qcs8w_lwg( - ComputeGraph* graph, - const vkapi::ShaderInfo& shader, - const GlobalWorkGrid& gwg, - const std::vector& args, - const std::vector& resize_args) { - (void)graph; - (void)shader; - (void)gwg; - (void)args; - (void)resize_args; - return LocalWorkGroup(64u, 1u, 1u); -} - // Custom global workgroup size function for linear_qcsnw_tiled GlobalWorkGrid linear_qcsnw_tiled_gwg( ComputeGraph* graph, @@ -178,81 +150,6 @@ void resize_linear_qcsnw_node( graph->virtual_resize(out, new_out_sizes); } -void add_linear_qcs8w_node( - ComputeGraph& graph, - const ValueRef mat1, - const ValueRef q_mat2_data, - const ValueRef scales_data, - const ValueRef out) { - auto viewFn = VK_GET_OP_FN("aten.view_copy.default"); - ValueRef mat1_W_packed = mat1; - ValueRef out_W_packed = out; - // Create temporary tensors to store the width packed versions of mat1 and out - TmpTensor mat1_tmp( - &graph, graph.sizes_of(mat1), graph.dtype_of(mat1), utils::kWidthPacked); - TmpTensor out_tmp( - &graph, graph.sizes_of(out), graph.dtype_of(out), utils::kWidthPacked); - if (!graph.is_buffer_storage(out) && - graph.packed_dim_of(mat1) != WHCN::kWidthDim) { - // Ensure mat1 is width packed - mat1_W_packed = mat1_tmp; - viewFn(graph, {mat1, graph.add_none(), mat1_W_packed}); - // Ensure out is packed correctly - out_W_packed = out_tmp; - } - ValueRef q_mat2 = prepack_standard_hw_transposed( - graph, q_mat2_data, graph.storage_type_of(out), utils::kWidthPacked); - ValueRef scales = prepack_standard( - graph, scales_data, graph.storage_type_of(out), utils::kWidthPacked); - - std::string kernel_name = "linear_qcs8w"; - kernel_name.reserve(kShaderNameReserve); - add_packed_dim_suffix(kernel_name, graph.packed_dim_of(mat1_W_packed)); - add_packed_dim_suffix(kernel_name, graph.packed_dim_of(q_mat2)); - add_dtype_suffix(kernel_name, graph.dtype_of(out_W_packed)); - add_storage_type_suffix(kernel_name, graph.storage_type_of(out_W_packed)); - - std::vector pcs; - if (graph.is_buffer_storage(out_W_packed)) { - pcs = { - graph.sizes_pc_of(out_W_packed), - graph.strides_pc_of(out_W_packed), - graph.sizes_pc_of(mat1_W_packed), - graph.strides_pc_of(mat1), - graph.strides_pc_of(q_mat2), - graph.strides_pc_of(scales), - graph.numel_pc_of(out_W_packed)}; - } else { - pcs = { - graph.logical_limits_pc_of(out_W_packed), - graph.sizes_pc_of(mat1_W_packed), - graph.sizes_pc_of(q_mat2)}; - } - - graph.execute_nodes().emplace_back(new DynamicDispatchNode( - graph, - VK_KERNEL_FROM_STR(kernel_name), - linear_qcs8w_gwg, - linear_qcs8w_lwg, - // Inputs and Outputs - {{out_W_packed, vkapi::MemoryAccessType::WRITE}, - {{mat1_W_packed, q_mat2, scales}, vkapi::MemoryAccessType::READ}}, - // Shader params buffers - {}, - // Push Constants - pcs, - // Specialization Constants - {}, - // Resize Args - {}, - // Resizing Logic - resize_linear_qcsnw_node)); - if (!graph.is_buffer_storage(out) && - graph.packed_dim_of(out) != WHCN::kWidthDim) { - viewFn(graph, {out_W_packed, graph.add_none(), out}); - } -} - void add_linear_qcsnw_tiled_node( ComputeGraph& graph, const bool use_coop_algorithm, @@ -293,6 +190,8 @@ void add_linear_qcsnw_tiled_node( kernel_name = use_coop_algorithm ? "linear_qcs4w_coop" : "linear_qcs4w_tiled"; } else { + // Unreachable: linear_qcs4w is the only caller of this function and it + // always passes quant_nbits == 4. Kept so the 4-bit path is not disturbed. kernel_name = use_coop_algorithm ? "linear_qcs8w_coop" : "linear_qcs8w_tiled"; } @@ -389,18 +288,6 @@ bool can_use_coop_impl(ComputeGraph& graph, const ValueRef mat1) { return (graph.size_at(-2, mat1) == 1); } -void weight_int8pack_mm( - ComputeGraph& graph, - const std::vector& args) { - check_linear_qcsnw_args(graph, 8, args[0], args[1], args[2], args[3]); - if (can_use_tiled_impl(graph, args[0], args[1], args[2], args[3])) { - bool use_coop_algorithm = can_use_coop_impl(graph, args[0]); - return add_linear_qcsnw_tiled_node( - graph, use_coop_algorithm, 8, args[0], args[1], args[2], args[3]); - } - return add_linear_qcs8w_node(graph, args[0], args[1], args[2], args[3]); -} - void linear_qcs4w(ComputeGraph& graph, const std::vector& args) { check_linear_qcsnw_args(graph, 4, args[0], args[1], args[2], args[3]); @@ -411,7 +298,9 @@ void linear_qcs4w(ComputeGraph& graph, const std::vector& args) { } REGISTER_OPERATORS { - VK_REGISTER_OP(aten._weight_int8pack_mm.default, weight_int8pack_mm); + // aten._weight_int8pack_mm is registered in QuantizedLinear.cpp, on the + // maintained weight-only quantized linear implementation. The 8-bit path + // here dispatched to linear_qcs8w_* shaders that no longer exist. VK_REGISTER_OP(et_vk.linear_qcs4w.default, linear_qcs4w); } From 74f553eaf7ad874c2fe4985c5204996cee38a140 Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 9 Sep 2026 11:22:00 -0700 Subject: [PATCH 110/190] Cortex-M: preserve calibration input layout (#22572) ### Summary Calibration observers use logical tensor values, so avoid converting calibration samples to channels-last. Cover the behavior in the MobileNet tests and clarify the documentation. Authored with Codex. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell --- backends/arm/scripts/aot_arm_compiler.py | 2 -- backends/cortex_m/test/models/test_mobilenet_v2.py | 5 +---- backends/cortex_m/test/models/test_mobilenet_v3.py | 5 +---- docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md | 3 +++ 4 files changed, 5 insertions(+), 10 deletions(-) diff --git a/backends/arm/scripts/aot_arm_compiler.py b/backends/arm/scripts/aot_arm_compiler.py index 81ca626a031..5c89c79a960 100644 --- a/backends/arm/scripts/aot_arm_compiler.py +++ b/backends/arm/scripts/aot_arm_compiler.py @@ -981,8 +981,6 @@ def _to_channels_last(x): calibration_samples = [example_inputs] for sample in calibration_samples: - if not args.cortex_m_explicit_layout: - sample = tuple(_to_channels_last(x) for x in sample) prepared(*sample) model_quant = convert_pt2e(prepared) diff --git a/backends/cortex_m/test/models/test_mobilenet_v2.py b/backends/cortex_m/test/models/test_mobilenet_v2.py index 9bc99e4bf2c..7a3bc9ecfe3 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v2.py +++ b/backends/cortex_m/test/models/test_mobilenet_v2.py @@ -37,10 +37,7 @@ } # Use larger sample set for calibration to get better quantization -calibration_samples = [ - (torch.randn(1, 3, 224, 224).to(memory_format=torch.channels_last),) - for _ in range(100) -] +calibration_samples = [(torch.randn(1, 3, 224, 224),) for _ in range(100)] test_cases = { "mobilenet_v2": McuTestCase( diff --git a/backends/cortex_m/test/models/test_mobilenet_v3.py b/backends/cortex_m/test/models/test_mobilenet_v3.py index 08633d54dd6..3a6f0a5004f 100644 --- a/backends/cortex_m/test/models/test_mobilenet_v3.py +++ b/backends/cortex_m/test/models/test_mobilenet_v3.py @@ -39,10 +39,7 @@ } # Use bigger sample set for calibration. -calibration_samples = [ - (torch.randn(1, 3, 232, 232).to(memory_format=torch.channels_last),) - for i in (range(100)) -] +calibration_samples = [(torch.randn(1, 3, 232, 232),) for _ in range(100)] test_cases = { "mobilenet_v3_small": McuTestCase( diff --git a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md index 9094ab96d51..47eb2fceeb7 100644 --- a/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md +++ b/docs/source/backends/arm-cortex-m/arm-cortex-m-overview.md @@ -100,6 +100,9 @@ quantized = convert_pt2e(prepared) quantized_exported_program = torch.export.export(quantized, (example_input,)) ``` +Calibration observes logical tensor values, so calibration inputs do not need +to use the same memory format as the export example. + ### 2. Lower to edge and apply Cortex-M passes Lower to the edge dialect with the backend's `EdgeCompileConfig`, then run the `CortexMPassManager` to replace quantized subgraphs with CMSIS-NN operator implementations: From 1e937f28d183a0266f1de2d4d1ec3c9120e134c9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sun, 30 Aug 2026 15:52:56 +0200 Subject: [PATCH 111/190] [ET-VK] Fix out-of-range broadcast size calculation for 0-dim tensors calculate_broadcasted_output_size() guards its loop with `i >= -out_sizes.size()`. That is size_t arithmetic: when both operands are 0-dimensional the bound evaluates to 0, `i` promotes to a huge unsigned value, the guard stays true, and the body runs `out_sizes.at(size() - 1)` on an empty vector, throwing std::out_of_range. Any binary op whose output is 0-dimensional therefore aborts at graph build. Whisper hits it in the log-mel normalisation, which makes its whole encode method unloadable on Vulkan. Hold the bound in a signed local so the loop is skipped when the output is 0-dimensional. Non-empty cases are unchanged: the previous unsigned wraparound happened to compare correctly for them. --- .../runtime/graph/ops/impl/utils/TensorUtils.cpp | 9 +++++++-- backends/vulkan/test/test_vulkan_delegate.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp index df83bdec7e0..90b923194ad 100644 --- a/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.cpp @@ -18,10 +18,15 @@ std::vector calculate_broadcasted_output_size( const std::vector& sizes1, const std::vector& sizes2) { std::vector out_sizes(std::max(sizes1.size(), sizes2.size())); + // ndim must be signed. `-out_sizes.size()` is size_t arithmetic, so when both + // inputs are 0-dimensional it evaluates to 0 while `i` promotes to a huge + // unsigned value, the guard stays true, and `out_sizes.at(size() - 1)` throws + // std::out_of_range instead of the loop being skipped. + const int64_t ndim = static_cast(out_sizes.size()); // Match the sizes in reverse because sizes are in NCHW order - for (int i = -1; i >= -out_sizes.size(); --i) { - out_sizes.at(out_sizes.size() + i) = + for (int64_t i = -1; i >= -ndim; --i) { + out_sizes.at(static_cast(ndim + i)) = std::max(utils::val_at(i, sizes1), utils::val_at(i, sizes2)); } diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 50ac763679a..84d9c88fbfc 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -1080,6 +1080,20 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_binary_op_zero_dim(self): + # Both operands, and therefore the output, are 0-dimensional. This is + # what a reduction to a scalar followed by arithmetic produces, e.g. the + # log-mel normalisation in Whisper's preprocessor. + class ZeroDimModule(torch.nn.Module): + def forward(self, x): + m = x.max() + return (m - (m - 1.0)).reshape(1) + + self.lower_module_and_test_output( + ZeroDimModule(), + (torch.randn(size=(64,), dtype=torch.float32),), + ) + @disable_test("layer norm compute shader not working with swiftshader") def test_vulkan_backend_native_layer_norm(self): class NativeLayerNormModule(torch.nn.Module): From 90c5a67a74db9ef0b3469cb38504996f845d8523 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sat, 29 Aug 2026 16:25:47 +0200 Subject: [PATCH 112/190] Vulkan: handle a single tensor output arriving as a list num_tensors_in_node() counts the tensors associated with a node rather than the nesting of meta["val"], so it returns 1 both for a bare FakeTensor and for a one-element list. OpRepSets treats that count as proof of the former and passes meta["val"] straight to filter_invalid_reprs(), which then does tensor_val.shape on a list: File "backends/vulkan/utils.py", line 1203, in filter_invalid_reprs extents = required_image_extents(tensor_val.shape, memory_layout) AttributeError: 'list' object has no attribute 'shape' Any op declared to return Tensor[] hits this whenever it happens to produce exactly one tensor. aten.split_with_sizes_copy.default is the case seen in the wild (RF-DETR); it aborts partitioning for the whole model instead of the node being reported unsupported. Unwrap the single element before filtering, matching what the multiple-output branch below already does per element. --- .../vulkan/test/test_vulkan_tensor_repr.py | 23 +++++++++++++++++++ backends/vulkan/utils.py | 9 +++++++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/backends/vulkan/test/test_vulkan_tensor_repr.py b/backends/vulkan/test/test_vulkan_tensor_repr.py index 5a0fc664c17..83d68195407 100644 --- a/backends/vulkan/test/test_vulkan_tensor_repr.py +++ b/backends/vulkan/test/test_vulkan_tensor_repr.py @@ -605,6 +605,29 @@ def test_unary_op_construction(self): self.assertEqual(op_repsets.primary_arg_idx, 0) self.assertTrue(op_repsets.sync_primary_io_repr) + def test_single_tensor_output_in_list_construction(self): + """An op declared to return Tensor[] that yields exactly one tensor. + + num_tensors_in_node() counts tensors rather than nesting, so such a + node reports 1 while meta["val"] is a one-element list rather than a + bare FakeTensor. + """ + arg = _make_tensor_arg_node((1, 3, 8, 8)) + node = _make_op_node( + target=torch.ops.aten.split_with_sizes_copy.default, + args=(arg, [3]), + output_val=[_make_fake_tensor((1, 3, 8, 8))], + ) + + op_repsets = OpRepSets( + TensorRepSetList(ANY_STORAGE), + TensorRepSetList(ANY_STORAGE), + node, + DEFAULT_TEXTURE_LIMITS, + ) + + self.assertFalse(op_repsets.any_is_empty()) + def test_binary_op_syncs_args(self): """When a single repset covers all inputs, sync_args_repr is True.""" op_repsets = self._make_binary_op() diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 84b901b6b6e..066ee78ea0d 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -1442,8 +1442,15 @@ def __init__( # noqa: C901 outs_repset_list = TensorRepSetList([]) common_out_repset = ANY_STORAGE_INCL_PACKED_INT8 if num_tensors_in_node(op_node) == 1: + out_val = op_node.meta["val"] + # num_tensors_in_node counts tensors, not nesting: an op declared + # to return Tensor[] still lands here when it happens to produce + # exactly one, and meta["val"] is then a one-element list rather + # than a bare FakeTensor. + if isinstance(out_val, (list, tuple)): + out_val = out_val[0] common_out_repset = filter_invalid_reprs( - op_node.meta["val"], outputs_repsets[0], texture_limits + out_val, outputs_repsets[0], texture_limits ) outs_repset_list.append(common_out_repset) # Multiple output tensors From f3dd4b4c5cd2a3f2fa7018688623656b282e0cfe Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:41:46 +0200 Subject: [PATCH 113/190] [ET-VK] Keep unused placeholders in the delegate's input list Fixes #22478. `VulkanBackend::execute` matches its `args` positionally against `ComputeGraph::inputs()`, but `process_placeholder_node` drops any placeholder with no users before it reaches `input_ids`. The delegate call still passes that argument, so the counts disagree and the model fails at every `execute()` with "Vulkan graph declares N inputs and M outputs, but the delegate call supplied K arguments". A placeholder can lose its only consumers to the passes that run inside `preprocess`, after partitioning has already fixed the call's argument list. A qwen3 0.6B export with attention kept on the host hits this: the final partition is handed 28 per-layer symints that nothing in the subgraph consumes by serialization time, and declares 2 inputs and 1 output against a 31 argument call. Params keep the existing skip, since they are serialized into the blob rather than passed at call time. Test: `test_unused_placeholder_is_still_declared_as_an_input` fails on `main` and passes with the change; the companion test covering used placeholders passes either way. `test_serialization.py`, `test_vulkan_compile_options.py` and `test_vulkan_passes.py` (17 tests) still pass. --- .../serialization/vulkan_graph_builder.py | 25 +++++--- .../vulkan/test/test_vulkan_graph_builder.py | 60 +++++++++++++++++++ 2 files changed, 78 insertions(+), 7 deletions(-) create mode 100644 backends/vulkan/test/test_vulkan_graph_builder.py diff --git a/backends/vulkan/serialization/vulkan_graph_builder.py b/backends/vulkan/serialization/vulkan_graph_builder.py index 46e01e701b1..e5dfde9d865 100644 --- a/backends/vulkan/serialization/vulkan_graph_builder.py +++ b/backends/vulkan/serialization/vulkan_graph_builder.py @@ -418,15 +418,26 @@ def get_or_create_value_for(self, arg: _Argument): raise RuntimeError(f"Cannot create value for arg of type {type(arg)}") def process_placeholder_node(self, node: Node) -> None: - # ignores any tensors that don't get used in any ops - if len(node.users) == 0: + # A non-param placeholder occupies a slot in the delegate call's + # argument list whether or not this graph goes on to use it, and + # VulkanBackend::execute matches `args` to graph inputs positionally. + # Dropping an unused one from input_ids desynchronises the two, and the + # runtime then rejects the call because it was handed more arguments + # than the graph declares inputs and outputs. That happens in practice + # when a placeholder's only consumers are folded away by the passes + # that run after partitioning, so the graph the partitioner tagged and + # the graph serialized here disagree about which inputs are live. + if is_param_node(self.program, node): + # Params are serialized into the blob rather than passed at call + # time, so an unused one costs nothing to skip. + if len(node.users) > 0: + self.create_node_value(node) return None ids = self.create_node_value(node) - if not is_param_node(self.program, node): - if isinstance(ids, int): - self.input_ids.append(ids) - else: - self.input_ids += ids + if isinstance(ids, int): + self.input_ids.append(ids) + else: + self.input_ids += ids def process_getitem_node(self, node: Node) -> None: # Find ValueList id from the collection node. diff --git a/backends/vulkan/test/test_vulkan_graph_builder.py b/backends/vulkan/test/test_vulkan_graph_builder.py new file mode 100644 index 00000000000..65afc3a2542 --- /dev/null +++ b/backends/vulkan/test/test_vulkan_graph_builder.py @@ -0,0 +1,60 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.vulkan.serialization.vulkan_graph_builder import VkGraphBuilder +from executorch.backends.vulkan.vulkan_preprocess import apply_passes +from executorch.exir import to_edge +from executorch.exir.backend.utils import DelegateMappingBuilder +from executorch.exir.passes import SpecPropPass + + +class TestVkGraphBuilderInputIds(unittest.TestCase): + """The serialized input list has to match the delegate call's arguments. + + VulkanBackend::execute walks `args` positionally against + ComputeGraph::inputs() and rejects the call when the counts disagree, so + every placeholder that the delegate call passes must appear in input_ids, + including ones this graph happens not to use. Unused placeholders are not + hypothetical: passes that run after partitioning can fold away a + placeholder's only consumers, leaving the argument list and the serialized + graph out of step. + """ + + def _build(self, module: torch.nn.Module, inputs) -> VkGraphBuilder: + edge = to_edge(torch.export.export(module, inputs, strict=True)) + # The builder reads node specs, which the backend's own preprocess + # populates before it gets here. + program = apply_passes(edge.exported_program(), [SpecPropPass()]) + builder = VkGraphBuilder( + program, DelegateMappingBuilder(generated_identifiers=True) + ) + builder.build_graph() + return builder + + def test_unused_placeholder_is_still_declared_as_an_input(self) -> None: + class UsesOnlyTheFirstInput(torch.nn.Module): + def forward(self, used, unused): + return used + used + + builder = self._build( + UsesOnlyTheFirstInput(), (torch.randn(2, 3), torch.randn(2, 3)) + ) + self.assertEqual(len(builder.input_ids), 2) + + def test_used_placeholders_are_declared_in_order(self) -> None: + class UsesBothInputs(torch.nn.Module): + def forward(self, first, second): + return first + second + + builder = self._build(UsesBothInputs(), (torch.randn(2, 3), torch.randn(2, 3))) + self.assertEqual(len(builder.input_ids), 2) + + +if __name__ == "__main__": + unittest.main() From a489948c40274c67e8a06266352f4c622dc45874 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sun, 30 Aug 2026 11:17:31 +0200 Subject: [PATCH 114/190] [ET-VK] Clamp the tanh argument in the gelu shader The gelu shader evaluates the tanh approximation with an unclamped argument. For an input of x the argument is sqrt(2/pi) * (x + 0.044715 * x^3) which grows cubically, so x = -13.24 already yields -93.4. A driver that evaluates tanh as (e^y - e^-y) / (e^y + e^-y) overflows fp32 at |y| > ~88 and returns inf/inf = NaN. On a Mali-G76 this makes the Whisper encoder emit NaN from conv2 for exactly the two activations below -13, and the first LayerNorm then propagates them across the whole tensor. The tanh op in this same file already clamps to +/-15 for this reason; gelu now does the same. The clamp is numerically free: 1 - tanh(15) is 1.9e-13, well below fp32 epsilon, so the result is bit-identical for every input that gets clamped. --- backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml index 0331a15fde6..fc70b54076b 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/unary_op.yaml @@ -31,7 +31,7 @@ unary_op: - NAME: exp OPERATOR: exp(X) - NAME: gelu - OPERATOR: 0.5 * X * (1 + tanh(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X))) + OPERATOR: 0.5 * X * (1 + tanh(clamp(sqrt(2 / 3.141593) * (X + 0.044715 * X * X * X), -15.0, 15.0))) - NAME: neg OPERATOR: -X - NAME: sigmoid From 43afe51f34ceb49ecdac0d1439d7e41fbe9563dd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Thu, 3 Sep 2026 09:20:55 +0200 Subject: [PATCH 115/190] Vulkan: do not partition batch norm on a non-4d input add_native_batch_norm_node() asserts VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor") on both the input and the output sizes, but the partitioner claims the node at any rank. A batch norm on rank-3 activations, which is every conv1d model, lowers cleanly and then aborts at execute time. Decline the node instead so it falls back. Split out of #22399, which is now scoped to constant_pad_nd. --- backends/vulkan/op_registry.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index eef73ed04a4..d04d91f32e9 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1756,6 +1756,22 @@ def register_embedding_q4gsw(): # ============================================================================= +def _check_batch_norm_is_4d(node: torch.fx.Node) -> bool: + """Only support batch norm on a 4d input. + + add_native_batch_norm_node() asserts + VK_CHECK_COND(in_sizes.size() == 4, "BatchNorm only support 4d tensor") on + both the input and the output, so partitioning a batch norm whose input is + not 4d yields a .pte that lowers cleanly and then aborts at execute time. + Any conv1d model reaches here with rank-3 activations. + """ + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + return False + val = input_node.meta.get("val") + return val is not None and val.dim() == 4 + + @update_features(exir_ops.edge.aten._native_batch_norm_legit_no_training.default) def register_native_batch_norm_legit_no_training(): return OpFeatures( @@ -1763,6 +1779,7 @@ def register_native_batch_norm_legit_no_training(): inputs_dtypes=utils.FP_T, supports_prepacking=True, supports_resize=True, + are_node_inputs_supported_fn=_check_batch_norm_is_4d, ) From 6a03b0ba4592158c34544ac5226b29592d1cff5a Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Wed, 9 Sep 2026 20:54:41 +0200 Subject: [PATCH 116/190] [ET-VK] Classify the conv2d method by weight shape in conv2d_local_wg_size (#22051) ### Summary Fixes #21942. `conv2d_local_wg_size()` picked the convolution method from the shader name alone, and the condition it used matched every conv2d shader: ```cpp if (kernel_name.find("conv2d_pw") != npos || (kernel_name.find("conv2d") != npos && kernel_name.find("conv_transpose2d") == npos)) { method = Conv2dMethod::Pointwise; } else { method = Conv2dMethod::SlidingWindow; } ``` The sliding window shader is itself named `conv2d`, so it matched and was labelled `Pointwise`, making the `SlidingWindow` branch unreachable for every conv2d variant. Only `conv_transpose2d` reached the `else`. The sibling `conv2d_global_wg_size()` directly above uses the identical outer name test but then disambiguates by inspecting the weight's spatial extent, so the two could disagree about the same dispatch: the global size computed as sliding window while the local size was computed as pointwise. This factors that classification into one function used by both, so they cannot drift apart again. ### Why this is a restoration, not a new heuristic The weight-shape check arrived with #13173, which introduced the tuned `{64 / y, y, 1}` local size for pointwise convolutions. Before that commit the dispatch used `create_local_wg_size(global_size)` for every method: ``` - const utils::uvec3 local_size = graph.create_local_wg_size(global_size); ``` The name test swept sliding window convolutions into the new pointwise size along with the pointwise ones, so this restores the local size they had before #13173. ### Scope of the behavior change Only sliding window conv2d changes. - **Pointwise** was classified correctly before and after, so it keeps the tuned size. - **Depthwise** is routed to `conv2d_dw_impl()` before this dispatch and never reaches either function. - **Transposed** resolves to the same non-pointwise branch as before (`conv_transpose2d` does not contain the substring `conv2d`, so it already fell through to `create_local_wg_size`). `create_local_wg_size()` is pure arithmetic on the global workgroup size, so the exact before/after can be computed on the host. Across 7452 realistic non-pointwise conv2d output shapes (spatial extents 1 to 512, 3 to 2048 output channels): - every case is 64 threads per group, before and after, so occupancy per dispatch is unchanged and nothing exceeds 64 - only the group shape differs, most commonly `{8, 8, 1}` to one of `{8, 4, 2}`, `{4, 8, 2}`, `{2, 8, 4}`, `{4, 2, 8}` For example, a 128x128x64 output has global size `{128, 128, 16}` and goes from `{8, 8, 1}` to `{8, 4, 2}`. Since sliding window convolutions were last benchmarked under `create_local_wg_size` before #13173, and the pointwise tuning in that commit was not aimed at them, I would expect this to be neutral to positive. I do not have perf numbers across GPUs, so if you would like this validated on a specific device before landing, say which and I will run it. ### Test plan `vulkan_backend` builds clean. The change is a refactor of method classification plus the restored branch; existing conv2d correctness coverage applies unchanged, since the local workgroup size affects scheduling and not results. cc @SS-JIA @manuelcandales @digantdesai @cbilgin Co-authored-by: Sicheng Stephen Jia --- .../runtime/graph/ops/impl/Convolution.cpp | 64 +++++++++++-------- 1 file changed, 37 insertions(+), 27 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp index 4e8ff50fdac..b839ac8a61a 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Convolution.cpp @@ -350,6 +350,38 @@ GlobalWorkGrid create_conv2d_gwg( } } +// Determines which convolution method a dispatch uses. +// +// Depthwise and transposed convolutions have shader names of their own, but +// the name alone cannot separate pointwise from sliding window: the sliding +// window shader is itself named "conv2d", and a pointwise convolution also +// takes that name when its weights are prepacked. Those two are therefore +// separated by the weight's spatial extent. Shared by the global and local +// workgroup size functions below so that the two cannot disagree about the +// same dispatch. +Conv2dMethod infer_conv2d_method_from_shader( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const ValueRef weight_data) { + const std::string& kernel_name = shader.kernel_name; + // Checked before the plain "conv2d" test below, which "conv2d_dw" and + // "conv2d_pw" would otherwise match too. + if (kernel_name.find("conv2d_dw") != std::string::npos) { + return Conv2dMethod::Depthwise; + } + if (kernel_name.find("conv2d_pw") != std::string::npos) { + return Conv2dMethod::Pointwise; + } + if (kernel_name.find("conv_transpose2d") != std::string::npos) { + return Conv2dMethod::Transposed; + } + const auto& weight_sizes = graph->get_tref(weight_data)->sizes; + if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { + return Conv2dMethod::Pointwise; + } + return Conv2dMethod::SlidingWindow; +} + // Custom global workgroup size function for conv2d GlobalWorkGrid conv2d_gwg( ComputeGraph* graph, @@ -359,23 +391,8 @@ GlobalWorkGrid conv2d_gwg( const ValueRef out = args.at(0).refs.at(0); const ValueRef weight_data = resize_args.at(0); - // Determine method from shader name - Conv2dMethod method; - if (shader.kernel_name.find("conv2d_pw") != std::string::npos || - (shader.kernel_name.find("conv2d") != std::string::npos && - shader.kernel_name.find("conv_transpose2d") == std::string::npos)) { - // Check if it's pointwise by examining weight sizes - const auto& weight_sizes = graph->get_tref(weight_data)->sizes; - if (weight_sizes.at(2) == 1 && weight_sizes.at(3) == 1) { - method = Conv2dMethod::Pointwise; - } else { - method = Conv2dMethod::SlidingWindow; - } - } else if (shader.kernel_name.find("conv_transpose2d") != std::string::npos) { - method = Conv2dMethod::Transposed; - } else { - method = Conv2dMethod::SlidingWindow; - } + const Conv2dMethod method = + infer_conv2d_method_from_shader(graph, shader, weight_data); // Determine stride_equals_dilation from shader name bool stride_equals_dilation = @@ -404,17 +421,10 @@ LocalWorkGroup conv2d_lwg( const std::vector& args, const std::vector& resize_args) { (void)args; - (void)resize_args; - // Determine method from shader name - Conv2dMethod method; - if (shader.kernel_name.find("conv2d_pw") != std::string::npos || - (shader.kernel_name.find("conv2d") != std::string::npos && - shader.kernel_name.find("conv_transpose2d") == std::string::npos)) { - method = Conv2dMethod::Pointwise; - } else { - method = Conv2dMethod::SlidingWindow; - } + const ValueRef weight_data = resize_args.at(0); + const Conv2dMethod method = + infer_conv2d_method_from_shader(graph, shader, weight_data); if (method == Conv2dMethod::Pointwise) { uint32_t lwg_y = 1; From e9025a912253e330fac44a0d5afa18b3f9411b88 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 1 Sep 2026 11:16:32 +0200 Subject: [PATCH 117/190] Vulkan: size texture-vs-buffer choice by the bound, not the trace hint filter_invalid_reprs() decides whether a tensor can live in a texture by comparing required_image_extents(tensor_val.shape) against the device limits. For a symbolic dimension that comparison resolves at the dim's *hint* -- the size of the example the model happened to be traced with -- rather than the maximum its exported range allows. So a model traced with a small example picks a texture that is only valid at that size, and silently produces wrong results once it runs larger. Nothing raises. Repro: the Supertonic TTS text encoder, dynamic T in [8, 512], always executed at T=512 on a Galaxy S26 Ultra (Adreno 840). Only the trace-time example length differs; the edge graphs are identical (552 nodes, same ops, same literal args, same symbolic shapes): traced at T= 64 cosine 0.537966 .pte 18147970 bytes traced at T=128 cosine 0.537966 .pte 18147970 bytes (identical file) traced at T=192 cosine 0.999414 .pte 18149890 bytes traced at T=512 cosine 0.999414 .pte 18149890 bytes (identical file) The emitted binaries cluster exactly on the correctness boundary, which is what identifies this as a lowering decision rather than a bad kernel. With this change the T=64 export emits the 18149890-byte artifact and scores 0.999414. Evaluating the bound instead also covers the unbounded case: a dim with no finite upper bound cannot be shown to fit any texture, so it falls back to buffer storage. Measured on Supertonic (CPU reference reproduced by the XNNPACK delegate at cosine 1.000000): text_encoder 0.406303 -> 0.999414 vector_estimator 0.994432 -> 0.999994 vocoder 0.016757 -> 0.999977 --- backends/vulkan/utils.py | 41 ++++++++++++++++++++++++++++++++++++---- 1 file changed, 37 insertions(+), 4 deletions(-) diff --git a/backends/vulkan/utils.py b/backends/vulkan/utils.py index 066ee78ea0d..c43f318ae00 100644 --- a/backends/vulkan/utils.py +++ b/backends/vulkan/utils.py @@ -1181,6 +1181,33 @@ def make_tensor_repset(tensor_repr: TensorRepr) -> TensorRepSet: raise RuntimeError(f"Unsupported storage type {tensor_repr.storage_type}") +def upper_bound_size(dim: Union[int, torch.SymInt]) -> Optional[int]: + """Largest value a (possibly symbolic) tensor dimension can take. + + Returns None if no finite bound is known. + + A symbolic dim compares against a limit using its *hint* -- the size of the + example input the model happened to be traced with -- not the maximum the + exported range allows. Sizing decisions must use the bound instead, or a + model traced with a small example will make a choice that is invalid once it + runs at a larger size. + """ + if not isinstance(dim, torch.SymInt): + return int(dim) + if not dim.node.expr.free_symbols: + return int(dim.node.expr) + shape_env = dim.node.shape_env + if shape_env is None: + return None + try: + upper = shape_env.bound_sympy(dim.node.expr).upper + except Exception: + return None + if upper is None or not upper.is_finite: + return None + return int(upper) + + def filter_invalid_reprs( tensor_val: FakeTensor, tensor_repset: TensorRepSet, @@ -1198,11 +1225,17 @@ def filter_invalid_reprs( can be used to produce a valid image texture for the given tensor (i.e. fits within texture limits). """ + # Size the texture by what the dimension CAN be, not by the example the + # model was traced with. An unbounded dim cannot be shown to fit, so it + # falls back to buffer storage. + bounds = [upper_bound_size(d) for d in tensor_val.shape] valid_texture_layouts = set() - for memory_layout in tensor_repset.valid_texture_layouts: - extents = required_image_extents(tensor_val.shape, memory_layout) - if extents_are_valid(extents, texture_limits): - valid_texture_layouts.add(memory_layout) + if all(b is not None for b in bounds): + max_shape = torch.Size(bounds) + for memory_layout in tensor_repset.valid_texture_layouts: + extents = required_image_extents(max_shape, memory_layout) + if extents_are_valid(extents, texture_limits): + valid_texture_layouts.add(memory_layout) # High dimensional tensors require buffer storage if len(tensor_val.shape) > 4: From db3390dc70042f80a06d9e63a37175b2c899308f Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 9 Sep 2026 12:01:26 -0700 Subject: [PATCH 118/190] Pin mypy CI Transformers to the examples version (#22654) ### Summary Match the Transformers 5.0.0rc1 pin in requirements-examples.txt and include the workflow in the pip cache key. This prevents upstream API removals from breaking mypy on unchanged model tests. ### Test plan Validated with lintrunner --all-files --take MYPY. Authored with OpenAI Codex. --- .github/workflows/lint.yml | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/lint.yml b/.github/workflows/lint.yml index 9b217d00292..f8a2008ff9e 100644 --- a/.github/workflows/lint.yml +++ b/.github/workflows/lint.yml @@ -42,6 +42,7 @@ jobs: python-version: '3.11' cache: 'pip' cache-dependency-path: | + .github/workflows/lint.yml requirements-lintrunner.txt torch_pin.py @@ -56,7 +57,8 @@ jobs: pip install lintrunner==0.12.7 lintrunner-adapters==0.14.1 pip install -r requirements-lintrunner.txt USE_CPP=0 pip install --no-build-isolation third-party/ao - pip install pytest numpy parameterized huggingface_hub transformers timm expecttest types-requests + # Match requirements-examples.txt for the model APIs checked by mypy. + pip install pytest numpy parameterized huggingface_hub "transformers==5.0.0rc1" timm expecttest types-requests - name: Generate mypy stubs for C++ bindings run: | From b3b8c9ab118db4f500c29170f90e63fc89166a8f Mon Sep 17 00:00:00 2001 From: Kiymet Akdemir <54183514+kiymetakdemir@users.noreply.github.com> Date: Wed, 9 Sep 2026 12:22:03 -0700 Subject: [PATCH 119/190] Add a batched-sequence KV cache to the eager reference (#22646) ### Summary Adds a third reference layout: a private history per sequence instead of one shared cell table. Attention splits the flat token axis into its declared spans and runs independently over each sequence's own history, so no sequence attends another and no cross-sequence mask is built. That means a step can be answered by more than one attention, so update_and_fetch returns a list of AttendSpec and the op attends each over the query tokens it covers. A spec now carries its own K/V and how many queries it answers; the offsets follow from the running total, and the op checks the specs cover the step. The single-history layouts return a one-element list and are behaviourally unchanged; most of the test diff is call sites reshaping around that. Read reference_cache.py first (AttendSpec, then BatchedSequenceReferenceCache), then update_and_attend.py for the loop, then the tests. --- backends/mlx/test/test_ops.py | 4 +- extension/llm/cache/reference_cache.py | 302 +++++++++-- extension/llm/cache/test_update_and_attend.py | 502 ++++++++++++++---- extension/llm/cache/update_and_attend.py | 12 +- 4 files changed, 687 insertions(+), 133 deletions(-) diff --git a/backends/mlx/test/test_ops.py b/backends/mlx/test/test_ops.py index 1f8459564c3..0d5c72af06b 100644 --- a/backends/mlx/test/test_ops.py +++ b/backends/mlx/test/test_ops.py @@ -8598,14 +8598,14 @@ def compute_expected_outputs(self, model, test_inputs): # an oracle cache installed for its duration. from executorch.extension.llm.cache.reference_cache import ( CacheConfig, - ContiguousReferenceCache, + SequenceReferenceCache, ) from executorch.extension.llm.cache.update_and_attend import REGISTRY key = f"{self.name}-oracle" REGISTRY.install( key, - ContiguousReferenceCache( + SequenceReferenceCache( CacheConfig( n_layers=self.n_layers, n_kv_heads=self.n_kv_heads, diff --git a/extension/llm/cache/reference_cache.py b/extension/llm/cache/reference_cache.py index 7abc5d5138a..c0e206811f5 100644 --- a/extension/llm/cache/reference_cache.py +++ b/extension/llm/cache/reference_cache.py @@ -20,9 +20,10 @@ The cache places K/V and returns the history plus an ``AttendSpec`` (a mask *semantic*). The attend mechanism (``attend`` below) is applied by the op/backend from that spec. -Two caches share the op: ``ContiguousReferenceCache`` (one sequence appended in -place) and ``CellReferenceCache`` (many sequences over a pool of per-token cells, -with sharing and eviction). Both store float KV. +Three caches share the op: ``SequenceReferenceCache`` (one sequence), +``BatchedSequenceReferenceCache`` (many private sequence caches), and +``CellReferenceCache`` (many sequences over a shared pool of per-token cells, +with sharing and eviction). All store float KV. """ from __future__ import annotations @@ -50,7 +51,17 @@ class MaskKind(Enum): @dataclass class AttendSpec: + """One attention: what to attend over, which queries do it, how to mask it. + + A step is answered with a list of these -- one for a cache holding a single + history, one per sequence for a cache holding a private history each. They + cover the query axis in order, so ``q_len`` alone places each. + """ + + k: torch.Tensor # [B, H_kv, total, head_dim] -- key history + v: torch.Tensor # [B, H_kv, total, v_head_dim] -- value history kind: MaskKind + q_len: int # query tokens this spec answers, following the one before it mask: Optional[torch.Tensor] = None # EXPLICIT only: bool, true = attend @@ -109,7 +120,7 @@ def policy_for(self, layer_id: int) -> LayerPolicy: @experimental( "update_and_attend KV cache is experimental and may change without notice." ) -class ContiguousReferenceCache: +class SequenceReferenceCache: """Per-layer contiguous float KV history for a single sequence.""" def __init__(self, config: CacheConfig): @@ -126,6 +137,36 @@ def __init__(self, config: CacheConfig): def used(self, layer_id: int) -> int: return self._used[layer_id] + def rewind(self, new_len: int) -> None: + """Drop everything from ``new_len`` on, in every layer. + + A windowed layer retains only its last ``window`` positions, so it + cannot go back further than that even though this reference keeps the + older ones -- the window is applied to the mask here and to the storage + in a byte layer, and a rewind past it would attend cells that layer no + longer holds. + """ + used = self._used[0] + if new_len < 0 or new_len > used: + raise ValueError(f"rewind to {new_len}: the history holds {used}") + floor = max( + ( + used - self.config.policy_for(layer_id).window + for layer_id in range(self.config.n_layers) + if self.config.policy_for(layer_id).window > 0 + ), + default=0, + ) + if new_len < floor: + raise ValueError( + f"rewind to {new_len}: a windowed layer retains only from {floor}" + ) + for layer_id in range(self.config.n_layers): + if self.config.sizing == CacheSizing.DYNAMIC: + self._k[layer_id] = self._k[layer_id][:, :, :new_len, :] + self._v[layer_id] = self._v[layer_id][:, :, :new_len, :] + self._used[layer_id] = new_len + def reset(self): self._used = [0] * self.config.n_layers if self.config.sizing == CacheSizing.DYNAMIC: @@ -144,8 +185,10 @@ def update_and_fetch( k: torch.Tensor, v: torch.Tensor, position: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, AttendSpec]: - """Place this step's K/V and return the full history + mask semantic. + ) -> List[AttendSpec]: + """Place this step's K/V and return what to attend over. + + One sequence, one history, so the list is always one long. Per the design, ``position`` is the cache's placement + masking input. This contiguous single-sequence cache appends at its used length, so the @@ -160,9 +203,8 @@ def update_and_fetch( position: ``[q_len, n_dims]`` int -- per-query-token positions. Returns: - ``(k_hist, v_hist, spec)`` -- history ``[B, H_kv, total, head_dim]`` / - ``[B, H_kv, total, v_head_dim]`` (``total`` = prior length + q_len) and - the AttendSpec mask semantic. + one AttendSpec over the whole history, ``total`` = prior length + + q_len. """ q_len = k.shape[-2] used = self._used[layer_id] @@ -188,10 +230,14 @@ def update_and_fetch( v_hist = self._v[layer_id] self._used[layer_id] = new_used - return k_hist, v_hist, self._spec(layer_id, q_len, new_used, k.device) + return [self._spec(layer_id, k_hist, v_hist, q_len)] def _spec( - self, layer_id: int, q_len: int, total: int, device: torch.device + self, + layer_id: int, + k_hist: torch.Tensor, + v_hist: torch.Tensor, + q_len: int, ) -> AttendSpec: """The mask semantic for q_len new cells at the tail of a total window. @@ -200,21 +246,25 @@ def _spec( ``i + total - q_len - window``. Whichever bound the fused kinds cannot express is what makes the step EXPLICIT. """ + total = k_hist.shape[-2] window = self.config.policy_for(layer_id).window windowed = 0 < window < total if q_len == 1 and not windowed: - return AttendSpec(kind=MaskKind.NONE) + return AttendSpec(k=k_hist, v=v_hist, kind=MaskKind.NONE, q_len=q_len) if q_len == total and not windowed: - return AttendSpec(kind=MaskKind.CAUSAL) + return AttendSpec(k=k_hist, v=v_hist, kind=MaskKind.CAUSAL, q_len=q_len) # torch's is_causal is upper-left and expresses no window, so the band # is handed back explicitly. + device = k_hist.device offsets = torch.arange(total, device=device) - torch.arange( q_len, device=device ).unsqueeze(-1) band = offsets <= total - q_len if windowed: band &= offsets > total - q_len - window - return AttendSpec(kind=MaskKind.EXPLICIT, mask=band) + return AttendSpec( + k=k_hist, v=v_hist, kind=MaskKind.EXPLICIT, q_len=q_len, mask=band + ) # A cell's owners are a bitset in a torch int64, so bit 63 (the sign bit) is out. @@ -242,7 +292,7 @@ def flatten_step( Returns: ``(tokens, positions, seq_ids, logits_indices)`` -- tokens concatenated on the token axis and ``positions`` (``[n_tok, 1]``) as model inputs, - ``seq_ids`` for ``begin_step``, and ``logits_indices`` selecting each + ``seq_ids`` for ``declare_step``, and ``logits_indices`` selecting each sequence's last token, the rows worth running the LM head on. """ tokens, positions, seq_ids, logits_indices = [], [], [], [] @@ -259,6 +309,180 @@ def flatten_step( ) +@dataclass(frozen=True) +class _SequenceSpan: + seq_id: int + start: int + length: int + + +@experimental( + "update_and_attend KV cache is experimental and may change without notice." +) +class BatchedSequenceReferenceCache: + """A private ``SequenceReferenceCache`` per sequence in a flat batch. + + Projections share one model forward over the flattened token axis. Attention + splits that axis into its declared sequence spans, runs independently over + each sequence's private history, then concatenates the outputs in input + order. No sequence attends another and no dense cross-sequence mask is built. + """ + + def __init__(self, config: CacheConfig): + if config.batch_size != 1: + raise ValueError( + "batched sequence cache is flat on the token axis: batch_size must be 1" + ) + self.config = config + self._sequences: Dict[int, SequenceReferenceCache] = {} + self._spans: List[_SequenceSpan] = [] + self._served: Set[int] = set() + self._declared = False + + def declare_step(self, seq_ids: Sequence[int]) -> None: + if not seq_ids: + raise ValueError("a step carries at least one token") + for seq_id in seq_ids: + self._check_seq_id(seq_id) + + spans: List[_SequenceSpan] = [] + start = 0 + while start < len(seq_ids): + seq_id = seq_ids[start] + end = start + 1 + while end < len(seq_ids) and seq_ids[end] == seq_id: + end += 1 + spans.append(_SequenceSpan(seq_id, start, end - start)) + start = end + + # capacity bounds the whole cache. Checked before anything is created so a refusal changes + # nothing. + held = sum(sequence.used(0) for sequence in self._sequences.values()) + if held + len(seq_ids) > self.config.capacity: + raise RuntimeError( + f"KV cache overflow: {held + len(seq_ids)} cells exceeds " + f"capacity {self.config.capacity}" + ) + + for span in spans: + if span.seq_id not in self._sequences: + self._sequences[span.seq_id] = SequenceReferenceCache(self.config) + + self._spans = spans + self._served.clear() + self._declared = True + + def update_and_fetch( + self, + layer_id: int, + k: torch.Tensor, + v: torch.Tensor, + position: torch.Tensor, + ) -> List[AttendSpec]: + """Place each span's K/V in its own sequence and return one spec each. + + The specs follow the declared spans, so they cover the query axis in + order and no sequence appears in another's window. + """ + if not self._declared: + raise RuntimeError( + "no step declared: declare_step must precede every forward" + ) + if layer_id in self._served: + raise RuntimeError( + f"layer {layer_id} served twice for one step: " + "declare_step must precede every forward" + ) + token_count = k.shape[-2] + if not position.shape[0] == token_count == v.shape[-2]: + raise ValueError("position, k, and v must have the same token count") + if token_count != sum(span.length for span in self._spans): + raise ValueError("the forward token count must match declare_step") + if position.shape[-1] != 1: + raise NotImplementedError( + "sequence placement needs one position per token, got " + f"{position.shape[-1]}" + ) + self._check_positions(layer_id, position.reshape(-1).tolist()) + + specs: List[AttendSpec] = [] + for span in self._spans: + end = span.start + span.length + sequence = self._sequences[span.seq_id] + specs.extend( + sequence.update_and_fetch( + layer_id, + k[:, :, span.start : end, :], + v[:, :, span.start : end, :], + position[span.start : end], + ) + ) + + self._served.add(layer_id) + return specs + + def reset(self) -> None: + self._sequences.clear() + self._spans.clear() + self._served.clear() + self._declared = False + + def seq_rm(self, seq_id: int, p0: int = 0, p1: Optional[int] = None) -> None: + """Drop seq_id's claim on positions [p0, p1); p1 = None runs to the end. + + A private contiguous history drops its tail but not its middle: it has + no per-token position map to reindex what would survive. Bounding + history from below is a layer policy, not a verb. + """ + self._check_seq_id(seq_id) + if p1 is not None: + raise NotImplementedError( + "a private history drops only its tail, so [p0, p1) with a " + "bounded end has nothing to reindex the remainder against" + ) + sequence = self._sequences.get(seq_id) + if sequence is not None: + if p0 == 0: + del self._sequences[seq_id] + else: + sequence.rewind(p0) + self._spans.clear() + self._served.clear() + self._declared = False + + def seq_len(self, seq_id: int) -> int: + self._check_seq_id(seq_id) + sequence = self._sequences.get(seq_id) + return sequence.used(0) if sequence is not None else 0 + + def _check_positions(self, layer_id: int, positions: List[int]) -> None: + """Every span continues its own sequence, from where that sequence ends. + + A private history appends at its used length and never reads + ``position``, so a step that declared the wrong one would still place + its tokens contiguously -- correct cells under the wrong names, and no + later step would notice. A sequence spanned twice in one step continues + across both. + """ + ends: Dict[int, int] = {} + for span in self._spans: + at = ends.get(span.seq_id, self._sequences[span.seq_id].used(layer_id)) + got = positions[span.start : span.start + span.length] + want = list(range(at, at + span.length)) + if got != want: + raise ValueError( + f"sequence {span.seq_id} holds {at} positions on layer " + f"{layer_id}: the step declares {got}, not {want}" + ) + ends[span.seq_id] = at + span.length + + @staticmethod + def _check_seq_id(seq_id: int) -> None: + # No upper bound: a sequence is a dict entry, not a bit in an owner set. + if seq_id < 0: + raise ValueError(f"seq_id must be non-negative, got {seq_id}") + + @dataclass class _CellStepPlan: """One step's allocation, shared by every layer of that forward. @@ -296,7 +520,7 @@ class CellReferenceCache: that, so the spec is always EXPLICIT. The batch is flat: tokens from every sequence sit on one axis with B = 1, - and sequence identity is supplied out-of-band. ``begin_step`` declares which + and sequence identity is supplied out-of-band. ``declare_step`` declares which sequence each of the next forward's tokens belongs to; the positions arrive with the forward itself, in the op's ``position`` tensor, so cells are allocated on the first layer of the step and memoized for the rest of it. @@ -332,7 +556,7 @@ def __init__(self, config: CacheConfig): for _ in range(config.n_layers) ] self._step_seq_ids: List[int] = [] - self._declared = False # set by begin_step, cleared by the step it authorizes + self._declared = False # set by declare_step, cleared by the step it authorizes self._plan: Optional[_CellStepPlan] = None self._served: Set[int] = set() @@ -354,7 +578,7 @@ def seq_len(self, seq_id: int) -> int: bit = 1 << seq_id return sum(1 for owners in self._owners if owners & bit) - def begin_step(self, seq_ids: Sequence[int]) -> None: + def declare_step(self, seq_ids: Sequence[int]) -> None: """Declare the sequence each of the next forward's tokens belongs to. Admission is decided here, before the forward: the token count is known @@ -430,17 +654,19 @@ def update_and_fetch( k: torch.Tensor, v: torch.Tensor, position: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, AttendSpec]: + ) -> List[AttendSpec]: """Scatter this step's K/V into its cells and return the read window. The first layer of a step allocates; the rest reuse that allocation, so the cells and the mask are computed once per forward, not once per - layer. Args are as ``ContiguousReferenceCache.update_and_fetch``. + layer. Every sequence reads the same window and the mask holds them + apart, so the list is always one long. Args are as + ``SequenceReferenceCache.update_and_fetch``. """ if layer_id in self._served: raise RuntimeError( f"layer {layer_id} served twice for one step: " - "begin_step must precede every forward" + "declare_step must precede every forward" ) if self._plan is None: self._plan = self._allocate(position) @@ -451,14 +677,15 @@ def update_and_fetch( cells = self._plan.cells self._k[layer_id][:, :, cells, :] = k.to(self.config.dtype) self._v[layer_id][:, :, cells, :] = v.to(self.config.dtype) - return ( - self._k[layer_id][:, :, :read_len, :], - self._v[layer_id][:, :, :read_len, :], + return [ AttendSpec( + k=self._k[layer_id][:, :, :read_len, :], + v=self._v[layer_id][:, :, :read_len, :], kind=MaskKind.EXPLICIT, + q_len=len(cells), mask=self._plan.mask_for(self.config.policy_for(layer_id).window), - ), - ) + ) + ] # -- internals ---------------------------------------------------------- @@ -467,18 +694,17 @@ def _allocate(self, position: torch.Tensor) -> _CellStepPlan: device = self._k[0].device if not self._declared: raise RuntimeError( - "no step declared: begin_step must precede every forward" + "no step declared: declare_step must precede every forward" ) self._declared = False # one declaration, one attempt at allocating it if position.shape[-1] != 1: raise NotImplementedError( - "cell placement needs one position per token, got " - f"{position.shape[-1]}" + f"cell placement needs one position per token, got {position.shape[-1]}" ) positions = position.reshape(-1).tolist() if len(positions) != len(self._step_seq_ids): raise ValueError( - f"begin_step declared {len(self._step_seq_ids)} tokens, " + f"declare_step declared {len(self._step_seq_ids)} tokens, " f"the forward carries {len(positions)}" ) cells = [ @@ -537,7 +763,7 @@ def _claim(self, pos: int, owners: int) -> int: self._owners[i] = owners self._used_end = max(self._used_end, i + 1) return i - raise RuntimeError("no free cell") # begin_step admitted the step + raise RuntimeError("no free cell") # declare_step admitted the step def _shrink(self): while self._used_end > 0 and self._pos[self._used_end - 1] < 0: @@ -546,7 +772,7 @@ def _shrink(self): def _invalidate_plan(self): # A mutated cell table leaves a built plan's cells and mask stale. The # step protocol state is deliberately left alone: a mutation must not - # disguise a forward that skipped begin_step. + # disguise a forward that skipped declare_step. self._plan = None @staticmethod @@ -563,8 +789,6 @@ def _in_range(pos: int, p0: int, p1: Optional[int]) -> bool: def attend( q: torch.Tensor, - k: torch.Tensor, - v: torch.Tensor, spec: AttendSpec, scale: float, out_dtype: torch.dtype, @@ -579,17 +803,17 @@ def attend( so a cache must declare EXPLICIT for a chunked or multi-turn step. Args (BHSD): - q: ``[B, H_q, q_len, head_dim]`` -- queries (already RoPE-rotated). - k: ``[B, H_kv, total, head_dim]`` -- key history. - v: ``[B, H_kv, total, v_head_dim]`` -- value history. - spec: mask semantic (NONE = attend all; CAUSAL = causal; EXPLICIT = the - spec's bool mask). + q: ``[B, H_q, q_len, head_dim]`` -- queries (already RoPE-rotated), the + ones this spec answers. + spec: the K/V history to attend over and its mask semantic (NONE = + attend all; CAUSAL = causal; EXPLICIT = the spec's bool mask). scale: attention softmax scale. out_dtype: output dtype. Returns: ``[B, H_q, q_len, v_head_dim]`` attention output, in ``out_dtype``. """ + k, v = spec.k, spec.v n_q_heads = q.shape[1] n_kv_heads = k.shape[1] if n_q_heads != n_kv_heads: diff --git a/extension/llm/cache/test_update_and_attend.py b/extension/llm/cache/test_update_and_attend.py index 96933a43b01..eed13f5b20c 100644 --- a/extension/llm/cache/test_update_and_attend.py +++ b/extension/llm/cache/test_update_and_attend.py @@ -11,15 +11,16 @@ from executorch.extension.llm.cache.reference_cache import ( attend, AttendSpec, + BatchedSequenceReferenceCache, CacheConfig, CacheSizing, CellReferenceCache, - ContiguousReferenceCache, flatten_step, LayerKind, LayerPolicy, MaskKind, MAX_SEQS, + SequenceReferenceCache, ) from executorch.extension.llm.cache.update_and_attend import REGISTRY, update_and_attend @@ -175,7 +176,7 @@ def test_prefill_matches_baseline(self): (CacheSizing.STATIC, seq_len), ]: with self.subTest(sizing=sizing): - cache = ContiguousReferenceCache(self._config(sizing, cap)) + cache = SequenceReferenceCache(self._config(sizing, cap)) REGISTRY.install(self.cache_key, cache) with REGISTRY.active(self.cache_key): out = ep.module()(x, _positions(0, seq_len), torch.arange(seq_len)) @@ -195,7 +196,7 @@ def test_incremental_decode_matches_baseline(self): (CacheSizing.STATIC, total), ]: with self.subTest(sizing=sizing): - cache = ContiguousReferenceCache(self._config(sizing, cap)) + cache = SequenceReferenceCache(self._config(sizing, cap)) REGISTRY.install(self.cache_key, cache) with REGISTRY.active(self.cache_key): ep_prefill.module()( @@ -221,7 +222,7 @@ def test_chunked_prefill_matches_baseline(self): ref = self.model.reference_forward(x, torch.arange(total)) ep = self._export(chunk) - cache = ContiguousReferenceCache(self._config(CacheSizing.DYNAMIC, total)) + cache = SequenceReferenceCache(self._config(CacheSizing.DYNAMIC, total)) REGISTRY.install(self.cache_key, cache) with REGISTRY.active(self.cache_key): for start in range(0, total, chunk): @@ -236,13 +237,35 @@ def test_chunked_prefill_matches_baseline(self): def test_static_overflow_raises(self): ep = self._export(seq_len=5) - cache = ContiguousReferenceCache(self._config(CacheSizing.STATIC, capacity=3)) + cache = SequenceReferenceCache(self._config(CacheSizing.STATIC, capacity=3)) REGISTRY.install(self.cache_key, cache) with self.assertRaises(RuntimeError), REGISTRY.active(self.cache_key): ep.module()( torch.randn(1, 5, self.hidden), _positions(0, 5), torch.arange(5) ) + def test_the_specs_must_cover_every_query_token(self): + # Each spec's queries are placed by the running total of the ones + # before it, so a cache that miscounts would attend the wrong slice + # rather than fail. Only this check separates the two. + class Miscounting: + def __init__(self, q_len): + self.q_len = q_len + + def update_and_fetch(self, layer_id, k, v, position): + return [AttendSpec(k=k, v=v, kind=MaskKind.NONE, q_len=self.q_len)] + + q = torch.randn(1, self.n_heads, 3, self.head_dim) + kv = torch.randn(1, self.n_kv_heads, 3, self.head_dim) + for q_len in (2, 4): # answering too few, and claiming too many + with self.subTest(q_len=q_len): + REGISTRY.install(self.cache_key, Miscounting(q_len)) + with self.assertRaisesRegex(ValueError, "of 3 query tokens"): + with REGISTRY.active(self.cache_key): + update_and_attend( + q, kv, kv, _positions(0, 3), 0, 0.125, torch.float32 + ) + def test_output_shape_uses_value_head_dim(self): # The output's last dim comes from v, which may differ from q's head dim # (e.g. MLA). Export (fake kernel only) and check the op node's meta. @@ -263,6 +286,295 @@ def forward(self, q, k, v, position): self.assertEqual(tuple(node.meta["val"].shape), (1, 4, 3, 5)) +class BatchedSequenceCacheTest(unittest.TestCase): + def setUp(self): + torch.manual_seed(0) + self.n_layers, self.hidden = 2, 16 + self.n_heads, self.n_kv_heads, self.head_dim = 4, 2, 8 + self.model = TinyAttentionModel( + self.n_layers, + self.hidden, + self.n_heads, + self.n_kv_heads, + self.head_dim, + 40, + ).eval() + self.cache_key = "batched-sequences" + + def tearDown(self): + REGISTRY.uninstall(self.cache_key) + + def _cache(self, capacity=16, layers=None): + cache = BatchedSequenceReferenceCache( + CacheConfig( + n_layers=self.n_layers, + n_kv_heads=self.n_kv_heads, + head_dim=self.head_dim, + capacity=capacity, + layers=[LayerPolicy.flat()] if layers is None else layers, + ) + ) + REGISTRY.install(self.cache_key, cache) + return cache + + def _step(self, cache, x, positions, seq_ids): + cache.declare_step(seq_ids) + with REGISTRY.active(self.cache_key): + return self.model(x, positions, torch.arange(x.shape[1])) + + def _attention_inputs(self, length): + return ( + torch.randn(1, self.n_heads, length, self.head_dim), + torch.randn(1, self.n_kv_heads, length, self.head_dim), + torch.randn(1, self.n_kv_heads, length, self.head_dim), + _positions(0, length), + ) + + def _attend(self, cache, inputs, layer_id=0): + # What the op does: fetch one spec per span, attend each over the query + # tokens it answers, rejoin. + q, k, v, positions = inputs + specs = cache.update_and_fetch(layer_id, k, v, positions) + outputs, start = [], 0 + for spec in specs: + end = start + spec.q_len + outputs.append( + attend( + q[:, :, start:end, :], + spec, + self.head_dim**-0.5, + torch.float32, + ) + ) + start = end + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=2) + + def test_single_span_matches_single_sequence(self): + x = torch.randn(1, 5, self.hidden) + out = self._step(self._cache(), x, _positions(0, 5), [3] * 5) + torch.testing.assert_close( + out, + self.model.reference_forward(x, torch.arange(5)), + atol=1e-4, + rtol=1e-4, + ) + + def test_multiple_sequence_spans_match_separate_runs(self): + a = torch.randn(1, 4, self.hidden) + b = torch.randn(1, 3, self.hidden) + tokens, positions, seq_ids, _ = flatten_step({2: (a, 0), 7: (b, 0)}) + + out = self._step(self._cache(), tokens, positions, seq_ids) + + torch.testing.assert_close( + out[:, :4], + self.model.reference_forward(a, torch.arange(4)), + atol=1e-4, + rtol=1e-4, + ) + torch.testing.assert_close( + out[:, 4:], + self.model.reference_forward(b, torch.arange(3)), + atol=1e-4, + rtol=1e-4, + ) + + def test_repeated_sequence_spans_preserve_input_order(self): + a = torch.randn(1, 3, self.hidden) + b = torch.randn(1, 1, self.hidden) + tokens = torch.cat([a[:, :2], b, a[:, 2:]], dim=1) + positions = torch.tensor([[0], [1], [0], [2]], dtype=torch.long) + out = self._step(self._cache(), tokens, positions, [2, 2, 7, 2]) + + a_out = self.model.reference_forward(a, torch.arange(3)) + b_out = self.model.reference_forward(b, torch.arange(1)) + torch.testing.assert_close(out[:, [0, 1, 3]], a_out, atol=1e-4, rtol=1e-4) + torch.testing.assert_close(out[:, 2:3], b_out, atol=1e-4, rtol=1e-4) + + def test_decode_continues_each_private_sequence(self): + a = torch.randn(1, 4, self.hidden) + b = torch.randn(1, 3, self.hidden) + cache = self._cache() + + tokens, positions, seq_ids, _ = flatten_step( + {2: (a[:, :3], 0), 7: (b[:, :2], 0)} + ) + self._step(cache, tokens, positions, seq_ids) + + tokens, positions, seq_ids, _ = flatten_step( + {2: (a[:, 3:], 3), 7: (b[:, 2:], 2)} + ) + out = self._step(cache, tokens, positions, seq_ids) + + torch.testing.assert_close( + out[:, 0], + self.model.reference_forward(a, torch.arange(4))[:, -1], + atol=1e-4, + rtol=1e-4, + ) + torch.testing.assert_close( + out[:, 1], + self.model.reference_forward(b, torch.arange(3))[:, -1], + atol=1e-4, + rtol=1e-4, + ) + + def test_a_span_must_continue_its_own_sequence(self): + # A private history appends at its length, so a wrong position would + # still land contiguously. Only this check separates the two. + cache = self._cache() + q, k, v, _ = self._attention_inputs(2) + cache.declare_step([5, 5]) + + with self.assertRaisesRegex(ValueError, r"declares \[1, 2\], not \[0, 1\]"): + self._attend(cache, (q, k, v, _positions(1, 2))) + self.assertEqual(cache.seq_len(5), 0) # a refusal writes nothing + + # Ascending is not enough; a span is a consecutive run. + gapped = torch.tensor([[0], [2]], dtype=torch.long) + with self.assertRaisesRegex(ValueError, r"declares \[0, 2\], not \[0, 1\]"): + self._attend(cache, (q, k, v, gapped)) + self.assertEqual(cache.seq_len(5), 0) + + # The declaration still stands, so the same layer can be retried. + self._attend(cache, (q, k, v, _positions(0, 2))) + self.assertEqual(cache.seq_len(5), 2) + + def test_a_sequence_spanned_twice_continues_across_both(self): + cache = self._cache() + q, k, v, _ = self._attention_inputs(3) + + # Tokens 0 and 2 are seq 4, token 1 is seq 9; seq 4's second span picks + # up where its first left off rather than at its prior length. + cache.declare_step([4, 9, 4]) + self._attend(cache, (q, k, v, torch.tensor([[0], [0], [1]]))) + self.assertEqual(cache.seq_len(4), 2) + self.assertEqual(cache.seq_len(9), 1) + + # The bad position is in the last span, so a per-span check would have + # written the first two before refusing. + cache.declare_step([4, 9, 4]) + with self.assertRaisesRegex(ValueError, r"holds 3 .*declares \[4\], not \[3\]"): + self._attend(cache, (q, k, v, torch.tensor([[2], [1], [4]]))) + self.assertEqual(cache.seq_len(4), 2) + self.assertEqual(cache.seq_len(9), 1) + + def test_requires_one_declared_step_per_forward(self): + cache = self._cache() + inputs = self._attention_inputs(1) + + with self.assertRaisesRegex(RuntimeError, "no step declared"): + self._attend(cache, inputs) + + cache.declare_step([2]) + self._attend(cache, inputs) + with self.assertRaisesRegex(RuntimeError, "served twice"): + self._attend(cache, inputs) + + def test_declaration_and_sequence_verbs_validate_ids(self): + cache = self._cache() + with self.assertRaisesRegex(ValueError, "at least one token"): + cache.declare_step([]) + + for call in ( + lambda: cache.declare_step([-1]), + lambda: cache.seq_rm(-1), + lambda: cache.seq_len(-1), + ): + with self.subTest(call=call), self.assertRaises(ValueError): + call() + + # Private histories are dict entries, so nothing caps the id. + cache.declare_step([9999]) + self.assertEqual(cache.seq_len(9999), 0) + + def test_capacity_is_the_pool_total_and_refusal_changes_nothing(self): + cache = self._cache(capacity=4) + with self.assertRaisesRegex(RuntimeError, "exceeds capacity"): + cache.declare_step([2] * 5) + self.assertEqual(cache.seq_len(2), 0) + + # Two sequences share the budget rather than each getting one. + a = torch.randn(1, 2, self.hidden) + b = torch.randn(1, 2, self.hidden) + tokens, positions, seq_ids, _ = flatten_step({2: (a, 0), 7: (b, 0)}) + self._step(cache, tokens, positions, seq_ids) + self.assertEqual(cache.seq_len(2), 2) + self.assertEqual(cache.seq_len(7), 2) + + # Either sequence is now blocked by what the other holds. + with self.assertRaisesRegex(RuntimeError, "exceeds capacity"): + cache.declare_step([2]) + self.assertEqual(cache.seq_len(2), 2) + self.assertEqual(cache.seq_len(7), 2) + + def test_forward_width_must_match_tensors_and_declaration(self): + one = self._attention_inputs(1) + two = self._attention_inputs(2) + cases = ( + ("position", [2], (one[0], one[1], one[2], two[3]), "same token count"), + ("q/k", [2, 2], (two[0], one[1], one[2], two[3]), "same token count"), + ("k/v", [2], (one[0], one[1], two[2], one[3]), "same token count"), + ("declaration", [2, 2], one, "must match declare_step"), + ) + for name, seq_ids, inputs, message in cases: + with self.subTest(name=name): + cache = self._cache() + cache.declare_step(seq_ids) + with self.assertRaisesRegex(ValueError, message): + self._attend(cache, inputs) + + def test_sequence_removal_invalidates_a_declared_step(self): + cache = self._cache() + cache.declare_step([2]) + cache.seq_rm(2) + + self.assertEqual(cache.seq_len(2), 0) + with self.assertRaisesRegex(RuntimeError, "no step declared"): + self._attend(cache, self._attention_inputs(1)) + + def test_seq_rm_truncates_or_drops_and_refuses_a_bounded_range(self): + cache = self._cache() + self._step(cache, torch.randn(1, 4, self.hidden), _positions(0, 4), [1] * 4) + self._step(cache, torch.randn(1, 2, self.hidden), _positions(0, 2), [6] * 2) + + with self.assertRaises(NotImplementedError): + cache.seq_rm(1, 0, 2) + self.assertEqual(cache.seq_len(1), 4) + + cache.seq_rm(1, 2) # keep positions 0..1 + self.assertEqual(cache.seq_len(1), 2) + self.assertEqual(cache.seq_len(6), 2) # its neighbour is untouched + + cache.seq_rm(1) # the whole sequence + self.assertEqual(cache.seq_len(1), 0) + self.assertEqual(cache.seq_len(6), 2) + + def test_rewinding_then_continuing_matches_an_unbroken_run(self): + x = torch.randn(1, 5, self.hidden) + ref = self.model.reference_forward(x, torch.arange(5)) + + cache = self._cache() + self._step(cache, x[:, :4], _positions(0, 4), [3] * 4) + cache.seq_rm(3, 2) # discard positions 2..3 + out = self._step(cache, x[:, 2:], _positions(2, 3), [3] * 3) + + torch.testing.assert_close(out, ref[:, 2:], atol=1e-4, rtol=1e-4) + + def test_rewind_refuses_to_grow_or_pass_a_window(self): + cache = self._cache(layers=[LayerPolicy.ring(2)]) + self._step(cache, torch.randn(1, 5, self.hidden), _positions(0, 5), [0] * 5) + + with self.assertRaisesRegex(ValueError, "the history holds 5"): + cache.seq_rm(0, 6) + # A windowed layer keeps only its last two positions, so 3 is the floor + # even though this reference still holds the older ones. + with self.assertRaisesRegex(ValueError, "retains only from 3"): + cache.seq_rm(0, 1) + cache.seq_rm(0, 3) + self.assertEqual(cache.seq_len(0), 3) + + class CellCacheTest(unittest.TestCase): # Many sequences over one pool of per-token cells, flat on the token axis. # The baseline throughout is the cacheless model: whatever a sequence would @@ -305,7 +617,7 @@ def _cache( def _step(self, cache, x, positions, seqs): """One forward carrying `x`, whose tokens have these positions/seqs.""" - cache.begin_step(seqs) + cache.declare_step(seqs) pos = torch.tensor(positions, dtype=torch.long).unsqueeze(-1) with REGISTRY.active(self.cache_key): return self.model(x, pos, torch.arange(x.shape[1])) @@ -325,7 +637,7 @@ def test_batched_sequences_match_separate_runs(self): # {seq_id: (tokens, start_pos)} -> the step's parallel arrays tokens, positions, seq_ids, _ = flatten_step({0: (a, 0), 1: (b, 0)}) - cache.begin_step(seq_ids) + cache.declare_step(seq_ids) with REGISTRY.active(self.cache_key): # every row, not one per sequence: each token is compared below out = self.model(tokens, positions, torch.arange(tokens.shape[1])) @@ -352,14 +664,14 @@ def test_batched_decode_continues_each_sequence(self): tokens, positions, seq_ids, logits_indices = flatten_step( {0: (a[:, :3], 0), 1: (b[:, :2], 0)} ) - cache.begin_step(seq_ids) + cache.declare_step(seq_ids) with REGISTRY.active(self.cache_key): self.model(tokens, positions, logits_indices) tokens, positions, seq_ids, logits_indices = flatten_step( {0: (a[:, 3:], 3), 1: (b[:, 2:], 2)} ) - cache.begin_step(seq_ids) + cache.declare_step(seq_ids) with REGISTRY.active(self.cache_key): out = self.model(tokens, positions, logits_indices) @@ -433,18 +745,18 @@ def test_fork_at_a_position_shares_only_the_prefix(self): def test_freeing_the_tail_shrinks_the_read_window(self): cache = self._cache() kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) - k, _, _ = cache.update_and_fetch(0, kv, kv, _positions(0, 4)) - self.assertEqual(k.shape[2], 4) # four cells held, so a window of four + cache.declare_step([0] * 4) + spec = cache.update_and_fetch(0, kv, kv, _positions(0, 4))[0] + self.assertEqual(spec.k.shape[2], 4) # four cells held, so a window of four cache.seq_rm(0) # frees all four, so used_end walks back to 0 self.assertEqual(cache.free_cells(), self.CAPACITY) # one token reclaims cell 0, so the window is its own single cell kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) - cache.begin_step([1]) - k, _, spec = cache.update_and_fetch(0, kv, kv, torch.tensor([[0]])) - self.assertEqual(k.shape[2], 1) # the window length is 1, not the old 4 + cache.declare_step([1]) + spec = cache.update_and_fetch(0, kv, kv, torch.tensor([[0]]))[0] + self.assertEqual(spec.k.shape[2], 1) # the window length is 1, not the old 4 self.assertEqual(spec.mask.shape[-1], 1) def test_seq_rm_over_a_range_frees_only_that_window(self): @@ -464,7 +776,7 @@ def test_every_verb_range_checks_the_seq_id(self): # much later as an overflow while building the mask. cache = self._cache() for call in ( - lambda: cache.begin_step([MAX_SEQS]), + lambda: cache.declare_step([MAX_SEQS]), lambda: cache.seq_cp(0, MAX_SEQS), lambda: cache.seq_cp(MAX_SEQS, 0), lambda: cache.seq_rm(MAX_SEQS), @@ -491,8 +803,8 @@ def test_layer_policy_rejects_a_mismatched_window(self): def test_window_narrows_each_query_without_crossing_sequences(self): cache = self._cache(layers=[LayerPolicy.ring(2)]) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) - spec = cache.update_and_fetch(0, kv, kv, _positions(0, 4))[2] + cache.declare_step([0] * 4) + spec = cache.update_and_fetch(0, kv, kv, _positions(0, 4))[0] # offsets[i][j] = j - i, so <= 0 is causal and > -2 keeps the newest # two: a band whose row 2 drops key 0, which plain causal would keep. @@ -501,9 +813,9 @@ def test_window_narrows_each_query_without_crossing_sequences(self): # a second sequence is bounded the same way, and still sees none of # the first's cells even though they are inside its window - cache.begin_step([1, 1]) + cache.declare_step([1, 1]) kv = torch.randn(1, self.n_kv_heads, 2, self.head_dim) - spec = cache.update_and_fetch(0, kv, kv, _positions(0, 2))[2] + spec = cache.update_and_fetch(0, kv, kv, _positions(0, 2))[0] expected = torch.zeros(2, 6, dtype=torch.bool) expected[0, 4] = expected[1, 4] = expected[1, 5] = True torch.testing.assert_close(spec.mask, expected) @@ -513,10 +825,10 @@ def test_layers_can_window_independently(self): # per policy, not per layer, so a mixed model costs one extra mask. cache = self._cache(layers=[LayerPolicy.flat(), LayerPolicy.ring(2)]) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) + cache.declare_step([0] * 4) pos = _positions(0, 4) - flat = cache.update_and_fetch(0, kv, kv, pos)[2].mask - windowed = cache.update_and_fetch(1, kv, kv, pos)[2].mask + flat = cache.update_and_fetch(0, kv, kv, pos)[0].mask + windowed = cache.update_and_fetch(1, kv, kv, pos)[0].mask offsets = torch.arange(4) - torch.arange(4).unsqueeze(-1) torch.testing.assert_close(flat, offsets <= 0) @@ -530,11 +842,11 @@ def test_layers_sharing_a_window_share_one_mask(self): n_layers=3, ) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) + cache.declare_step([0] * 4) pos = _positions(0, 4) - flat = cache.update_and_fetch(0, kv, kv, pos)[2].mask - first = cache.update_and_fetch(1, kv, kv, pos)[2].mask - second = cache.update_and_fetch(2, kv, kv, pos)[2].mask + flat = cache.update_and_fetch(0, kv, kv, pos)[0].mask + first = cache.update_and_fetch(1, kv, kv, pos)[0].mask + second = cache.update_and_fetch(2, kv, kv, pos)[0].mask self.assertIs(first, second) self.assertIsNot(flat, first) @@ -543,22 +855,25 @@ def test_windowed_decode_attends_only_the_retained_cells(self): window = 2 cache = self._cache(layers=[LayerPolicy.ring(window)]) kv = torch.randn(1, self.n_kv_heads, 4, self.head_dim) - cache.begin_step([0] * 4) + cache.declare_step([0] * 4) cache.update_and_fetch(0, kv, kv, _positions(0, 4)) - cache.begin_step([0]) + cache.declare_step([0]) kv = torch.randn(1, self.n_kv_heads, 1, self.head_dim) - k, v, spec = cache.update_and_fetch(0, kv, kv, _positions(4, 1)) + spec = cache.update_and_fetch(0, kv, kv, _positions(4, 1))[0] q = torch.randn(1, self.n_heads, 1, self.head_dim) scale = self.head_dim**-0.5 torch.testing.assert_close( - attend(q, k, v, spec, scale, torch.float32), + attend(q, spec, scale, torch.float32), attend( # the last `window` cells of its sequence, unmasked q, - k[:, :, -window:, :], - v[:, :, -window:, :], - AttendSpec(kind=MaskKind.NONE), + AttendSpec( + k=spec.k[:, :, -window:, :], + v=spec.v[:, :, -window:, :], + kind=MaskKind.NONE, + q_len=1, + ), scale, torch.float32, ), @@ -568,7 +883,7 @@ def test_admission_fails_before_the_forward(self): cache = self._cache(capacity=4) self.assertFalse(cache.can_extend(5)) with self.assertRaises(RuntimeError): - cache.begin_step([0] * 5) + cache.declare_step([0] * 5) def test_step_protocol_is_enforced(self): cache = self._cache() @@ -576,17 +891,17 @@ def test_step_protocol_is_enforced(self): pos = torch.tensor([[0]]) with self.assertRaises(ValueError): # a step with no tokens - cache.begin_step([]) + cache.declare_step([]) - cache.begin_step([0, 0]) # declares two tokens, forward carries one + cache.declare_step([0, 0]) # declares two tokens, forward carries one with self.assertRaises(ValueError): cache.update_and_fetch(0, kv, kv, pos) with self.assertRaises(RuntimeError): # the failed attempt still cleared it cache.update_and_fetch(0, kv, kv, torch.tensor([[0], [1]])) - cache.begin_step([0]) + cache.declare_step([0]) cache.update_and_fetch(0, kv, kv, pos) - with self.assertRaises(RuntimeError): # a second step, no begin_step + with self.assertRaises(RuntimeError): # a second step, no declare_step cache.update_and_fetch(0, kv, kv, pos) def test_growth_keeps_cell_indices_and_bytes(self): @@ -595,18 +910,19 @@ def test_growth_keeps_cell_indices_and_bytes(self): # move history without anything noticing. cache = self._cache() first = torch.randn(1, self.n_kv_heads, 2, self.head_dim) - cache.begin_step([0, 0]) - k, _, _ = cache.update_and_fetch(0, first, first, torch.tensor([[0], [1]])) - self.assertEqual(k.shape[2], 2) # a short session reserves a short pool + cache.declare_step([0, 0]) + spec = cache.update_and_fetch(0, first, first, torch.tensor([[0], [1]]))[0] + self.assertEqual(spec.k.shape[2], 2) # a short session reserves a short pool rest = torch.randn(1, self.n_kv_heads, 6, self.head_dim) - cache.begin_step([0] * 6) - k, v, _ = cache.update_and_fetch( + cache.declare_step([0] * 6) + spec = cache.update_and_fetch( 0, rest, rest, torch.tensor([[p] for p in range(2, 8)]) - ) - self.assertEqual(k.shape[2], 8) - torch.testing.assert_close(k[:, :, :2, :], first) # cells 0,1 unmoved - torch.testing.assert_close(v[:, :, 2:, :], rest) + )[0] + self.assertEqual(spec.k.shape[2], 8) + # cells 0,1 unmoved + torch.testing.assert_close(spec.k[:, :, :2, :], first) + torch.testing.assert_close(spec.v[:, :, 2:, :], rest) def test_sizings_agree(self): x = torch.randn(1, 5, self.hidden) @@ -616,14 +932,14 @@ def test_sizings_agree(self): ] torch.testing.assert_close(out[0], out[1]) - def test_a_verb_does_not_hide_a_missing_begin_step(self): + def test_a_verb_does_not_hide_a_missing_declare_step(self): # A sequence verb drops the memoized plan, which must not be mistaken # for the start of a step -- that would silently reuse the previous # step's sequence assignment for the new tokens. cache = self._cache() kv = torch.randn(1, self.n_kv_heads, 2, self.head_dim) pos = torch.tensor([[0], [0]]) - cache.begin_step([0, 1]) + cache.declare_step([0, 1]) cache.update_and_fetch(0, kv, kv, pos) cache.seq_rm(2) # any verb; a no-op here beyond dropping the plan @@ -633,18 +949,18 @@ def test_a_verb_does_not_hide_a_missing_begin_step(self): cache.update_and_fetch(1, kv, kv, pos) -class ContiguousSpecTest(unittest.TestCase): +class SequenceSpecTest(unittest.TestCase): # Which mask semantic the cache declares for each shape of step. def setUp(self): torch.manual_seed(0) - self.cache = ContiguousReferenceCache( + self.cache = SequenceReferenceCache( CacheConfig(n_layers=1, n_kv_heads=2, head_dim=4, capacity=8) ) def _update(self, start, q_len): kv = torch.randn(1, 2, q_len, 4) - return self.cache.update_and_fetch(0, kv, kv, _positions(start, q_len))[2] + return self.cache.update_and_fetch(0, kv, kv, _positions(start, q_len))[0] def test_decode_is_unmasked(self): self.assertEqual(self._update(0, 1).kind, MaskKind.NONE) @@ -674,7 +990,7 @@ def setUp(self): self.scale = self.DIM**-0.5 def _cache(self, policy, sizing=CacheSizing.DYNAMIC): - return ContiguousReferenceCache( + return SequenceReferenceCache( CacheConfig( n_layers=1, n_kv_heads=self.HEADS, @@ -689,50 +1005,51 @@ def _update(self, cache, n, start): kv = torch.randn(1, self.HEADS, n, self.DIM) return cache.update_and_fetch(0, kv, kv, _positions(start, n)) - def _attend(self, q, k, v, spec): - return attend(q, k, v, spec, self.scale, torch.float32) + def _attend(self, q, spec): + return attend(q, spec, self.scale, torch.float32) + + def _unmasked(self, q, k, v): + # The same queries over a hand-picked window, for the spec to match. + spec = AttendSpec(k=k, v=v, kind=MaskKind.NONE, q_len=q.shape[-2]) + return attend(q, spec, self.scale, torch.float32) def test_decode_attends_only_the_window(self): window = 3 cache = self._cache(LayerPolicy.ring(window)) self._update(cache, 5, 0) - k, v, spec = self._update(cache, 1, 5) + spec = self._update(cache, 1, 5)[0] q = torch.randn(1, self.HEADS, 1, self.DIM) torch.testing.assert_close( - self._attend(q, k, v, spec), - self._attend( # the last `window` cells, unmasked - q, - k[:, :, -window:, :], - v[:, :, -window:, :], - AttendSpec(kind=MaskKind.NONE), + self._attend(q, spec), + self._unmasked( # the last `window` cells, unmasked + q, spec.k[:, :, -window:, :], spec.v[:, :, -window:, :] ), ) def test_each_prefill_query_attends_its_own_window(self): window = 2 cache = self._cache(LayerPolicy.ring(window)) - k, v, spec = self._update(cache, 4, 0) + spec = self._update(cache, 4, 0)[0] self.assertEqual(spec.kind, MaskKind.EXPLICIT) q = torch.randn(1, self.HEADS, 4, self.DIM) - out = self._attend(q, k, v, spec) + out = self._attend(q, spec) for i in range(4): # query at position i sees (i - window, i] lo = max(0, i - window + 1) torch.testing.assert_close( out[:, :, i : i + 1, :], - self._attend( + self._unmasked( q[:, :, i : i + 1, :], - k[:, :, lo : i + 1, :], - v[:, :, lo : i + 1, :], - AttendSpec(kind=MaskKind.NONE), + spec.k[:, :, lo : i + 1, :], + spec.v[:, :, lo : i + 1, :], ), ) def test_layers_can_window_independently(self): # gemma-style: only some layers are windowed, so one step yields two # different semantics from the same cache. - cache = ContiguousReferenceCache( + cache = SequenceReferenceCache( CacheConfig( n_layers=2, n_kv_heads=self.HEADS, @@ -743,8 +1060,8 @@ def test_layers_can_window_independently(self): ) kv = torch.randn(1, self.HEADS, 4, self.DIM) pos = _positions(0, 4) - flat = cache.update_and_fetch(0, kv, kv, pos)[2] - windowed = cache.update_and_fetch(1, kv, kv, pos)[2] + flat = cache.update_and_fetch(0, kv, kv, pos)[0] + windowed = cache.update_and_fetch(1, kv, kv, pos)[0] self.assertEqual(flat.kind, MaskKind.CAUSAL) self.assertEqual(windowed.kind, MaskKind.EXPLICIT) @@ -757,7 +1074,7 @@ def test_windowed_continuation_bounds_the_band_at_both_ends(self): window = 2 cache = self._cache(LayerPolicy.ring(window)) self._update(cache, 4, 0) - _, _, spec = self._update(cache, 3, 4) + spec = self._update(cache, 3, 4)[0] self.assertEqual(spec.kind, MaskKind.EXPLICIT) q_len, total = 3, 7 @@ -773,8 +1090,8 @@ def test_window_equal_to_history_stays_fused(self): # later the band appears. window = 4 cache = self._cache(LayerPolicy.ring(window)) - self.assertEqual(self._update(cache, window, 0)[2].kind, MaskKind.CAUSAL) - self.assertEqual(self._update(cache, 1, window)[2].kind, MaskKind.EXPLICIT) + self.assertEqual(self._update(cache, window, 0)[0].kind, MaskKind.CAUSAL) + self.assertEqual(self._update(cache, 1, window)[0].kind, MaskKind.EXPLICIT) def test_static_sizing_windows_like_dynamic(self): # STATIC writes into a preallocated buffer and slices it; the window is @@ -782,13 +1099,13 @@ def test_static_sizing_windows_like_dynamic(self): window = 2 torch.manual_seed(0) static = self._cache(LayerPolicy.ring(window), sizing=CacheSizing.STATIC) - sk, sv, s_spec = self._update(static, 4, 0) + s_spec = self._update(static, 4, 0)[0] torch.manual_seed(0) dynamic = self._cache(LayerPolicy.ring(window)) - dk, dv, d_spec = self._update(dynamic, 4, 0) + d_spec = self._update(dynamic, 4, 0)[0] - torch.testing.assert_close(sk, dk) - torch.testing.assert_close(sv, dv) + torch.testing.assert_close(s_spec.k, d_spec.k) + torch.testing.assert_close(s_spec.v, d_spec.v) self.assertEqual(s_spec.kind, d_spec.kind) torch.testing.assert_close(s_spec.mask, d_spec.mask) @@ -796,8 +1113,8 @@ def test_window_wider_than_history_stays_fused(self): # Nothing to bound from below, so the window must not force a mask. for policy in (LayerPolicy.flat(), LayerPolicy.ring(64)): cache = self._cache(policy) - self.assertEqual(self._update(cache, 4, 0)[2].kind, MaskKind.CAUSAL) - self.assertEqual(self._update(cache, 1, 4)[2].kind, MaskKind.NONE) + self.assertEqual(self._update(cache, 4, 0)[0].kind, MaskKind.CAUSAL) + self.assertEqual(self._update(cache, 1, 4)[0].kind, MaskKind.NONE) class AttendExplicitTest(unittest.TestCase): @@ -812,16 +1129,21 @@ def setUp(self): self.v = torch.randn(1, 2, self.total, self.head_dim) self.scale = self.head_dim**-0.5 - def _attend(self, spec, k=None, v=None): - k = self.k if k is None else k - v = self.v if v is None else v - return attend(self.q, k, v, spec, self.scale, torch.float32) + def _attend(self, kind, mask=None, k=None, v=None): + spec = AttendSpec( + k=self.k if k is None else k, + v=self.v if v is None else v, + kind=kind, + q_len=self.q_len, + mask=mask, + ) + return attend(self.q, spec, self.scale, torch.float32) def test_causal_rejects_a_non_square_window(self): # torch's is_causal is upper-left, so it cannot serve a continuation; # a cache must declare EXPLICIT there rather than CAUSAL. with self.assertRaises(ValueError): - self._attend(AttendSpec(kind=MaskKind.CAUSAL)) + self._attend(MaskKind.CAUSAL) def test_explicit_attends_the_true_cells(self): # Polarity: masking to cells {0, 2} must equal attending over just those @@ -830,11 +1152,11 @@ def test_explicit_attends_the_true_cells(self): mask = torch.zeros(self.q_len, self.total, dtype=torch.bool) mask[:, keep] = True torch.testing.assert_close( - self._attend(AttendSpec(kind=MaskKind.EXPLICIT, mask=mask)), + self._attend(MaskKind.EXPLICIT, mask=mask), self._attend( - AttendSpec(kind=MaskKind.NONE), - self.k.index_select(2, keep), - self.v.index_select(2, keep), + MaskKind.NONE, + k=self.k.index_select(2, keep), + v=self.v.index_select(2, keep), ), ) diff --git a/extension/llm/cache/update_and_attend.py b/extension/llm/cache/update_and_attend.py index 9feaabc4038..227f0aebbe6 100644 --- a/extension/llm/cache/update_and_attend.py +++ b/extension/llm/cache/update_and_attend.py @@ -95,8 +95,16 @@ def update_and_attend( Returns: ``[B, H_q, q_len, v_head_dim]`` attention output, in ``out_dtype``. """ - k_hist, v_hist, spec = REGISTRY.current().update_and_fetch(layer_id, k, v, position) - return attend(q, k_hist, v_hist, spec, scale, out_dtype) + specs = REGISTRY.current().update_and_fetch(layer_id, k, v, position) + outputs = [] + start = 0 + for spec in specs: + end = start + spec.q_len + outputs.append(attend(q[:, :, start:end, :], spec, scale, out_dtype)) + start = end + if start != q.shape[-2]: + raise ValueError(f"the cache answered {start} of {q.shape[-2]} query tokens") + return outputs[0] if len(outputs) == 1 else torch.cat(outputs, dim=2) @update_and_attend.register_fake From 75c2b301c1388b19604bb5e3d5d68c3b70e66b9a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Sun, 30 Aug 2026 10:20:44 +0200 Subject: [PATCH 120/190] Vulkan: add the missing int32 -> uint8 view_convert variant The view_convert combo lists generate uint8 -> int32 but not the reverse, so a graph containing an int32 -> uint8 view_convert lowers cleanly and then aborts at run time: Could not find ShaderInfo with name view_convert_buffer_int32_uint8 (ShaderRegistry.cpp:54, (it != listings_.end()) is false) The partitioner has no visibility into which dtype pairs were generated, so this surfaces as a crash on device rather than the node being left on host. sentence-transformers/all-MiniLM-L6-v2 hits it through its attention mask on any graph lowered with VulkanPartitioner. view_convert is a plain elementwise OUT_T(in) cast, so the pair needs nothing beyond being listed. Verified on a Galaxy S26 Ultra (Adreno 840): with this variant generated, all-MiniLM-L6-v2 runs and its embedding matches the eager reference to cosine 0.99999726 (max abs deviation 4.0e-04), bit-identical across 90 in-process replays. --- backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml | 1 + backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml | 1 + 2 files changed, 2 insertions(+) diff --git a/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml b/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml index 200d58e1217..a770d986fd5 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/view_convert_buffer.yaml @@ -18,6 +18,7 @@ view_convert_buffer: - parameter_values: [uint8, float] - parameter_values: [uint8, half] - parameter_values: [uint8, int32] + - parameter_values: [int32, uint8] - parameter_values: [float, int32] - parameter_values: [float, half] - parameter_values: [half, float] diff --git a/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml b/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml index 47e9c43ee24..227a18b6ccf 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml +++ b/backends/vulkan/runtime/graph/ops/glsl/view_convert_texture.yaml @@ -18,6 +18,7 @@ view_convert_texture: - parameter_values: [uint8, float] - parameter_values: [uint8, half] - parameter_values: [uint8, int32] + - parameter_values: [int32, uint8] - parameter_values: [float, int32] - parameter_values: [float, half] - parameter_values: [half, float] From eadb4c3a93d8874b6f83d52f3cb6ce493325f2b1 Mon Sep 17 00:00:00 2001 From: Gasoonjia Date: Wed, 9 Sep 2026 12:58:13 -0700 Subject: [PATCH 121/190] [cuda backend] support multi-SM AOTI PTEs (#22198) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ### Summary Add one current `ETCUDAFQN0` metadata format that stores target-specific CUDA AOTI shared libraries with one shared FQN weight manifest. At runtime, CUDA first selects an exact-SM regular variant. If no exact regular variant exists, it may use the PTX from one explicitly designated fallback PTE when that PTX is compatible with the current GPU. The fallback is no longer inferred from the lowest regular target SM. Add a PTE/PTD merge tool with two input classes: - One or more regular PTEs provide exact-SM native variants. Their merged metadata sets `ptx_compute=0`, so they are never considered as PTX fallbacks. Regular exports should use `CompileSpec("cuda_include_ptx", b"OFF")` to avoid embedding redundant PTX. - At most one fallback PTE provides the PTX-capable variant. It is marked `fallback_only`, is never selected as an exact native variant, and is considered only when the merged PTE has no regular variant for the current SM. The merge validates the ExecuTorch program, CUDA delegate layout, non-CUDA named data, weight metadata, and streamed PTD weight hashes. CUDA PTE merging is NVIDIA-only and explicitly rejects ROCm. ### Merged artifact layout For an A100 (`sm80`) regular export, an RTX 5090 (`sm120`) regular export, and one PTX fallback export, the merged artifact is structured as follows: ```text merged.pte ├── program and CUDA delegates │ └── `ETCUDAFQN0` CUDA variant metadata (per delegate) │ ├── regular: target_sm=80, ptx_compute=0 -> sm80 source SO key │ ├── regular: target_sm=120, ptx_compute=0 -> sm120 source SO key │ └── fallback_only: target_sm=, │ ptx_compute= -> fallback source SO key ├── named data │ ├── sm80 source SO key -> actual AOTI SO containing the prebuilt sm80 cubin │ ├── sm120 source SO key -> actual AOTI SO containing the prebuilt sm120 cubin │ └── fallback source SO key -> actual PTX-bearing fallback AOTI SO └── one shared FQN weight manifest merged.ptd └── shared external tensor data reused from the regular base export ``` The tool copies each AOTI SO atomically rather than rewriting its ELF/fatbin. The regular/fallback distinction is encoded in the merged metadata: only regular variants participate in exact-SM native selection, while the fallback SO contributes PTX only to fallback selection. After merging, the CLI prints a provenance table mapping every stored SM cubin and the fallback PTX entry to its source PTE. This PR was authored with OpenAI Codex. ### Test plan The primary validation is a real cross-hardware CUDA CI pipeline spanning A100 (`sm80`) and A10G (`sm86`): 1. Export an A100 native PTE and an A10G native PTE independently on their respective GPUs, both with PTX disabled and isolated Inductor caches. 2. Export a separate PTX fallback PTE on A100 with the A10G shared-memory constraint. 3. Verify byte-for-byte that the A100 native PTD, A10G native PTD, fallback PTD, and final merged PTD are identical. 4. On A10G, verify the A10G native PTE succeeds, the A100 native-only PTE is rejected with no compatible native/PTX variant, and the merged PTE succeeds. 5. On A100, verify the A100 native PTE succeeds, the A10G native-only PTE is rejected with no compatible native/PTX variant, and the same merged PTE succeeds. Additional Python coverage exercises the single `ETCUDAFQN0` metadata format, untargeted ROCm metadata, explicit fallback-only encoding, PTE/PTD merging, shared-weight validation, duplicate regular SM rejection, provenance reporting, and ROCm merge rejection. C++ coverage exercises metadata parsing, exact regular selection, PTX fallback selection, incompatibility, and duplicate targets. Locally ran the 17 targeted CUDA metadata/merge tests, Python lint/format checks, workflow YAML parsing, and `git diff --check`. --- .github/workflows/cuda.yml | 162 ++++- backends/cuda/BUCK | 16 + backends/cuda/cuda_backend.py | 34 +- backends/cuda/cuda_weight_collector.py | 190 ++++- backends/cuda/merge_ptes.py | 675 ++++++++++++++++++ backends/cuda/runtime/cuda_backend.cpp | 41 +- backends/cuda/runtime/cuda_weight_cache.cpp | 87 ++- backends/cuda/runtime/cuda_weight_cache.h | 17 +- .../runtime/test/test_cuda_weight_cache.cpp | 181 ++++- backends/cuda/tests/test_cuda_export.py | 27 + backends/cuda/tests/test_cuda_partitioner.py | 11 +- .../cuda/tests/test_cuda_weight_metadata.py | 130 ++++ backends/cuda/tests/test_merge_ptes.py | 387 ++++++++++ examples/cuda/README.md | 39 + examples/cuda/scripts/export.py | 23 +- 15 files changed, 1985 insertions(+), 35 deletions(-) create mode 100644 backends/cuda/merge_ptes.py create mode 100644 backends/cuda/tests/test_cuda_weight_metadata.py create mode 100644 backends/cuda/tests/test_merge_ptes.py diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index 553e6ac721b..f3705104ded 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -160,7 +160,7 @@ jobs: done export-cuda-target-smem-cross-arch: - name: export-cuda-target-smem-cross-arch-a100 + name: export-cuda-multi-arch-a100 needs: [changed-files, run-decision] if: | contains(needs.changed-files.outputs.changed-files, 'backends/cuda') || @@ -178,15 +178,13 @@ jobs: gpu-arch-version: "13.0" use-custom-docker-registry: false submodules: recursive - upload-artifact: cuda-target-smem-cross-arch + upload-artifact: cuda-multi-arch-a100 ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux PYTHON_EXECUTABLE=python ./install_executorch.sh export LD_LIBRARY_PATH="/opt/conda/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" - # A10G (sm_86) supports 99 KiB of opt-in shared memory per block. - export ET_CUDA_TARGET_SMEM_BYTES=101376 # The A100 exporter has an AVX-512 host CPU while the A10G runner has # AVX2. Keep the generated AOTI host wrapper compatible with both. export ATEN_CPU_CAPABILITY=avx2 @@ -199,16 +197,33 @@ jobs: assert capability == (8, 0), f"Expected A100 sm_80 exporter, got {capability}" PY + # Export an sm80 native PTE with no PTX. This artifact must fail on the + # A10G because its metadata has neither an sm86 native variant nor PTX. + model_dir="${RUNNER_ARTIFACT_DIR}/native-sm80/linear" + mkdir -p "${model_dir}" + TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-native-sm80" \ + python -m examples.cuda.scripts.export \ + --model_name=linear \ + --output_dir="${model_dir}" \ + --cuda_include_ptx=OFF \ + --seed=0 + + # Export the explicit PTX fallback separately. A10G supports 99 KiB of + # opt-in shared memory per block, so constrain the A100 export to it. for model in linear sdpa; do - model_dir="${RUNNER_ARTIFACT_DIR}/${model}" + model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}" mkdir -p "${model_dir}" - python -m examples.cuda.scripts.export \ + ET_CUDA_TARGET_SMEM_BYTES=101376 \ + TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-fallback-sm80-${model}" \ + python -m examples.cuda.scripts.export \ --model_name="${model}" \ - --output_dir "${model_dir}" + --output_dir="${model_dir}" \ + --cuda_include_ptx=ON \ + --seed=0 done test-cuda-target-smem-cross-arch: - name: test-cuda-target-smem-cross-arch-a10g + name: merge-and-test-cuda-multi-arch-a10g needs: [export-cuda-target-smem-cross-arch] if: needs.export-cuda-target-smem-cross-arch.result == 'success' uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main @@ -222,7 +237,8 @@ jobs: gpu-arch-version: "13.0" use-custom-docker-registry: false submodules: recursive - download-artifact: cuda-target-smem-cross-arch + download-artifact: cuda-multi-arch-a100 + upload-artifact: cuda-multi-arch-merged ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux @@ -243,20 +259,141 @@ jobs: ) PY + # Independently export an sm86 native PTE with no PTX. A fixed seed + # makes its external weights byte-identical to both A100 exports. + model_dir="${RUNNER_ARTIFACT_DIR}/native-sm86/linear" + mkdir -p "${model_dir}" + TORCHINDUCTOR_CACHE_DIR="${RUNNER_TEMP}/inductor-native-sm86" \ + python -m examples.cuda.scripts.export \ + --model_name=linear \ + --output_dir="${model_dir}" \ + --cuda_include_ptx=OFF \ + --seed=0 + + sm80_ptd="${RUNNER_ARTIFACT_DIR}/native-sm80/linear/aoti_cuda_blob.ptd" + sm86_ptd="${RUNNER_ARTIFACT_DIR}/native-sm86/linear/aoti_cuda_blob.ptd" + fallback_ptd="${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/aoti_cuda_blob.ptd" + cmp "${sm80_ptd}" "${sm86_ptd}" + cmp "${sm80_ptd}" "${fallback_ptd}" + sha256sum "${sm80_ptd}" "${sm86_ptd}" "${fallback_ptd}" + + merged_dir="${RUNNER_ARTIFACT_DIR}/merged" + mkdir -p "${merged_dir}" + python -m executorch.backends.cuda.merge_ptes \ + --input-pte "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ + --input-pte "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ + --input-ptd "${sm80_ptd}" \ + --input-ptd "${sm86_ptd}" \ + --fallback-pte "${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/linear.pte" \ + --fallback-ptd "${fallback_ptd}" \ + --output-pte "${merged_dir}/linear.pte" \ + --output-ptd "${merged_dir}/aoti_cuda_blob.ptd" \ + | tee "${merged_dir}/provenance.txt" + cmp "${sm80_ptd}" "${merged_dir}/aoti_cuda_blob.ptd" + cmake -DCMAKE_BUILD_TYPE=Release \ -DEXECUTORCH_BUILD_CUDA=ON \ -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ -DPYTHON_EXECUTABLE=python \ -Bcmake-out . cmake --build cmake-out --target executor_runner -j4 for model in linear sdpa; do - model_dir="${RUNNER_ARTIFACT_DIR}/${model}" + model_dir="${RUNNER_ARTIFACT_DIR}/fallback-sm80/${model}" ./cmake-out/executor_runner \ --model_path "${model_dir}/${model}.pte" \ --data_path "${model_dir}/aoti_cuda_blob.ptd" done + # The local native PTE and merged PTE must run on sm86. + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ + --data_path "${sm86_ptd}" + ./cmake-out/executor_runner \ + --model_path "${merged_dir}/linear.pte" \ + --data_path "${merged_dir}/aoti_cuda_blob.ptd" + + # The sm80 native-only PTE must be rejected on sm86. + set +e + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ + --data_path "${sm80_ptd}" \ + >"${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" 2>&1 + native_status=$? + set -e + cat "${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" + test "${native_status}" -ne 0 + grep -q "no native or PTX variant compatible with sm86" \ + "${RUNNER_ARTIFACT_DIR}/native-sm80-on-sm86.log" + + test-cuda-merged-pte-a100: + name: test-cuda-merged-pte-a100 + needs: [test-cuda-target-smem-cross-arch] + if: needs.test-cuda-target-smem-cross-arch.result == 'success' + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read + with: + timeout: 90 + runner: mt-l-x86iavx512-11-125-a100 + gpu-arch-type: cuda + gpu-arch-version: "13.0" + use-custom-docker-registry: false + submodules: recursive + download-artifact: cuda-multi-arch-merged + ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} + script: | + set -eux + + PYTHON_EXECUTABLE=python ./install_executorch.sh + export LD_LIBRARY_PATH="/opt/conda/lib${LD_LIBRARY_PATH:+:$LD_LIBRARY_PATH}" + + python - <<'PY' + import torch + + capability = torch.cuda.get_device_capability() + assert capability == (8, 0), f"Expected A100 sm_80 runner, got {capability}" + PY + + cmake -DCMAKE_BUILD_TYPE=Release \ + -DEXECUTORCH_BUILD_CUDA=ON \ + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON \ + -DEXECUTORCH_ENABLE_LOGGING=ON \ + -DPYTHON_EXECUTABLE=python \ + -Bcmake-out . + cmake --build cmake-out --target executor_runner -j4 + + sm80_ptd="${RUNNER_ARTIFACT_DIR}/native-sm80/linear/aoti_cuda_blob.ptd" + sm86_ptd="${RUNNER_ARTIFACT_DIR}/native-sm86/linear/aoti_cuda_blob.ptd" + fallback_ptd="${RUNNER_ARTIFACT_DIR}/fallback-sm80/linear/aoti_cuda_blob.ptd" + merged_ptd="${RUNNER_ARTIFACT_DIR}/merged/aoti_cuda_blob.ptd" + cmp "${sm80_ptd}" "${sm86_ptd}" + cmp "${sm80_ptd}" "${fallback_ptd}" + cmp "${sm80_ptd}" "${merged_ptd}" + + # The local native PTE and merged PTE must run on sm80. + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm80/linear/linear.pte" \ + --data_path "${sm80_ptd}" + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/merged/linear.pte" \ + --data_path "${merged_ptd}" + + # The sm86 native-only PTE must be rejected on sm80. + set +e + ./cmake-out/executor_runner \ + --model_path "${RUNNER_ARTIFACT_DIR}/native-sm86/linear/linear.pte" \ + --data_path "${sm86_ptd}" \ + >"${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" 2>&1 + native_status=$? + set -e + cat "${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" + test "${native_status}" -ne 0 + grep -q "no native or PTX variant compatible with sm80" \ + "${RUNNER_ARTIFACT_DIR}/native-sm86-on-sm80.log" + unittest-cuda: name: unittest-cuda needs: [changed-files, run-decision] @@ -370,6 +507,11 @@ jobs: conda install -y -c conda-forge 'libstdcxx-ng>=12' export LD_LIBRARY_PATH=/opt/conda/lib:$LD_LIBRARY_PATH + python -m pytest \ + backends/cuda/tests/test_cuda_weight_metadata.py \ + backends/cuda/tests/test_merge_ptes.py \ + -v -o "addopts=" + cmake --preset llm-release-cuda -DEXECUTORCH_BUILD_TESTS=ON cmake --build cmake-out --target test_cuda_allocator test_cuda_mutable_state test_cuda_weight_cache -j$(nproc) ctest --test-dir cmake-out -R test_cuda_allocator --output-on-failure -V diff --git a/backends/cuda/BUCK b/backends/cuda/BUCK index 3fe33473472..4602978fad4 100644 --- a/backends/cuda/BUCK +++ b/backends/cuda/BUCK @@ -109,6 +109,22 @@ fbcode_target( ], ) +fbcode_target( + _kind = runtime.python_library, + name = "merge_ptes", + srcs = [ + "merge_ptes.py", + ], + visibility = ["PUBLIC"], + deps = [ + ":cuda_backend", + "//caffe2:torch", + "//executorch/exir/_serialize:lib", + "//executorch/exir:schema", + "//executorch/extension/flat_tensor/serialize:serialize", + ], +) + fbcode_target( _kind = runtime.python_library, name = "cuda_partitioner", diff --git a/backends/cuda/cuda_backend.py b/backends/cuda/cuda_backend.py index 89e1b5ad79a..4db99b6b040 100644 --- a/backends/cuda/cuda_backend.py +++ b/backends/cuda/cuda_backend.py @@ -519,6 +519,14 @@ def _setup_cuda_environment_for_fatbin() -> bool: except Exception: return False + @classmethod + def _should_include_ptx(cls, compile_specs: List[CompileSpec]) -> bool: + include_ptx = True + for spec in compile_specs: + if spec.key == "cuda_include_ptx": + include_ptx = _on_off_compile_spec_value(spec) + return include_ptx and cls._setup_cuda_environment_for_fatbin() + @classmethod def save_data_externally(cls) -> bool: """ @@ -539,7 +547,29 @@ def _preprocess_with_weight_collector( result = super().preprocess(edge_program, compile_specs) if capture.artifact is None: raise RuntimeError("CUDA AOTI did not return a structured Weights output") - collector.add_preprocess_result(result, capture.artifact, cls.get_device_name()) + target_sm = None + ptx_compute = 0 + if torch.version.hip is None: + from torch._inductor.codegen.cuda.compile_utils import ( + _nvcc_arch_as_compile_option, + ) + + compiled_arch = _nvcc_arch_as_compile_option() + target_arch = compiled_arch.removesuffix("a") + if not target_arch.isdigit(): + raise RuntimeError(f"Unsupported CUDA architecture: {compiled_arch}") + target_sm = int(target_arch) + # Architecture-accelerated PTX targets such as compute_90a and + # compute_120a are not forward-compatible fallback images. + if cls._should_include_ptx(compile_specs) and compiled_arch.isdigit(): + ptx_compute = int(compiled_arch) + collector.add_preprocess_result( + result, + capture.artifact, + cls.get_device_name(), + target_sm=target_sm, + ptx_compute=ptx_compute, + ) return result @classmethod @@ -762,7 +792,7 @@ def get_aoti_compile_options( # Configure CUDA environment variables based on detected version - emit_multi_arch_kernel = CudaBackend._setup_cuda_environment_for_fatbin() + emit_multi_arch_kernel = cls._should_include_ptx(compile_specs) # Base options for all platforms options: Dict[str, typing.Any] = { diff --git a/backends/cuda/cuda_weight_collector.py b/backends/cuda/cuda_weight_collector.py index ac47f7c7e2c..08806c2bc68 100644 --- a/backends/cuda/cuda_weight_collector.py +++ b/backends/cuda/cuda_weight_collector.py @@ -22,7 +22,7 @@ from executorch.exir.tensor import scalar_type_enum -CUDA_WEIGHT_CACHE_MAGIC = b"ETCUDAFQN3" +CUDA_AOTI_METADATA_MAGIC = b"ETCUDAFQN0" AOTI_DEVICE_TYPE_CPU = 0 AOTI_DEVICE_TYPE_CUDA = 1 @@ -46,6 +46,24 @@ class CudaWeightArtifact: storages: Dict[str, FileBackedData] +@dataclass(frozen=True) +class CudaAotiVariant: + """One AOTI shared library and its CUDA runtime-selection metadata.""" + + target_sm: int + ptx_compute: int + so_blob_key: str + fallback_only: bool = False + + +@dataclass(frozen=True) +class CudaAotiMetadata: + """CUDA AOTI native and fallback variants sharing one weight manifest.""" + + variants: List[CudaAotiVariant] + entries: List[CudaWeightEntry] + + @dataclass class _CudaWeightCapture: collector: "CudaWeightCollector" @@ -114,18 +132,75 @@ def _is_aoti_library_local_fqn(fqn: str) -> bool: return fqn.startswith("_tensor_constant") -def encode_cuda_weight_metadata( - so_blob_key: str, entries: List[CudaWeightEntry] +def _validate_cuda_aoti_variant( + variant: CudaAotiVariant, has_fallback: bool, regular_sms: set[int] +) -> None: + if variant.target_sm < 0: + raise ValueError(f"Invalid CUDA target SM: {variant.target_sm}") + if variant.ptx_compute < 0 or variant.ptx_compute > variant.target_sm: + raise ValueError( + f"Invalid PTX compute target {variant.ptx_compute} for sm{variant.target_sm}" + ) + if not variant.so_blob_key: + raise ValueError("CUDA AOTI variant is missing its shared-object key") + if variant.fallback_only: + if variant.ptx_compute == 0: + raise ValueError("CUDA fallback variant must contain PTX") + return + if variant.target_sm in regular_sms: + raise ValueError(f"Duplicate CUDA target SM: {variant.target_sm}") + regular_sms.add(variant.target_sm) + if has_fallback and variant.ptx_compute != 0: + raise ValueError("Regular CUDA variants cannot advertise PTX fallback") + + +def _validate_cuda_aoti_variants( + variants: List[CudaAotiVariant], has_fallback: bool +) -> None: + untargeted = [variant for variant in variants if variant.target_sm == 0] + if untargeted: + if len(variants) != 1 or untargeted[0].ptx_compute or has_fallback: + raise ValueError( + "Untargeted CUDA AOTI metadata requires one non-fallback variant" + ) + if not untargeted[0].so_blob_key: + raise ValueError("CUDA AOTI variant is missing its shared-object key") + return + if sum(variant.fallback_only for variant in variants) > 1: + raise ValueError("CUDA AOTI metadata supports only one fallback variant") + if len(variants) > 1 and any( + variant.ptx_compute and not variant.fallback_only for variant in variants + ): + raise ValueError( + "Multi-variant CUDA AOTI metadata requires an explicit PTX fallback" + ) + regular_sms: set[int] = set() + for variant in variants: + _validate_cuda_aoti_variant(variant, has_fallback, regular_sms) + + +def encode_cuda_aoti_metadata( + variants: List[CudaAotiVariant], entries: List[CudaWeightEntry] ) -> bytes: - """Encode the per-method FQN-to-tensor metadata consumed by CUDA runtime.""" - output = bytearray(CUDA_WEIGHT_CACHE_MAGIC) + """Encode CUDA AOTI variants followed by one shared weight manifest.""" + if not variants: + raise ValueError("CUDA AOTI metadata requires at least one variant") + + has_fallback = any(variant.fallback_only for variant in variants) + _validate_cuda_aoti_variants(variants, has_fallback) + output = bytearray(CUDA_AOTI_METADATA_MAGIC) def write_string(value: str) -> None: encoded = value.encode("utf-8") output.extend(struct.pack(" None: return bytes(output) +class _MetadataReader: + def __init__(self, data: bytes) -> None: + self._data = memoryview(data) + self._offset = 0 + + def read(self, size: int) -> memoryview: + end = self._offset + size + if size < 0 or end > len(self._data): + raise ValueError("Truncated CUDA AOTI metadata") + value = self._data[self._offset : end] + self._offset = end + return value + + def unpack(self, format: str) -> Tuple[Any, ...]: + size = struct.calcsize(format) + return struct.unpack(format, self.read(size)) + + def read_string(self) -> str: + (size,) = self.unpack(" None: + if self._offset != len(self._data): + raise ValueError("CUDA AOTI metadata contains trailing bytes") + + +def _decode_weight_entries(reader: _MetadataReader) -> List[CudaWeightEntry]: + (num_entries,) = reader.unpack(" 1 << 20: + raise ValueError(f"CUDA AOTI metadata has too many weights: {num_entries}") + + entries = [] + for _ in range(num_entries): + fqn = reader.read_string() + storage_key = reader.read_string() + storage_nbytes, dtype, device_type, storage_offset, ndim = reader.unpack( + " 64: + raise ValueError("CUDA AOTI metadata contains an invalid weight entry") + sizes = reader.unpack(f"<{ndim}q") if ndim else () + strides = reader.unpack(f"<{ndim}q") if ndim else () + if storage_offset < 0 or any(value < 0 for value in sizes + strides): + raise ValueError("CUDA AOTI metadata contains invalid tensor metadata") + entries.append( + CudaWeightEntry( + fqn=fqn, + storage_key=storage_key, + storage_nbytes=storage_nbytes, + dtype=dtype, + device_type=device_type, + storage_offset=storage_offset, + sizes=tuple(sizes), + strides=tuple(strides), + ) + ) + return entries + + +def decode_cuda_aoti_metadata(data: bytes) -> CudaAotiMetadata: + """Decode CUDA AOTI variant and shared-weight metadata.""" + reader = _MetadataReader(data) + magic = bytes(reader.read(len(CUDA_AOTI_METADATA_MAGIC))) + if magic != CUDA_AOTI_METADATA_MAGIC: + raise ValueError("Unrecognized CUDA AOTI metadata") + + variants = [] + (num_variants,) = reader.unpack(" 256: + raise ValueError( + f"CUDA AOTI metadata has invalid variant count: {num_variants}" + ) + for _ in range(num_variants): + target_sm, ptx_compute, flags = reader.unpack(" value`` store for all methods.""" @@ -311,6 +482,8 @@ def add_preprocess_result( result: PreprocessResult, artifact: CudaWeightArtifact, device_name: str, + target_sm: Optional[int] = None, + ptx_compute: int = 0, ) -> None: if result.data_store_output is None: raise RuntimeError("CUDA AOTI preprocess returned no named data") @@ -344,8 +517,9 @@ def add_preprocess_result( self._add_weight(entry, data, external_tag) serialized_entries.append(entry) - result.processed_bytes = encode_cuda_weight_metadata( - so_blob_key, serialized_entries + result.processed_bytes = encode_cuda_aoti_metadata( + [CudaAotiVariant(target_sm or 0, ptx_compute, so_blob_key)], + serialized_entries, ) self._results.append(result) diff --git a/backends/cuda/merge_ptes.py b/backends/cuda/merge_ptes.py new file mode 100644 index 00000000000..2a7212e608a --- /dev/null +++ b/backends/cuda/merge_ptes.py @@ -0,0 +1,675 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import argparse +import copy +import hashlib +import os +import shutil +import tempfile +from dataclasses import dataclass, replace +from pathlib import Path +from typing import Dict, Iterable, List, Optional, Sequence, Tuple + +import torch + +from executorch.backends.cuda.cuda_weight_collector import ( + CudaAotiMetadata, + CudaAotiVariant, + CudaWeightEntry, + decode_cuda_aoti_metadata, + encode_cuda_aoti_metadata, +) +from executorch.exir._serialize._cord import Cord +from executorch.exir._serialize._named_data_store import ( + NamedDataStore, + NamedDataStoreOutput, +) +from executorch.exir._serialize._program import ( + deserialize_pte_binary, + PTEFile, + serialize_pte_binary, +) +from executorch.exir.schema import ( + BackendDelegateDataReference, + BackendDelegateInlineData, + DataLocation, + Program, +) +from executorch.extension.flat_tensor.serialize.serialize import ( + _deserialize_to_flat_tensor, + FlatTensorHeader, +) + + +CUDA_BACKEND_ID = "CudaBackend" + + +@dataclass(frozen=True) +class CudaPteInput: + """One native CUDA export and its optional external tensor data.""" + + pte_path: Path + ptd_path: Optional[Path] + + +@dataclass(frozen=True) +class CudaPteProvenance: + delegate: Tuple[str, int] + kind: str + target_sm: int + ptx_compute: int + source_pte: Path + + +@dataclass(frozen=True) +class CudaPteMergeResult: + pte: Cord + provenance: Tuple[CudaPteProvenance, ...] + + +@dataclass(frozen=True) +class _PtdBlob: + offset: int + size: int + + +class _PtdIndex: + def __init__(self, path: Path) -> None: + self.path = path + with path.open("rb") as source: + prefix = source.read(8 + FlatTensorHeader.EXPECTED_LENGTH) + header = FlatTensorHeader.from_bytes(prefix[8:]) + if not header.is_valid(): + raise ValueError(f"Invalid PTD header in {path}") + source.seek(0) + flatbuffer_size = header.flatbuffer_offset + header.flatbuffer_size + flat_tensor = _deserialize_to_flat_tensor(source.read(flatbuffer_size)) + + file_size = path.stat().st_size + self._blobs: Dict[str, _PtdBlob] = {} + for named_data in flat_tensor.named_data: + if named_data.key in self._blobs: + raise ValueError(f"PTD contains duplicate key {named_data.key!r}") + if named_data.segment_index >= len(flat_tensor.segments): + raise ValueError( + f"PTD key {named_data.key!r} has an invalid segment index" + ) + segment = flat_tensor.segments[named_data.segment_index] + offset = header.segment_base_offset + segment.offset + if offset + segment.size > file_size: + raise ValueError(f"PTD key {named_data.key!r} extends past end of file") + self._blobs[named_data.key] = _PtdBlob(offset, segment.size) + self._digests: Dict[str, bytes] = {} + + def keys(self) -> set[str]: + return set(self._blobs) + + def size(self, key: str) -> int: + try: + return self._blobs[key].size + except KeyError as error: + raise ValueError(f"PTD {self.path} does not contain key {key!r}") from error + + def sha256(self, key: str) -> bytes: + digest = self._digests.get(key) + if digest is not None: + return digest + try: + blob = self._blobs[key] + except KeyError as error: + raise ValueError(f"PTD {self.path} does not contain key {key!r}") from error + + hasher = hashlib.sha256() + remaining = blob.size + with self.path.open("rb") as source: + source.seek(blob.offset) + while remaining: + chunk = source.read(min(8 * 1024 * 1024, remaining)) + if not chunk: + raise ValueError(f"PTD {self.path} is truncated at key {key!r}") + hasher.update(chunk) + remaining -= len(chunk) + digest = hasher.digest() + self._digests[key] = digest + return digest + + +@dataclass +class _CudaDelegate: + identity: Tuple[str, int] + metadata: CudaAotiMetadata + + +@dataclass +class _Artifact: + source: CudaPteInput + pte: PTEFile + pte_named_data: Dict[str, bytes] + ptd: Optional[_PtdIndex] + delegates: List[_CudaDelegate] + + def blob_size(self, key: str) -> int: + data = self.pte_named_data.get(key) + if data is not None: + return len(data) + if self.ptd is not None: + return self.ptd.size(key) + raise ValueError(f"{self.source.pte_path} does not contain named data {key!r}") + + def blob_sha256(self, key: str) -> bytes: + data = self.pte_named_data.get(key) + if data is not None: + return hashlib.sha256(data).digest() + if self.ptd is not None: + return self.ptd.sha256(key) + raise ValueError(f"{self.source.pte_path} does not contain named data {key!r}") + + +def _named_data_bytes(output: Optional[NamedDataStoreOutput]) -> Dict[str, bytes]: + if output is None: + return {} + return { + key: bytes(output.buffers[entry.buffer_index]) + for key, entry in output.pte_data.items() + } + + +def _delegate_payload(program: Program, delegate) -> bytes: + if delegate.processed.location != DataLocation.INLINE: + raise ValueError("PTE deserialization did not restore delegate data inline") + try: + return bytes(program.backend_delegate_data[delegate.processed.index].data) + except IndexError as error: + raise ValueError("CUDA delegate references invalid processed data") from error + + +def _load_artifact(source: CudaPteInput) -> _Artifact: + pte = deserialize_pte_binary(source.pte_path.read_bytes()) + delegates = [] + for plan in pte.program.execution_plan: + for delegate_index, delegate in enumerate(plan.delegates): + if delegate.id != CUDA_BACKEND_ID: + continue + metadata = decode_cuda_aoti_metadata( + _delegate_payload(pte.program, delegate) + ) + if metadata.variants[0].target_sm == 0: + raise ValueError( + f"Untargeted CUDA metadata in {source.pte_path} cannot be merged" + ) + delegates.append(_CudaDelegate((plan.name, delegate_index), metadata)) + if not delegates: + raise ValueError(f"{source.pte_path} contains no CUDA delegates") + ptd = _PtdIndex(source.ptd_path) if source.ptd_path is not None else None + return _Artifact(source, pte, _named_data_bytes(pte.named_data), ptd, delegates) + + +def _normalized_program(program: Program) -> Program: + normalized = copy.deepcopy(program) + payloads = [] + for plan in normalized.execution_plan: + for delegate in plan.delegates: + payload = _delegate_payload(normalized, delegate) + if delegate.id == CUDA_BACKEND_ID: + payload = b"CUDA_AOTI_VARIANTS" + delegate.compile_specs = [ + spec + for spec in delegate.compile_specs + if spec.key != "cuda_include_ptx" + ] + delegate.processed = BackendDelegateDataReference( + location=DataLocation.INLINE, index=len(payloads) + ) + payloads.append(BackendDelegateInlineData(data=payload)) + normalized.backend_delegate_data = payloads + return normalized + + +def _entry_map(entries: Iterable[CudaWeightEntry]) -> Dict[str, CudaWeightEntry]: + result = {} + for entry in entries: + if entry.fqn in result: + raise ValueError(f"Duplicate CUDA weight FQN {entry.fqn!r}") + result[entry.fqn] = entry + return result + + +def _entry_without_storage_key(entry: CudaWeightEntry) -> Tuple[object, ...]: + return ( + entry.fqn, + entry.storage_nbytes, + entry.dtype, + entry.device_type, + entry.storage_offset, + entry.sizes, + entry.strides, + ) + + +def _validate_shared_weights( + reference: _Artifact, + reference_metadata: CudaAotiMetadata, + candidate: _Artifact, + candidate_metadata: CudaAotiMetadata, + identity: Tuple[str, int], +) -> None: + reference_entries = _entry_map(reference_metadata.entries) + candidate_entries = _entry_map(candidate_metadata.entries) + if reference_entries.keys() != candidate_entries.keys(): + raise ValueError(f"CUDA weights differ for delegate {identity}: FQN mismatch") + + for fqn, reference_entry in reference_entries.items(): + candidate_entry = candidate_entries[fqn] + if _entry_without_storage_key(reference_entry) != _entry_without_storage_key( + candidate_entry + ): + raise ValueError( + f"CUDA weight metadata differs for delegate {identity}, FQN {fqn!r}" + ) + if ( + reference.blob_size(reference_entry.storage_key) + != reference_entry.storage_nbytes + ): + raise ValueError( + f"CUDA weight {fqn!r} has an invalid size in {reference.source.pte_path}" + ) + if ( + candidate.blob_size(candidate_entry.storage_key) + != candidate_entry.storage_nbytes + ): + raise ValueError( + f"CUDA weight {fqn!r} has an invalid size in {candidate.source.pte_path}" + ) + if reference.blob_sha256(reference_entry.storage_key) != candidate.blob_sha256( + candidate_entry.storage_key + ): + raise ValueError( + f"CUDA weight content differs for delegate {identity}, FQN {fqn!r}" + ) + + +def _cuda_so_keys(artifact: _Artifact) -> set[str]: + return { + variant.so_blob_key + for delegate in artifact.delegates + for variant in delegate.metadata.variants + } + + +def _cuda_weight_keys(artifact: _Artifact) -> set[str]: + return { + entry.storage_key + for delegate in artifact.delegates + for entry in delegate.metadata.entries + } + + +def _validate_programs(reference: _Artifact, candidate: _Artifact) -> None: + if _normalized_program(reference.pte.program) != _normalized_program( + candidate.pte.program + ): + raise ValueError( + f"ExecuTorch programs differ between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + if reference.pte.mutable_data != candidate.pte.mutable_data: + raise ValueError( + f"Mutable program data differs between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + + reference_non_cuda = { + key: value + for key, value in reference.pte_named_data.items() + if key not in _cuda_so_keys(reference) + and key not in _cuda_weight_keys(reference) + } + candidate_non_cuda = { + key: value + for key, value in candidate.pte_named_data.items() + if key not in _cuda_so_keys(candidate) + and key not in _cuda_weight_keys(candidate) + } + if reference_non_cuda != candidate_non_cuda: + raise ValueError( + f"Non-CUDA named data differs between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + + reference_external = ( + reference.ptd.keys() - _cuda_weight_keys(reference) + if reference.ptd is not None + else set() + ) + candidate_external = ( + candidate.ptd.keys() - _cuda_weight_keys(candidate) + if candidate.ptd is not None + else set() + ) + if reference_external != candidate_external: + raise ValueError( + f"Non-CUDA external data differs between {reference.source.pte_path} and " + f"{candidate.source.pte_path}" + ) + for key in reference_external: + if reference.blob_size(key) != candidate.blob_size( + key + ) or reference.blob_sha256(key) != candidate.blob_sha256(key): + raise ValueError( + f"External data {key!r} differs between " + f"{reference.source.pte_path} and {candidate.source.pte_path}" + ) + + +def _compact_delegate_data(program: Program) -> None: + payloads = [] + for plan in program.execution_plan: + for delegate in plan.delegates: + payload = _delegate_payload(program, delegate) + delegate.processed = BackendDelegateDataReference( + location=DataLocation.INLINE, index=len(payloads) + ) + payloads.append(BackendDelegateInlineData(data=payload)) + program.backend_delegate_data = payloads + + +def _merge_delegate_variants( + regular_artifacts: Sequence[_Artifact], + regular_delegates: Sequence[Dict[Tuple[str, int], CudaAotiMetadata]], + fallback_artifact: Optional[_Artifact], + fallback_delegates: Optional[Dict[Tuple[str, int], CudaAotiMetadata]], + reference_metadata: CudaAotiMetadata, + identity: Tuple[str, int], + merged_store: NamedDataStore, + provenance: List[CudaPteProvenance], +) -> List[CudaAotiVariant]: + variants = [] + target_sms = set() + reference = regular_artifacts[0] + for artifact, delegates in zip(regular_artifacts, regular_delegates): + metadata = delegates[identity] + _validate_shared_weights( + reference, reference_metadata, artifact, metadata, identity + ) + for variant in metadata.variants: + if variant.target_sm in target_sms: + raise ValueError(f"Duplicate CUDA target sm{variant.target_sm}") + target_sms.add(variant.target_sm) + try: + so_data = artifact.pte_named_data[variant.so_blob_key] + except KeyError as error: + raise ValueError( + f"{artifact.source.pte_path} does not contain CUDA SO " + f"{variant.so_blob_key!r}" + ) from error + merged_store.add_named_data(variant.so_blob_key, so_data) + variants.append(replace(variant, ptx_compute=0, fallback_only=False)) + provenance.append( + CudaPteProvenance( + delegate=identity, + kind="cubin", + target_sm=variant.target_sm, + ptx_compute=0, + source_pte=artifact.source.pte_path, + ) + ) + + variants.sort(key=lambda variant: variant.target_sm) + if fallback_artifact is not None: + assert fallback_delegates is not None + metadata = fallback_delegates[identity] + _validate_shared_weights( + reference, reference_metadata, fallback_artifact, metadata, identity + ) + fallback_variants = [ + variant for variant in metadata.variants if variant.ptx_compute + ] + if len(fallback_variants) != 1: + raise ValueError( + f"Fallback PTE {fallback_artifact.source.pte_path} must contain " + f"exactly one PTX-capable variant for delegate {identity}" + ) + fallback = replace(fallback_variants[0], fallback_only=True) + try: + so_data = fallback_artifact.pte_named_data[fallback.so_blob_key] + except KeyError as error: + raise ValueError( + f"{fallback_artifact.source.pte_path} does not contain CUDA SO " + f"{fallback.so_blob_key!r}" + ) from error + merged_store.add_named_data(fallback.so_blob_key, so_data) + variants.append(fallback) + provenance.append( + CudaPteProvenance( + delegate=identity, + kind="ptx-fallback", + target_sm=fallback.target_sm, + ptx_compute=fallback.ptx_compute, + source_pte=fallback_artifact.source.pte_path, + ) + ) + return variants + + +def _load_merge_artifacts( + inputs: Sequence[CudaPteInput], fallback: Optional[CudaPteInput] +) -> Tuple[List[_Artifact], Optional[_Artifact], List[Tuple[str, int]]]: + regular_artifacts = [_load_artifact(source) for source in inputs] + fallback_artifact = _load_artifact(fallback) if fallback is not None else None + artifacts = [*regular_artifacts] + if fallback_artifact is not None: + artifacts.append(fallback_artifact) + + reference = regular_artifacts[0] + reference_identities = [delegate.identity for delegate in reference.delegates] + for candidate in artifacts[1:]: + _validate_programs(reference, candidate) + candidate_identities = [delegate.identity for delegate in candidate.delegates] + if candidate_identities != reference_identities: + raise ValueError( + f"CUDA delegate layout differs between {reference.source.pte_path} " + f"and {candidate.source.pte_path}" + ) + return regular_artifacts, fallback_artifact, reference_identities + + +def _prepare_merged_output(reference: _Artifact) -> Tuple[Program, NamedDataStore]: + merged_program = copy.deepcopy(reference.pte.program) + for plan in merged_program.execution_plan: + for delegate in plan.delegates: + if delegate.id == CUDA_BACKEND_ID: + delegate.compile_specs = [ + spec + for spec in delegate.compile_specs + if spec.key != "cuda_include_ptx" + ] + + merged_store = NamedDataStore() + if reference.pte.named_data is not None: + merged_store.merge_named_data_store(reference.pte.named_data) + return merged_program, merged_store + + +def merge_cuda_pte_files_with_provenance( + inputs: Sequence[CudaPteInput], fallback: Optional[CudaPteInput] = None +) -> CudaPteMergeResult: + """Merge exact-SM CUDA exports and an optional PTX-only fallback source.""" + if torch.version.hip is not None: + raise RuntimeError( + "CUDA PTE merging supports only NVIDIA CUDA and is not supported on ROCm" + ) + if not inputs: + raise ValueError("At least one regular CUDA PTE input is required") + if len(inputs) + int(fallback is not None) < 2: + raise ValueError("At least two CUDA PTE inputs are required") + regular_artifacts, fallback_artifact, reference_identities = _load_merge_artifacts( + inputs, fallback + ) + reference = regular_artifacts[0] + merged_program, merged_store = _prepare_merged_output(reference) + + regular_delegates = [ + {delegate.identity: delegate.metadata for delegate in artifact.delegates} + for artifact in regular_artifacts + ] + fallback_delegates = ( + { + delegate.identity: delegate.metadata + for delegate in fallback_artifact.delegates + } + if fallback_artifact is not None + else None + ) + expected_variants = None + provenance: List[CudaPteProvenance] = [] + for identity_index, identity in enumerate(reference_identities): + reference_metadata = reference.delegates[identity_index].metadata + variants = _merge_delegate_variants( + regular_artifacts, + regular_delegates, + fallback_artifact, + fallback_delegates, + reference_metadata, + identity, + merged_store, + provenance, + ) + current_variants = tuple( + (variant.target_sm, variant.ptx_compute, variant.fallback_only) + for variant in variants + ) + if expected_variants is None: + expected_variants = current_variants + elif current_variants != expected_variants: + raise ValueError( + f"CUDA target variants differ across delegates at {identity}" + ) + merged_payload = encode_cuda_aoti_metadata(variants, reference_metadata.entries) + plan_name, delegate_index = identity + plan = next( + plan for plan in merged_program.execution_plan if plan.name == plan_name + ) + delegate = plan.delegates[delegate_index] + delegate.processed = BackendDelegateDataReference( + location=DataLocation.INLINE, + index=len(merged_program.backend_delegate_data), + ) + merged_program.backend_delegate_data.append( + BackendDelegateInlineData(data=merged_payload) + ) + + _compact_delegate_data(merged_program) + return CudaPteMergeResult( + pte=serialize_pte_binary( + PTEFile( + program=merged_program, + mutable_data=reference.pte.mutable_data, + named_data=merged_store.get_named_data_store_output(), + ), + extract_delegate_segments=True, + ), + provenance=tuple(provenance), + ) + + +def merge_cuda_pte_files( + inputs: Sequence[CudaPteInput], fallback: Optional[CudaPteInput] = None +) -> Cord: + return merge_cuda_pte_files_with_provenance(inputs, fallback).pte + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=("Merge exact-SM CUDA PTEs with an optional explicit PTX fallback") + ) + parser.add_argument( + "--input-pte", + action="append", + required=True, + type=Path, + help=( + "Regular CUDA PTE contributing exact-SM native cubins; the first " + "input supplies common data" + ), + ) + parser.add_argument( + "--input-ptd", + action="append", + type=Path, + default=[], + help="PTD paired by position with --input-pte", + ) + parser.add_argument( + "--fallback-pte", + type=Path, + help="Single CUDA PTE contributing only the PTX runtime fallback", + ) + parser.add_argument( + "--fallback-ptd", + type=Path, + help="PTD paired with --fallback-pte", + ) + parser.add_argument("--output-pte", required=True, type=Path) + parser.add_argument("--output-ptd", type=Path) + return parser.parse_args() + + +def main() -> None: + """Command-line entry point for CUDA PTE merging.""" + args = _parse_args() + if args.input_ptd and len(args.input_ptd) != len(args.input_pte): + raise ValueError("--input-ptd must be provided once per --input-pte") + if args.input_ptd and args.output_ptd is None: + raise ValueError("--output-ptd is required when --input-ptd is provided") + if args.fallback_ptd is not None and args.fallback_pte is None: + raise ValueError("--fallback-ptd requires --fallback-pte") + sources = [ + CudaPteInput( + pte_path=pte_path, + ptd_path=args.input_ptd[index] if args.input_ptd else None, + ) + for index, pte_path in enumerate(args.input_pte) + ] + fallback = ( + CudaPteInput( + pte_path=args.fallback_pte, + ptd_path=args.fallback_ptd, + ) + if args.fallback_pte is not None + else None + ) + result = merge_cuda_pte_files_with_provenance(sources, fallback) + args.output_pte.parent.mkdir(parents=True, exist_ok=True) + with tempfile.NamedTemporaryFile( + dir=args.output_pte.parent, prefix=f".{args.output_pte.name}.", delete=False + ) as temporary: + temporary_path = Path(temporary.name) + result.pte.write_to_file(temporary) + os.replace(temporary_path, args.output_pte) + + if args.output_ptd is not None: + if not args.input_ptd: + raise ValueError("--output-ptd requires --input-ptd") + args.output_ptd.parent.mkdir(parents=True, exist_ok=True) + if args.input_ptd[0].resolve() != args.output_ptd.resolve(): + shutil.copyfile(args.input_ptd[0], args.output_ptd) + + print("Merged CUDA code provenance:") + print("delegate\tkind\ttarget\tsource PTE") + for entry in result.provenance: + delegate = f"{entry.delegate[0]}[{entry.delegate[1]}]" + if entry.kind == "cubin": + target = f"sm{entry.target_sm}" + source = str(entry.source_pte) + else: + target = f"compute_{entry.ptx_compute} (source sm{entry.target_sm})" + source = f"{entry.source_pte} [fallback]" + print(f"{delegate}\t{entry.kind}\t{target}\t{source}") + + +if __name__ == "__main__": + main() diff --git a/backends/cuda/runtime/cuda_backend.cpp b/backends/cuda/runtime/cuda_backend.cpp index 0f63f662c83..3e0c3175818 100644 --- a/backends/cuda/runtime/cuda_backend.cpp +++ b/backends/cuda/runtime/cuda_backend.cpp @@ -322,7 +322,46 @@ class ET_EXPERIMENTAL CudaBackend final CudaWeightCache::parse( processed->data(), processed->size(), fqn_weights), "Malformed CUDA FQN weight metadata"); - so_blob_key = fqn_weights.so_blob_key; + size_t variant_index = 0; + bool uses_ptx_fallback = false; + uint32_t current_sm = 0; + if (!(fqn_weights.variants.size() == 1 && + fqn_weights.variants[0].target_sm == 0)) { +#if defined(EXECUTORCH_USE_HIP) + ET_LOG( + Error, + "Multi-SM CUDA AOTI metadata is not supported by the ROCm runtime"); + return Error::NotSupported; +#else + int device_index = 0; + cudaDeviceProp device_properties{}; + ET_CUDA_CHECK_OR_RETURN_ERROR(cudaGetDevice(&device_index)); + ET_CUDA_CHECK_OR_RETURN_ERROR( + cudaGetDeviceProperties(&device_properties, device_index)); + current_sm = static_cast( + device_properties.major * 10 + device_properties.minor); + ET_CHECK_OK_OR_RETURN_ERROR( + CudaWeightCache::select_variant( + fqn_weights, current_sm, variant_index, uses_ptx_fallback), + "Failed to select a CUDA AOTI variant for sm%u", + current_sm); +#endif + } + const auto& variant = fqn_weights.variants[variant_index]; + so_blob_key = variant.so_blob_key; + if (variant.target_sm == 0) { + ET_LOG(Info, "Selected untargeted CUDA AOTI variant"); + } else if (uses_ptx_fallback) { + ET_LOG( + Info, + "Selected sm%u CUDA AOTI PTX fallback (compute_%u) for sm%u", + variant.target_sm, + variant.ptx_compute, + current_sm); + } else { + ET_LOG( + Info, "Selected native sm%u CUDA AOTI variant", variant.target_sm); + } } else { ET_CHECK_OK_OR_RETURN_ERROR( executorch::backends::aoti::resolve_blob_keys( diff --git a/backends/cuda/runtime/cuda_weight_cache.cpp b/backends/cuda/runtime/cuda_weight_cache.cpp index 09b4bd76312..5a2e465f243 100644 --- a/backends/cuda/runtime/cuda_weight_cache.cpp +++ b/backends/cuda/runtime/cuda_weight_cache.cpp @@ -152,9 +152,51 @@ Error CudaWeightCache::parse( } MetadataReader reader(data, size); - if (!reader.skip(kFormatMagicSize) || - !reader.read_string(metadata.so_blob_key) || - metadata.so_blob_key.empty()) { + if (!reader.skip(kFormatMagicSize)) { + return Error::InvalidProgram; + } + + metadata.variants.clear(); + uint32_t num_variants = 0; + constexpr uint32_t kMaxVariants = 256; + if (!reader.read_u32(num_variants) || num_variants == 0 || + num_variants > kMaxVariants) { + return Error::InvalidProgram; + } + metadata.variants.reserve(num_variants); + std::unordered_set target_sms; + bool found_fallback = false; + bool regular_has_ptx = false; + for (uint32_t index = 0; index < num_variants; ++index) { + Variant variant; + uint32_t flags = 0; + if (!reader.read_u32(variant.target_sm) || + !reader.read_u32(variant.ptx_compute) || + variant.ptx_compute > variant.target_sm || !reader.read_u32(flags) || + (flags & ~1U) != 0 || !reader.read_string(variant.so_blob_key) || + variant.so_blob_key.empty()) { + return Error::InvalidProgram; + } + variant.fallback_only = (flags & 1U) != 0; + if (variant.target_sm == 0) { + if (num_variants != 1 || variant.ptx_compute != 0 || + variant.fallback_only) { + return Error::InvalidProgram; + } + } else if (variant.fallback_only) { + if (variant.ptx_compute == 0 || found_fallback) { + return Error::InvalidProgram; + } + found_fallback = true; + } else { + if (!target_sms.emplace(variant.target_sm).second) { + return Error::InvalidProgram; + } + regular_has_ptx |= variant.ptx_compute != 0; + } + metadata.variants.push_back(std::move(variant)); + } + if (num_variants > 1 && regular_has_ptx) { return Error::InvalidProgram; } @@ -202,6 +244,45 @@ Error CudaWeightCache::parse( return reader.empty() ? Error::Ok : Error::InvalidProgram; } +Error CudaWeightCache::select_variant( + const Metadata& metadata, + uint32_t current_sm, + size_t& variant_index, + bool& uses_ptx_fallback) { + ET_CHECK_OR_RETURN_ERROR( + !metadata.variants.empty(), InvalidProgram, "CUDA AOTI has no variants"); + + if (metadata.variants.size() == 1 && metadata.variants[0].target_sm == 0) { + variant_index = 0; + uses_ptx_fallback = false; + return Error::Ok; + } + + for (size_t index = 0; index < metadata.variants.size(); ++index) { + if (!metadata.variants[index].fallback_only && + metadata.variants[index].target_sm == current_sm) { + variant_index = index; + uses_ptx_fallback = false; + return Error::Ok; + } + } + + for (size_t index = 0; index < metadata.variants.size(); ++index) { + const Variant& variant = metadata.variants[index]; + if (variant.ptx_compute != 0 && variant.ptx_compute <= current_sm && + (variant.fallback_only || metadata.variants.size() == 1)) { + variant_index = index; + uses_ptx_fallback = true; + return Error::Ok; + } + } + ET_CHECK_OR_RETURN_ERROR( + false, + NotSupported, + "CUDA AOTI has no native or PTX variant compatible with sm%u", + current_sm); +} + Error CudaWeightCache::validate_view(const Entry& entry) { uint64_t item_size = 0; switch (static_cast(entry.dtype)) { diff --git a/backends/cuda/runtime/cuda_weight_cache.h b/backends/cuda/runtime/cuda_weight_cache.h index eb58ebc8fb1..ee77f95bcd7 100644 --- a/backends/cuda/runtime/cuda_weight_cache.h +++ b/backends/cuda/runtime/cuda_weight_cache.h @@ -24,9 +24,16 @@ namespace executorch::backends::cuda { class CudaWeightCache final { public: - static constexpr char kFormatMagic[] = "ETCUDAFQN3"; + static constexpr char kFormatMagic[] = "ETCUDAFQN0"; static constexpr size_t kFormatMagicSize = sizeof(kFormatMagic) - 1; + struct Variant { + uint32_t target_sm{0}; + uint32_t ptx_compute{0}; + std::string so_blob_key; + bool fallback_only{false}; + }; + struct Entry { std::string fqn; std::string storage_key; @@ -39,7 +46,7 @@ class CudaWeightCache final { }; struct Metadata { - std::string so_blob_key; + std::vector variants; std::vector entries; }; @@ -48,6 +55,12 @@ class CudaWeightCache final { static runtime::Error parse(const void* data, size_t size, Metadata& metadata); + static runtime::Error select_variant( + const Metadata& metadata, + uint32_t current_sm, + size_t& variant_index, + bool& uses_ptx_fallback); + runtime::Error load( CudaDelegateHandle* handle, const runtime::NamedDataMap* named_data_map, diff --git a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp index d16e7143058..daa0a41bf0f 100644 --- a/backends/cuda/runtime/test/test_cuda_weight_cache.cpp +++ b/backends/cuda/runtime/test/test_cuda_weight_cache.cpp @@ -43,6 +43,10 @@ std::vector serialized_metadata( cuda::CudaWeightCache::kFormatMagic, cuda::CudaWeightCache::kFormatMagic + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(output, 1); // variants + append_u32(output, 0); // untargeted (ROCm) + append_u32(output, 0); // no PTX + append_u32(output, 0); // regular append_string(output, "so-key"); append_u32(output, 1); // entries append_string(output, "model.weight"); @@ -59,9 +63,60 @@ std::vector serialized_metadata( return output; } +std::vector serialized_multi_arch_metadata() { + std::vector output( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(output, 3); // variants + append_u32(output, 80); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm80-so"); + append_u32(output, 90); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm90-so"); + append_u32(output, 120); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm120-so"); + append_u32(output, 1); // entries + append_string(output, "model.weight"); + append_string(output, "storage-key"); + append_u64(output, 24); + append_u32(output, 6); + append_u32(output, 1); + append_u64(output, 0); + append_u32(output, 2); + append_u64(output, 2); + append_u64(output, 3); + append_u64(output, 3); + append_u64(output, 1); + return output; +} + +std::vector serialized_fallback_metadata() { + std::vector output( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(output, 2); // variants + append_u32(output, 80); + append_u32(output, 0); + append_u32(output, 0); // regular + append_string(output, "sm80-so"); + append_u32(output, 80); + append_u32(output, 80); + append_u32(output, 1); // fallback only + append_string(output, "fallback-so"); + append_u32(output, 0); // entries + return output; +} + } // namespace -TEST(CudaWeightCacheTest, LegacyPayloadIsNotMisdetected) { +TEST(CudaWeightCacheTest, RawAotiPayloadIsNotMisdetected) { const std::string legacy = "so-key\nweights-key"; EXPECT_FALSE( cuda::CudaWeightCache::is_serialized(legacy.data(), legacy.size())); @@ -73,7 +128,9 @@ TEST(CudaWeightCacheTest, ParsesSerializedMetadata) { ASSERT_EQ( cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), Error::Ok); - ASSERT_EQ(metadata.so_blob_key, "so-key"); + ASSERT_EQ(metadata.variants.size(), 1u); + EXPECT_EQ(metadata.variants[0].target_sm, 0u); + EXPECT_EQ(metadata.variants[0].so_blob_key, "so-key"); ASSERT_EQ(metadata.entries.size(), 1u); const auto& entry = metadata.entries[0]; EXPECT_EQ(entry.fqn, "model.weight"); @@ -85,6 +142,126 @@ TEST(CudaWeightCacheTest, ParsesSerializedMetadata) { EXPECT_EQ(entry.strides, (std::vector{3, 1})); } +TEST(CudaWeightCacheTest, ParsesAndSelectsMultiArchMetadata) { + const std::vector bytes = serialized_multi_arch_metadata(); + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.variants.size(), 3u); + EXPECT_EQ(metadata.variants[0].target_sm, 80u); + EXPECT_EQ(metadata.variants[2].so_blob_key, "sm120-so"); + + size_t variant_index = 0; + bool uses_ptx_fallback = false; + EXPECT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 120, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 2u); + EXPECT_FALSE(uses_ptx_fallback); + + EXPECT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 100, variant_index, uses_ptx_fallback), + Error::NotSupported); +} + +TEST(CudaWeightCacheTest, SelectsPtxFromPortableSingleVariant) { + cuda::CudaWeightCache::Metadata metadata; + metadata.variants = {{80, 80, "sm80-so"}}; + size_t variant_index = 0; + bool uses_ptx_fallback = false; + ASSERT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 100, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 0u); + EXPECT_TRUE(uses_ptx_fallback); +} + +TEST(CudaWeightCacheTest, FallbackOnlyVariantNeverWinsNativeMatch) { + const std::vector bytes = serialized_fallback_metadata(); + cuda::CudaWeightCache::Metadata metadata; + ASSERT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::Ok); + ASSERT_EQ(metadata.variants.size(), 2u); + EXPECT_FALSE(metadata.variants[0].fallback_only); + EXPECT_TRUE(metadata.variants[1].fallback_only); + + size_t variant_index = 0; + bool uses_ptx_fallback = false; + ASSERT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 80, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 0u); + EXPECT_FALSE(uses_ptx_fallback); + + ASSERT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 90, variant_index, uses_ptx_fallback), + Error::Ok); + EXPECT_EQ(variant_index, 1u); + EXPECT_TRUE(uses_ptx_fallback); +} + +TEST(CudaWeightCacheTest, RejectsWhenNoVariantIsCompatible) { + cuda::CudaWeightCache::Metadata metadata; + metadata.variants = {{120, 0, "sm120-so"}}; + size_t variant_index = 0; + bool uses_ptx_fallback = false; + EXPECT_EQ( + cuda::CudaWeightCache::select_variant( + metadata, 90, variant_index, uses_ptx_fallback), + Error::NotSupported); +} + +TEST(CudaWeightCacheTest, RejectsDuplicateMultiArchTarget) { + std::vector bytes( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(bytes, 2); + append_u32(bytes, 80); + append_u32(bytes, 0); + append_u32(bytes, 0); + append_string(bytes, "first-so"); + append_u32(bytes, 80); + append_u32(bytes, 0); + append_u32(bytes, 0); + append_string(bytes, "second-so"); + append_u32(bytes, 0); + + cuda::CudaWeightCache::Metadata metadata; + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} + +TEST(CudaWeightCacheTest, RejectsImplicitPtxFallbackInMultiVariantMetadata) { + std::vector bytes( + cuda::CudaWeightCache::kFormatMagic, + cuda::CudaWeightCache::kFormatMagic + + cuda::CudaWeightCache::kFormatMagicSize); + append_u32(bytes, 2); + append_u32(bytes, 80); + append_u32(bytes, 80); + append_u32(bytes, 0); + append_string(bytes, "sm80-so"); + append_u32(bytes, 120); + append_u32(bytes, 0); + append_u32(bytes, 0); + append_string(bytes, "sm120-so"); + append_u32(bytes, 0); + + cuda::CudaWeightCache::Metadata metadata; + EXPECT_EQ( + cuda::CudaWeightCache::parse(bytes.data(), bytes.size(), metadata), + Error::InvalidProgram); +} + TEST(CudaWeightCacheTest, RejectsTruncationAndTrailingData) { std::vector bytes = serialized_metadata(); cuda::CudaWeightCache::Metadata metadata; diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index eda5e46de41..7b45e677da2 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -133,6 +133,33 @@ def test_target_smem_context_only_patches_exact_triton_limit(self): self.assertIs(triton_compiler.max_shared_mem, local_max_shared_mem) + def test_cuda_include_ptx_compile_spec(self): + with mock.patch.object( + CudaBackend, "_setup_cuda_environment_for_fatbin", return_value=True + ): + options = CudaBackend.get_aoti_compile_options( + [CompileSpec(key="cuda_include_ptx", value=b"ON")] + ) + + self.assertTrue(options["aot_inductor.emit_multi_arch_kernel"]) + + def test_cuda_include_ptx_off_disables_multi_arch_kernel(self): + with mock.patch.object( + CudaBackend, "_setup_cuda_environment_for_fatbin" + ) as setup_fatbin: + options = CudaBackend.get_aoti_compile_options( + [CompileSpec(key="cuda_include_ptx", value=b"OFF")] + ) + + setup_fatbin.assert_not_called() + self.assertFalse(options["aot_inductor.emit_multi_arch_kernel"]) + + def test_invalid_cuda_include_ptx_compile_spec(self): + with self.assertRaisesRegex(ValueError, "Invalid cuda_include_ptx"): + CudaBackend.get_aoti_compile_options( + [CompileSpec(key="cuda_include_ptx", value=b"MAYBE")] + ) + class TestCudaExport(unittest.TestCase): """Test CUDA export functionality for various operations using to_edge_transform_and_lower.""" diff --git a/backends/cuda/tests/test_cuda_partitioner.py b/backends/cuda/tests/test_cuda_partitioner.py index 153828b2cb0..9a1fafaab68 100644 --- a/backends/cuda/tests/test_cuda_partitioner.py +++ b/backends/cuda/tests/test_cuda_partitioner.py @@ -21,9 +21,10 @@ from executorch.backends.cuda.cuda_weight_collector import ( AOTI_DEVICE_TYPE_CPU, AOTI_DEVICE_TYPE_CUDA, - CUDA_WEIGHT_CACHE_MAGIC, + CUDA_AOTI_METADATA_MAGIC, + CudaAotiVariant, CudaWeightCollector, - encode_cuda_weight_metadata, + encode_cuda_aoti_metadata, ) from executorch.exir._serialize._cord import FileBackedData from executorch.exir._serialize._named_data_store import NamedDataStore @@ -95,8 +96,10 @@ def test_weights_are_materialized_as_independent_storages(self) -> None: artifact.storages[artifact.entries[1].storage_key].to_bytes(), ) - metadata = encode_cuda_weight_metadata("so-key", artifact.entries) - self.assertTrue(metadata.startswith(CUDA_WEIGHT_CACHE_MAGIC)) + metadata = encode_cuda_aoti_metadata( + [CudaAotiVariant(0, 0, "so-key")], artifact.entries + ) + self.assertTrue(metadata.startswith(CUDA_AOTI_METADATA_MAGIC)) self.assertIn(b"first", metadata) self.assertIn(b"second", metadata) for storage in artifact.storages.values(): diff --git a/backends/cuda/tests/test_cuda_weight_metadata.py b/backends/cuda/tests/test_cuda_weight_metadata.py new file mode 100644 index 00000000000..d7fb1d93179 --- /dev/null +++ b/backends/cuda/tests/test_cuda_weight_metadata.py @@ -0,0 +1,130 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CUDA, + CUDA_AOTI_METADATA_MAGIC, + CudaAotiVariant, + CudaWeightEntry, + decode_cuda_aoti_metadata, + encode_cuda_aoti_metadata, +) + + +class TestCudaWeightMetadata(unittest.TestCase): + @staticmethod + def _entry() -> CudaWeightEntry: + return CudaWeightEntry( + fqn="model.weight", + storage_key="cuda_fqn_weight:cuda:model.weight", + storage_nbytes=24, + dtype=6, + device_type=AOTI_DEVICE_TYPE_CUDA, + storage_offset=0, + sizes=(2, 3), + strides=(3, 1), + ) + + def test_targeted_metadata_has_shared_weights(self) -> None: + entry = self._entry() + encoded = encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(120, 0, "sm120-so"), + ], + [entry], + ) + self.assertTrue(encoded.startswith(CUDA_AOTI_METADATA_MAGIC)) + decoded = decode_cuda_aoti_metadata(encoded) + self.assertEqual( + decoded.variants, + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(120, 0, "sm120-so"), + ], + ) + self.assertEqual(decoded.entries, [entry]) + + def test_metadata_rejects_duplicate_target(self) -> None: + with self.assertRaisesRegex(ValueError, "Duplicate CUDA target SM"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 0, "first"), + CudaAotiVariant(80, 0, "second"), + ], + [self._entry()], + ) + + def test_fallback_metadata_allows_matching_regular_target(self) -> None: + entry = self._entry() + encoded = encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(80, 80, "fallback-so", fallback_only=True), + ], + [entry], + ) + self.assertTrue(encoded.startswith(CUDA_AOTI_METADATA_MAGIC)) + decoded = decode_cuda_aoti_metadata(encoded) + self.assertEqual( + decoded.variants, + [ + CudaAotiVariant(80, 0, "sm80-so"), + CudaAotiVariant(80, 80, "fallback-so", fallback_only=True), + ], + ) + + def test_fallback_metadata_rejects_multiple_fallbacks(self) -> None: + with self.assertRaisesRegex(ValueError, "only one fallback"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 80, "first", fallback_only=True), + CudaAotiVariant(75, 75, "second", fallback_only=True), + ], + [self._entry()], + ) + + def test_multi_variant_metadata_rejects_implicit_ptx_fallback(self) -> None: + with self.assertRaisesRegex(ValueError, "explicit PTX fallback"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(80, 80, "sm80-so"), + CudaAotiVariant(120, 0, "sm120-so"), + ], + [self._entry()], + ) + + def test_untargeted_metadata_for_rocm(self) -> None: + encoded = encode_cuda_aoti_metadata( + [CudaAotiVariant(0, 0, "rocm-so")], [self._entry()] + ) + self.assertTrue(encoded.startswith(CUDA_AOTI_METADATA_MAGIC)) + decoded = decode_cuda_aoti_metadata(encoded) + self.assertEqual(decoded.variants, [CudaAotiVariant(0, 0, "rocm-so")]) + self.assertEqual(decoded.entries, [self._entry()]) + + def test_untargeted_metadata_cannot_mix_with_targeted_variants(self) -> None: + with self.assertRaisesRegex(ValueError, "requires one non-fallback variant"): + encode_cuda_aoti_metadata( + [ + CudaAotiVariant(0, 0, "rocm-so"), + CudaAotiVariant(80, 0, "sm80-so"), + ], + [self._entry()], + ) + + def test_metadata_rejects_trailing_data(self) -> None: + encoded = encode_cuda_aoti_metadata( + [CudaAotiVariant(80, 80, "sm80-so")], [self._entry()] + ) + with self.assertRaisesRegex(ValueError, "trailing bytes"): + decode_cuda_aoti_metadata(encoded + b"\0") + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/cuda/tests/test_merge_ptes.py b/backends/cuda/tests/test_merge_ptes.py new file mode 100644 index 00000000000..7872f066321 --- /dev/null +++ b/backends/cuda/tests/test_merge_ptes.py @@ -0,0 +1,387 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import hashlib +import io +import tempfile +import unittest +from contextlib import redirect_stdout +from pathlib import Path +from unittest.mock import patch + +from executorch.backends.cuda.cuda_weight_collector import ( + AOTI_DEVICE_TYPE_CUDA, + CudaAotiVariant, + CudaWeightEntry, + decode_cuda_aoti_metadata, + encode_cuda_aoti_metadata, +) +from executorch.backends.cuda.merge_ptes import ( + CudaPteInput, + main as merge_main, + merge_cuda_pte_files, +) +from executorch.exir._serialize._named_data_store import NamedDataStore +from executorch.exir._serialize._program import ( + deserialize_pte_binary, + PTEFile, + serialize_pte_binary, +) +from executorch.exir._serialize.data_serializer import DataEntry, DataPayload +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.schema import ( + BackendDelegate, + BackendDelegateDataReference, + BackendDelegateInlineData, + ContainerMetadata, + DataLocation, + ExecutionPlan, + Program, + SubsegmentOffsets, +) +from executorch.extension.flat_tensor.serialize.serialize import FlatTensorSerializer + + +@patch("executorch.backends.cuda.merge_ptes.torch.version.hip", None) +class TestMergeCudaPtes(unittest.TestCase): + def _write_artifact( + self, + directory: Path, + target_sm: int, + so_data: bytes, + weight_data: bytes, + *, + fqn: str = "model.weight", + weight_key: str = "cuda_fqn_weight:cuda:model.weight", + ptx_compute: int | None = None, + ) -> CudaPteInput: + if ptx_compute is None: + ptx_compute = target_sm + so_key = hashlib.sha256(so_data).hexdigest() + "_so_blob" + entry = CudaWeightEntry( + fqn=fqn, + storage_key=weight_key, + storage_nbytes=len(weight_data), + dtype=1, + device_type=AOTI_DEVICE_TYPE_CUDA, + storage_offset=0, + sizes=(len(weight_data),), + strides=(1,), + ) + metadata = encode_cuda_aoti_metadata( + [CudaAotiVariant(target_sm, ptx_compute, so_key)], [entry] + ) + compile_specs = [CompileSpec("method_name", b"forward")] + compile_specs.append( + CompileSpec("cuda_include_ptx", b"ON" if ptx_compute else b"OFF") + ) + delegate = BackendDelegate( + id="CudaBackend", + processed=BackendDelegateDataReference(DataLocation.INLINE, 0), + compile_specs=compile_specs, + ) + program = Program( + version=0, + execution_plan=[ + ExecutionPlan( + name="forward", + container_meta_type=ContainerMetadata("", ""), + values=[], + inputs=[], + outputs=[], + chains=[], + operators=[], + delegates=[delegate], + non_const_buffer_sizes=[], + ) + ], + constant_buffer=[], + backend_delegate_data=[BackendDelegateInlineData(metadata)], + segments=[], + constant_segment=SubsegmentOffsets(0, []), + ) + store = NamedDataStore() + store.add_named_data(so_key, so_data) + pte_path = directory / "model.pte" + with pte_path.open("wb") as output: + serialize_pte_binary( + PTEFile( + program=program, named_data=store.get_named_data_store_output() + ), + extract_delegate_segments=True, + ).write_to_file(output) + + ptd_path = directory / "aoti_cuda_blob.ptd" + serializer = FlatTensorSerializer() + with ptd_path.open("wb") as output: + serializer.serialize( + DataPayload( + buffers=[weight_data], + named_data={weight_key: DataEntry(0, 1, None)}, + ) + ).write_to_file(output) + return CudaPteInput( + pte_path=pte_path, + ptd_path=ptd_path, + ) + + def test_merges_variants_and_keeps_one_weight_manifest(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"shared-weight"), + self._write_artifact(sm120, 120, b"sm120-so", b"shared-weight"), + ] + + merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) + delegate = merged.program.execution_plan[0].delegates[0] + payload = merged.program.backend_delegate_data[ + delegate.processed.index + ].data + metadata = decode_cuda_aoti_metadata(payload) + self.assertEqual( + [variant.target_sm for variant in metadata.variants], [80, 120] + ) + self.assertEqual( + [variant.ptx_compute for variant in metadata.variants], [0, 0] + ) + self.assertFalse( + any(variant.fallback_only for variant in metadata.variants) + ) + self.assertEqual(len(metadata.entries), 1) + self.assertEqual(metadata.entries[0].fqn, "model.weight") + self.assertEqual( + set(merged.named_data.pte_data), + { + hashlib.sha256(b"sm80-so").hexdigest() + "_so_blob", + hashlib.sha256(b"sm120-so").hexdigest() + "_so_blob", + }, + ) + self.assertNotIn( + "cuda_include_ptx", + { + spec.key + for spec in merged.program.execution_plan[0] + .delegates[0] + .compile_specs + }, + ) + + def test_uses_ptx_only_from_explicit_fallback(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + fallback_dir = root / "fallback" + sm80.mkdir() + sm120.mkdir() + fallback_dir.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"weight"), + self._write_artifact( + sm120, + 120, + b"sm120-so", + b"weight", + ptx_compute=0, + ), + ] + fallback = self._write_artifact(fallback_dir, 75, b"fallback-so", b"weight") + + merged = deserialize_pte_binary( + bytes(merge_cuda_pte_files(inputs, fallback)) + ) + delegate = merged.program.execution_plan[0].delegates[0] + metadata = decode_cuda_aoti_metadata( + merged.program.backend_delegate_data[delegate.processed.index].data + ) + self.assertEqual( + [variant.target_sm for variant in metadata.variants], [80, 120, 75] + ) + self.assertEqual( + [variant.ptx_compute for variant in metadata.variants], [0, 0, 75] + ) + self.assertEqual( + [variant.fallback_only for variant in metadata.variants], + [False, False, True], + ) + + def test_preserves_no_ptx_fallback(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"weight", ptx_compute=0), + self._write_artifact( + sm120, + 120, + b"sm120-so", + b"weight", + ptx_compute=0, + ), + ] + + merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) + delegate = merged.program.execution_plan[0].delegates[0] + metadata = decode_cuda_aoti_metadata( + merged.program.backend_delegate_data[delegate.processed.index].data + ) + self.assertEqual( + [variant.ptx_compute for variant in metadata.variants], [0, 0] + ) + + def test_rejects_different_weight_content(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact(sm80, 80, b"sm80-so", b"first-weight"), + self._write_artifact(sm120, 120, b"sm120-so", b"other-weight"), + ] + with self.assertRaisesRegex(ValueError, "weight content differs"): + merge_cuda_pte_files(inputs) + + def test_library_local_weight_keys_are_normalized_to_base(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + sm80.mkdir() + sm120.mkdir() + inputs = [ + self._write_artifact( + sm80, + 80, + b"sm80-so", + b"constant", + fqn="_tensor_constant0", + weight_key="cuda_fqn_weight:cuda:sm80-so:_tensor_constant0", + ), + self._write_artifact( + sm120, + 120, + b"sm120-so", + b"constant", + fqn="_tensor_constant0", + weight_key="cuda_fqn_weight:cuda:sm120-so:_tensor_constant0", + ), + ] + + merged = deserialize_pte_binary(bytes(merge_cuda_pte_files(inputs))) + delegate = merged.program.execution_plan[0].delegates[0] + metadata = decode_cuda_aoti_metadata( + merged.program.backend_delegate_data[delegate.processed.index].data + ) + self.assertEqual( + metadata.entries[0].storage_key, + "cuda_fqn_weight:cuda:sm80-so:_tensor_constant0", + ) + + def test_rejects_duplicate_target_sm(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + first = root / "first" + second = root / "second" + first.mkdir() + second.mkdir() + inputs = [ + self._write_artifact(first, 80, b"first-so", b"weight"), + self._write_artifact(second, 80, b"second-so", b"weight"), + ] + with self.assertRaisesRegex(ValueError, "Duplicate CUDA target sm80"): + merge_cuda_pte_files(inputs) + + def test_rejects_fallback_without_ptx(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + regular_dir = root / "regular" + fallback_dir = root / "fallback" + regular_dir.mkdir() + fallback_dir.mkdir() + regular = self._write_artifact(regular_dir, 80, b"sm80-so", b"weight") + fallback = self._write_artifact( + fallback_dir, + 75, + b"fallback-so", + b"weight", + ptx_compute=0, + ) + + with self.assertRaisesRegex(ValueError, "exactly one PTX-capable"): + merge_cuda_pte_files([regular], fallback) + + def test_rejects_rocm(self) -> None: + with patch( + "executorch.backends.cuda.merge_ptes.torch.version.hip", "6.3" + ), self.assertRaisesRegex(RuntimeError, "only NVIDIA CUDA"): + merge_cuda_pte_files([]) + + def test_cli_writes_merged_pte_and_reuses_base_ptd(self) -> None: + with tempfile.TemporaryDirectory() as temporary: + root = Path(temporary) + sm80 = root / "sm80" + sm120 = root / "sm120" + fallback_dir = root / "fallback" + output = root / "output" + sm80.mkdir() + sm120.mkdir() + fallback_dir.mkdir() + first = self._write_artifact(sm80, 80, b"sm80-so", b"weight") + second = self._write_artifact(sm120, 120, b"sm120-so", b"weight") + fallback = self._write_artifact(fallback_dir, 75, b"fallback-so", b"weight") + self.assertIsNotNone(first.ptd_path) + output_pte = output / "model.pte" + output_ptd = output / "aoti_cuda_blob.ptd" + + with patch( + "sys.argv", + [ + "merge_ptes", + "--input-pte", + str(first.pte_path), + "--input-pte", + str(second.pte_path), + "--input-ptd", + str(first.ptd_path), + "--input-ptd", + str(second.ptd_path), + "--fallback-pte", + str(fallback.pte_path), + "--fallback-ptd", + str(fallback.ptd_path), + "--output-pte", + str(output_pte), + "--output-ptd", + str(output_ptd), + ], + ), redirect_stdout(io.StringIO()) as stdout: + merge_main() + + self.assertTrue(output_pte.is_file()) + assert first.ptd_path is not None + self.assertEqual(output_ptd.read_bytes(), first.ptd_path.read_bytes()) + report = stdout.getvalue() + self.assertIn(f"cubin\tsm80\t{first.pte_path}", report) + self.assertIn(f"cubin\tsm120\t{second.pte_path}", report) + self.assertIn( + f"ptx-fallback\tcompute_75 (source sm75)\t" + f"{fallback.pte_path} [fallback]", + report, + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/cuda/README.md b/examples/cuda/README.md index a5421edb035..76831505be5 100644 --- a/examples/cuda/README.md +++ b/examples/cuda/README.md @@ -35,3 +35,42 @@ installed PyTorch build does not list in `torch.cuda.get_arch_list()`. The example emits `amd_triton.pte` and `aoti_cuda_blob.ptd`. It uses a fresh Inductor cache and fails unless it finds generated Triton source there, and it checks that the `.pte` embeds a code object for the architecture it compiled for. + +## Merge native NVIDIA GPU exports + +CUDA AOTI exports record their compiled target SM. Exports of the same program +and weights can be combined so the runtime selects an exactly matching native +AOTI library, or uses one explicitly designated PTX fallback. + +```bash +python -m executorch.backends.cuda.merge_ptes \ + --input-pte a100/model.pte \ + --input-pte rtx5090/model.pte \ + --input-ptd a100/aoti_cuda_blob.ptd \ + --input-ptd rtx5090/aoti_cuda_blob.ptd \ + --fallback-pte portable/model.pte \ + --fallback-ptd portable/aoti_cuda_blob.ptd \ + --output-pte merged/model.pte \ + --output-ptd merged/aoti_cuda_blob.ptd +``` + +The inputs must come from the same ExecuTorch program and contain identical +weights. Each regular `--input-pte` contributes only exact-SM native cubins; +any PTX capability in a regular input is ignored. At most one +`--fallback-pte` may be provided, and it must contain exactly one PTX-capable +variant. The runtime uses it only when no regular input provides a native cubin +for the current SM. The output PTD reuses one validated copy of the weights. +After merging, the tool prints every native SM and PTX fallback together with +its source PTE. + +Export every regular input with PTX disabled: + +```python +CompileSpec("cuda_include_ptx", b"OFF") +``` + +Export the fallback with PTX enabled and with any portability constraints, such +as a shared-memory limit, required by its target GPU set. AOTI host code and its +CUDA fatbin are linked into one shared library, so the merge step does not +rewrite ELF sections. Instead, the merged metadata makes regular libraries +native-only and marks the fallback library as PTX-only for runtime selection. diff --git a/examples/cuda/scripts/export.py b/examples/cuda/scripts/export.py index ee4390bf938..8eb9c651b24 100644 --- a/examples/cuda/scripts/export.py +++ b/examples/cuda/scripts/export.py @@ -19,6 +19,7 @@ from executorch.examples.models.model_factory import EagerModelFactory from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower +from executorch.exir.backend.compile_spec_schema import CompileSpec from executorch.extension.export_util.utils import save_pte_program @@ -51,6 +52,16 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--generate_etrecord", action=argparse.BooleanOptionalAction) parser.add_argument("--save_processed_bytes", action=argparse.BooleanOptionalAction) + parser.add_argument( + "--cuda_include_ptx", + choices=("ON", "OFF"), + help="Explicitly enable or disable PTX in the exported CUDA AOTI library", + ) + parser.add_argument( + "--seed", + type=int, + help="Seed model initialization and example input generation", + ) args = parser.parse_args() return args @@ -73,6 +84,9 @@ def main(): f"Available models are {list(MODEL_NAME_TO_MODEL.keys())}." ) + if args.seed is not None: + torch.manual_seed(args.seed) + ( model, example_args, @@ -87,9 +101,12 @@ def main(): dynamic_shapes=dynamic_shapes, ) - partitioner = CudaPartitioner( - [CudaBackend.generate_method_name_compile_spec(args.model_name)] - ) + compile_specs = [CudaBackend.generate_method_name_compile_spec(args.model_name)] + if args.cuda_include_ptx is not None: + compile_specs.append( + CompileSpec("cuda_include_ptx", args.cuda_include_ptx.encode()) + ) + partitioner = CudaPartitioner(compile_specs) et_prog = to_edge_transform_and_lower( exported_programs, From 04f900291eabf68a3eb501328a3a31c5c2285c39 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= Date: Tue, 1 Sep 2026 14:46:29 +0200 Subject: [PATCH 122/190] [ET-VK] Fix squeeze_copy of the outermost dim under dynamic shapes add_squeeze_copy_dims_node() skips dim 0 and falls back to add_clone_node(). resize_clone_node() only propagates sizes when input and output have the same dim count, which a squeeze never does, so the output keeps the extents it was built with. With static shapes that is invisible. With dynamic shapes the output holds its upper-bound extents while consumers read it at the real size, so the copy lands in the wrong places and roughly half the output comes back zeroed -- silently, with no error. Route dim 0 through the permute path like every other squeeze dim; resize_permute_node() already has an explicit branch for the rank-reducing case. Repro: any model that ends up with torch.cat(list(x), -1) over a rank-4 tensor with a dynamic dim. The unbind lowers to slice_copy plus squeeze_copy.dims, and the second slice comes back zeroed for every extent below the bound. Reduced to a 15-line case: y[1:2] is correct while y[1:2].squeeze(0) returns exactly half zeros (cosine 0.704 = sqrt of 0.5 against the reference), correct only at the bound. Verified on a Galaxy S26 Ultra (Adreno 840): the reduced case goes from 0.704 to 1.000000 at extents 200, 500 and 1000, and a TTS model whose classifier-free-guidance batch is built this way goes from cosine 0.36 to 0.99993 against its CPU reference. --- .../vulkan/runtime/graph/ops/impl/Squeeze.cpp | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp index fce7600a035..a78456538bf 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp @@ -23,21 +23,27 @@ void add_squeeze_copy_dims_node( const ValueRef out) { const int64_t in_dim = graph.dim_of(in); const std::vector in_sizes = graph.sizes_of(in); - const std::vector out_sizes = graph.sizes_of(in); const std::vector dims = graph.extract_int_or_symint_list(dims_ref); std::vector squeeze_dims; - // Filter out edge cases that we don't need squeeze: - // 1. The size of squeeze dim is larger than 1. - // 2. Squeeze outter most dim - // For these cases, just pass input to output via clone. + // Filter out the edge case that we don't need to squeeze: the size of the + // squeeze dim is larger than 1. For that case, just pass input to output via + // clone. + // + // Note that the outermost dim must NOT be excluded here. Routing it to + // add_clone_node() leaves the output unresized at runtime, because + // resize_clone_node() only propagates sizes when input and output have the + // same dim count -- which is never true for a squeeze. Under dynamic shapes + // the output then keeps its upper-bound extents while consumers read it at + // the real size, silently producing wrong values. add_permute_node()'s + // resize function handles the rank-reducing case explicitly. for (int i = 0; i < dims.size(); ++i) { // adjust negative dims int64_t dim_val = dims.at(i); if (dim_val < 0) { dim_val += in_dim; } - if (dims.at(i) != 0 && in_sizes.at(dim_val) == 1) { + if (in_sizes.at(dim_val) == 1) { squeeze_dims.push_back(dim_val); } } From 5b7718754e2ab4d634dc8e7ad837209345307983 Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Wed, 9 Sep 2026 13:24:29 -0700 Subject: [PATCH 123/190] Add batch runner example and refactor single sequence runner to use text stream (#22627) Co-authored-by: kiymetakdemir --- .github/workflows/mlx.yml | 36 +- backends/mlx/examples/llm/CMakeLists.txt | 32 +- backends/mlx/examples/llm/run_llm_batched.cpp | 439 ++++++++++++++++++ backends/mlx/examples/llm/run_llm_hf.cpp | 185 +++----- backends/mlx/examples/llm/runner_utils.h | 109 +++++ backends/mlx/test/mlx_cell_cache_test.cpp | 10 +- extension/llm/cache/cell_cache.cpp | 42 +- extension/llm/cache/cell_cache.h | 9 +- extension/llm/cache/test/cache_test.cpp | 21 +- 9 files changed, 725 insertions(+), 158 deletions(-) create mode 100644 backends/mlx/examples/llm/run_llm_batched.cpp create mode 100644 backends/mlx/examples/llm/runner_utils.h diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index 7977b234589..e762110693f 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -1037,10 +1037,13 @@ jobs: ${CONDA_RUN} cmake --build cmake-out/backends/mlx/examples/llm \ -j$(( $(sysctl -n hw.ncpu) - 1 )) RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_hf - if [ ! -x "${RUNNER}" ]; then - echo "Failed: runner not found at ${RUNNER}" - exit 1 - fi + BATCHED_RUNNER=cmake-out/backends/mlx/examples/llm/mlx_run_llm_batched + for binary in "${RUNNER}" "${BATCHED_RUNNER}"; do + if [ ! -x "${binary}" ]; then + echo "Failed: runner not found at ${binary}" + exit 1 + fi + done echo "::endgroup::" echo "::group::Install LLM requirements" @@ -1085,3 +1088,28 @@ jobs: exit 1 fi echo "::endgroup::" + + echo "::group::Run ${MODEL_NAME} batched off-graph inference" + ${BATCHED_RUNNER} \ + --pte /tmp/${MODEL_NAME}_offgraph.pte \ + --tokenizer "${TOKENIZER}" \ + --chat "${CHAT}" \ + --max-session-tokens 1024 \ + --max-new-tokens 50 \ + --max-decode-sequences 2 \ + --out-prefix /tmp/${MODEL_NAME}_batched \ + "What is the capital of France?" \ + "What color is grass?" + for check in "0:Paris" "1:green"; do + index="${check%%:*}" + expected="${check#*:}" + output="/tmp/${MODEL_NAME}_batched_${index}.txt" + if grep -iq "${expected}" "${output}"; then + echo "Success: '${expected}' found in ${output}" + else + echo "Failed: Expected '${expected}' not found in ${output}" + cat "${output}" + exit 1 + fi + done + echo "::endgroup::" diff --git a/backends/mlx/examples/llm/CMakeLists.txt b/backends/mlx/examples/llm/CMakeLists.txt index fa041b7129a..b8ef4981bce 100644 --- a/backends/mlx/examples/llm/CMakeLists.txt +++ b/backends/mlx/examples/llm/CMakeLists.txt @@ -35,18 +35,10 @@ executorch_target_link_options_shared_lib(executorch) set(gflags_DIR ${CMAKE_CURRENT_BINARY_DIR}/../../../../third-party/gflags) find_package(gflags REQUIRED) -set(link_libraries - executorch - extension_module - extension_tensor - extension_llm_cache - extension_llm_runner - extension_llm_sampler - gflags -) +set(link_libraries executorch extension_llm_runner gflags) if(NOT TARGET mlxdelegate) - message(FATAL_ERROR "mlx_run_llm_hf requires the MLX backend (mlxdelegate)") + message(FATAL_ERROR "MLX LLM runners require the MLX backend (mlxdelegate)") endif() list(APPEND link_libraries mlxdelegate mlx) executorch_target_link_options_shared_lib(mlxdelegate) @@ -67,7 +59,27 @@ target_include_directories( ) target_link_libraries(mlx_run_llm_hf PUBLIC ${link_libraries}) +add_executable(mlx_run_llm_batched run_llm_batched.cpp) +target_include_directories( + mlx_run_llm_batched PUBLIC ${_common_include_directories} ${_json_include} + ${_flatbuffers_include} +) +target_link_libraries( + mlx_run_llm_batched + PUBLIC executorch + extension_llm_batching_module + extension_llm_runner + gflags + mlxdelegate + mlx + tokenizers::tokenizers +) +if(TARGET optimized_native_cpu_ops_lib) + target_link_libraries(mlx_run_llm_batched PUBLIC optimized_native_cpu_ops_lib) +endif() + # The copy helper is gated on EXECUTORCH_BUILD_MLX, which the installed config # does not set; reaching here means mlxdelegate exists. set(EXECUTORCH_BUILD_MLX ON) executorch_target_copy_mlx_metallib(mlx_run_llm_hf) +executorch_target_copy_mlx_metallib(mlx_run_llm_batched) diff --git a/backends/mlx/examples/llm/run_llm_batched.cpp b/backends/mlx/examples/llm/run_llm_batched.cpp new file mode 100644 index 00000000000..92d9e6cd7e7 --- /dev/null +++ b/backends/mlx/examples/llm/run_llm_batched.cpp @@ -0,0 +1,439 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +// Sample application demonstrating continuous batching and streaming output +// for independent prompts submitted from one thread. +// +// Required flags are --pte and --tokenizer. Each remaining positional argument +// is a prompt, and generated text is streamed to +// _.txt. + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include + +DEFINE_string(pte, "", "Path to the .pte exported with --use-offgraph-cache"); +DEFINE_string(tokenizer, "", "Path to a supported tokenizer file"); +DEFINE_string(out_prefix, "gen", "Output files are _.txt"); +DEFINE_int32(max_session_tokens, 2048, "Maximum tokens retained per session"); +DEFINE_string( + kv_storage_dtype, + "bf16", + "KV storage dtype: bf16, fp16, or fp32"); +DEFINE_int32( + kv_initial_capacity, + -1, + "Initial cache pool capacity; -1 keeps the cache default"); +DEFINE_int32(max_new_tokens, 128, "Maximum generated tokens per prompt"); +DEFINE_int32(flush_every, 8, "Flush each output file every N generated tokens"); +DEFINE_int32( + max_decode_sequences, + 32, + "Maximum decode sequences admitted to one batch"); +DEFINE_double(temperature, 0.0, "Sampling temperature; 0 is greedy"); +DEFINE_double(top_p, 1.0, "Nucleus sampling probability"); +DEFINE_int32(top_k, 0, "Top-k sampling limit; 0 disables it"); +DEFINE_uint64(seed, 42, "Per-generation sampling seed"); +DEFINE_bool(metrics, true, "Print per-generation and engine reports"); +DEFINE_string( + chat, + "llama3", + "Chat template: llama3, gemma, gemma4, or 0 for raw text"); + +namespace batching = ::executorch::extension::llm::batching; +using ::executorch::backends::mlx::examples::llm::resolve_stop_tokens; +using ::executorch::backends::mlx::examples::llm::StopTokens; +using ::executorch::backends::mlx::examples::llm::storage_dtype; +using ::executorch::backends::mlx::examples::llm::wrap_turn; +using ::executorch::extension::Module; +using ::executorch::extension::llm::TextStream; +using ::executorch::runtime::Error; +using ::executorch::runtime::Result; + +namespace { + +struct Emitter { + Emitter( + const tokenizers::Tokenizer& tokenizer, + batching::Token previous, + const std::string& path, + std::size_t flush_every) + : file(path, std::ios::binary), + flush_every(flush_every), + stream( + tokenizer, + [this](const std::string& piece) { + file.write( + piece.data(), static_cast(piece.size())); + }, + previous) {} + + void append(batching::Token token) { + if (stream.append(token) != Error::Ok || !file) { + throw std::runtime_error("failed to decode or write output"); + } + if (++tokens_since_flush == flush_every) { + file.flush(); + tokens_since_flush = 0; + if (!file) { + throw std::runtime_error("failed to flush output"); + } + } + } + + void finish() { + stream.flush(); + file.flush(); + if (!file) { + throw std::runtime_error("failed to flush output"); + } + } + + std::ofstream file; + const std::size_t flush_every; + std::size_t tokens_since_flush = 0; + TextStream stream; +}; + +struct JobResult { + std::optional session; + batching::GenerationHandle handle; + std::optional reason; + std::optional metrics; + std::string message; + std::string output_path; + + bool failed() const { + return !reason || *reason == batching::FinishReason::Cancelled || + *reason == batching::FinishReason::Failed; + } +}; + +const char* reason_name(const std::optional& reason) { + if (!reason) { + return "never started"; + } + switch (*reason) { + case batching::FinishReason::StopToken: + return "stop token"; + case batching::FinishReason::NewTokenLimit: + return "token limit"; + case batching::FinishReason::Cancelled: + return "cancelled"; + case batching::FinishReason::Failed: + return "failed"; + } + return "unknown"; +} + +Result> optional_const_int( + Module& module, + const char* name) { + const auto methods = module.method_names(); + if (!methods.ok()) { + return methods.error(); + } + if (methods->count(name) == 0) { + return std::optional{}; + } + const auto result = module.execute(name); + if (!result.ok()) { + return result.error(); + } + if (result->size() != 1 || !result->at(0).isInt()) { + return Error::InvalidProgram; + } + return std::optional{result->at(0).toInt()}; +} + +void submit_prompt( + batching::Session& session, + const tokenizers::Tokenizer& tokenizer, + const std::string& prompt, + const std::vector& stop_tokens, + JobResult& result) { + try { + std::string wrapped; + if (!wrap_turn(FLAGS_chat, prompt, true, wrapped)) { + result.message = "unknown --chat template: " + FLAGS_chat; + return; + } + auto encoded = tokenizer.encode(wrapped, FLAGS_chat == "0" ? 1 : 0, 0); + if (!encoded.ok() || encoded->empty()) { + result.message = "could not encode prompt"; + return; + } + // Reserve the full generation budget so an admitted job is never shortened. + if (encoded->size() > + static_cast( + FLAGS_max_session_tokens - FLAGS_max_new_tokens)) { + result.message = + "prompt plus --max_new_tokens exceeds --max_session_tokens"; + return; + } + + auto emitter = std::make_shared( + tokenizer, + encoded->back(), + result.output_path, + static_cast(FLAGS_flush_every)); + if (!emitter->file) { + result.message = "could not open output file"; + return; + } + + batching::GenConfig config; + config.max_new_tokens = FLAGS_max_new_tokens; + config.sampling.temperature = static_cast(FLAGS_temperature); + config.sampling.top_p = static_cast(FLAGS_top_p); + config.sampling.top_k = FLAGS_top_k; + config.stop_tokens = stop_tokens; + config.seed = FLAGS_seed; + + result.handle = session.generate_async( + std::move(*encoded), + std::move(config), + [emitter](const batching::GenerationUpdate& update) { + std::size_t count = update.tokens.size(); + if (update.finish_reason == batching::FinishReason::StopToken && + count > 0) { + --count; + } + for (std::size_t i = 0; i < count; ++i) { + emitter->append(update.tokens[i]); + } + if (update.finish_reason) { + emitter->finish(); + } + }); + } catch (const std::exception& error) { + result.message = error.what(); + } +} + +} // namespace + +int main(int argc, char** argv) { + gflags::ParseCommandLineFlags(&argc, &argv, true); + const std::vector prompts(argv + 1, argv + argc); + + if (FLAGS_pte.empty() || FLAGS_tokenizer.empty() || prompts.empty()) { + std::cerr << "usage: " << argv[0] + << " --pte model.pte --tokenizer tokenizer-file \"prompt\" [...]" + << std::endl; + return 1; + } + if (FLAGS_max_session_tokens <= 0 || FLAGS_max_new_tokens <= 0 || + FLAGS_max_decode_sequences <= 0 || FLAGS_flush_every <= 0) { + std::cerr + << "session, generation, decode, and flush limits must be positive" + << std::endl; + return 1; + } + if (FLAGS_temperature < 0.0 || FLAGS_top_p <= 0.0 || FLAGS_top_p > 1.0 || + FLAGS_top_k < 0) { + std::cerr << "invalid sampling parameters" << std::endl; + return 1; + } + if (FLAGS_kv_initial_capacity < -1) { + std::cerr << "--kv_initial_capacity must be -1 or non-negative" + << std::endl; + return 1; + } + if (prompts.size() > + static_cast(std::numeric_limits::max())) { + std::cerr << "too many prompts" << std::endl; + return 1; + } + const int kv_dtype = storage_dtype(FLAGS_kv_storage_dtype); + if (kv_dtype < 0) { + std::cerr << "--kv_storage_dtype must be bf16, fp16, or fp32" << std::endl; + return 1; + } + if (FLAGS_max_new_tokens > FLAGS_max_session_tokens) { + std::cerr << "--max_new_tokens exceeds --max_session_tokens" << std::endl; + return 1; + } + + auto tokenizer = + ::executorch::extension::llm::load_tokenizer(FLAGS_tokenizer); + if (!tokenizer) { + std::cerr << "could not load tokenizer: " << FLAGS_tokenizer << std::endl; + return 1; + } + + auto module = std::make_unique(FLAGS_pte); + if (module->load() != Error::Ok) { + std::cerr << "could not load " << FLAGS_pte << std::endl; + return 1; + } + + const auto model_max_context = optional_const_int(*module, "get_max_ctx_len"); + if (!model_max_context.ok()) { + std::cerr << "could not read get_max_ctx_len" << std::endl; + return 1; + } + if (*model_max_context && + (**model_max_context <= 0 || + FLAGS_max_session_tokens > **model_max_context)) { + std::cerr << "--max_session_tokens " << FLAGS_max_session_tokens + << " exceeds the model context limit " << **model_max_context + << std::endl; + return 1; + } + + StopTokens resolved_stop_tokens; + if (!resolve_stop_tokens( + *tokenizer, *module, FLAGS_chat, resolved_stop_tokens)) { + std::cerr << "could not resolve stop tokens for --chat=" << FLAGS_chat + << std::endl; + return 1; + } + const std::vector stop_tokens( + resolved_stop_tokens.ids.begin(), resolved_stop_tokens.ids.end()); + + auto executor = batching::ModuleExecutor::create( + std::move(module), + static_cast(prompts.size()), + FLAGS_max_session_tokens, + kv_dtype, + FLAGS_kv_initial_capacity); + if (!executor.ok()) { + std::cerr << "could not create executor: " + << ::executorch::runtime::to_string(executor.error()) + << std::endl; + return 1; + } + + const std::size_t width = (*executor)->preferred_batch_tokens(); + const std::size_t decode_slots = + static_cast(FLAGS_max_decode_sequences); + if (width == 0) { + std::cerr << "the model's forward token input has no usable width (its " + "traced seq_len dimension is 0); re-export it with a dynamic " + "token dimension" + << std::endl; + return 1; + } + if (decode_slots >= width) { + std::cerr << "--max_decode_sequences " << decode_slots + << " leaves no room for prefill in a " << width + << "-token forward" << std::endl; + return 1; + } + if (decode_slots > width / 4) { + std::cerr << "warning: --max_decode_sequences " << decode_slots + << " leaves only " << width - decode_slots + << " prefill tokens of a " << width << "-token forward" + << std::endl; + } + + auto scheduler = batching::DecodeFirstScheduler::create( + width, decode_slots, width - decode_slots); + if (!scheduler) { + std::cerr << "the scheduler refused those limits" << std::endl; + return 1; + } + + batching::Runner runner(**executor, std::move(scheduler)); + std::vector results(prompts.size()); + // Fail before opening sessions if any output path cannot be created. + for (std::size_t i = 0; i < results.size(); ++i) { + results[i].output_path = + FLAGS_out_prefix + "_" + std::to_string(i) + ".txt"; + std::ofstream output( + results[i].output_path, std::ios::binary | std::ios::trunc); + if (!output) { + runner.shutdown(); + std::cerr << "could not create " << results[i].output_path << std::endl; + return 1; + } + } + + std::vector>> session_futures; + session_futures.reserve(prompts.size()); + for (std::size_t i = 0; i < prompts.size(); ++i) { + session_futures.push_back(runner.open_session_async()); + } + for (std::size_t i = 0; i < prompts.size(); ++i) { + results[i].session = session_futures[i].get(); + } + for (std::size_t i = 0; i < prompts.size(); ++i) { + if (!results[i].session) { + results[i].message = "could not open session"; + continue; + } + submit_prompt( + *results[i].session, *tokenizer, prompts[i], stop_tokens, results[i]); + } + for (JobResult& result : results) { + if (!result.handle.valid()) { + continue; + } + result.handle.wait(); + result.metrics = result.handle.metrics(); + result.reason = result.handle.finish_reason(); + result.message = result.handle.error_message(); + } + + runner.shutdown(); + const batching::EngineMetrics engine = runner.metrics(); + + if (FLAGS_metrics) { + std::cout << "\n"; + for (std::size_t i = 0; i < results.size(); ++i) { + if (results[i].metrics) { + std::cout << "[" << i << "] " + << batching::format_report(*results[i].metrics); + } + } + std::cout << "\n" << batching::format_report(engine); + } + + std::size_t failures = 0; + for (const JobResult& result : results) { + failures += result.failed() ? 1 : 0; + } + if (failures > 0) { + std::cout << "\nfailures:\n"; + for (std::size_t i = 0; i < results.size(); ++i) { + if (!results[i].failed()) { + continue; + } + std::cout << " [" << i << "] " << results[i].output_path << ": " + << reason_name(results[i].reason); + if (!results[i].message.empty()) { + std::cout << ": " << results[i].message; + } + std::cout << "\n"; + } + } + + std::cout << prompts.size() - failures << "/" << prompts.size() + << " generations completed" << std::endl; + return failures == 0 ? 0 : 1; +} diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp index 1872b9a8a33..925c1360cf3 100644 --- a/backends/mlx/examples/llm/run_llm_hf.cpp +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -32,6 +32,7 @@ #include #include #include +#include #include #include #include @@ -42,14 +43,14 @@ #include #include -#include +#include + #include #include #include #include #include #include -#include #include DEFINE_string(pte, "", "Model .pte file."); @@ -100,11 +101,14 @@ DEFINE_bool( false, "Run once before measuring, to absorb JIT and pool growth."); +using ::executorch::backends::mlx::examples::llm::resolve_stop_tokens; +using ::executorch::backends::mlx::examples::llm::StopTokens; +using ::executorch::backends::mlx::examples::llm::storage_dtype; +using ::executorch::backends::mlx::examples::llm::wrap_turn; using ::executorch::extension::make_tensor_ptr; using ::executorch::extension::Module; -using ::executorch::extension::TensorPtr; +using ::executorch::extension::llm::TextStream; using ::executorch::runtime::Error; -using ::executorch::runtime::EValue; namespace cache = ::executorch::extension::llm::cache; @@ -142,20 +146,6 @@ bool parse_int_list( return true; } -int storage_dtype(const std::string& name) { - using S = ::executorch::runtime::etensor::ScalarType; - if (name == "bf16") { - return static_cast(S::BFloat16); - } - if (name == "fp16") { - return static_cast(S::Half); - } - if (name == "fp32") { - return static_cast(S::Float); - } - return -1; -} - // Constant methods the export publishes (get_n_caches and friends). They carry // no delegate, so reading them only needs the program loaded -- which is what // lets the cache be built before forward's backend init consumes its key. @@ -292,33 +282,6 @@ void print_cache_summary(const cache::CacheConfig& cfg) { std::cout << std::endl; } -// One user turn wrapped in the model's instruct template. Returns false for an -// unknown template name. The leading BOS belongs to the first turn only, so a -// continuing conversation passes with_bos=false. -bool wrap_turn( - const std::string& chat, - const std::string& prompt, - bool with_bos, - std::string& out) { - if (chat == "0") { - out = prompt; - } else if (chat == "llama3") { - out = std::string(with_bos ? "<|begin_of_text|>" : "") + - "<|start_header_id|>user<|end_header_id|>\n\n" + prompt + - "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; - } else if (chat == "gemma") { - out = std::string(with_bos ? "" : "") + "user\n" + - prompt + "\nmodel\n"; - } else if (chat == "gemma4") { - // Gemma 4 renamed the turn markers; its own is also the eos. - out = std::string(with_bos ? "" : "") + "<|turn>user\n" + prompt + - "\n<|turn>model\n"; - } else { - return false; - } - return true; -} - } // namespace int main(int argc, char** argv) { @@ -341,6 +304,12 @@ int main(int argc, char** argv) { "[--kv-max-capacity N for off-graph models]\n"; return 1; } + if (warmup && kv_capacity <= 0) { + std::cerr << "--warmup requires an off-graph cache selected with " + "--kv_max_capacity" + << std::endl; + return 1; + } try { // The shared loader sniffs the format, so --tokenizer takes any of the @@ -373,6 +342,15 @@ int main(int argc, char** argv) { return 1; } const int prefill_chunk = static_cast(*published_prefill_chunk); + StopTokens stop_tokens; + if (!resolve_stop_tokens(*tokenizer, module, chat, stop_tokens)) { + std::cerr << "Could not resolve stop tokens for --chat=" << chat + << std::endl; + return 1; + } + auto write_text = [](const std::string& text) { + std::cout << text << std::flush; + }; // Everything past load_method is identical for both model kinds; only // setup differs. ctl is null for an in-graph model, which owns its cache @@ -399,47 +377,8 @@ int main(int argc, char** argv) { std::cout << "[mem] after load : " << mem_at_load << " MiB" << std::endl; - // Encode. HFTokenizer maps special-token markers in the string to their - // ids, so the template's <|...|> tokens encode correctly; it already - // carries <|begin_of_text|>, so pass bos=0 to avoid a doubled BOS. - std::string enc_input; - if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { - std::cerr << "Unknown --chat template: " << chat - << " (expected llama3, gemma, gemma4, or 0)" << std::endl; - return 1; - } - // The template carries its own BOS, so only a raw prompt asks for one. - const int8_t bos = chat == "0" ? 1 : 0; - auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); - if (!enc.ok()) { - std::cerr << "Encode failed" << std::endl; - return 1; - } - std::vector tokens = std::move(*enc); - const int prompt_len = static_cast(tokens.size()); - - // End-of-text from the model's metadata when it publishes any, else the - // tokenizer's. The turn-end token is ours: it depends on --chat, which - // the .pte knows nothing about. - std::unordered_set stop_ids = - ::executorch::extension::llm::get_eos_ids(tokenizer.get(), &module); - std::optional turn_end_id; - if (chat != "0") { - const char* turn_end = chat == "llama3" ? "<|eot_id|>" - : chat == "gemma4" ? "" - : ""; - if (auto eot = tokenizer->piece_to_id(turn_end); eot.ok()) { - turn_end_id = static_cast(*eot); - stop_ids.insert(*eot); - } - } - auto is_stop = [&](int64_t t) { - for (uint64_t s : stop_ids) { - if (t == static_cast(s)) { - return true; - } - } - return false; + auto is_stop = [&](int64_t token) { + return stop_tokens.ids.count(static_cast(token)) != 0; }; // One Sampler for the whole run, as the shared runner does: constructing @@ -494,7 +433,6 @@ int main(int argc, char** argv) { std::cerr << "--interactive requires --kv-max-capacity\n"; return 1; } - auto* control = ctl; std::cout << "Multi-turn chat. /reset clears, /undo drops the last turn, " "/undo N drops N tokens, /quit exits.\n"; @@ -506,7 +444,7 @@ int main(int argc, char** argv) { break; } if (line == "/reset") { - control->clear(); + ctl->clear(); position = turn_start = 0; std::cout << "[cleared]\n"; continue; @@ -523,7 +461,7 @@ int main(int argc, char** argv) { continue; } } - if (control->rewind(static_cast(target))) { + if (ctl->rewind(static_cast(target))) { position = target; turn_start = std::min(turn_start, position); std::cout << "[rewound to " << position << "]\n"; @@ -540,8 +478,8 @@ int main(int argc, char** argv) { std::string turn; wrap_turn(chat, line, /*with_bos=*/position == 0, turn); auto te = tokenizer->encode(turn, /*bos=*/chat == "0" ? 1 : 0, 0); - if (!te.ok()) { - std::cerr << "Encode failed\n"; + if (!te.ok() || te->empty()) { + std::cerr << "Encode failed or produced no tokens\n"; continue; } const int n = static_cast(te->size()); @@ -549,15 +487,15 @@ int main(int argc, char** argv) { // whole max_new budget up front would report "full" with most of the // cache still free. Generation is then clamped to the room that // remains. - if (!control->can_extend(n + 1)) { - std::cout << "[cache full: " << position << "/" - << control->capacity() << ", turn " << n << " tokens" - << (control->can_extend(1) ? "" : ", length at capacity") + if (!ctl->can_extend(n + 1)) { + std::cout << "[cache full: " << position << "/" << ctl->capacity() + << ", turn " << n << " tokens" + << (ctl->can_extend(1) ? "" : ", length at capacity") << ", use /reset]\n"; continue; } const int budget = std::min( - max_new, control->capacity() - static_cast(position) - n); + max_new, ctl->capacity() - static_cast(position) - n); turn_start = position; std::vector tin(te->begin(), te->end()), tpos; @@ -567,27 +505,28 @@ int main(int argc, char** argv) { int64_t next = prefill(tin, tpos); position += n; - uint64_t prev = te->back(); + TextStream text_stream(*tokenizer, write_text, te->back()); for (int i = 0; i < budget && !is_stop(next); ++i) { - if (auto piece = - tokenizer->decode(prev, static_cast(next)); - piece.ok()) { - std::cout << *piece << std::flush; + if (text_stream.append(static_cast(next)) != Error::Ok) { + text_stream.flush(); + std::cerr << "Failed to decode generated token" << std::endl; + return 1; } - prev = static_cast(next); next = step({next}, {position}); ++position; } + text_stream.flush(); // The turn-end token stops generation, so it is neither printed nor // fed back -- but the next turn opens without closing this one, and // an unterminated assistant turn compounds over a session. Commit it, // at the cost of one extra step per turn. - if (turn_end_id && next == *turn_end_id && control->can_extend(1)) { + if (stop_tokens.turn_end_id && + static_cast(next) == *stop_tokens.turn_end_id && + ctl->can_extend(1)) { step({next}, {position}); ++position; } - std::cout << "\n[" << position << "/" << control->capacity() - << " tokens" + std::cout << "\n[" << position << "/" << ctl->capacity() << " tokens" << (budget < max_new ? ", generation capped by capacity" : "") << "]\n"; @@ -595,13 +534,24 @@ int main(int argc, char** argv) { return 0; } + std::string enc_input; + if (!wrap_turn(chat, prompt, /*with_bos=*/true, enc_input)) { + std::cerr << "Unknown --chat template: " << chat + << " (expected llama3, gemma, gemma4, or 0)" << std::endl; + return 1; + } + const int8_t bos = chat == "0" ? 1 : 0; + auto enc = tokenizer->encode(enc_input, bos, /*eos=*/0); + if (!enc.ok() || enc->empty()) { + std::cerr << "Encode failed or produced no tokens" << std::endl; + return 1; + } + std::vector tokens = std::move(*enc); + const int prompt_len = static_cast(tokens.size()); std::vector ids(tokens.begin(), tokens.end()), prefill_pos; for (int i = 0; i < prompt_len; ++i) { prefill_pos.push_back(i); } - auto ms = [](auto a, auto b) { - return std::chrono::duration(b - a).count(); - }; // Sequence length against the configured ceiling, with what MLX actually // holds for it. Pools start at initial_capacity and grow by doubling, so // the bytes lag the token count in steps; bf16 storage (kv_dtype 15) @@ -640,24 +590,25 @@ int main(int argc, char** argv) { std::cout << "\n"; // blank line before the streamed generation } - uint64_t prev = tokens.back(); + TextStream::Sink sink; + if (measured) { + sink = write_text; + } + TextStream text_stream(*tokenizer, std::move(sink), tokens.back()); int generated = 0; for (int i = 0; i < max_new; ++i) { if (is_stop(next)) { break; } - if (measured) { - if (auto piece = - tokenizer->decode(prev, static_cast(next)); - piece.ok()) { - ::executorch::extension::llm::safe_printf(piece->c_str()); - fflush(stdout); - } + if (text_stream.append(static_cast(next)) != Error::Ok) { + text_stream.flush(); + std::cerr << "Failed to decode generated token" << std::endl; + return 1; } - prev = static_cast(next); ++generated; next = step({next}, {prompt_len + i}); } + text_stream.flush(); stats.inference_end_ms = ::executorch::extension::llm::time_in_ms(); if (measured) { std::cout << "\n\n"; // close the generation line + blank separator diff --git a/backends/mlx/examples/llm/runner_utils.h b/backends/mlx/examples/llm/runner_utils.h new file mode 100644 index 00000000000..6e8ac83c1d9 --- /dev/null +++ b/backends/mlx/examples/llm/runner_utils.h @@ -0,0 +1,109 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch { +namespace backends { +namespace mlx { +namespace examples { +namespace llm { + +struct StopTokens { + std::unordered_set ids; + std::optional turn_end_id; +}; + +inline const char* turn_end_piece(const std::string& chat) { + if (chat == "llama3") { + return "<|eot_id|>"; + } + if (chat == "gemma") { + return ""; + } + if (chat == "gemma4") { + return ""; + } + return nullptr; +} + +inline bool resolve_stop_tokens( + tokenizers::Tokenizer& tokenizer, + ::executorch::extension::Module& module, + const std::string& chat, + StopTokens& out) { + out.ids = ::executorch::extension::llm::get_eos_ids(&tokenizer, &module); + out.turn_end_id.reset(); + if (chat == "0") { + return true; + } + const char* piece = turn_end_piece(chat); + if (piece == nullptr) { + return false; + } + auto id = tokenizer.piece_to_id(piece); + if (!id.ok()) { + return false; + } + out.turn_end_id = *id; + out.ids.insert(*id); + return true; +} + +inline bool wrap_turn( + const std::string& chat, + const std::string& prompt, + bool with_bos, + std::string& out) { + if (chat == "0") { + out = prompt; + } else if (chat == "llama3") { + out = std::string(with_bos ? "<|begin_of_text|>" : "") + + "<|start_header_id|>user<|end_header_id|>\n\n" + prompt + + "<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n"; + } else if (chat == "gemma") { + out = std::string(with_bos ? "" : "") + "user\n" + + prompt + "\nmodel\n"; + } else if (chat == "gemma4") { + out = std::string(with_bos ? "" : "") + "<|turn>user\n" + prompt + + "\n<|turn>model\n"; + } else { + return false; + } + return true; +} + +inline int storage_dtype(const std::string& name) { + using ScalarType = ::executorch::runtime::etensor::ScalarType; + if (name == "bf16") { + return static_cast(ScalarType::BFloat16); + } + if (name == "fp16") { + return static_cast(ScalarType::Half); + } + if (name == "fp32") { + return static_cast(ScalarType::Float); + } + return -1; +} + +} // namespace llm +} // namespace examples +} // namespace mlx +} // namespace backends +} // namespace executorch diff --git a/backends/mlx/test/mlx_cell_cache_test.cpp b/backends/mlx/test/mlx_cell_cache_test.cpp index 6e68ba79ef5..f808965d1a5 100644 --- a/backends/mlx/test/mlx_cell_cache_test.cpp +++ b/backends/mlx/test/mlx_cell_cache_test.cpp @@ -224,8 +224,9 @@ TEST_F(MLXCellCacheTest, StorageDtypeDiffersCastsOnWrite) { EXPECT_TRUE(allclose(spec.K, k, 1e-2f)); } -// The step verbs are a contract: no declaration, a miscounted call, a repeated -// layer and a position a sequence already holds are all refused. +// The step verbs are a contract: no declaration, a miscounted call, and a +// position a sequence already holds are refused. A repeated layer with the same +// tokens (a KV-shared donor re-serving) is served again idempotently. TEST_F(MLXCellCacheTest, IllFormedStepsThrow) { using namespace ::mlx::core; MLXCellCache c(flat_config(32, 1, H, D, kHalf)); @@ -239,8 +240,9 @@ TEST_F(MLXCellCacheTest, IllFormedStepsThrow) { EXPECT_ANY_THROW(c.update_and_fetch(1, {0, 1}, k, v, s)); // no such layer c.update_and_fetch(0, {0, 1}, k, v, s); - EXPECT_ANY_THROW( - c.update_and_fetch(0, {0, 1}, k, v, s)); // layer served twice + // A KV-shared layer re-serves its donor's id with the same tokens; the repeat + // is idempotent and returns the same step rather than throwing. + EXPECT_NO_THROW(c.update_and_fetch(0, {0, 1}, k, v, s)); EXPECT_TRUE(c.declare_step({a})); array k1 = randn(1), v1 = randn(1); diff --git a/extension/llm/cache/cell_cache.cpp b/extension/llm/cache/cell_cache.cpp index ceaa31a6686..43866357baf 100644 --- a/extension/llm/cache/cell_cache.cpp +++ b/extension/llm/cache/cell_cache.cpp @@ -19,8 +19,7 @@ namespace cache { CellCache::CellCache(const CacheConfig& cfg) : capacity_(cfg.capacity), pos_(cfg.capacity, -1), - owners_(cfg.capacity, 0), - served_(cfg.n_layers, false) { + owners_(cfg.capacity, 0) { assert(valid(cfg)); // One window per layer, from the same per-layer config the sequence cache // reads. Layers agreeing on a window share a step. @@ -53,7 +52,6 @@ void CellCache::clear() { declared_ = false; step_seq_ids_.clear(); step_pos_.clear(); - std::fill(served_.begin(), served_.end(), false); invalidate_steps(); } @@ -71,7 +69,6 @@ bool CellCache::declare_step(const std::vector& seq_ids) { step_seq_ids_ = seq_ids; declared_ = true; invalidate_steps(); - std::fill(served_.begin(), served_.end(), false); return true; } @@ -155,25 +152,32 @@ int CellCache::used_end() const { const CellStep* CellCache::place_step(int layer, const int32_t* positions, int length) { - if (layer < 0 || layer >= static_cast(windows_.size()) || - served_[layer]) { - return nullptr; // out of range, or a forward that skipped declare_step + if (layer < 0 || layer >= static_cast(windows_.size())) { + return nullptr; // layer out of range } - if (!placed_) { - if (!declared_ || length != static_cast(step_seq_ids_.size())) { - return nullptr; // no declaration, or a token count disagreeing with it - } - if (!extends(positions, length)) { - return nullptr; // nothing mutated yet, so the step can be re-placed - } - step_pos_.assign(positions, positions + length); - if (!place()) { + if (placed_) { + // Re-serve within the placed forward. Every layer of a forward places the + // same tokens, and a KV-shared layer re-serves its donor's id, so a repeat + // with the same positions returns the same step and claims no new cells. + // Different positions mean a new step that never declared, still refused. + if (length != static_cast(step_pos_.size()) || + !std::equal(positions, positions + length, step_pos_.begin())) { return nullptr; } - declared_ = false; // one declaration, one placement - placed_ = true; + return &step_for(windows_[layer]); + } + if (!declared_ || length != static_cast(step_seq_ids_.size())) { + return nullptr; // no declaration, or a token count disagreeing with it + } + if (!extends(positions, length)) { + return nullptr; // a position a sequence already holds + } + step_pos_.assign(positions, positions + length); + if (!place()) { + return nullptr; // out of cells } - served_[layer] = true; + declared_ = false; // one declaration, one placement + placed_ = true; return &step_for(windows_[layer]); } diff --git a/extension/llm/cache/cell_cache.h b/extension/llm/cache/cell_cache.h index a307c72b1cd..256a43739c5 100644 --- a/extension/llm/cache/cell_cache.h +++ b/extension/llm/cache/cell_cache.h @@ -40,8 +40,12 @@ struct ET_EXPERIMENTAL CellStep { // every later layer reuses that placement. `layer` selects the window, which // decides the kind and mask, so a step is per policy and memoized for the // forward. The returned step is owned by the cache and valid until the next -// verb. nullptr = no declaration, a token count disagreeing with it, a position -// a sequence already holds, a layer out of range, or a layer served twice. +// verb. A layer may be served more than once per forward -- a KV-shared layer +// re-serves its donor's id -- provided the repeat passes the same positions; it +// returns the same step and claims no new cells. nullptr = no declaration, a +// token count disagreeing with it, a position a sequence already holds, a layer +// out of range, or a re-serve whose positions differ (a step that never +// declared). class ET_EXPERIMENTAL CellStepper { public: static constexpr const char* kFaceName = "et.cache.CellStepper"; @@ -148,7 +152,6 @@ class ET_EXPERIMENTAL CellCache : public Cache, std::vector step_seq_ids_; // set by declare_step std::vector step_pos_; // set when the step is placed std::vector cells_; // the step's placement, shared by every layer - std::vector served_; // layers this step has already answered std::vector windows_; // per layer; 0 = keeps all history // window -> step, memoized per forward. Node-based is required: a step // handed to one layer must survive another layer's insert. diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index fda446412a7..654db1a5b72 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -532,7 +532,7 @@ TEST_F(CacheTest, CellPlacementIsSharedByEveryLayerOfTheStep) { const auto* first = c.place(0, args.positions); // layer 0 places the cells ASSERT_NE(first, nullptr); EXPECT_EQ(c.place(1, args.positions), first); // later layers reuse them - EXPECT_EQ(c.place(0, args.positions), nullptr); // asking twice is a new step + EXPECT_EQ(c.place(0, args.positions), first); // re-serving repeats the step EXPECT_EQ(c.cache.free_cells(), 14); // placed once, not once per layer } @@ -776,3 +776,22 @@ TEST_F(CacheTest, CellClearReturnsEveryCell) { EXPECT_EQ(c.ctl->next_pos(s0), 0); // the sequence is gone EXPECT_EQ(c.place(0, {0}), nullptr); // and the step went with it } + +TEST_F(CacheTest, KvSharedLayerReservesIdempotently) { + // A KV-shared layer re-serves its donor's id with the same tokens: the repeat + // returns the donor's step and claims no new cells. A re-serve with different + // tokens is a new step that never declared, and is refused. + Cells c(16, {flat_layer(), flat_layer()}); + const int32_t s0 = c.seq_new(); + const auto args = flatten_step({{s0, 0, 3}}); + ASSERT_TRUE(c.ctl->declare_step(args.seq_ids)); + + const auto* donor = c.place(0, args.positions); + ASSERT_NE(donor, nullptr); + const int free_after_place = c.cache.free_cells(); + + EXPECT_EQ(c.place(0, args.positions), donor); // same tokens -> same step + EXPECT_EQ(c.cache.free_cells(), free_after_place); // no new cells claimed + + EXPECT_EQ(c.place(0, {7, 8, 9}), nullptr); // different tokens, never declared +} From ad2de690ce06dc46eaec72fb07d3e9d2172f4869 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Wed, 9 Sep 2026 14:08:18 -0700 Subject: [PATCH 124/190] Optimize CPU SDPA for ring attention (#22358) ## Summary - skip fully masked KV blocks and trim masked boundaries in CPU custom SDPA - split wrapped ring windows around large masked gaps while preserving additive-mask semantics - size local KV caches as sliding window plus the maximum in-flight prefill chunk, capped by full context - support regular, custom-op, and quantized ring caches --- examples/models/llama/attention.py | 44 ++- examples/models/llama/export_llama_lib.py | 5 + .../source_transformation/attention_sink.py | 122 +++++-- .../source_transformation/custom_kv_cache.py | 116 +++++-- .../test_attention_sink.py | 73 ++-- .../llama/tests/test_replace_kv_cache.py | 14 +- .../models/llama/tests/test_ring_attention.py | 81 ++++- .../models/llama/tests/test_ring_kv_cache.py | 33 +- extension/llm/custom_ops/op_sdpa_impl.h | 322 ++++++++++++++---- extension/llm/custom_ops/op_sdpa_test.cpp | 165 +++++++++ 10 files changed, 794 insertions(+), 181 deletions(-) diff --git a/examples/models/llama/attention.py b/examples/models/llama/attention.py index 98a838bfa30..9980da91958 100644 --- a/examples/models/llama/attention.py +++ b/examples/models/llama/attention.py @@ -262,6 +262,23 @@ def calculate_positions_and_update_indices(self, input_pos: torch.Tensor, seq_le return indices +def _get_ring_cache_size( + max_context_length: int, + window_size: int, + max_seq_len: Optional[int] = None, +) -> int: + """Size an SWA cache for one retained window plus one in-flight chunk.""" + assert window_size > 0, "Sliding-window size must be positive" + if max_seq_len is None: + max_seq_len = max_context_length + assert max_seq_len > 0, "Maximum sequence length must be positive" + assert window_size <= max_context_length, ( + f"Sliding-window size ({window_size}) cannot exceed the full context " + f"length ({max_context_length})" + ) + return min(max_context_length, window_size + max_seq_len) + + class RingKVCache(KVCache): def __init__( self, @@ -271,10 +288,17 @@ def __init__( head_dim: int, enable_dynamic_shape: bool, dtype=torch.float32, + *, + window_size: int, + max_seq_len: Optional[int] = None, ): - self.window_size = max_context_length """ - Reason why we want the kv cache size to be twice the context length: + The cache needs room for the retained sliding window and the current + prefill chunk. Its size is window_size + max_seq_len, capped by the + full-context cache. If max_seq_len is omitted, the full context length + is used. + + Reason why a cache larger than the sliding window is needed: Sliding window attention without ringbuffer pos 0 1 2 3 4 5 6 7 8 9 10 0 x 0 0 0 0 0 0 0 0 0 0 @@ -307,20 +331,28 @@ def __init__( So not having kept 2, 3 and 4 in cache means we will have divergent behavior. Worst case of this would have been when update it equal to the length of the cache. like in our case pos = 5 seq len = 4. - Thus we need to have a cache that is larger. How much larger, as much as - the sliding window size. So twice the max_context_length. + Thus we need a cache larger than the sliding window by enough space for + the largest in-flight input chunk. How would that have helped. Lets see. At pos = 5 our cache would have [0, 1, 2, 3, 4, NA, NA, NA] After cache update we would have [8, 1, 2, 3, 4, 5, 6, 7]. We kicked out token at pos = 0. However, the current step still has access to [pos - sliding_window_size, pos] tokens. - + To make sure we dont over attend, i.e. we dont have pos = 5 to attend to pos = 1, mask calculaton has to account for the sliding window size. """ + self.window_size = window_size + self.full_context_length = max_context_length + self.max_seq_len = ( + max_context_length if max_seq_len is None else int(max_seq_len) + ) + ring_cache_size = _get_ring_cache_size( + max_context_length, window_size, self.max_seq_len + ) super().__init__( max_batch_size, - max_context_length * 2, + ring_cache_size, n_heads, head_dim, enable_dynamic_shape, diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 07240b11d8c..074f8e24ea5 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -969,6 +969,7 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager: preq_group_size=llm_config.base.preq_group_size, preq_embedding_quantize=llm_config.base.preq_embedding_quantize, local_global_attention=llm_config.model.local_global_attention, + max_seq_len=llm_config.export.max_seq_length, use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear, use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding, quantize_with_hqq=llm_config.quantization.use_hqq, @@ -2064,6 +2065,7 @@ def _get_source_transforms( # noqa preq_group_size: Optional[int] = None, preq_embedding_quantize: Optional[str] = None, local_global_attention: Optional[List[int]] = None, + max_seq_len: Optional[int] = None, use_torchao_kernels_linear: bool = False, use_torchao_kernels_tied_embedding: bool = False, quantize_with_hqq: bool = True, @@ -2104,6 +2106,8 @@ def _get_source_transforms( # noqa preq_mode: Pre-quantization mode. preq_group_size: Pre-quantization group size. preq_embedding_quantize: Pre-quantization embedding quantize. + max_seq_len: Largest input chunk accepted by local-attention layers. If + omitted, each layer's full context length is used. Returns: A list of transformation functions. @@ -2280,6 +2284,7 @@ def _get_source_transforms( # noqa partial( replace_kv_cache_with_ring_kv_cache, layer_sizes=local_global_attention, + max_seq_len=max_seq_len, ) ) diff --git a/examples/models/llama/source_transformation/attention_sink.py b/examples/models/llama/source_transformation/attention_sink.py index c2d548fa798..c76ec7514a1 100644 --- a/examples/models/llama/source_transformation/attention_sink.py +++ b/examples/models/llama/source_transformation/attention_sink.py @@ -18,13 +18,36 @@ _create_causal_mask_for_ring_buffer, AttentionMHA, KVCache, - RingKVCache, ) from executorch.examples.models.llama.model_args import ModelArgs from executorch.examples.models.llama.rope import Rope from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter +def _get_attention_sink_cache_size( + max_context_length: int, + window_size: int, + sink_size: int, + max_seq_len: Optional[int] = None, +) -> int: + """Size a sink cache for fixed sinks, one window, and one input chunk.""" + assert sink_size >= 0, "Attention sink size must be non-negative" + assert sink_size < max_context_length, ( + f"Attention sink size ({sink_size}) must be smaller than the full " + f"context length ({max_context_length})" + ) + assert window_size > 0, "Sliding-window size must be positive" + assert sink_size + window_size <= max_context_length, ( + f"Attention sink size ({sink_size}) plus sliding-window size " + f"({window_size}) cannot exceed the full context length " + f"({max_context_length})" + ) + if max_seq_len is None: + max_seq_len = max_context_length + assert max_seq_len > 0, "Maximum sequence length must be positive" + return sink_size + window_size + max_seq_len + + class RopeWithAttentionSink(Rope): """ Rope subclass for Attention Sink models. @@ -37,10 +60,11 @@ class RopeWithAttentionSink(Rope): - Window tokens (pos >= sink_size): wrapped into ring buffer range [sink_size, sink_size + ring_size) via modulo - The ring buffer is 2x window_size for write-ahead headroom, not to keep the - live window contiguous -- it can span a wrap. Across a wrap two positions - remap to a difference that is not their true distance, so RoPE preserves - relative distance only within a wrap. + The ring buffer holds one retained window plus one maximum-size input + chunk. It is larger than the live window for write-ahead headroom, not to + keep the live window contiguous -- it can span a wrap. Across a wrap two + positions remap to a difference that is not their true distance, so RoPE + preserves relative distance only within a wrap. """ def __init__( @@ -52,7 +76,26 @@ def __init__( super().__init__(params) self.window_size = window_size self.sink_size = sink_size - self.ring_size = window_size * 2 + max_seq_len = ( + params.max_context_len + if getattr(params, "max_seq_len", None) is None + else int(params.max_seq_len) + ) + cache_size = _get_attention_sink_cache_size( + params.max_context_len, + window_size, + sink_size, + max_seq_len, + ) + self.ring_size = cache_size - sink_size + if self.freqs_cos.size(0) < cache_size: + freqs_cos, freqs_sin = self.precompute_freqs_cis( + self.params.head_dim, + cache_size, + self.params.rope_freq_base, + ) + self.freqs_cos = freqs_cos + self.freqs_sin = freqs_sin def _remap_input_pos(self, input_pos: torch.Tensor) -> torch.Tensor: """Remap positions: sink tokens stay, window tokens wrap in ring buffer.""" @@ -133,7 +176,7 @@ class CachePositionsManagerWithSink(nn.Module): For sink_size=0: behaves exactly like original CachePositionsManager. For sink_size>0: sink tokens go to fixed positions, rest uses ring buffer. - IMPORTANT: cache_size should be the actual cache dimension size (2x window for ring buffer). + IMPORTANT: cache_size is the actual cache dimension, including sink slots. """ def __init__(self, cache_size: int, sink_size: int = 0): @@ -191,7 +234,7 @@ class KVCacheWithAttentionSink(KVCache): Uses a ring buffer approach for the sliding window portion while keeping the first sink_size tokens fixed. This avoids dynamic shape operations. - Cache layout: [sink: 0 to sink_size-1] [ring_buffer: sink_size to sink_size + window_size*2 - 1] + Cache layout: [fixed sink tokens] [ring buffer for one window + one chunk] """ def __init__( @@ -202,16 +245,27 @@ def __init__( rope: RopeWithAttentionSink, window_size: int, sink_size: int, + max_context_length: int, + max_seq_len: Optional[int] = None, max_batch_size: int = 1, dtype=torch.float32, ): - # Total cache size is sink_size + window_size * 2. - # The ring buffer needs 2x the window size because at the moment a new - # token is written, the previous window_size tokens must still be readable - # (they haven't been overwritten yet). With only 1x, writing a new entry - # would immediately evict the oldest visible token, leaving fewer than - # window_size tokens available for attention. - total_cache_size = sink_size + window_size * 2 + self.full_context_length = max_context_length + self.max_seq_len = ( + max_context_length if max_seq_len is None else int(max_seq_len) + ) + # Keep the fixed sinks separate from the ring space needed for the + # retained window and the largest in-flight input chunk. + total_cache_size = _get_attention_sink_cache_size( + max_context_length, + window_size, + sink_size, + self.max_seq_len, + ) + assert rope.ring_size == total_cache_size - sink_size, ( + f"RoPE ring size ({rope.ring_size}) must match the KV-cache ring " + f"size ({total_cache_size - sink_size})" + ) super().__init__( max_batch_size=max_batch_size, max_context_length=total_cache_size, @@ -310,6 +364,7 @@ def _replace_attention( rope_with_attention_sink: RopeWithAttentionSink, sink_size: int, window_size: int, + max_seq_len: Optional[int], ): for _, child_module in module._modules.items(): if len(list(child_module.children())) > 0: # pyre-ignore [16] @@ -318,32 +373,23 @@ def _replace_attention( rope_with_attention_sink=rope_with_attention_sink, sink_size=sink_size, window_size=window_size, + max_seq_len=max_seq_len, ) if isinstance(child_module, AttentionMHA): kv_cache = child_module.kv_cache - if sink_size == 0: - # No sink tokens needed — use standard RingKVCache directly - child_module.kv_cache = RingKVCache( - kv_cache.max_batch_size, - window_size, # RingKVCache expects user-provided window size - kv_cache.n_heads, - kv_cache.head_dim, - kv_cache.enable_dynamic_shape, - kv_cache.k_cache.dtype, - ) - else: - kv_cache_with_attention_sink = KVCacheWithAttentionSink( - n_heads=kv_cache.n_heads, - head_dim=kv_cache.head_dim, - enable_dynamic_shape=kv_cache.enable_dynamic_shape, - rope=rope_with_attention_sink, - max_batch_size=kv_cache.max_batch_size, - window_size=window_size, - sink_size=sink_size, - dtype=kv_cache.k_cache.dtype, - ) - child_module.kv_cache = kv_cache_with_attention_sink + child_module.kv_cache = KVCacheWithAttentionSink( + n_heads=kv_cache.n_heads, + head_dim=kv_cache.head_dim, + enable_dynamic_shape=kv_cache.enable_dynamic_shape, + rope=rope_with_attention_sink, + max_batch_size=kv_cache.max_batch_size, + window_size=window_size, + sink_size=sink_size, + max_context_length=kv_cache.max_context_length, + max_seq_len=max_seq_len, + dtype=kv_cache.k_cache.dtype, + ) # Don't replace forward - let the original AttentionMHA.forward handle it # since our KVCache has is_ring_buffer=True, it will use the ring buffer mask @@ -365,11 +411,13 @@ def enable_attention_sink( window_size=window_size, sink_size=sink_size, ) + max_seq_len = getattr(params, "max_seq_len", None) _replace_rope(module, rope_with_attention_sink) _replace_attention( module=module, rope_with_attention_sink=rope_with_attention_sink, sink_size=sink_size, window_size=window_size, + max_seq_len=max_seq_len, ) return module diff --git a/examples/models/llama/source_transformation/custom_kv_cache.py b/examples/models/llama/source_transformation/custom_kv_cache.py index dbaac9accf4..7fb0adfb923 100644 --- a/examples/models/llama/source_transformation/custom_kv_cache.py +++ b/examples/models/llama/source_transformation/custom_kv_cache.py @@ -14,6 +14,7 @@ import torch.nn as nn from executorch.examples.models.llama.attention import ( _create_causal_mask_for_ring_buffer, + _get_ring_cache_size, CachePositionsManager, KVCache, RingKVCache, @@ -661,17 +662,19 @@ def _replace_kv_cache_with_custom_kv_cache(module): if sdpa is not None and hasattr(sdpa, "use_attention_mask"): sdpa.use_attention_mask = True elif isinstance(child, RingKVCache): - # RingKVCache (e.g., from attention sink with sink_size=0) needs - # CustomRingKVCache, not plain CustomKVCache + # Preserve ring-buffer sizing and masking when converting a + # local-attention cache to the custom update op. setattr( module, name, CustomRingKVCache( child.max_batch_size, - child.window_size, + child.full_context_length, child.n_heads, child.head_dim, dtype=child.k_cache.dtype, + window_size=child.window_size, + max_seq_len=child.max_seq_len, ), ) sdpa = getattr(module, "SDPA", None) @@ -707,11 +710,20 @@ def __init__( cache_type: QuantizedCacheType = QuantizedCacheType.AffineSymmetric, use_custom_update_cache_op: bool = False, return_float_values: bool = True, + *, + window_size: int, + max_seq_len: Optional[int] = None, ): - # Look at attention.py for explanation on why max_context_length * 2 + self.full_context_length = max_context_length + self.max_seq_len = ( + max_context_length if max_seq_len is None else int(max_seq_len) + ) + ring_cache_size = _get_ring_cache_size( + max_context_length, window_size, self.max_seq_len + ) super().__init__( max_batch_size, - max_context_length * 2, + ring_cache_size, n_heads, head_dim, cache_type, @@ -720,7 +732,7 @@ def __init__( ) self.cache_positions_manager = CachePositionsManager(self.max_context_length) self.is_ring_buffer = True - self.window_size = max_context_length + self.window_size = window_size def create_causal_mask_for_ring_buffer(self, start_pos, seq_len): cache_positions = self.cache_positions_manager.cache_positions @@ -742,7 +754,7 @@ def update(self, input_pos, k_val, v_val): seq_len = k_val.transpose(1, 2).size(1) assert seq_len <= self.k_cache.size( 1 - ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(2)})" + ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(1)})" indices = self.cache_positions_manager.calculate_positions_and_update_indices( input_pos, seq_len ) @@ -755,19 +767,22 @@ def from_quantized_kv_cache( cls, kv_cache, sliding_window_size, + max_seq_len, ): assert isinstance( kv_cache, QuantizedKVCache ), "For QuantizedRingKVCache expect QuantizedKVCache as input kv_cache" - max_batch_size, _, n_heads, head_dim = kv_cache.k_cache.shape + max_batch_size, max_context_length, n_heads, head_dim = kv_cache.k_cache.shape return cls( max_batch_size, - sliding_window_size, + max_context_length, n_heads, head_dim, kv_cache.cache_type, kv_cache.use_custom_update_cache_op, kv_cache.return_float_values, + window_size=sliding_window_size, + max_seq_len=max_seq_len, ) @@ -783,19 +798,30 @@ class CustomKVCacheWithAttentionSink(CustomKVCache): def __init__( self, max_batch_size, + max_context_length, n_heads, head_dim, window_size, sink_size, + max_seq_len=None, dtype=torch.float32, ): - # Total cache size: sink slots + ring buffer (2x window for wrap safety) - total_cache_size = sink_size + window_size * 2 - super().__init__(max_batch_size, total_cache_size, n_heads, head_dim, dtype) from executorch.examples.models.llama.source_transformation.attention_sink import ( + _get_attention_sink_cache_size, CachePositionsManagerWithSink, ) + self.full_context_length = max_context_length + self.max_seq_len = ( + max_context_length if max_seq_len is None else int(max_seq_len) + ) + total_cache_size = _get_attention_sink_cache_size( + max_context_length, + window_size, + sink_size, + self.max_seq_len, + ) + super().__init__(max_batch_size, total_cache_size, n_heads, head_dim, dtype) self.cache_positions_manager = CachePositionsManagerWithSink( total_cache_size, sink_size ) @@ -846,10 +872,12 @@ def from_kv_cache_with_attention_sink(cls, kv_cache): max_batch_size, n_heads, _, head_dim = kv_cache.k_cache.shape return cls( max_batch_size, + kv_cache.full_context_length, n_heads, head_dim, kv_cache.window_size, kv_cache.sink_size, + max_seq_len=kv_cache.max_seq_len, dtype=kv_cache.k_cache.dtype, ) @@ -862,14 +890,21 @@ def __init__( n_heads, head_dim, dtype=torch.float32, + *, + window_size: int, + max_seq_len: Optional[int] = None, ): - # Look at attention.py for explanation on why max_context_length * 2 - super().__init__( - max_batch_size, max_context_length * 2, n_heads, head_dim, dtype + self.full_context_length = max_context_length + self.max_seq_len = ( + max_context_length if max_seq_len is None else int(max_seq_len) + ) + ring_cache_size = _get_ring_cache_size( + max_context_length, window_size, self.max_seq_len ) + super().__init__(max_batch_size, ring_cache_size, n_heads, head_dim, dtype) self.cache_positions_manager = CachePositionsManager(self.max_context_length) self.is_ring_buffer = True - self.window_size = max_context_length + self.window_size = window_size def create_causal_mask_for_ring_buffer(self, start_pos, seq_len): cache_positions = self.cache_positions_manager.cache_positions @@ -891,7 +926,7 @@ def update(self, input_pos, k_val, v_val): seq_len = k_val.transpose(1, 2).size(1) assert seq_len <= self.k_cache.size( 1 - ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(2)})" + ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(1)})" indices = self.cache_positions_manager.calculate_positions_and_update_indices( input_pos, seq_len ) @@ -904,21 +939,24 @@ def from_custom_kv_cache( cls, kv_cache, sliding_window_size, + max_seq_len, ): - max_batch_size, n_heads, _, head_dim = kv_cache.k_cache.shape - if isinstance(kv_cache, CustomKVCache): - # If replacing custom kv cache, then the shape is [B, S, H, D] - max_batch_size, _, n_heads, head_dim = kv_cache.k_cache.shape + # CustomKVCache storage is [B, S, H, D]. + max_batch_size, max_context_length, n_heads, head_dim = kv_cache.k_cache.shape return cls( max_batch_size, - sliding_window_size, + max_context_length, n_heads, head_dim, dtype=kv_cache.k_cache.dtype, + window_size=sliding_window_size, + max_seq_len=max_seq_len, ) -def _replace_kv_cache_with_ring_kv_cache(attention, layer_size): +def _replace_kv_cache_with_ring_kv_cache( + attention, layer_size: int, max_seq_len: Optional[int] +): sliding_window_size = layer_size assert ( getattr(attention, "kv_cache", None) is not None @@ -927,23 +965,28 @@ def _replace_kv_cache_with_ring_kv_cache(attention, layer_size): if isinstance(kv_cache, KVCache): attention.kv_cache = RingKVCache( kv_cache.max_batch_size, - sliding_window_size, + kv_cache.max_context_length, kv_cache.n_heads, kv_cache.head_dim, kv_cache.enable_dynamic_shape, kv_cache.k_cache.dtype, + window_size=sliding_window_size, + max_seq_len=max_seq_len, ) elif isinstance(kv_cache, CustomKVCache): attention.kv_cache = CustomRingKVCache.from_custom_kv_cache( - kv_cache, layer_size + kv_cache, layer_size, max_seq_len ) elif isinstance(kv_cache, QuantizedKVCache): attention.kv_cache = QuantizedRingKVCache.from_quantized_kv_cache( - kv_cache, layer_size + kv_cache, layer_size, max_seq_len ) -def replace_kv_cache_with_ring_kv_cache(module, layer_sizes): +def replace_kv_cache_with_ring_kv_cache( + module, layer_sizes, max_seq_len: Optional[int] = None +): + """Replace local-layer caches with window-plus-in-flight ring caches.""" # This is needed to ensure that custom ops are registered from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 @@ -959,6 +1002,7 @@ def replace_kv_cache_with_ring_kv_cache(module, layer_sizes): logging.info( f"Applying local sliding window attention with following pattern {layer_sizes}." ) + logged_full_context_cache = False assert len(layer_sizes) == len( module.layers ), f"Length of layer sizes {len(layer_sizes)} must match the number of layers in the module {len(module.layers)}." @@ -970,7 +1014,23 @@ def replace_kv_cache_with_ring_kv_cache(module, layer_sizes): getattr(transformer_block, "attention", None) is not None ), f"Transfomer block must have attention module. Transformer block {transformer_block}" attention = transformer_block.attention - _replace_kv_cache_with_ring_kv_cache(attention, sliding_window_size) + full_context_length = attention.kv_cache.max_context_length + effective_max_seq_len = ( + full_context_length if max_seq_len is None else max_seq_len + ) + if ( + not logged_full_context_cache + and sliding_window_size + effective_max_seq_len >= full_context_length + ): + logging.info( + "Local attention KV caches will use the full context length; " + "set max_seq_length below max_context_length - window_size " + "to retain KV-cache memory savings." + ) + logged_full_context_cache = True + _replace_kv_cache_with_ring_kv_cache( + attention, sliding_window_size, max_seq_len + ) # if attention's sdpa is custom sdpa then we have to make sure # it is not doing causal attention if "SDPACustom" in attention.SDPA.__class__.__name__: diff --git a/examples/models/llama/source_transformation/test_attention_sink.py b/examples/models/llama/source_transformation/test_attention_sink.py index c4105338f0e..f5fb0c4d063 100644 --- a/examples/models/llama/source_transformation/test_attention_sink.py +++ b/examples/models/llama/source_transformation/test_attention_sink.py @@ -69,7 +69,10 @@ def setUp(self) -> None: # Ring top is 20. The table is deliberately longer, so a slice running # past the ring still lands on real rows instead of going out of bounds. self.params = ModelArgs( - use_kv_cache=True, enable_dynamic_shape=True, max_context_len=64 + use_kv_cache=True, + enable_dynamic_shape=True, + max_context_len=64, + max_seq_len=self.WINDOW_SIZE, ) self.rope = RopeWithAttentionSink( params=self.params, @@ -226,14 +229,15 @@ def setUp(self): use_kv_cache=True, enable_dynamic_shape=True, max_context_len=256, + max_seq_len=self.window_size, ) self.rope = RopeWithAttentionSink( params=self.params, window_size=self.window_size, sink_size=self.sink_size, ) - # Total cache size = sink_size + window_size * 2 = 4 + 56 = 60 - self.cache_size = self.sink_size + self.window_size * 2 + # Total cache size = sink_size + window_size + max_seq_len = 60. + self.cache_size = self.sink_size + self.window_size + self.params.max_seq_len self.kv_cache = KVCacheWithAttentionSink( n_heads=self.params.n_heads, head_dim=self.params.head_dim, @@ -242,6 +246,8 @@ def setUp(self): max_batch_size=self.max_batch_size, window_size=self.window_size, sink_size=self.sink_size, + max_context_length=self.params.max_context_len, + max_seq_len=self.params.max_seq_len, dtype=self.dtype, ) @@ -372,11 +378,15 @@ def test_causal_mask_blocks_future(self): ) def test_no_sink_degenerates_to_ring_buffer(self, sink_size): """With sink_size=0, behavior should match a plain ring buffer.""" + window_size = 100 params = ModelArgs( - use_kv_cache=True, enable_dynamic_shape=True, max_context_len=256 + use_kv_cache=True, + enable_dynamic_shape=True, + max_context_len=128, + max_seq_len=64, ) rope = RopeWithAttentionSink( - params=params, window_size=self.window_size, sink_size=0 + params=params, window_size=window_size, sink_size=0 ) cache = KVCacheWithAttentionSink( n_heads=params.n_heads, @@ -384,11 +394,16 @@ def test_no_sink_degenerates_to_ring_buffer(self, sink_size): enable_dynamic_shape=params.enable_dynamic_shape, rope=rope, max_batch_size=1, - window_size=self.window_size, + window_size=window_size, sink_size=0, + max_context_length=params.max_context_len, + max_seq_len=params.max_seq_len, dtype=self.dtype, ) - cache_size = self.window_size * 2 # 56 + cache_size = window_size + params.max_seq_len + self.assertEqual(cache_size, 164) + self.assertEqual(rope.ring_size, cache_size) + self.assertEqual(cache.max_context_length, cache_size) # Fill and wrap k_init, v_init = self._rand_kv(cache_size) @@ -434,7 +449,10 @@ def _build_model(self, args, sink_size, window_size, use_custom_sdpa=False): model = construct_transformer(args) model = enable_attention_sink( - model, params=args, sink_size=sink_size, window_size=window_size + model, + params=args, + sink_size=sink_size, + window_size=window_size, ) if use_custom_sdpa: @@ -502,12 +520,12 @@ def test_beyond_context_window_basic(self): """Generate tokens well beyond the KV cache size using standard SDPA.""" sink_size = 4 window_size = 16 - # KV cache size = sink_size + window_size * 2 = 36 + # KV cache size = sink_size + window_size + max_seq_len = 52 # max_context_len = 128 (for RoPE table) args = self._make_args(max_context_len=128) model = self._build_model(args, sink_size, window_size, use_custom_sdpa=False) - # Generate 80 tokens — well beyond KV cache size of 36 + # Generate 80 tokens — beyond the KV cache size of 52 outputs = self._run_generation(model, args, num_tokens=80) self.assertEqual(len(outputs), 77) # 1 prefill + 76 decode steps @@ -520,10 +538,13 @@ def test_beyond_max_context_len(self): """Generate tokens beyond max_context_len with RoPE position remapping.""" sink_size = 4 window_size = 16 - # KV cache size = 36, max_context_len = 64 + # With max_seq_len omitted, the cache is capped at max_context_len. # Generate 100 tokens — well beyond max_context_len args = self._make_args(max_context_len=64) + args.max_seq_len = None model = self._build_model(args, sink_size, window_size, use_custom_sdpa=False) + cache = model.layers[0].attention.kv_cache + self.assertEqual(cache.max_context_length, 84) outputs = self._run_generation(model, args, num_tokens=100) @@ -537,10 +558,9 @@ def test_beyond_max_context_len(self): def test_chunked_prefill_across_the_ring_wrap(self): """Chunked prefill where a chunk spans the ring wrap. - sink_size=4, window_size=16, so the ring is slots [4, 36). Feeding 5 - tokens at a time puts chunk starts at 0, 5, ..., 95. The chunk at 35 - covers positions 35..39 and needs rows 35, 4, 5, 6, 7; the chunk at 65 - covers 65..69 and needs rows 33, 34, 35, 4, 5. Both span the wrap. + sink_size=4, window_size=16, and max_seq_len=48, so the ring is slots + [4, 68). Feeding 40 tokens at a time exceeds the old 2x-window ring and + crosses the new ring boundary on the second chunk. The other beyond-context-window tests decode one token at a time, and a chunk of one can never span the wrap however far the position runs, so @@ -554,27 +574,32 @@ def test_chunked_prefill_across_the_ring_wrap(self): """ sink_size = 4 window_size = 16 - args = self._make_args(max_context_len=64) + chunk_size = 40 + args = self._make_args(max_context_len=128) + args.max_seq_len = 48 torch.manual_seed(0) model = self._build_model(args, sink_size, window_size) - tokens = torch.randint(0, args.vocab_size, (1, 100)) + tokens = torch.randint(0, args.vocab_size, (1, 80)) - chunked = self._feed_in_chunks(copy.deepcopy(model), tokens, chunk_size=5) + chunked = self._feed_in_chunks( + copy.deepcopy(model), tokens, chunk_size=chunk_size + ) one_at_a_time = self._feed_in_chunks(copy.deepcopy(model), tokens, chunk_size=1) - self.assertEqual(len(chunked), 20) - self.assertEqual(len(one_at_a_time), 100) + self.assertEqual(len(chunked), 2) + self.assertEqual(len(one_at_a_time), 80) # generate_full_logits is off, so each call returns logits for its last # position only. How the input was chunked must not change the result: - # chunk i ends on the same token as one_at_a_time[5 * i + 4]. + # chunk i ends on the same token as the corresponding decode call. for i, out in enumerate(chunked): self.assertTrue(torch.isfinite(out).all(), f"chunk {i} is not finite") torch.testing.assert_close( out, - one_at_a_time[5 * i + 4], - msg=lambda m, i=i: f"chunk {i}, positions {5 * i}..{5 * i + 4}, " + one_at_a_time[chunk_size * (i + 1) - 1], + msg=lambda m, i=i: f"chunk {i}, positions {chunk_size * i}.." + f"{chunk_size * (i + 1) - 1}, " f"disagrees with feeding the same tokens one at a time:\n{m}", ) @@ -599,7 +624,7 @@ def test_beyond_context_window_custom_sdpa(self): found_custom_cache, "Expected CustomKVCacheWithAttentionSink in model" ) - # Generate 80 tokens — well beyond KV cache size of 36 + # Generate 80 tokens — beyond the KV cache size of 52 outputs = self._run_generation(model, args, num_tokens=80) self.assertEqual(len(outputs), 77) diff --git a/examples/models/llama/tests/test_replace_kv_cache.py b/examples/models/llama/tests/test_replace_kv_cache.py index 383fe6d6882..a4595b4311b 100644 --- a/examples/models/llama/tests/test_replace_kv_cache.py +++ b/examples/models/llama/tests/test_replace_kv_cache.py @@ -89,13 +89,14 @@ def test_replace_kv_cache_with_ring_kv_cache(self): # Replace KVCache with RingKVCache layer_sizes = [8] # Sliding window size for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes, max_seq_len=4) # Verify that KVCache has been replaced with RingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, RingKVCache) # Verify that the sliding window size is set correctly self.assertEqual(model.layers[0].attention.kv_cache.window_size, layer_sizes[0]) + self.assertEqual(model.layers[0].attention.kv_cache.k_cache.size(2), 12) def test_replace_custom_kv_cache_with_custom_ring_kv_cache(self): """Test replacing CustomKVCache with CustomRingKVCache.""" @@ -112,10 +113,11 @@ def test_replace_custom_kv_cache_with_custom_ring_kv_cache(self): # Replace CustomKVCache with CustomRingKVCache layer_sizes = [8] # Sliding window size for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes, max_seq_len=4) # Verify that CustomKVCache has been replaced with CustomRingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, CustomRingKVCache) + self.assertEqual(model.layers[0].attention.kv_cache.k_cache.size(1), 12) def test_replace_quantized_kv_cache_with_quantized_ring_kv_cache(self): """Test replacing QuantizedKVCache with QuantizedRingKVCache.""" @@ -134,10 +136,11 @@ def test_replace_quantized_kv_cache_with_quantized_ring_kv_cache(self): # Replace QuantizedKVCache with QuantizedRingKVCache layer_sizes = [8] # Sliding window size for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes, max_seq_len=4) # Verify that QuantizedKVCache has been replaced with QuantizedRingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, QuantizedRingKVCache) + self.assertEqual(model.layers[0].attention.kv_cache.k_cache.size(1), 12) def test_replace_static_quantized_kv_cache(self): """Test replacing KVCache with static-qparams int8 KV storage.""" @@ -287,6 +290,7 @@ def test_static_quantized_kv_cache_rejects_specialized_cache(self): self.n_kv_heads, self.head_dim, self.enable_dynamic_shape, + window_size=self.max_context_len, ) model = self._create_mock_model([attention]) @@ -305,7 +309,9 @@ def test_multiple_layers_with_different_window_sizes(self): # Replace KVCache with RingKVCache with different window sizes layer_sizes = [4, 8, 16] # Different sliding window sizes for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes) + replace_kv_cache_with_ring_kv_cache( + model, layer_sizes, max_seq_len=self.max_context_len + ) # Verify that each layer has the correct window size self.assertIsInstance(model.layers[0].attention.kv_cache, RingKVCache) diff --git a/examples/models/llama/tests/test_ring_attention.py b/examples/models/llama/tests/test_ring_attention.py index ae440e00e47..5ab696bc5f5 100644 --- a/examples/models/llama/tests/test_ring_attention.py +++ b/examples/models/llama/tests/test_ring_attention.py @@ -84,7 +84,10 @@ def _create_baseline_attention( return attention def _create_ring_attention( - self, attention, kv_cache_type: KVCacheType = KVCacheType.REGULAR + self, + attention, + max_seq_len, + kv_cache_type: KVCacheType = KVCacheType.REGULAR, ): """Create attention with ring buffer KV cache.""" assert self.sliding_window is not None @@ -98,25 +101,82 @@ def _create_ring_attention( baseline_attention.kv_cache = QuantizedRingKVCache.from_quantized_kv_cache( baseline_attention.kv_cache, self.sliding_window, + max_seq_len, ) elif isinstance(baseline_attention.kv_cache, CustomKVCache): # Replace CustomKVCache with CustomRingKVCache baseline_attention.kv_cache = CustomRingKVCache.from_custom_kv_cache( baseline_attention.kv_cache, self.sliding_window, + max_seq_len, ) else: # Replace regular KVCache with RingKVCache baseline_attention.kv_cache = RingKVCache( self.args.max_batch_size, - self.sliding_window, + self.args.max_context_len, self.n_kv_heads, self.head_dim, self.args.enable_dynamic_shape, self.dtype, + window_size=self.sliding_window, + max_seq_len=max_seq_len, ) return baseline_attention + def test_continuation_prefill_larger_than_window( + self, kv_cache_type: KVCacheType = KVCacheType.REGULAR + ): + """A W+C cache preserves history when the incoming chunk exceeds W.""" + self.sliding_window = 4 + chunk_size = 6 + baseline_attn = self._create_baseline_attention(12, kv_cache_type) + ring_attn = self._create_ring_attention( + baseline_attn, chunk_size, kv_cache_type + ) + + self.assertEqual(ring_attn.kv_cache.max_context_length, 10) + + with torch.nn.attention.sdpa_kernel( + [SDPBackend.FLASH_ATTENTION] + ), torch.no_grad(): + for pos in (0, chunk_size): + x = torch.randn( + (self.batch_size, chunk_size, self.dim), dtype=self.dtype + ) + input_pos = torch.tensor([pos], dtype=torch.long) + freqs_cos, freqs_sin = self.rope.get_freqs(input_pos, chunk_size) + + baseline_out, _ = baseline_attn.forward( + x, freqs_cos, freqs_sin, input_pos=input_pos + ) + ring_out, _ = ring_attn.forward( + x, freqs_cos, freqs_sin, input_pos=input_pos + ) + + tolerance = 1e-6 if kv_cache_type != KVCacheType.REGULAR else 1e-7 + self.assertTrue( + torch.allclose( + baseline_out, + ring_out, + rtol=tolerance, + atol=tolerance, + ), + f"Outputs differ at position {pos}", + ) + + def test_continuation_prefill_larger_than_window_quantized(self): + self._run_test_with_kv_cache_type( + self.test_continuation_prefill_larger_than_window, + KVCacheType.QUANTIZED, + ) + + def test_continuation_prefill_larger_than_window_custom(self): + self._run_test_with_kv_cache_type( + self.test_continuation_prefill_larger_than_window, + KVCacheType.CUSTOM, + ) + def _create_sliding_window_mask(self, seq_len, context_len, window_size): """Create a sliding window mask for the baseline.""" mask = torch.full((seq_len, context_len), float("-inf"), dtype=self.dtype) @@ -140,7 +200,7 @@ def test_single_token_processing( seq_len = 10 self.sliding_window = 4 baseline_attn = self._create_baseline_attention(seq_len, kv_cache_type) - ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) + ring_attn = self._create_ring_attention(baseline_attn, 1, kv_cache_type) # Process tokens one by one with torch.nn.attention.sdpa_kernel( @@ -199,7 +259,7 @@ def test_sliding_window_attention( baseline_attn = self._create_baseline_attention(seq_len, kv_cache_type) # Create ring attention with sliding window size - ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) + ring_attn = self._create_ring_attention(baseline_attn, 1, kv_cache_type) # Process tokens one by one with torch.nn.attention.sdpa_kernel( @@ -251,7 +311,7 @@ def test_ring_buffer_wrapping( ) # Create ring attention with sliding window size - ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) + ring_attn = self._create_ring_attention(baseline_attn, 1, kv_cache_type) # Process enough tokens to cause wrapping seq_len = 1 @@ -277,13 +337,12 @@ def test_ring_buffer_wrapping( f"Outputs differ at position {pos}", ) - # After processing 8 tokens with window size 4, the ring buffer should have wrapped around + # With W=3 and max_seq_len=1, the physical cache has 4 slots and wraps. # Check the cache positions to verify wrapping cache_positions = ring_attn.kv_cache.cache_positions_manager.cache_positions - # The cache positions should contain the most recent 4 positions (4, 5, 6, 7) - # mapped to the ring buffer indices - expected_positions = torch.tensor([6, 7, 2, 3, 4, 5], dtype=torch.long) + # Positions 4 through 7 occupy physical slots 0 through 3. + expected_positions = torch.tensor([4, 5, 6, 7], dtype=torch.long) self.assertTrue( torch.all(cache_positions == expected_positions), @@ -316,7 +375,9 @@ def test_large_context_with_sliding_window( baseline_attn = self._create_baseline_attention(seq_len, kv_cache_type) # Create ring attention with sliding window size - ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) + ring_attn = self._create_ring_attention( + baseline_attn, max(token_lens), kv_cache_type + ) pos = 0 with torch.nn.attention.sdpa_kernel( diff --git a/examples/models/llama/tests/test_ring_kv_cache.py b/examples/models/llama/tests/test_ring_kv_cache.py index f923a1bf2bf..561ab117078 100644 --- a/examples/models/llama/tests/test_ring_kv_cache.py +++ b/examples/models/llama/tests/test_ring_kv_cache.py @@ -14,7 +14,9 @@ class TestRingKVCache(unittest.TestCase): def setUp(self): # Common test parameters self.max_batch_size = 2 - self.max_context_length = 8 + self.max_context_length = 16 + self.window_size = 8 + self.max_seq_len = 8 self.n_heads = 4 self.head_dim = 16 self.enable_dynamic_shape = True @@ -32,7 +34,7 @@ def test_dynamic_kv_cache_update_on_cuda(self): self._require_usable_cuda() cache = KVCache( max_batch_size=1, - max_context_length=self.max_context_length, + max_context_length=self.window_size, n_heads=self.n_heads, head_dim=self.head_dim, enable_dynamic_shape=True, @@ -60,6 +62,8 @@ def test_ring_cache_positions_and_mask_on_cuda(self): head_dim=self.head_dim, enable_dynamic_shape=True, dtype=self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ).cuda() input_pos = torch.tensor([0], dtype=torch.long, device="cuda") seq_len = 3 @@ -103,8 +107,11 @@ def test_basic_update(self): self.head_dim, self.enable_dynamic_shape, self.dtype, + window_size=self.window_size, ) + self.assertEqual(cache.max_seq_len, self.max_context_length) + # Create input tensors input_pos = torch.tensor([0], dtype=torch.long) seq_len = 3 @@ -129,7 +136,7 @@ def test_basic_update(self): self.assertTrue(torch.all(v_out[:, :, i] == 2.0)) # Check that the rest of the cache is still zeros - for i in range(seq_len, self.max_context_length): + for i in range(seq_len, self.window_size): self.assertTrue(torch.all(k_out[:, :, i] == 0.0)) self.assertTrue(torch.all(v_out[:, :, i] == 0.0)) @@ -153,6 +160,8 @@ def test_ring_buffer_wrapping(self): self.head_dim, self.enable_dynamic_shape, self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ) # Create input tensors for first update @@ -216,6 +225,8 @@ def test_multiple_updates(self): self.head_dim, self.enable_dynamic_shape, self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ) # First update @@ -346,6 +357,8 @@ def test_edge_case_input_pos_zero(self): self.head_dim, self.enable_dynamic_shape, self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ) # Create input tensors @@ -374,7 +387,7 @@ def test_edge_case_input_pos_zero(self): self.assertTrue(torch.all(v_out[:, :, 0] == 12.0)) # Check that the rest of the cache is still zeros - for i in range(1, self.max_context_length): + for i in range(1, self.window_size): self.assertTrue(torch.all(k_out[:, :, i] == 0.0)) self.assertTrue(torch.all(v_out[:, :, i] == 0.0)) @@ -390,7 +403,7 @@ def test_edge_case_input_pos_zero(self): ) def test_edge_case_exceeding_context_length(self): - """Test the edge case where input_pos + seq_len > max_context_length.""" + """Test the edge case where input_pos + seq_len exceeds cache capacity.""" cache = RingKVCache( self.max_batch_size, self.max_context_length, @@ -398,6 +411,8 @@ def test_edge_case_exceeding_context_length(self): self.head_dim, self.enable_dynamic_shape, self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ) # Create input tensors @@ -462,6 +477,8 @@ def test_original_indices_tracking(self): self.head_dim, self.enable_dynamic_shape, self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ) # First update at position 10 (will be mapped to position 10 in the ring buffer) @@ -481,7 +498,7 @@ def test_original_indices_tracking(self): # Check that cache_positions correctly tracks the original indices # For input_pos=10 and seq_len=4, the original indices should be 10, 11, 12, 13 - # These map to positions 10, 11, 12, 13 in the ring buffer (since max_context_length=8 but buffer size is 16) + # These map directly to positions 10, 11, 12, 13 in the 16-slot cache. # Note that positions 0-9 are 0 because in actual ring # updates those positions would have been updated for start_pos = 0. # So CachePositionsManager thinks they are updated because start_pos > (0-9) @@ -532,6 +549,8 @@ def test_non_dynamic_shape(self): self.head_dim, enable_dynamic_shape=False, dtype=self.dtype, + window_size=self.window_size, + max_seq_len=self.max_seq_len, ) # Create input tensors @@ -561,6 +580,6 @@ def test_non_dynamic_shape(self): self.assertTrue(torch.all(v_out[:, :, i] == 16.0)) # Check that the rest of the cache is still zeros - for i in range(seq_len, self.max_context_length): + for i in range(seq_len, self.window_size): self.assertTrue(torch.all(k_out[:, :, i] == 0.0)) self.assertTrue(torch.all(v_out[:, :, i] == 0.0)) diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index f6ed378ec03..fac940cf3ee 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -16,8 +16,6 @@ // @lint-ignore CLANGTIDY facebook-unused-include-check #include -#include - #ifdef ET_USE_THREADPOOL #include #include @@ -990,6 +988,193 @@ void cpu_flash_attention( scalar_t* buf_reduced_data = is_reduced_type ? reinterpret_cast(buf_reduced) : nullptr; + // An explicit mask is shared across batches and heads. Precompute up to two + // useful K/V intervals for every (query tile, K/V tile) pair. Ring attention + // can mask large portions of its backing cache with -inf; discovering those + // ranges once lets every head skip fully masked tiles and trim the boundary + // tiles before either GEMM. Arbitrary additive masks retain their existing + // behavior because only values that are exactly -inf are excluded. + struct MaskBlockRanges { + int64_t first_begin; + int64_t first_end; + int64_t second_begin; + int64_t second_end; + bool first_mask_is_zero; + bool second_mask_is_zero; + }; + const int64_t num_kv_blocks = (kvSize - 1) / kvSplitSize + 1; + MaskBlockRanges* mask_block_ranges = nullptr; + std::unique_ptr allocated_mask_block_ranges; + uint8_t* column_states_by_thread = nullptr; + std::unique_ptr allocated_column_states; + bool use_mask_ranges = false; + if (has_attn_mask) { + const int64_t num_mask_block_ranges = qSlice * num_kv_blocks; + const int64_t mask_block_ranges_bytes = + num_mask_block_ranges * sizeof(MaskBlockRanges); + Result mask_block_ranges_scratch = + ctx.allocate_temp(mask_block_ranges_bytes, 64); + if (!mask_block_ranges_scratch.ok()) { + allocated_mask_block_ranges = + std::make_unique(mask_block_ranges_bytes); + mask_block_ranges = + reinterpret_cast(allocated_mask_block_ranges.get()); + } else { + mask_block_ranges = + reinterpret_cast(mask_block_ranges_scratch.get()); + } + + // Bit 0 means at least one query row can attend to the column. Bit 1 + // means every query row has a zero additive mask for the column. + const int64_t column_states_bytes = kvSplitSize * num_thread + qSlice; + Result column_states_scratch = + ctx.allocate_temp(column_states_bytes, 64); + if (!column_states_scratch.ok()) { + allocated_column_states = std::make_unique(column_states_bytes); + column_states_by_thread = + reinterpret_cast(allocated_column_states.get()); + } else { + column_states_by_thread = + reinterpret_cast(column_states_scratch.get()); + } + + const accum_t neg_inf = -std::numeric_limits::infinity(); + uint8_t* useful_mask_ranges = + column_states_by_thread + kvSplitSize * num_thread; + auto find_mask_ranges = [&](int64_t begin, int64_t end) { + const int64_t thread_index = torch::executor::get_thread_num(); + uint8_t* column_states = + column_states_by_thread + thread_index * kvSplitSize; + for (int64_t q_block = begin; q_block < end; ++q_block) { + bool found_useful_range = false; + const int64_t query_begin = q_block * qSplitSize; + const int64_t query_end = std::min(query_begin + qSplitSize, qSize); + const int64_t causal_end = + is_causal ? std::min(start_pos + query_end, kvSize) : kvSize; + for (int64_t kv_block = 0; kv_block < num_kv_blocks; ++kv_block) { + const int64_t block_begin = kv_block * kvSplitSize; + const int64_t block_end = + std::min(std::min(block_begin + kvSplitSize, kvSize), causal_end); + const int64_t range_index = q_block * num_kv_blocks + kv_block; + auto& ranges = mask_block_ranges[range_index]; + ranges.first_begin = block_begin; + ranges.first_end = block_begin; + ranges.second_begin = block_begin; + ranges.second_end = block_begin; + ranges.first_mask_is_zero = false; + ranges.second_mask_is_zero = false; + if (block_begin >= block_end) { + continue; + } + + const int64_t block_size = block_end - block_begin; + std::fill(column_states, column_states + block_size, uint8_t{2}); + + for (int64_t row = query_begin; row < query_end; ++row) { + const accum_t* mask_row = mask_data + row * mStrideM; + for (int64_t col = block_begin; col < block_end; ++col) { + const accum_t mask_value = mask_row[col]; + auto& state = column_states[col - block_begin]; + state |= mask_value != neg_inf; + if (mask_value != static_cast(0)) { + state &= uint8_t{1}; + } + } + } + + int64_t active_begin = block_begin; + while (active_begin < block_end && + (column_states[active_begin - block_begin] & uint8_t{1}) == + 0) { + ++active_begin; + } + int64_t active_end = block_end; + while (active_end > active_begin && + (column_states[active_end - 1 - block_begin] & uint8_t{1}) == + 0) { + --active_end; + } + ranges.first_begin = active_begin; + ranges.first_end = active_end; + ranges.second_begin = block_end; + ranges.second_end = block_end; + if (active_begin >= active_end) { + found_useful_range = true; + continue; + } + + // A wrapped ring window has at most one interior gap. custom_sdpa + // also accepts arbitrary additive masks, which may contain several + // gaps, so find the largest one and split around it when it is large + // enough to repay the extra pair of GEMM calls. Smaller gaps stay + // represented by their existing -inf values inside one bounding + // interval. + int64_t largest_gap_begin = active_begin; + int64_t largest_gap_end = active_begin; + int64_t gap_begin = active_begin; + while (gap_begin < active_end) { + while (gap_begin < active_end && + (column_states[gap_begin - block_begin] & uint8_t{1}) != 0) { + ++gap_begin; + } + int64_t gap_end = gap_begin; + while (gap_end < active_end && + (column_states[gap_end - block_begin] & uint8_t{1}) == 0) { + ++gap_end; + } + if (gap_end - gap_begin > largest_gap_end - largest_gap_begin) { + largest_gap_begin = gap_begin; + largest_gap_end = gap_end; + } + gap_begin = gap_end; + } + + constexpr int64_t min_gap_to_split = 64; + if (largest_gap_end - largest_gap_begin >= min_gap_to_split) { + ranges.first_end = largest_gap_begin; + ranges.second_begin = largest_gap_end; + ranges.second_end = active_end; + } + + auto mask_range_is_zero = [&](int64_t range_begin, + int64_t range_end) { + for (int64_t col = range_begin; col < range_end; ++col) { + if ((column_states[col - block_begin] & uint8_t{2}) == 0) { + return false; + } + } + return true; + }; + ranges.first_mask_is_zero = + mask_range_is_zero(ranges.first_begin, ranges.first_end); + if (ranges.second_begin < ranges.second_end) { + ranges.second_mask_is_zero = + mask_range_is_zero(ranges.second_begin, ranges.second_end); + } + + const int64_t retained_size = ranges.first_end - ranges.first_begin + + ranges.second_end - ranges.second_begin; + if (retained_size < block_size || ranges.first_mask_is_zero || + ranges.second_mask_is_zero) { + found_useful_range = true; + } + } + useful_mask_ranges[q_block] = found_useful_range; + } + }; + const bool mask_ranges_computed = + torch::executor::parallel_for(0, qSlice, 1, find_mask_ranges); + ET_KERNEL_CHECK_MSG( + ctx, + mask_ranges_computed, + Internal, + , + "parallel_for failed while precomputing attention mask ranges"); + for (int64_t q_block = 0; q_block < qSlice; ++q_block) { + use_mask_ranges |= useful_mask_ranges[q_block]; + } + } + auto compute_lambda = [&](int64_t begin, int64_t end) { int64_t i = 0, j = 0, k = 0; data_index_init(begin, i, batchSize, j, num_head, k, qSlice); @@ -1045,11 +1230,38 @@ void cpu_flash_attention( is_causal ? std::min(m + start_pos + qBlockSize, kvSize) : kvSize; int64_t m_start_pos = m + start_pos; auto j_kv = j / num_reps; - fill_stub(dst_data, static_cast(0), qSplitSize * headSize); - for (int64_t n = 0; n < num_keys; n += kvSplitSize) { - int64_t kvBlockSize = std::min(kvSplitSize, kvSize - n); - // Calculate scale * q @ k.T - fill_stub(qk_data, static_cast(0), qSplitSize * kvSplitSize); + fill_stub(dst_data, static_cast(0), qBlockSize * headSize); + bool has_processed_kv = false; + const int64_t num_ranges = use_mask_ranges + ? 2 * num_kv_blocks + : (num_keys + kvSplitSize - 1) / kvSplitSize; + for (int64_t range = 0; range < num_ranges; ++range) { + int64_t kvBlockStart; + int64_t kvBlockEnd; + bool range_mask_is_zero = false; + if (use_mask_ranges) { + const int64_t kv_block = range / 2; + const bool use_second_range = range % 2 != 0; + const auto& ranges = mask_block_ranges[k * num_kv_blocks + kv_block]; + if (use_second_range) { + kvBlockStart = ranges.second_begin; + kvBlockEnd = ranges.second_end; + range_mask_is_zero = ranges.second_mask_is_zero; + } else { + kvBlockStart = ranges.first_begin; + kvBlockEnd = ranges.first_end; + range_mask_is_zero = ranges.first_mask_is_zero; + } + } else { + kvBlockStart = range * kvSplitSize; + kvBlockEnd = std::min(kvBlockStart + kvSplitSize, kvSize); + } + kvBlockEnd = std::min(kvBlockEnd, num_keys); + if (kvBlockStart >= kvBlockEnd) { + continue; + } + const int64_t kvBlockSize = kvBlockEnd - kvBlockStart; + const bool apply_attn_mask = has_attn_mask && !range_mask_is_zero; const void* q_sub_matrix_data_ptr; const void* k_sub_matrix_data_ptr; @@ -1058,12 +1270,14 @@ void cpu_flash_attention( const int8_t* q_zero_points_ptr = nullptr; const int8_t* k_zero_points_ptr = nullptr; int64_t q_offset = i * qStrideB + j * qStrideH + m * qStrideM; - int64_t k_offset = i * kStrideB + j_kv * kStrideH + n * kStrideN; + int64_t k_offset = + i * kStrideB + j_kv * kStrideH + kvBlockStart * kStrideN; if (is_quantized_sdpa) { int64_t q_quant_params_offset = i * q_quant_params_StrideB + j * q_quant_params_StrideH + m * q_quant_params_StrideM; int64_t k_quant_params_offset = i * k_quant_params_StrideB + - j_kv * k_quant_params_StrideH + n * k_quant_params_StrideN; + j_kv * k_quant_params_StrideH + + kvBlockStart * k_quant_params_StrideN; q_scales_ptr = q_scales.value().const_data_ptr() + q_quant_params_offset; k_scales_ptr = @@ -1106,57 +1320,31 @@ void cpu_flash_attention( (widen_qk && qBlockSize >= kMinQBlockForWidenedQK) ? widen_ptr : nullptr); - // There are 4 cases that is_causal has to cover to fill - // not-attendable-position with -inf - /* 1. Everything is attended to. This happens when m_start_pos > n + - kvSplitSize e.g m_pos [8:15] and n_pos [0:7]. Since you must attend to - all previous tokens matrix is full - + + + + + + + + - + + + + + + + + - + + + + + + + + - + + + + + + + + - + + + + + + + + - + + + + + + + + - + + + + + + + + - 2. Everything is not attended to. However only some tokens at the - beginning dont attend to everything. This happens when m_start_pos <= n - + kvSplitSize but m_start_pos + qBlockSize > n + kvSplitSize m_start_pos - = 8 qBlockSize = 8 n = 4 kvSplitSize = 8 For example m_pos [8:15] but - n_pos is [4:11] - + + + + + - - - - + + + + + + - - - + + + + + + + - - + + + + + + + + - + + + + + + + + - + + + + + + + + - + + + + + + + + - + + + + + + + + - 3. In this case only last few tokens have something to attend to. - This happens when m_start_pos < n and m_start_pos + qBlockSize >= n and - m_start_pos + qBlockSize <= n + kvSplitSize m_start_pos = 8 qBlockSize = - 8 n = 13 kvSplitSize = 8 For example m_pos [8:15] but n_pos is [13:20] - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - + - - - - - - - - + + - - - - - - - + + + - - - - - - 4. In this no tokens attend to anything, but we dont really have to - take care of this case because the loop for (int64_t n = 0; n < - num_keys; n += kvSplitSize) will exit before that. - */ - if (is_causal && m_start_pos <= n + kvSplitSize) { - // For this fn to work k_split_size > q_split_size + // Apply causal masking relative to the retained KV block. These are + // the two overlap configurations; a KV block wholly before the new + // query needs no causal masking. Rows are new-query tokens, columns + // are KV-block tokens, '+' is attendable, and '-' is causally masked: + // + // New query begins midway through KV block: Tail of new query lies in + // KV block: + // + + + - - - - - - - - - + // + + + + - - - - - - - - + // + + + + + - + - - - - - + // + + + + + + + + - - - - + // + + + + + + + + + - - - + // + // Each row may attend through its own logical position, so last_col + // is the number of keys in [kvBlockStart, kvBlockEnd) that are not in + // its future. + if (is_causal && m_start_pos < kvBlockEnd) { for (int32_t row = 0; - row < qBlockSize && (m_start_pos + row < n + (kvSplitSize - 1)); + row < qBlockSize && (m_start_pos + row < kvBlockEnd - 1); ++row) { - // When last_col is 0, it means that the entire row is not attended - // to because m_pos is smaller than n_pos. So everything in n is for - // future. - int64_t last_col = - n > (m_start_pos + row) ? 0 : row + m_start_pos + 1 - n; + // When last_col is 0, it means that the entire row is not + // attended to because the range begins after the query position. + int64_t last_col = kvBlockStart > (m_start_pos + row) + ? 0 + : row + m_start_pos + 1 - kvBlockStart; accum_t* row_ptr = qk_data + row * kvBlockSize; fill_stub( row_ptr + last_col, @@ -1167,7 +1355,7 @@ void cpu_flash_attention( // Update attention weights with attention mask // And apply scaling factor // qk <- qk * scaling + attn_mask - if (has_attn_mask) { + if (apply_attn_mask) { for (int64_t row = 0; row < qBlockSize; ++row) { vec::map2( [scaling_factor](Vec x, Vec y) { @@ -1176,14 +1364,14 @@ void cpu_flash_attention( qk_data + row * kvBlockSize, qk_data + row * kvBlockSize, mask_data + i * mStrideB + j * mStrideH + (m + row) * mStrideM + - n, + kvBlockStart, kvBlockSize); } } // Update coefficients with Softmax accum_t tmp_max = 0, tmp_sum = 0, exp_tmp = 0; for (int64_t row = 0; row < qBlockSize; ++row) { - if (has_attn_mask) { + if (apply_attn_mask) { // max per row tmp_max = vec::reduce_all( [](Vec& x, Vec& y) { return vec::maximum(x, y); }, @@ -1220,7 +1408,7 @@ void cpu_flash_attention( // max[row] <- max qk_max_data[row] = tmp_max; // dst <- dst * exp_tmp - if (n > 0) { + if (has_processed_kv) { vec::map( [exp_tmp](Vec x) { return x * Vec(exp_tmp); }, dst_data + row * headSize, @@ -1233,10 +1421,12 @@ void cpu_flash_attention( const void* v_sub_matrix_data_ptr; const float* v_scales_ptr = nullptr; const int8_t* v_zero_points_ptr = nullptr; - int64_t v_offset = i * vStrideB + j_kv * vStrideH + n * vStrideN; + int64_t v_offset = + i * vStrideB + j_kv * vStrideH + kvBlockStart * vStrideN; if (is_quantized_sdpa) { int64_t v_quant_params_offset = i * v_quant_params_StrideB + - j_kv * v_quant_params_StrideH + n * v_quant_params_StrideN; + j_kv * v_quant_params_StrideH + + kvBlockStart * v_quant_params_StrideN; v_scales_ptr = v_scales.value().const_data_ptr() + v_quant_params_offset; v_zero_points_ptr = v_zero_points.value().const_data_ptr() + @@ -1287,10 +1477,12 @@ void cpu_flash_attention( vStrideN, dst_data, headSize, - n == 0 ? static_cast(0) : static_cast(1), + has_processed_kv ? static_cast(1) + : static_cast(0), buf_qdq_ptr, widen_v, use_fp32_qk_weights); + has_processed_kv = true; } // dst <- dst / sum[row] // reorder MHA output with strides diff --git a/extension/llm/custom_ops/op_sdpa_test.cpp b/extension/llm/custom_ops/op_sdpa_test.cpp index 0572fd7ff80..84fdb2e3966 100644 --- a/extension/llm/custom_ops/op_sdpa_test.cpp +++ b/extension/llm/custom_ops/op_sdpa_test.cpp @@ -33,6 +33,26 @@ executorch::aten::Tensor op_scaled_dot_product_attention( context, query, key, value, attn_mask, dropout_p, is_causal, scale, out); } +executorch::aten::Tensor op_custom_sdpa( + const executorch::aten::Tensor& query, + const executorch::aten::Tensor& key, + const executorch::aten::Tensor& value, + const std::optional& attn_mask, + executorch::aten::Tensor& out) { + executorch::runtime::KernelRuntimeContext context{}; + return torch::executor::native::custom_sdpa_out( + context, + query, + key, + value, + /*start_pos=*/0, + attn_mask, + /*dropout_p=*/0.0, + /*is_causal=*/false, + /*scale=*/std::nullopt, + out); +} + std::tuple op_gated_delta_rule( executorch::runtime::KernelRuntimeContext& context, @@ -336,6 +356,151 @@ TEST(OpScaledDotProductAttentionTest, CorrectnessTest_105) { EXPECT_TENSOR_CLOSE_WITH_TOL(ret, ret_expected, 1e-4, 1e-4); } +TEST(OpScaledDotProductAttentionTest, SparseMaskRangesMatchReference) { + TensorFactory tfFloat; + + constexpr int32_t q_size = 64; + constexpr int32_t kv_size = 1024; + std::vector key_values(kv_size, 0.0f); + std::vector value_values(kv_size); + std::vector mask_values( + q_size * kv_size, -std::numeric_limits::infinity()); + std::vector expected_values(q_size); + + for (int32_t i = 0; i < kv_size; ++i) { + value_values[i] = static_cast(i); + } + // The first 32-row query tile has staggered active columns separated by a + // 38-column gap, which is too small to split. This exercises the per-column + // union across query rows while retaining the detailed mask inside the range. + for (int32_t row = 0; row < 32; ++row) { + mask_values[row * kv_size + 10 + row] = 0.0f; + mask_values[row * kv_size + 80 + row] = 0.0f; + expected_values[row] = 45.0f + row; + } + // The second query tile uses KV block 1 and has a gap larger than the split + // threshold. The zero first range and additive second range exercise both + // mask paths while also checking query-tile indexing for k > 0. + for (int32_t row = 32; row < q_size; ++row) { + mask_values[row * kv_size + 600] = 0.0f; + mask_values[row * kv_size + 700] = 1.0986122886681098f; // log(3) + expected_values[row] = 675.0f; + } + + // custom_sdpa uses [batch, sequence, heads, head_dim]. + auto query = tfFloat.zeros({1, q_size, 1, 1}); + auto key = tfFloat.make({1, kv_size, 1, 1}, key_values); + auto value = tfFloat.make({1, kv_size, 1, 1}, value_values); + auto attn_mask = tfFloat.make({q_size, kv_size}, mask_values); + auto out = tfFloat.zeros({1, q_size, 1, 1}); + + auto result = op_custom_sdpa(query, key, value, attn_mask, out); + + auto expected = tfFloat.make({1, q_size, 1, 1}, expected_values); + EXPECT_TENSOR_CLOSE_WITH_TOL(result, expected, 1e-5, 1e-5); +} + +TEST( + OpScaledDotProductAttentionTest, + CausalSparseMaskRangesCoverGqaAndReducedPrecision) { + TensorFactory tfHalf; + TensorFactory tfFloat; + + constexpr int32_t q_size = 256; + constexpr int32_t kv_size = 1024; + constexpr int32_t num_query_heads = 2; + std::vector value_values(kv_size); + std::vector mask_values( + q_size * kv_size, -std::numeric_limits::infinity()); + std::vector expected_values(num_query_heads * q_size); + + for (int32_t col = 0; col < kv_size; ++col) { + value_values[col] = static_cast(col); + } + for (int32_t row = 0; row < q_size; ++row) { + mask_values[row * kv_size] = 0.0f; + for (int32_t col = 100; col < q_size; ++col) { + mask_values[row * kv_size + col] = 0.0f; + } + const float expected = + row < 100 ? 0.0f : (100.0f + row) * (row - 99) / (2.0f * (row - 98)); + for (int32_t head = 0; head < num_query_heads; ++head) { + expected_values[head * q_size + row] = + static_cast(expected); + } + } + + auto query = tfHalf.zeros({1, num_query_heads, q_size, 1}); + auto key = tfHalf.zeros({1, 1, kv_size, 1}); + auto value = tfHalf.make({1, 1, kv_size, 1}, value_values); + auto attn_mask = tfFloat.make({q_size, kv_size}, mask_values); + auto out = tfHalf.zeros({1, num_query_heads, q_size, 1}); + + auto result = op_scaled_dot_product_attention( + query, + key, + value, + attn_mask, + /*dropout_p=*/0.0, + /*is_causal=*/true, + /*scale=*/std::nullopt, + out); + + auto expected = tfHalf.make({1, num_query_heads, q_size, 1}, expected_values); + EXPECT_TENSOR_CLOSE_WITH_TOL(result, expected, 1e-2, 1e-2); +} + +TEST(OpScaledDotProductAttentionTest, QuantizedSparseMaskTrimmedBoundary) { + TensorFactory tfChar; + TensorFactory tfFloat; + + constexpr int32_t kv_size = 1024; + std::vector value_values(kv_size); + std::vector mask_values( + kv_size, -std::numeric_limits::infinity()); + for (int32_t col = 0; col < kv_size; ++col) { + value_values[col] = static_cast(col % 100); + } + for (int32_t col = 100; col < 200; ++col) { + mask_values[col] = 0.0f; + } + + auto query = tfChar.zeros({1, 1, 1, 1}); + auto key = tfChar.zeros({1, kv_size, 1, 1}); + auto value = tfChar.make({1, kv_size, 1, 1}, value_values); + auto query_zero_points = tfChar.zeros({1, 1, 1, 1}); + auto key_zero_points = tfChar.zeros({1, kv_size, 1, 1}); + auto value_zero_points = tfChar.zeros({1, kv_size, 1, 1}); + auto query_scales = tfFloat.ones({1, 1, 1, 1}); + auto key_scales = tfFloat.ones({1, kv_size, 1, 1}); + auto value_scales = tfFloat.ones({1, kv_size, 1, 1}); + auto attn_mask = tfFloat.make({1, kv_size}, mask_values); + auto out = tfFloat.zeros({1, 1, 1, 1}); + + executorch::runtime::KernelRuntimeContext context{}; + auto result = torch::executor::native::custom_quantized_sdpa_out( + context, + query, + key, + value, + /*start_pos=*/0, + attn_mask, + /*dropout_p=*/0.0, + /*is_causal=*/false, + /*scale=*/std::nullopt, + query_zero_points, + query_scales, + key_zero_points, + key_scales, + value_zero_points, + value_scales, + /*is_seq_at_dim_2=*/false, + out); + + auto expected = tfFloat.make({1, 1, 1, 1}, {49.5f}); + EXPECT_TENSOR_CLOSE_WITH_TOL(result, expected, 1e-4, 1e-4); +} + TEST(OpScaledDotProductAttentionTest, CorrectnessTest_11) { TensorFactory tfFloat; From e61d3d34adf2846261be59c728e823098a6a7de4 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Wed, 9 Sep 2026 14:54:01 -0700 Subject: [PATCH 125/190] Publish the CUDA trains PyTorch still offers, skip the ones it drops (#22652) ## Problem The nightly CUDA wheel build stopped publishing anything after PyTorch removed its CUDA 12.6 nightlies. ExecuTorch still listed `cu126` as a version to publish, and the release check treated a listed-but-absent train as a hard error, so it failed the whole job at the matrix-filter step. That blocked `cu130` and `cu132` from publishing too, even though their builds were fine. The result: no CUDA nightly wheel at all for several days, while the CPU wheel kept publishing normally. ## Fix Make the check degrade per train instead of all-or-nothing. A version listed in the policy is published only when the shared matrix generator still offers it. When PyTorch stops shipping a train, its rows simply do not appear, the release skips it with a note, and the other trains publish as usual. If PyTorch ships that train again later, it returns on its own with no code change here. The release still fails loudly in the two cases that are real problems: - an empty matrix, which would otherwise read as a green check for a build that never happened - a train that is offered but comes through missing some Python versions, which would ship that train incomplete Example, with PyTorch no longer offering cu126: ``` before: filter exits 1, publishes nothing after: filter exits 0, publishes cu130 and cu132, skips cu126 with a note ``` ## Test plan Extended the filter unit tests. A fully absent train is now asserted to be skipped while the others publish, plus a case for the exact cu126 drop. The incomplete-train and empty-matrix cases are asserted to still exit nonzero. Full suite: 21 passed, 2 subtests. Also ran the filter against today's real matrix (cu130 and cu132 only) and confirmed it exits 0 and publishes both. --- .ci/scripts/tests/test_filter_cuda_matrix.py | 51 +++++++++++++++----- .github/scripts/filter_cuda_matrix.py | 38 +++++++++------ 2 files changed, 61 insertions(+), 28 deletions(-) diff --git a/.ci/scripts/tests/test_filter_cuda_matrix.py b/.ci/scripts/tests/test_filter_cuda_matrix.py index d9b39024e8e..03a849df81f 100644 --- a/.ci/scripts/tests/test_filter_cuda_matrix.py +++ b/.ci/scripts/tests/test_filter_cuda_matrix.py @@ -191,18 +191,14 @@ def test_unparseable_matrix_exits_nonzero(self): FILTER.main(argv) self.assertNotEqual(raised.exception.code, 0) - def test_absent_train_exits_nonzero(self): - # A supported train the generator offers nothing for would publish no wheel at all. + def test_absent_train_is_skipped_not_fatal(self): + # A supported train the generator offers nothing for is one PyTorch stopped shipping. The + # release skips it and publishes the rest, so one dropped train cannot take the others down. # - # Patching the supported list rather than deleting rows, because deleting every row for one - # train also creates missing combinations, so both gates fire and the test cannot tell which - # one it exercised. Adding an extra supported train makes it absent while every offered - # combination stays complete. - # These two gates cannot be separated by input: any matrix leaving a train absent also - # leaves every combination for that train missing, so the later gate always catches what the - # earlier one would. Measured. So each gate gets its own case, and the case asserts on the - # message rather than only on a nonzero exit, which is the only way to tell them apart. + # Offering every train but the last leaves that train absent while every offered combination + # stays complete, which is exactly the shape of an upstream drop. offered = FILTER.SUPPORTED_CUDA_VERSIONS[:-1] + dropped = FILTER.SUPPORTED_CUDA_VERSIONS[-1] matrix = { "include": [ {"python_version": python, "desired_cuda": cuda} @@ -210,14 +206,43 @@ def test_absent_train_exits_nonzero(self): for cuda in offered ] } - message = self._exit_message(matrix) - self.assertIn("publish no wheel for that CUDA version", message) + emitted = _emitted(_run(matrix)) + published = sorted({row["desired_cuda"] for row in emitted["include"]}) + self.assertEqual(published, sorted(offered)) + self.assertNotIn(dropped, published) + + def test_dropped_train_still_publishes_the_others(self): + # The exact upstream drop this resilience is for: PyTorch stops shipping cu126, the generator + # offers only cu130 and cu132, and the release must still publish those two rather than fail + # because cu126 is gone. Skips the case cleanly if the policy no longer lists cu126. + if "cu126" not in FILTER.SUPPORTED_CUDA_VERSIONS: + self.skipTest("cu126 is not a published train") + survivors = [c for c in FILTER.SUPPORTED_CUDA_VERSIONS if c != "cu126"] + matrix = { + "include": [ + {"python_version": python, "desired_cuda": cuda} + for python in FILTER.SUPPORTED_PYTHON_VERSIONS + for cuda in survivors + ] + } + emitted = _emitted(_run(matrix)) + published = sorted({row["desired_cuda"] for row in emitted["include"]}) + self.assertEqual(published, sorted(survivors)) + self.assertNotIn("cu126", published) + # Every survivor keeps all its pythons, so what publishes is complete, just narrower. + self.assertEqual( + len(emitted["include"]), + len(survivors) * len(FILTER.SUPPORTED_PYTHON_VERSIONS), + ) def test_missing_combination_exits_nonzero(self): + # A train that IS offered but missing one python is a real break, not an upstream drop: the + # release would ship that train incomplete. Deleting one row from a full matrix leaves its + # train present, so this exercises the incomplete-train gate rather than the skip above. matrix = _full_matrix() del matrix["include"][0] message = self._exit_message(matrix) - self.assertIn("combination(s) produced no row", message) + self.assertIn("incomplete train", message) def test_jetpack_not_published_exits_nonzero(self): # Refused explicitly rather than allowed to fall through to an empty result, so the reason a diff --git a/.github/scripts/filter_cuda_matrix.py b/.github/scripts/filter_cuda_matrix.py index f7ac3a9debf..d0ec2ab076a 100644 --- a/.github/scripts/filter_cuda_matrix.py +++ b/.github/scripts/filter_cuda_matrix.py @@ -39,7 +39,7 @@ # not on that list is rejected whether or not it appears here. DISABLED_PYTHON_VERSIONS: List[str] = ["3.13t", "3.14t", "3.15", "3.15t"] -# CUDA versions to publish. +# CUDA versions to publish, when the generator offers them. # # Chosen so that every consumer row can find a matching wheel rather than by what is # convenient to verify. A delegate built against one of these has to be able to depend on an @@ -50,6 +50,12 @@ # cu130 the generator's stable choice, and the default for accelerator consumers # cu132 the newest, which consumers building against a current TensorRT need # +# A version listed here is published only when the shared generator still offers it. When +# PyTorch stops shipping a CUDA train, its rows simply do not appear and the release skips it, +# rather than failing the whole build. So a train PyTorch drops (as it did with cu126) costs +# only that train, and a train PyTorch restores returns here with no edit. The release still +# fails if a train that IS offered comes through incomplete, which is a real build break. +# # cu132 is included because omitting it would leave a published consumer row with no # ExecuTorch wheel to pair with. It is executable on a device one minor behind, since CUDA # minor versions are compatible, so a cu132 wheel has been run end to end on a CUDA 13.0 @@ -198,31 +204,33 @@ def main(argv: List[str]) -> None: # blind to a python that disappeared from every supported train. The generator lives in another # repository and its axes move independently of what this policy promises to publish. built = {(item["python_version"], item["desired_cuda"]) for item in items} - # A train that produced no row at all is missing for every python, so reporting it per python - # would read as a python problem. Named on its own instead, and first, because the per-pair - # report below would otherwise bury it. - absent_trains = sorted( - set(SUPPORTED_CUDA_VERSIONS) - {cuda for _, cuda in built} - ) + built_trains = {cuda for _, cuda in built} + # A train the generator offered nothing for is one PyTorch stopped shipping, not a build + # break here. Skip it and publish the rest, so one dropped train cannot take the others + # down with it. When PyTorch dropped CUDA 12.6, failing here also blocked cu130 and cu132 + # from publishing, which is the opposite of what a consumer needs. The train returns on its + # own if PyTorch ships it again, with no edit here. + absent_trains = sorted(set(SUPPORTED_CUDA_VERSIONS) - built_trains) if absent_trains: print( - f"this policy publishes {SUPPORTED_CUDA_VERSIONS}, but the generator offered no row " - f"this filter could keep for {absent_trains}, so a release would publish no wheel for " - "that CUDA version at all", + f"the generator offered no row for {absent_trains}, so they are skipped this run; " + f"publishing {sorted(built_trains)}", file=sys.stderr, ) - sys.exit(1) + # A train that IS offered but missing some python versions is a real break, not an upstream + # drop: the release would ship an incomplete train, fewer wheels than promised for a version + # that is otherwise present. Checked only against the trains actually offered, so a fully + # absent train handled above does not also trip this and read as a python problem. missing = sorted( f"{python}/{cuda}" for python in SUPPORTED_PYTHON_VERSIONS - for cuda in SUPPORTED_CUDA_VERSIONS + for cuda in built_trains if (python, cuda) not in built ) if missing: print( - f"this policy publishes {SUPPORTED_CUDA_VERSIONS} for each of " - f"{SUPPORTED_PYTHON_VERSIONS}, but {len(missing)} combination(s) produced no row, so a " - f"release would publish no wheel for them: {missing}", + f"a published train is missing some of {SUPPORTED_PYTHON_VERSIONS}, so a release " + f"would ship an incomplete train: {missing}", file=sys.stderr, ) sys.exit(1) From 25f155c32ddd2798bf8c59525e9dfb4b1e4f75ab Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 9 Sep 2026 15:09:43 -0700 Subject: [PATCH 126/190] Fix renamed compiler test source in `BUCK` Summary: D119200814 replaced `test_neutron_converter_manager.py` with `test_neutron_compiler_manager.py` but left the old source path in the mirrored `BUCK` files. This makes `test_neutron_converter_manager` fail to build with a missing-file error. Update both source paths to the renamed test file, preserving the existing target name. Authored with Codex. Differential Revision: D119340967 Pull Request resolved: https://github.com/pytorch/executorch/pull/22649 --- backends/nxp/tests/BUCK | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backends/nxp/tests/BUCK b/backends/nxp/tests/BUCK index a3a5adb215c..3d991e1eef0 100644 --- a/backends/nxp/tests/BUCK +++ b/backends/nxp/tests/BUCK @@ -148,7 +148,7 @@ fbcode_target(_kind = python_pytest, fbcode_target(_kind = python_pytest, name = "test_neutron_converter_manager", srcs = [ - "generic_tests/test_neutron_converter_manager.py", + "generic_tests/test_neutron_compiler_manager.py", ], deps = [ "//executorch/backends/nxp:neutron_sdk", From a33ac10541b9d2de5f95bde4f16cfe6c2fbff1fc Mon Sep 17 00:00:00 2001 From: RJ Ascani Date: Wed, 9 Sep 2026 16:16:51 -0700 Subject: [PATCH 127/190] Revert "Optimize CPU SDPA for ring attention" (#22660) Reverts pytorch/executorch#22358 Broke: https://github.com/pytorch/executorch/actions/runs/34405204388/job/102646426869 https://github.com/pytorch/executorch/actions/runs/34405204293/job/102646421578 --- examples/models/llama/attention.py | 44 +-- examples/models/llama/export_llama_lib.py | 5 - .../source_transformation/attention_sink.py | 122 ++----- .../source_transformation/custom_kv_cache.py | 116 ++----- .../test_attention_sink.py | 73 ++-- .../llama/tests/test_replace_kv_cache.py | 14 +- .../models/llama/tests/test_ring_attention.py | 81 +---- .../models/llama/tests/test_ring_kv_cache.py | 33 +- extension/llm/custom_ops/op_sdpa_impl.h | 322 ++++-------------- extension/llm/custom_ops/op_sdpa_test.cpp | 165 --------- 10 files changed, 181 insertions(+), 794 deletions(-) diff --git a/examples/models/llama/attention.py b/examples/models/llama/attention.py index 9980da91958..98a838bfa30 100644 --- a/examples/models/llama/attention.py +++ b/examples/models/llama/attention.py @@ -262,23 +262,6 @@ def calculate_positions_and_update_indices(self, input_pos: torch.Tensor, seq_le return indices -def _get_ring_cache_size( - max_context_length: int, - window_size: int, - max_seq_len: Optional[int] = None, -) -> int: - """Size an SWA cache for one retained window plus one in-flight chunk.""" - assert window_size > 0, "Sliding-window size must be positive" - if max_seq_len is None: - max_seq_len = max_context_length - assert max_seq_len > 0, "Maximum sequence length must be positive" - assert window_size <= max_context_length, ( - f"Sliding-window size ({window_size}) cannot exceed the full context " - f"length ({max_context_length})" - ) - return min(max_context_length, window_size + max_seq_len) - - class RingKVCache(KVCache): def __init__( self, @@ -288,17 +271,10 @@ def __init__( head_dim: int, enable_dynamic_shape: bool, dtype=torch.float32, - *, - window_size: int, - max_seq_len: Optional[int] = None, ): + self.window_size = max_context_length """ - The cache needs room for the retained sliding window and the current - prefill chunk. Its size is window_size + max_seq_len, capped by the - full-context cache. If max_seq_len is omitted, the full context length - is used. - - Reason why a cache larger than the sliding window is needed: + Reason why we want the kv cache size to be twice the context length: Sliding window attention without ringbuffer pos 0 1 2 3 4 5 6 7 8 9 10 0 x 0 0 0 0 0 0 0 0 0 0 @@ -331,28 +307,20 @@ def __init__( So not having kept 2, 3 and 4 in cache means we will have divergent behavior. Worst case of this would have been when update it equal to the length of the cache. like in our case pos = 5 seq len = 4. - Thus we need a cache larger than the sliding window by enough space for - the largest in-flight input chunk. + Thus we need to have a cache that is larger. How much larger, as much as + the sliding window size. So twice the max_context_length. How would that have helped. Lets see. At pos = 5 our cache would have [0, 1, 2, 3, 4, NA, NA, NA] After cache update we would have [8, 1, 2, 3, 4, 5, 6, 7]. We kicked out token at pos = 0. However, the current step still has access to [pos - sliding_window_size, pos] tokens. - + To make sure we dont over attend, i.e. we dont have pos = 5 to attend to pos = 1, mask calculaton has to account for the sliding window size. """ - self.window_size = window_size - self.full_context_length = max_context_length - self.max_seq_len = ( - max_context_length if max_seq_len is None else int(max_seq_len) - ) - ring_cache_size = _get_ring_cache_size( - max_context_length, window_size, self.max_seq_len - ) super().__init__( max_batch_size, - ring_cache_size, + max_context_length * 2, n_heads, head_dim, enable_dynamic_shape, diff --git a/examples/models/llama/export_llama_lib.py b/examples/models/llama/export_llama_lib.py index 074f8e24ea5..07240b11d8c 100644 --- a/examples/models/llama/export_llama_lib.py +++ b/examples/models/llama/export_llama_lib.py @@ -969,7 +969,6 @@ def _prepare_for_llama_export(llm_config: LlmConfig) -> LLMEdgeManager: preq_group_size=llm_config.base.preq_group_size, preq_embedding_quantize=llm_config.base.preq_embedding_quantize, local_global_attention=llm_config.model.local_global_attention, - max_seq_len=llm_config.export.max_seq_length, use_torchao_kernels_linear=llm_config.backend.torchao.use_torchao_kernels_linear, use_torchao_kernels_tied_embedding=llm_config.backend.torchao.use_torchao_kernels_tied_embedding, quantize_with_hqq=llm_config.quantization.use_hqq, @@ -2065,7 +2064,6 @@ def _get_source_transforms( # noqa preq_group_size: Optional[int] = None, preq_embedding_quantize: Optional[str] = None, local_global_attention: Optional[List[int]] = None, - max_seq_len: Optional[int] = None, use_torchao_kernels_linear: bool = False, use_torchao_kernels_tied_embedding: bool = False, quantize_with_hqq: bool = True, @@ -2106,8 +2104,6 @@ def _get_source_transforms( # noqa preq_mode: Pre-quantization mode. preq_group_size: Pre-quantization group size. preq_embedding_quantize: Pre-quantization embedding quantize. - max_seq_len: Largest input chunk accepted by local-attention layers. If - omitted, each layer's full context length is used. Returns: A list of transformation functions. @@ -2284,7 +2280,6 @@ def _get_source_transforms( # noqa partial( replace_kv_cache_with_ring_kv_cache, layer_sizes=local_global_attention, - max_seq_len=max_seq_len, ) ) diff --git a/examples/models/llama/source_transformation/attention_sink.py b/examples/models/llama/source_transformation/attention_sink.py index c76ec7514a1..c2d548fa798 100644 --- a/examples/models/llama/source_transformation/attention_sink.py +++ b/examples/models/llama/source_transformation/attention_sink.py @@ -18,36 +18,13 @@ _create_causal_mask_for_ring_buffer, AttentionMHA, KVCache, + RingKVCache, ) from executorch.examples.models.llama.model_args import ModelArgs from executorch.examples.models.llama.rope import Rope from torchao.quantization.quant_api import _replace_with_custom_fn_if_matches_filter -def _get_attention_sink_cache_size( - max_context_length: int, - window_size: int, - sink_size: int, - max_seq_len: Optional[int] = None, -) -> int: - """Size a sink cache for fixed sinks, one window, and one input chunk.""" - assert sink_size >= 0, "Attention sink size must be non-negative" - assert sink_size < max_context_length, ( - f"Attention sink size ({sink_size}) must be smaller than the full " - f"context length ({max_context_length})" - ) - assert window_size > 0, "Sliding-window size must be positive" - assert sink_size + window_size <= max_context_length, ( - f"Attention sink size ({sink_size}) plus sliding-window size " - f"({window_size}) cannot exceed the full context length " - f"({max_context_length})" - ) - if max_seq_len is None: - max_seq_len = max_context_length - assert max_seq_len > 0, "Maximum sequence length must be positive" - return sink_size + window_size + max_seq_len - - class RopeWithAttentionSink(Rope): """ Rope subclass for Attention Sink models. @@ -60,11 +37,10 @@ class RopeWithAttentionSink(Rope): - Window tokens (pos >= sink_size): wrapped into ring buffer range [sink_size, sink_size + ring_size) via modulo - The ring buffer holds one retained window plus one maximum-size input - chunk. It is larger than the live window for write-ahead headroom, not to - keep the live window contiguous -- it can span a wrap. Across a wrap two - positions remap to a difference that is not their true distance, so RoPE - preserves relative distance only within a wrap. + The ring buffer is 2x window_size for write-ahead headroom, not to keep the + live window contiguous -- it can span a wrap. Across a wrap two positions + remap to a difference that is not their true distance, so RoPE preserves + relative distance only within a wrap. """ def __init__( @@ -76,26 +52,7 @@ def __init__( super().__init__(params) self.window_size = window_size self.sink_size = sink_size - max_seq_len = ( - params.max_context_len - if getattr(params, "max_seq_len", None) is None - else int(params.max_seq_len) - ) - cache_size = _get_attention_sink_cache_size( - params.max_context_len, - window_size, - sink_size, - max_seq_len, - ) - self.ring_size = cache_size - sink_size - if self.freqs_cos.size(0) < cache_size: - freqs_cos, freqs_sin = self.precompute_freqs_cis( - self.params.head_dim, - cache_size, - self.params.rope_freq_base, - ) - self.freqs_cos = freqs_cos - self.freqs_sin = freqs_sin + self.ring_size = window_size * 2 def _remap_input_pos(self, input_pos: torch.Tensor) -> torch.Tensor: """Remap positions: sink tokens stay, window tokens wrap in ring buffer.""" @@ -176,7 +133,7 @@ class CachePositionsManagerWithSink(nn.Module): For sink_size=0: behaves exactly like original CachePositionsManager. For sink_size>0: sink tokens go to fixed positions, rest uses ring buffer. - IMPORTANT: cache_size is the actual cache dimension, including sink slots. + IMPORTANT: cache_size should be the actual cache dimension size (2x window for ring buffer). """ def __init__(self, cache_size: int, sink_size: int = 0): @@ -234,7 +191,7 @@ class KVCacheWithAttentionSink(KVCache): Uses a ring buffer approach for the sliding window portion while keeping the first sink_size tokens fixed. This avoids dynamic shape operations. - Cache layout: [fixed sink tokens] [ring buffer for one window + one chunk] + Cache layout: [sink: 0 to sink_size-1] [ring_buffer: sink_size to sink_size + window_size*2 - 1] """ def __init__( @@ -245,27 +202,16 @@ def __init__( rope: RopeWithAttentionSink, window_size: int, sink_size: int, - max_context_length: int, - max_seq_len: Optional[int] = None, max_batch_size: int = 1, dtype=torch.float32, ): - self.full_context_length = max_context_length - self.max_seq_len = ( - max_context_length if max_seq_len is None else int(max_seq_len) - ) - # Keep the fixed sinks separate from the ring space needed for the - # retained window and the largest in-flight input chunk. - total_cache_size = _get_attention_sink_cache_size( - max_context_length, - window_size, - sink_size, - self.max_seq_len, - ) - assert rope.ring_size == total_cache_size - sink_size, ( - f"RoPE ring size ({rope.ring_size}) must match the KV-cache ring " - f"size ({total_cache_size - sink_size})" - ) + # Total cache size is sink_size + window_size * 2. + # The ring buffer needs 2x the window size because at the moment a new + # token is written, the previous window_size tokens must still be readable + # (they haven't been overwritten yet). With only 1x, writing a new entry + # would immediately evict the oldest visible token, leaving fewer than + # window_size tokens available for attention. + total_cache_size = sink_size + window_size * 2 super().__init__( max_batch_size=max_batch_size, max_context_length=total_cache_size, @@ -364,7 +310,6 @@ def _replace_attention( rope_with_attention_sink: RopeWithAttentionSink, sink_size: int, window_size: int, - max_seq_len: Optional[int], ): for _, child_module in module._modules.items(): if len(list(child_module.children())) > 0: # pyre-ignore [16] @@ -373,23 +318,32 @@ def _replace_attention( rope_with_attention_sink=rope_with_attention_sink, sink_size=sink_size, window_size=window_size, - max_seq_len=max_seq_len, ) if isinstance(child_module, AttentionMHA): kv_cache = child_module.kv_cache - child_module.kv_cache = KVCacheWithAttentionSink( - n_heads=kv_cache.n_heads, - head_dim=kv_cache.head_dim, - enable_dynamic_shape=kv_cache.enable_dynamic_shape, - rope=rope_with_attention_sink, - max_batch_size=kv_cache.max_batch_size, - window_size=window_size, - sink_size=sink_size, - max_context_length=kv_cache.max_context_length, - max_seq_len=max_seq_len, - dtype=kv_cache.k_cache.dtype, - ) + if sink_size == 0: + # No sink tokens needed — use standard RingKVCache directly + child_module.kv_cache = RingKVCache( + kv_cache.max_batch_size, + window_size, # RingKVCache expects user-provided window size + kv_cache.n_heads, + kv_cache.head_dim, + kv_cache.enable_dynamic_shape, + kv_cache.k_cache.dtype, + ) + else: + kv_cache_with_attention_sink = KVCacheWithAttentionSink( + n_heads=kv_cache.n_heads, + head_dim=kv_cache.head_dim, + enable_dynamic_shape=kv_cache.enable_dynamic_shape, + rope=rope_with_attention_sink, + max_batch_size=kv_cache.max_batch_size, + window_size=window_size, + sink_size=sink_size, + dtype=kv_cache.k_cache.dtype, + ) + child_module.kv_cache = kv_cache_with_attention_sink # Don't replace forward - let the original AttentionMHA.forward handle it # since our KVCache has is_ring_buffer=True, it will use the ring buffer mask @@ -411,13 +365,11 @@ def enable_attention_sink( window_size=window_size, sink_size=sink_size, ) - max_seq_len = getattr(params, "max_seq_len", None) _replace_rope(module, rope_with_attention_sink) _replace_attention( module=module, rope_with_attention_sink=rope_with_attention_sink, sink_size=sink_size, window_size=window_size, - max_seq_len=max_seq_len, ) return module diff --git a/examples/models/llama/source_transformation/custom_kv_cache.py b/examples/models/llama/source_transformation/custom_kv_cache.py index 7fb0adfb923..dbaac9accf4 100644 --- a/examples/models/llama/source_transformation/custom_kv_cache.py +++ b/examples/models/llama/source_transformation/custom_kv_cache.py @@ -14,7 +14,6 @@ import torch.nn as nn from executorch.examples.models.llama.attention import ( _create_causal_mask_for_ring_buffer, - _get_ring_cache_size, CachePositionsManager, KVCache, RingKVCache, @@ -662,19 +661,17 @@ def _replace_kv_cache_with_custom_kv_cache(module): if sdpa is not None and hasattr(sdpa, "use_attention_mask"): sdpa.use_attention_mask = True elif isinstance(child, RingKVCache): - # Preserve ring-buffer sizing and masking when converting a - # local-attention cache to the custom update op. + # RingKVCache (e.g., from attention sink with sink_size=0) needs + # CustomRingKVCache, not plain CustomKVCache setattr( module, name, CustomRingKVCache( child.max_batch_size, - child.full_context_length, + child.window_size, child.n_heads, child.head_dim, dtype=child.k_cache.dtype, - window_size=child.window_size, - max_seq_len=child.max_seq_len, ), ) sdpa = getattr(module, "SDPA", None) @@ -710,20 +707,11 @@ def __init__( cache_type: QuantizedCacheType = QuantizedCacheType.AffineSymmetric, use_custom_update_cache_op: bool = False, return_float_values: bool = True, - *, - window_size: int, - max_seq_len: Optional[int] = None, ): - self.full_context_length = max_context_length - self.max_seq_len = ( - max_context_length if max_seq_len is None else int(max_seq_len) - ) - ring_cache_size = _get_ring_cache_size( - max_context_length, window_size, self.max_seq_len - ) + # Look at attention.py for explanation on why max_context_length * 2 super().__init__( max_batch_size, - ring_cache_size, + max_context_length * 2, n_heads, head_dim, cache_type, @@ -732,7 +720,7 @@ def __init__( ) self.cache_positions_manager = CachePositionsManager(self.max_context_length) self.is_ring_buffer = True - self.window_size = window_size + self.window_size = max_context_length def create_causal_mask_for_ring_buffer(self, start_pos, seq_len): cache_positions = self.cache_positions_manager.cache_positions @@ -754,7 +742,7 @@ def update(self, input_pos, k_val, v_val): seq_len = k_val.transpose(1, 2).size(1) assert seq_len <= self.k_cache.size( 1 - ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(1)})" + ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(2)})" indices = self.cache_positions_manager.calculate_positions_and_update_indices( input_pos, seq_len ) @@ -767,22 +755,19 @@ def from_quantized_kv_cache( cls, kv_cache, sliding_window_size, - max_seq_len, ): assert isinstance( kv_cache, QuantizedKVCache ), "For QuantizedRingKVCache expect QuantizedKVCache as input kv_cache" - max_batch_size, max_context_length, n_heads, head_dim = kv_cache.k_cache.shape + max_batch_size, _, n_heads, head_dim = kv_cache.k_cache.shape return cls( max_batch_size, - max_context_length, + sliding_window_size, n_heads, head_dim, kv_cache.cache_type, kv_cache.use_custom_update_cache_op, kv_cache.return_float_values, - window_size=sliding_window_size, - max_seq_len=max_seq_len, ) @@ -798,30 +783,19 @@ class CustomKVCacheWithAttentionSink(CustomKVCache): def __init__( self, max_batch_size, - max_context_length, n_heads, head_dim, window_size, sink_size, - max_seq_len=None, dtype=torch.float32, ): + # Total cache size: sink slots + ring buffer (2x window for wrap safety) + total_cache_size = sink_size + window_size * 2 + super().__init__(max_batch_size, total_cache_size, n_heads, head_dim, dtype) from executorch.examples.models.llama.source_transformation.attention_sink import ( - _get_attention_sink_cache_size, CachePositionsManagerWithSink, ) - self.full_context_length = max_context_length - self.max_seq_len = ( - max_context_length if max_seq_len is None else int(max_seq_len) - ) - total_cache_size = _get_attention_sink_cache_size( - max_context_length, - window_size, - sink_size, - self.max_seq_len, - ) - super().__init__(max_batch_size, total_cache_size, n_heads, head_dim, dtype) self.cache_positions_manager = CachePositionsManagerWithSink( total_cache_size, sink_size ) @@ -872,12 +846,10 @@ def from_kv_cache_with_attention_sink(cls, kv_cache): max_batch_size, n_heads, _, head_dim = kv_cache.k_cache.shape return cls( max_batch_size, - kv_cache.full_context_length, n_heads, head_dim, kv_cache.window_size, kv_cache.sink_size, - max_seq_len=kv_cache.max_seq_len, dtype=kv_cache.k_cache.dtype, ) @@ -890,21 +862,14 @@ def __init__( n_heads, head_dim, dtype=torch.float32, - *, - window_size: int, - max_seq_len: Optional[int] = None, ): - self.full_context_length = max_context_length - self.max_seq_len = ( - max_context_length if max_seq_len is None else int(max_seq_len) - ) - ring_cache_size = _get_ring_cache_size( - max_context_length, window_size, self.max_seq_len + # Look at attention.py for explanation on why max_context_length * 2 + super().__init__( + max_batch_size, max_context_length * 2, n_heads, head_dim, dtype ) - super().__init__(max_batch_size, ring_cache_size, n_heads, head_dim, dtype) self.cache_positions_manager = CachePositionsManager(self.max_context_length) self.is_ring_buffer = True - self.window_size = window_size + self.window_size = max_context_length def create_causal_mask_for_ring_buffer(self, start_pos, seq_len): cache_positions = self.cache_positions_manager.cache_positions @@ -926,7 +891,7 @@ def update(self, input_pos, k_val, v_val): seq_len = k_val.transpose(1, 2).size(1) assert seq_len <= self.k_cache.size( 1 - ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(1)})" + ), f"Update sequence length({seq_len}) for kv cache must be smaller than the cache size({self.k_cache.size(2)})" indices = self.cache_positions_manager.calculate_positions_and_update_indices( input_pos, seq_len ) @@ -939,24 +904,21 @@ def from_custom_kv_cache( cls, kv_cache, sliding_window_size, - max_seq_len, ): - # CustomKVCache storage is [B, S, H, D]. - max_batch_size, max_context_length, n_heads, head_dim = kv_cache.k_cache.shape + max_batch_size, n_heads, _, head_dim = kv_cache.k_cache.shape + if isinstance(kv_cache, CustomKVCache): + # If replacing custom kv cache, then the shape is [B, S, H, D] + max_batch_size, _, n_heads, head_dim = kv_cache.k_cache.shape return cls( max_batch_size, - max_context_length, + sliding_window_size, n_heads, head_dim, dtype=kv_cache.k_cache.dtype, - window_size=sliding_window_size, - max_seq_len=max_seq_len, ) -def _replace_kv_cache_with_ring_kv_cache( - attention, layer_size: int, max_seq_len: Optional[int] -): +def _replace_kv_cache_with_ring_kv_cache(attention, layer_size): sliding_window_size = layer_size assert ( getattr(attention, "kv_cache", None) is not None @@ -965,28 +927,23 @@ def _replace_kv_cache_with_ring_kv_cache( if isinstance(kv_cache, KVCache): attention.kv_cache = RingKVCache( kv_cache.max_batch_size, - kv_cache.max_context_length, + sliding_window_size, kv_cache.n_heads, kv_cache.head_dim, kv_cache.enable_dynamic_shape, kv_cache.k_cache.dtype, - window_size=sliding_window_size, - max_seq_len=max_seq_len, ) elif isinstance(kv_cache, CustomKVCache): attention.kv_cache = CustomRingKVCache.from_custom_kv_cache( - kv_cache, layer_size, max_seq_len + kv_cache, layer_size ) elif isinstance(kv_cache, QuantizedKVCache): attention.kv_cache = QuantizedRingKVCache.from_quantized_kv_cache( - kv_cache, layer_size, max_seq_len + kv_cache, layer_size ) -def replace_kv_cache_with_ring_kv_cache( - module, layer_sizes, max_seq_len: Optional[int] = None -): - """Replace local-layer caches with window-plus-in-flight ring caches.""" +def replace_kv_cache_with_ring_kv_cache(module, layer_sizes): # This is needed to ensure that custom ops are registered from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 @@ -1002,7 +959,6 @@ def replace_kv_cache_with_ring_kv_cache( logging.info( f"Applying local sliding window attention with following pattern {layer_sizes}." ) - logged_full_context_cache = False assert len(layer_sizes) == len( module.layers ), f"Length of layer sizes {len(layer_sizes)} must match the number of layers in the module {len(module.layers)}." @@ -1014,23 +970,7 @@ def replace_kv_cache_with_ring_kv_cache( getattr(transformer_block, "attention", None) is not None ), f"Transfomer block must have attention module. Transformer block {transformer_block}" attention = transformer_block.attention - full_context_length = attention.kv_cache.max_context_length - effective_max_seq_len = ( - full_context_length if max_seq_len is None else max_seq_len - ) - if ( - not logged_full_context_cache - and sliding_window_size + effective_max_seq_len >= full_context_length - ): - logging.info( - "Local attention KV caches will use the full context length; " - "set max_seq_length below max_context_length - window_size " - "to retain KV-cache memory savings." - ) - logged_full_context_cache = True - _replace_kv_cache_with_ring_kv_cache( - attention, sliding_window_size, max_seq_len - ) + _replace_kv_cache_with_ring_kv_cache(attention, sliding_window_size) # if attention's sdpa is custom sdpa then we have to make sure # it is not doing causal attention if "SDPACustom" in attention.SDPA.__class__.__name__: diff --git a/examples/models/llama/source_transformation/test_attention_sink.py b/examples/models/llama/source_transformation/test_attention_sink.py index f5fb0c4d063..c4105338f0e 100644 --- a/examples/models/llama/source_transformation/test_attention_sink.py +++ b/examples/models/llama/source_transformation/test_attention_sink.py @@ -69,10 +69,7 @@ def setUp(self) -> None: # Ring top is 20. The table is deliberately longer, so a slice running # past the ring still lands on real rows instead of going out of bounds. self.params = ModelArgs( - use_kv_cache=True, - enable_dynamic_shape=True, - max_context_len=64, - max_seq_len=self.WINDOW_SIZE, + use_kv_cache=True, enable_dynamic_shape=True, max_context_len=64 ) self.rope = RopeWithAttentionSink( params=self.params, @@ -229,15 +226,14 @@ def setUp(self): use_kv_cache=True, enable_dynamic_shape=True, max_context_len=256, - max_seq_len=self.window_size, ) self.rope = RopeWithAttentionSink( params=self.params, window_size=self.window_size, sink_size=self.sink_size, ) - # Total cache size = sink_size + window_size + max_seq_len = 60. - self.cache_size = self.sink_size + self.window_size + self.params.max_seq_len + # Total cache size = sink_size + window_size * 2 = 4 + 56 = 60 + self.cache_size = self.sink_size + self.window_size * 2 self.kv_cache = KVCacheWithAttentionSink( n_heads=self.params.n_heads, head_dim=self.params.head_dim, @@ -246,8 +242,6 @@ def setUp(self): max_batch_size=self.max_batch_size, window_size=self.window_size, sink_size=self.sink_size, - max_context_length=self.params.max_context_len, - max_seq_len=self.params.max_seq_len, dtype=self.dtype, ) @@ -378,15 +372,11 @@ def test_causal_mask_blocks_future(self): ) def test_no_sink_degenerates_to_ring_buffer(self, sink_size): """With sink_size=0, behavior should match a plain ring buffer.""" - window_size = 100 params = ModelArgs( - use_kv_cache=True, - enable_dynamic_shape=True, - max_context_len=128, - max_seq_len=64, + use_kv_cache=True, enable_dynamic_shape=True, max_context_len=256 ) rope = RopeWithAttentionSink( - params=params, window_size=window_size, sink_size=0 + params=params, window_size=self.window_size, sink_size=0 ) cache = KVCacheWithAttentionSink( n_heads=params.n_heads, @@ -394,16 +384,11 @@ def test_no_sink_degenerates_to_ring_buffer(self, sink_size): enable_dynamic_shape=params.enable_dynamic_shape, rope=rope, max_batch_size=1, - window_size=window_size, + window_size=self.window_size, sink_size=0, - max_context_length=params.max_context_len, - max_seq_len=params.max_seq_len, dtype=self.dtype, ) - cache_size = window_size + params.max_seq_len - self.assertEqual(cache_size, 164) - self.assertEqual(rope.ring_size, cache_size) - self.assertEqual(cache.max_context_length, cache_size) + cache_size = self.window_size * 2 # 56 # Fill and wrap k_init, v_init = self._rand_kv(cache_size) @@ -449,10 +434,7 @@ def _build_model(self, args, sink_size, window_size, use_custom_sdpa=False): model = construct_transformer(args) model = enable_attention_sink( - model, - params=args, - sink_size=sink_size, - window_size=window_size, + model, params=args, sink_size=sink_size, window_size=window_size ) if use_custom_sdpa: @@ -520,12 +502,12 @@ def test_beyond_context_window_basic(self): """Generate tokens well beyond the KV cache size using standard SDPA.""" sink_size = 4 window_size = 16 - # KV cache size = sink_size + window_size + max_seq_len = 52 + # KV cache size = sink_size + window_size * 2 = 36 # max_context_len = 128 (for RoPE table) args = self._make_args(max_context_len=128) model = self._build_model(args, sink_size, window_size, use_custom_sdpa=False) - # Generate 80 tokens — beyond the KV cache size of 52 + # Generate 80 tokens — well beyond KV cache size of 36 outputs = self._run_generation(model, args, num_tokens=80) self.assertEqual(len(outputs), 77) # 1 prefill + 76 decode steps @@ -538,13 +520,10 @@ def test_beyond_max_context_len(self): """Generate tokens beyond max_context_len with RoPE position remapping.""" sink_size = 4 window_size = 16 - # With max_seq_len omitted, the cache is capped at max_context_len. + # KV cache size = 36, max_context_len = 64 # Generate 100 tokens — well beyond max_context_len args = self._make_args(max_context_len=64) - args.max_seq_len = None model = self._build_model(args, sink_size, window_size, use_custom_sdpa=False) - cache = model.layers[0].attention.kv_cache - self.assertEqual(cache.max_context_length, 84) outputs = self._run_generation(model, args, num_tokens=100) @@ -558,9 +537,10 @@ def test_beyond_max_context_len(self): def test_chunked_prefill_across_the_ring_wrap(self): """Chunked prefill where a chunk spans the ring wrap. - sink_size=4, window_size=16, and max_seq_len=48, so the ring is slots - [4, 68). Feeding 40 tokens at a time exceeds the old 2x-window ring and - crosses the new ring boundary on the second chunk. + sink_size=4, window_size=16, so the ring is slots [4, 36). Feeding 5 + tokens at a time puts chunk starts at 0, 5, ..., 95. The chunk at 35 + covers positions 35..39 and needs rows 35, 4, 5, 6, 7; the chunk at 65 + covers 65..69 and needs rows 33, 34, 35, 4, 5. Both span the wrap. The other beyond-context-window tests decode one token at a time, and a chunk of one can never span the wrap however far the position runs, so @@ -574,32 +554,27 @@ def test_chunked_prefill_across_the_ring_wrap(self): """ sink_size = 4 window_size = 16 - chunk_size = 40 - args = self._make_args(max_context_len=128) - args.max_seq_len = 48 + args = self._make_args(max_context_len=64) torch.manual_seed(0) model = self._build_model(args, sink_size, window_size) - tokens = torch.randint(0, args.vocab_size, (1, 80)) + tokens = torch.randint(0, args.vocab_size, (1, 100)) - chunked = self._feed_in_chunks( - copy.deepcopy(model), tokens, chunk_size=chunk_size - ) + chunked = self._feed_in_chunks(copy.deepcopy(model), tokens, chunk_size=5) one_at_a_time = self._feed_in_chunks(copy.deepcopy(model), tokens, chunk_size=1) - self.assertEqual(len(chunked), 2) - self.assertEqual(len(one_at_a_time), 80) + self.assertEqual(len(chunked), 20) + self.assertEqual(len(one_at_a_time), 100) # generate_full_logits is off, so each call returns logits for its last # position only. How the input was chunked must not change the result: - # chunk i ends on the same token as the corresponding decode call. + # chunk i ends on the same token as one_at_a_time[5 * i + 4]. for i, out in enumerate(chunked): self.assertTrue(torch.isfinite(out).all(), f"chunk {i} is not finite") torch.testing.assert_close( out, - one_at_a_time[chunk_size * (i + 1) - 1], - msg=lambda m, i=i: f"chunk {i}, positions {chunk_size * i}.." - f"{chunk_size * (i + 1) - 1}, " + one_at_a_time[5 * i + 4], + msg=lambda m, i=i: f"chunk {i}, positions {5 * i}..{5 * i + 4}, " f"disagrees with feeding the same tokens one at a time:\n{m}", ) @@ -624,7 +599,7 @@ def test_beyond_context_window_custom_sdpa(self): found_custom_cache, "Expected CustomKVCacheWithAttentionSink in model" ) - # Generate 80 tokens — beyond the KV cache size of 52 + # Generate 80 tokens — well beyond KV cache size of 36 outputs = self._run_generation(model, args, num_tokens=80) self.assertEqual(len(outputs), 77) diff --git a/examples/models/llama/tests/test_replace_kv_cache.py b/examples/models/llama/tests/test_replace_kv_cache.py index a4595b4311b..383fe6d6882 100644 --- a/examples/models/llama/tests/test_replace_kv_cache.py +++ b/examples/models/llama/tests/test_replace_kv_cache.py @@ -89,14 +89,13 @@ def test_replace_kv_cache_with_ring_kv_cache(self): # Replace KVCache with RingKVCache layer_sizes = [8] # Sliding window size for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes, max_seq_len=4) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes) # Verify that KVCache has been replaced with RingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, RingKVCache) # Verify that the sliding window size is set correctly self.assertEqual(model.layers[0].attention.kv_cache.window_size, layer_sizes[0]) - self.assertEqual(model.layers[0].attention.kv_cache.k_cache.size(2), 12) def test_replace_custom_kv_cache_with_custom_ring_kv_cache(self): """Test replacing CustomKVCache with CustomRingKVCache.""" @@ -113,11 +112,10 @@ def test_replace_custom_kv_cache_with_custom_ring_kv_cache(self): # Replace CustomKVCache with CustomRingKVCache layer_sizes = [8] # Sliding window size for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes, max_seq_len=4) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes) # Verify that CustomKVCache has been replaced with CustomRingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, CustomRingKVCache) - self.assertEqual(model.layers[0].attention.kv_cache.k_cache.size(1), 12) def test_replace_quantized_kv_cache_with_quantized_ring_kv_cache(self): """Test replacing QuantizedKVCache with QuantizedRingKVCache.""" @@ -136,11 +134,10 @@ def test_replace_quantized_kv_cache_with_quantized_ring_kv_cache(self): # Replace QuantizedKVCache with QuantizedRingKVCache layer_sizes = [8] # Sliding window size for each layer - replace_kv_cache_with_ring_kv_cache(model, layer_sizes, max_seq_len=4) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes) # Verify that QuantizedKVCache has been replaced with QuantizedRingKVCache self.assertIsInstance(model.layers[0].attention.kv_cache, QuantizedRingKVCache) - self.assertEqual(model.layers[0].attention.kv_cache.k_cache.size(1), 12) def test_replace_static_quantized_kv_cache(self): """Test replacing KVCache with static-qparams int8 KV storage.""" @@ -290,7 +287,6 @@ def test_static_quantized_kv_cache_rejects_specialized_cache(self): self.n_kv_heads, self.head_dim, self.enable_dynamic_shape, - window_size=self.max_context_len, ) model = self._create_mock_model([attention]) @@ -309,9 +305,7 @@ def test_multiple_layers_with_different_window_sizes(self): # Replace KVCache with RingKVCache with different window sizes layer_sizes = [4, 8, 16] # Different sliding window sizes for each layer - replace_kv_cache_with_ring_kv_cache( - model, layer_sizes, max_seq_len=self.max_context_len - ) + replace_kv_cache_with_ring_kv_cache(model, layer_sizes) # Verify that each layer has the correct window size self.assertIsInstance(model.layers[0].attention.kv_cache, RingKVCache) diff --git a/examples/models/llama/tests/test_ring_attention.py b/examples/models/llama/tests/test_ring_attention.py index 5ab696bc5f5..ae440e00e47 100644 --- a/examples/models/llama/tests/test_ring_attention.py +++ b/examples/models/llama/tests/test_ring_attention.py @@ -84,10 +84,7 @@ def _create_baseline_attention( return attention def _create_ring_attention( - self, - attention, - max_seq_len, - kv_cache_type: KVCacheType = KVCacheType.REGULAR, + self, attention, kv_cache_type: KVCacheType = KVCacheType.REGULAR ): """Create attention with ring buffer KV cache.""" assert self.sliding_window is not None @@ -101,82 +98,25 @@ def _create_ring_attention( baseline_attention.kv_cache = QuantizedRingKVCache.from_quantized_kv_cache( baseline_attention.kv_cache, self.sliding_window, - max_seq_len, ) elif isinstance(baseline_attention.kv_cache, CustomKVCache): # Replace CustomKVCache with CustomRingKVCache baseline_attention.kv_cache = CustomRingKVCache.from_custom_kv_cache( baseline_attention.kv_cache, self.sliding_window, - max_seq_len, ) else: # Replace regular KVCache with RingKVCache baseline_attention.kv_cache = RingKVCache( self.args.max_batch_size, - self.args.max_context_len, + self.sliding_window, self.n_kv_heads, self.head_dim, self.args.enable_dynamic_shape, self.dtype, - window_size=self.sliding_window, - max_seq_len=max_seq_len, ) return baseline_attention - def test_continuation_prefill_larger_than_window( - self, kv_cache_type: KVCacheType = KVCacheType.REGULAR - ): - """A W+C cache preserves history when the incoming chunk exceeds W.""" - self.sliding_window = 4 - chunk_size = 6 - baseline_attn = self._create_baseline_attention(12, kv_cache_type) - ring_attn = self._create_ring_attention( - baseline_attn, chunk_size, kv_cache_type - ) - - self.assertEqual(ring_attn.kv_cache.max_context_length, 10) - - with torch.nn.attention.sdpa_kernel( - [SDPBackend.FLASH_ATTENTION] - ), torch.no_grad(): - for pos in (0, chunk_size): - x = torch.randn( - (self.batch_size, chunk_size, self.dim), dtype=self.dtype - ) - input_pos = torch.tensor([pos], dtype=torch.long) - freqs_cos, freqs_sin = self.rope.get_freqs(input_pos, chunk_size) - - baseline_out, _ = baseline_attn.forward( - x, freqs_cos, freqs_sin, input_pos=input_pos - ) - ring_out, _ = ring_attn.forward( - x, freqs_cos, freqs_sin, input_pos=input_pos - ) - - tolerance = 1e-6 if kv_cache_type != KVCacheType.REGULAR else 1e-7 - self.assertTrue( - torch.allclose( - baseline_out, - ring_out, - rtol=tolerance, - atol=tolerance, - ), - f"Outputs differ at position {pos}", - ) - - def test_continuation_prefill_larger_than_window_quantized(self): - self._run_test_with_kv_cache_type( - self.test_continuation_prefill_larger_than_window, - KVCacheType.QUANTIZED, - ) - - def test_continuation_prefill_larger_than_window_custom(self): - self._run_test_with_kv_cache_type( - self.test_continuation_prefill_larger_than_window, - KVCacheType.CUSTOM, - ) - def _create_sliding_window_mask(self, seq_len, context_len, window_size): """Create a sliding window mask for the baseline.""" mask = torch.full((seq_len, context_len), float("-inf"), dtype=self.dtype) @@ -200,7 +140,7 @@ def test_single_token_processing( seq_len = 10 self.sliding_window = 4 baseline_attn = self._create_baseline_attention(seq_len, kv_cache_type) - ring_attn = self._create_ring_attention(baseline_attn, 1, kv_cache_type) + ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) # Process tokens one by one with torch.nn.attention.sdpa_kernel( @@ -259,7 +199,7 @@ def test_sliding_window_attention( baseline_attn = self._create_baseline_attention(seq_len, kv_cache_type) # Create ring attention with sliding window size - ring_attn = self._create_ring_attention(baseline_attn, 1, kv_cache_type) + ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) # Process tokens one by one with torch.nn.attention.sdpa_kernel( @@ -311,7 +251,7 @@ def test_ring_buffer_wrapping( ) # Create ring attention with sliding window size - ring_attn = self._create_ring_attention(baseline_attn, 1, kv_cache_type) + ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) # Process enough tokens to cause wrapping seq_len = 1 @@ -337,12 +277,13 @@ def test_ring_buffer_wrapping( f"Outputs differ at position {pos}", ) - # With W=3 and max_seq_len=1, the physical cache has 4 slots and wraps. + # After processing 8 tokens with window size 4, the ring buffer should have wrapped around # Check the cache positions to verify wrapping cache_positions = ring_attn.kv_cache.cache_positions_manager.cache_positions - # Positions 4 through 7 occupy physical slots 0 through 3. - expected_positions = torch.tensor([4, 5, 6, 7], dtype=torch.long) + # The cache positions should contain the most recent 4 positions (4, 5, 6, 7) + # mapped to the ring buffer indices + expected_positions = torch.tensor([6, 7, 2, 3, 4, 5], dtype=torch.long) self.assertTrue( torch.all(cache_positions == expected_positions), @@ -375,9 +316,7 @@ def test_large_context_with_sliding_window( baseline_attn = self._create_baseline_attention(seq_len, kv_cache_type) # Create ring attention with sliding window size - ring_attn = self._create_ring_attention( - baseline_attn, max(token_lens), kv_cache_type - ) + ring_attn = self._create_ring_attention(baseline_attn, kv_cache_type) pos = 0 with torch.nn.attention.sdpa_kernel( diff --git a/examples/models/llama/tests/test_ring_kv_cache.py b/examples/models/llama/tests/test_ring_kv_cache.py index 561ab117078..f923a1bf2bf 100644 --- a/examples/models/llama/tests/test_ring_kv_cache.py +++ b/examples/models/llama/tests/test_ring_kv_cache.py @@ -14,9 +14,7 @@ class TestRingKVCache(unittest.TestCase): def setUp(self): # Common test parameters self.max_batch_size = 2 - self.max_context_length = 16 - self.window_size = 8 - self.max_seq_len = 8 + self.max_context_length = 8 self.n_heads = 4 self.head_dim = 16 self.enable_dynamic_shape = True @@ -34,7 +32,7 @@ def test_dynamic_kv_cache_update_on_cuda(self): self._require_usable_cuda() cache = KVCache( max_batch_size=1, - max_context_length=self.window_size, + max_context_length=self.max_context_length, n_heads=self.n_heads, head_dim=self.head_dim, enable_dynamic_shape=True, @@ -62,8 +60,6 @@ def test_ring_cache_positions_and_mask_on_cuda(self): head_dim=self.head_dim, enable_dynamic_shape=True, dtype=self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ).cuda() input_pos = torch.tensor([0], dtype=torch.long, device="cuda") seq_len = 3 @@ -107,11 +103,8 @@ def test_basic_update(self): self.head_dim, self.enable_dynamic_shape, self.dtype, - window_size=self.window_size, ) - self.assertEqual(cache.max_seq_len, self.max_context_length) - # Create input tensors input_pos = torch.tensor([0], dtype=torch.long) seq_len = 3 @@ -136,7 +129,7 @@ def test_basic_update(self): self.assertTrue(torch.all(v_out[:, :, i] == 2.0)) # Check that the rest of the cache is still zeros - for i in range(seq_len, self.window_size): + for i in range(seq_len, self.max_context_length): self.assertTrue(torch.all(k_out[:, :, i] == 0.0)) self.assertTrue(torch.all(v_out[:, :, i] == 0.0)) @@ -160,8 +153,6 @@ def test_ring_buffer_wrapping(self): self.head_dim, self.enable_dynamic_shape, self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ) # Create input tensors for first update @@ -225,8 +216,6 @@ def test_multiple_updates(self): self.head_dim, self.enable_dynamic_shape, self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ) # First update @@ -357,8 +346,6 @@ def test_edge_case_input_pos_zero(self): self.head_dim, self.enable_dynamic_shape, self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ) # Create input tensors @@ -387,7 +374,7 @@ def test_edge_case_input_pos_zero(self): self.assertTrue(torch.all(v_out[:, :, 0] == 12.0)) # Check that the rest of the cache is still zeros - for i in range(1, self.window_size): + for i in range(1, self.max_context_length): self.assertTrue(torch.all(k_out[:, :, i] == 0.0)) self.assertTrue(torch.all(v_out[:, :, i] == 0.0)) @@ -403,7 +390,7 @@ def test_edge_case_input_pos_zero(self): ) def test_edge_case_exceeding_context_length(self): - """Test the edge case where input_pos + seq_len exceeds cache capacity.""" + """Test the edge case where input_pos + seq_len > max_context_length.""" cache = RingKVCache( self.max_batch_size, self.max_context_length, @@ -411,8 +398,6 @@ def test_edge_case_exceeding_context_length(self): self.head_dim, self.enable_dynamic_shape, self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ) # Create input tensors @@ -477,8 +462,6 @@ def test_original_indices_tracking(self): self.head_dim, self.enable_dynamic_shape, self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ) # First update at position 10 (will be mapped to position 10 in the ring buffer) @@ -498,7 +481,7 @@ def test_original_indices_tracking(self): # Check that cache_positions correctly tracks the original indices # For input_pos=10 and seq_len=4, the original indices should be 10, 11, 12, 13 - # These map directly to positions 10, 11, 12, 13 in the 16-slot cache. + # These map to positions 10, 11, 12, 13 in the ring buffer (since max_context_length=8 but buffer size is 16) # Note that positions 0-9 are 0 because in actual ring # updates those positions would have been updated for start_pos = 0. # So CachePositionsManager thinks they are updated because start_pos > (0-9) @@ -549,8 +532,6 @@ def test_non_dynamic_shape(self): self.head_dim, enable_dynamic_shape=False, dtype=self.dtype, - window_size=self.window_size, - max_seq_len=self.max_seq_len, ) # Create input tensors @@ -580,6 +561,6 @@ def test_non_dynamic_shape(self): self.assertTrue(torch.all(v_out[:, :, i] == 16.0)) # Check that the rest of the cache is still zeros - for i in range(seq_len, self.window_size): + for i in range(seq_len, self.max_context_length): self.assertTrue(torch.all(k_out[:, :, i] == 0.0)) self.assertTrue(torch.all(v_out[:, :, i] == 0.0)) diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index fac940cf3ee..f6ed378ec03 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -16,6 +16,8 @@ // @lint-ignore CLANGTIDY facebook-unused-include-check #include +#include + #ifdef ET_USE_THREADPOOL #include #include @@ -988,193 +990,6 @@ void cpu_flash_attention( scalar_t* buf_reduced_data = is_reduced_type ? reinterpret_cast(buf_reduced) : nullptr; - // An explicit mask is shared across batches and heads. Precompute up to two - // useful K/V intervals for every (query tile, K/V tile) pair. Ring attention - // can mask large portions of its backing cache with -inf; discovering those - // ranges once lets every head skip fully masked tiles and trim the boundary - // tiles before either GEMM. Arbitrary additive masks retain their existing - // behavior because only values that are exactly -inf are excluded. - struct MaskBlockRanges { - int64_t first_begin; - int64_t first_end; - int64_t second_begin; - int64_t second_end; - bool first_mask_is_zero; - bool second_mask_is_zero; - }; - const int64_t num_kv_blocks = (kvSize - 1) / kvSplitSize + 1; - MaskBlockRanges* mask_block_ranges = nullptr; - std::unique_ptr allocated_mask_block_ranges; - uint8_t* column_states_by_thread = nullptr; - std::unique_ptr allocated_column_states; - bool use_mask_ranges = false; - if (has_attn_mask) { - const int64_t num_mask_block_ranges = qSlice * num_kv_blocks; - const int64_t mask_block_ranges_bytes = - num_mask_block_ranges * sizeof(MaskBlockRanges); - Result mask_block_ranges_scratch = - ctx.allocate_temp(mask_block_ranges_bytes, 64); - if (!mask_block_ranges_scratch.ok()) { - allocated_mask_block_ranges = - std::make_unique(mask_block_ranges_bytes); - mask_block_ranges = - reinterpret_cast(allocated_mask_block_ranges.get()); - } else { - mask_block_ranges = - reinterpret_cast(mask_block_ranges_scratch.get()); - } - - // Bit 0 means at least one query row can attend to the column. Bit 1 - // means every query row has a zero additive mask for the column. - const int64_t column_states_bytes = kvSplitSize * num_thread + qSlice; - Result column_states_scratch = - ctx.allocate_temp(column_states_bytes, 64); - if (!column_states_scratch.ok()) { - allocated_column_states = std::make_unique(column_states_bytes); - column_states_by_thread = - reinterpret_cast(allocated_column_states.get()); - } else { - column_states_by_thread = - reinterpret_cast(column_states_scratch.get()); - } - - const accum_t neg_inf = -std::numeric_limits::infinity(); - uint8_t* useful_mask_ranges = - column_states_by_thread + kvSplitSize * num_thread; - auto find_mask_ranges = [&](int64_t begin, int64_t end) { - const int64_t thread_index = torch::executor::get_thread_num(); - uint8_t* column_states = - column_states_by_thread + thread_index * kvSplitSize; - for (int64_t q_block = begin; q_block < end; ++q_block) { - bool found_useful_range = false; - const int64_t query_begin = q_block * qSplitSize; - const int64_t query_end = std::min(query_begin + qSplitSize, qSize); - const int64_t causal_end = - is_causal ? std::min(start_pos + query_end, kvSize) : kvSize; - for (int64_t kv_block = 0; kv_block < num_kv_blocks; ++kv_block) { - const int64_t block_begin = kv_block * kvSplitSize; - const int64_t block_end = - std::min(std::min(block_begin + kvSplitSize, kvSize), causal_end); - const int64_t range_index = q_block * num_kv_blocks + kv_block; - auto& ranges = mask_block_ranges[range_index]; - ranges.first_begin = block_begin; - ranges.first_end = block_begin; - ranges.second_begin = block_begin; - ranges.second_end = block_begin; - ranges.first_mask_is_zero = false; - ranges.second_mask_is_zero = false; - if (block_begin >= block_end) { - continue; - } - - const int64_t block_size = block_end - block_begin; - std::fill(column_states, column_states + block_size, uint8_t{2}); - - for (int64_t row = query_begin; row < query_end; ++row) { - const accum_t* mask_row = mask_data + row * mStrideM; - for (int64_t col = block_begin; col < block_end; ++col) { - const accum_t mask_value = mask_row[col]; - auto& state = column_states[col - block_begin]; - state |= mask_value != neg_inf; - if (mask_value != static_cast(0)) { - state &= uint8_t{1}; - } - } - } - - int64_t active_begin = block_begin; - while (active_begin < block_end && - (column_states[active_begin - block_begin] & uint8_t{1}) == - 0) { - ++active_begin; - } - int64_t active_end = block_end; - while (active_end > active_begin && - (column_states[active_end - 1 - block_begin] & uint8_t{1}) == - 0) { - --active_end; - } - ranges.first_begin = active_begin; - ranges.first_end = active_end; - ranges.second_begin = block_end; - ranges.second_end = block_end; - if (active_begin >= active_end) { - found_useful_range = true; - continue; - } - - // A wrapped ring window has at most one interior gap. custom_sdpa - // also accepts arbitrary additive masks, which may contain several - // gaps, so find the largest one and split around it when it is large - // enough to repay the extra pair of GEMM calls. Smaller gaps stay - // represented by their existing -inf values inside one bounding - // interval. - int64_t largest_gap_begin = active_begin; - int64_t largest_gap_end = active_begin; - int64_t gap_begin = active_begin; - while (gap_begin < active_end) { - while (gap_begin < active_end && - (column_states[gap_begin - block_begin] & uint8_t{1}) != 0) { - ++gap_begin; - } - int64_t gap_end = gap_begin; - while (gap_end < active_end && - (column_states[gap_end - block_begin] & uint8_t{1}) == 0) { - ++gap_end; - } - if (gap_end - gap_begin > largest_gap_end - largest_gap_begin) { - largest_gap_begin = gap_begin; - largest_gap_end = gap_end; - } - gap_begin = gap_end; - } - - constexpr int64_t min_gap_to_split = 64; - if (largest_gap_end - largest_gap_begin >= min_gap_to_split) { - ranges.first_end = largest_gap_begin; - ranges.second_begin = largest_gap_end; - ranges.second_end = active_end; - } - - auto mask_range_is_zero = [&](int64_t range_begin, - int64_t range_end) { - for (int64_t col = range_begin; col < range_end; ++col) { - if ((column_states[col - block_begin] & uint8_t{2}) == 0) { - return false; - } - } - return true; - }; - ranges.first_mask_is_zero = - mask_range_is_zero(ranges.first_begin, ranges.first_end); - if (ranges.second_begin < ranges.second_end) { - ranges.second_mask_is_zero = - mask_range_is_zero(ranges.second_begin, ranges.second_end); - } - - const int64_t retained_size = ranges.first_end - ranges.first_begin + - ranges.second_end - ranges.second_begin; - if (retained_size < block_size || ranges.first_mask_is_zero || - ranges.second_mask_is_zero) { - found_useful_range = true; - } - } - useful_mask_ranges[q_block] = found_useful_range; - } - }; - const bool mask_ranges_computed = - torch::executor::parallel_for(0, qSlice, 1, find_mask_ranges); - ET_KERNEL_CHECK_MSG( - ctx, - mask_ranges_computed, - Internal, - , - "parallel_for failed while precomputing attention mask ranges"); - for (int64_t q_block = 0; q_block < qSlice; ++q_block) { - use_mask_ranges |= useful_mask_ranges[q_block]; - } - } - auto compute_lambda = [&](int64_t begin, int64_t end) { int64_t i = 0, j = 0, k = 0; data_index_init(begin, i, batchSize, j, num_head, k, qSlice); @@ -1230,38 +1045,11 @@ void cpu_flash_attention( is_causal ? std::min(m + start_pos + qBlockSize, kvSize) : kvSize; int64_t m_start_pos = m + start_pos; auto j_kv = j / num_reps; - fill_stub(dst_data, static_cast(0), qBlockSize * headSize); - bool has_processed_kv = false; - const int64_t num_ranges = use_mask_ranges - ? 2 * num_kv_blocks - : (num_keys + kvSplitSize - 1) / kvSplitSize; - for (int64_t range = 0; range < num_ranges; ++range) { - int64_t kvBlockStart; - int64_t kvBlockEnd; - bool range_mask_is_zero = false; - if (use_mask_ranges) { - const int64_t kv_block = range / 2; - const bool use_second_range = range % 2 != 0; - const auto& ranges = mask_block_ranges[k * num_kv_blocks + kv_block]; - if (use_second_range) { - kvBlockStart = ranges.second_begin; - kvBlockEnd = ranges.second_end; - range_mask_is_zero = ranges.second_mask_is_zero; - } else { - kvBlockStart = ranges.first_begin; - kvBlockEnd = ranges.first_end; - range_mask_is_zero = ranges.first_mask_is_zero; - } - } else { - kvBlockStart = range * kvSplitSize; - kvBlockEnd = std::min(kvBlockStart + kvSplitSize, kvSize); - } - kvBlockEnd = std::min(kvBlockEnd, num_keys); - if (kvBlockStart >= kvBlockEnd) { - continue; - } - const int64_t kvBlockSize = kvBlockEnd - kvBlockStart; - const bool apply_attn_mask = has_attn_mask && !range_mask_is_zero; + fill_stub(dst_data, static_cast(0), qSplitSize * headSize); + for (int64_t n = 0; n < num_keys; n += kvSplitSize) { + int64_t kvBlockSize = std::min(kvSplitSize, kvSize - n); + // Calculate scale * q @ k.T + fill_stub(qk_data, static_cast(0), qSplitSize * kvSplitSize); const void* q_sub_matrix_data_ptr; const void* k_sub_matrix_data_ptr; @@ -1270,14 +1058,12 @@ void cpu_flash_attention( const int8_t* q_zero_points_ptr = nullptr; const int8_t* k_zero_points_ptr = nullptr; int64_t q_offset = i * qStrideB + j * qStrideH + m * qStrideM; - int64_t k_offset = - i * kStrideB + j_kv * kStrideH + kvBlockStart * kStrideN; + int64_t k_offset = i * kStrideB + j_kv * kStrideH + n * kStrideN; if (is_quantized_sdpa) { int64_t q_quant_params_offset = i * q_quant_params_StrideB + j * q_quant_params_StrideH + m * q_quant_params_StrideM; int64_t k_quant_params_offset = i * k_quant_params_StrideB + - j_kv * k_quant_params_StrideH + - kvBlockStart * k_quant_params_StrideN; + j_kv * k_quant_params_StrideH + n * k_quant_params_StrideN; q_scales_ptr = q_scales.value().const_data_ptr() + q_quant_params_offset; k_scales_ptr = @@ -1320,31 +1106,57 @@ void cpu_flash_attention( (widen_qk && qBlockSize >= kMinQBlockForWidenedQK) ? widen_ptr : nullptr); - // Apply causal masking relative to the retained KV block. These are - // the two overlap configurations; a KV block wholly before the new - // query needs no causal masking. Rows are new-query tokens, columns - // are KV-block tokens, '+' is attendable, and '-' is causally masked: - // - // New query begins midway through KV block: Tail of new query lies in - // KV block: - // + + + - - - - - - - - - - // + + + + - - - - - - - - - // + + + + + - + - - - - - - // + + + + + + + + - - - - - // + + + + + + + + + - - - - // - // Each row may attend through its own logical position, so last_col - // is the number of keys in [kvBlockStart, kvBlockEnd) that are not in - // its future. - if (is_causal && m_start_pos < kvBlockEnd) { + // There are 4 cases that is_causal has to cover to fill + // not-attendable-position with -inf + /* 1. Everything is attended to. This happens when m_start_pos > n + + kvSplitSize e.g m_pos [8:15] and n_pos [0:7]. Since you must attend to + all previous tokens matrix is full + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 2. Everything is not attended to. However only some tokens at the + beginning dont attend to everything. This happens when m_start_pos <= n + + kvSplitSize but m_start_pos + qBlockSize > n + kvSplitSize m_start_pos + = 8 qBlockSize = 8 n = 4 kvSplitSize = 8 For example m_pos [8:15] but + n_pos is [4:11] + + + + + + - - - + + + + + + + - - + + + + + + + + - + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + 3. In this case only last few tokens have something to attend to. + This happens when m_start_pos < n and m_start_pos + qBlockSize >= n and + m_start_pos + qBlockSize <= n + kvSplitSize m_start_pos = 8 qBlockSize = + 8 n = 13 kvSplitSize = 8 For example m_pos [8:15] but n_pos is [13:20] + - - - - - - - - + - - - - - - - - + - - - - - - - - + - - - - - - - - + - - - - - - - - + + - - - - - - - + + + - - - - - - + + + + - - - - - + 4. In this no tokens attend to anything, but we dont really have to + take care of this case because the loop for (int64_t n = 0; n < + num_keys; n += kvSplitSize) will exit before that. + */ + if (is_causal && m_start_pos <= n + kvSplitSize) { + // For this fn to work k_split_size > q_split_size for (int32_t row = 0; - row < qBlockSize && (m_start_pos + row < kvBlockEnd - 1); + row < qBlockSize && (m_start_pos + row < n + (kvSplitSize - 1)); ++row) { - // When last_col is 0, it means that the entire row is not - // attended to because the range begins after the query position. - int64_t last_col = kvBlockStart > (m_start_pos + row) - ? 0 - : row + m_start_pos + 1 - kvBlockStart; + // When last_col is 0, it means that the entire row is not attended + // to because m_pos is smaller than n_pos. So everything in n is for + // future. + int64_t last_col = + n > (m_start_pos + row) ? 0 : row + m_start_pos + 1 - n; accum_t* row_ptr = qk_data + row * kvBlockSize; fill_stub( row_ptr + last_col, @@ -1355,7 +1167,7 @@ void cpu_flash_attention( // Update attention weights with attention mask // And apply scaling factor // qk <- qk * scaling + attn_mask - if (apply_attn_mask) { + if (has_attn_mask) { for (int64_t row = 0; row < qBlockSize; ++row) { vec::map2( [scaling_factor](Vec x, Vec y) { @@ -1364,14 +1176,14 @@ void cpu_flash_attention( qk_data + row * kvBlockSize, qk_data + row * kvBlockSize, mask_data + i * mStrideB + j * mStrideH + (m + row) * mStrideM + - kvBlockStart, + n, kvBlockSize); } } // Update coefficients with Softmax accum_t tmp_max = 0, tmp_sum = 0, exp_tmp = 0; for (int64_t row = 0; row < qBlockSize; ++row) { - if (apply_attn_mask) { + if (has_attn_mask) { // max per row tmp_max = vec::reduce_all( [](Vec& x, Vec& y) { return vec::maximum(x, y); }, @@ -1408,7 +1220,7 @@ void cpu_flash_attention( // max[row] <- max qk_max_data[row] = tmp_max; // dst <- dst * exp_tmp - if (has_processed_kv) { + if (n > 0) { vec::map( [exp_tmp](Vec x) { return x * Vec(exp_tmp); }, dst_data + row * headSize, @@ -1421,12 +1233,10 @@ void cpu_flash_attention( const void* v_sub_matrix_data_ptr; const float* v_scales_ptr = nullptr; const int8_t* v_zero_points_ptr = nullptr; - int64_t v_offset = - i * vStrideB + j_kv * vStrideH + kvBlockStart * vStrideN; + int64_t v_offset = i * vStrideB + j_kv * vStrideH + n * vStrideN; if (is_quantized_sdpa) { int64_t v_quant_params_offset = i * v_quant_params_StrideB + - j_kv * v_quant_params_StrideH + - kvBlockStart * v_quant_params_StrideN; + j_kv * v_quant_params_StrideH + n * v_quant_params_StrideN; v_scales_ptr = v_scales.value().const_data_ptr() + v_quant_params_offset; v_zero_points_ptr = v_zero_points.value().const_data_ptr() + @@ -1477,12 +1287,10 @@ void cpu_flash_attention( vStrideN, dst_data, headSize, - has_processed_kv ? static_cast(1) - : static_cast(0), + n == 0 ? static_cast(0) : static_cast(1), buf_qdq_ptr, widen_v, use_fp32_qk_weights); - has_processed_kv = true; } // dst <- dst / sum[row] // reorder MHA output with strides diff --git a/extension/llm/custom_ops/op_sdpa_test.cpp b/extension/llm/custom_ops/op_sdpa_test.cpp index 84fdb2e3966..0572fd7ff80 100644 --- a/extension/llm/custom_ops/op_sdpa_test.cpp +++ b/extension/llm/custom_ops/op_sdpa_test.cpp @@ -33,26 +33,6 @@ executorch::aten::Tensor op_scaled_dot_product_attention( context, query, key, value, attn_mask, dropout_p, is_causal, scale, out); } -executorch::aten::Tensor op_custom_sdpa( - const executorch::aten::Tensor& query, - const executorch::aten::Tensor& key, - const executorch::aten::Tensor& value, - const std::optional& attn_mask, - executorch::aten::Tensor& out) { - executorch::runtime::KernelRuntimeContext context{}; - return torch::executor::native::custom_sdpa_out( - context, - query, - key, - value, - /*start_pos=*/0, - attn_mask, - /*dropout_p=*/0.0, - /*is_causal=*/false, - /*scale=*/std::nullopt, - out); -} - std::tuple op_gated_delta_rule( executorch::runtime::KernelRuntimeContext& context, @@ -356,151 +336,6 @@ TEST(OpScaledDotProductAttentionTest, CorrectnessTest_105) { EXPECT_TENSOR_CLOSE_WITH_TOL(ret, ret_expected, 1e-4, 1e-4); } -TEST(OpScaledDotProductAttentionTest, SparseMaskRangesMatchReference) { - TensorFactory tfFloat; - - constexpr int32_t q_size = 64; - constexpr int32_t kv_size = 1024; - std::vector key_values(kv_size, 0.0f); - std::vector value_values(kv_size); - std::vector mask_values( - q_size * kv_size, -std::numeric_limits::infinity()); - std::vector expected_values(q_size); - - for (int32_t i = 0; i < kv_size; ++i) { - value_values[i] = static_cast(i); - } - // The first 32-row query tile has staggered active columns separated by a - // 38-column gap, which is too small to split. This exercises the per-column - // union across query rows while retaining the detailed mask inside the range. - for (int32_t row = 0; row < 32; ++row) { - mask_values[row * kv_size + 10 + row] = 0.0f; - mask_values[row * kv_size + 80 + row] = 0.0f; - expected_values[row] = 45.0f + row; - } - // The second query tile uses KV block 1 and has a gap larger than the split - // threshold. The zero first range and additive second range exercise both - // mask paths while also checking query-tile indexing for k > 0. - for (int32_t row = 32; row < q_size; ++row) { - mask_values[row * kv_size + 600] = 0.0f; - mask_values[row * kv_size + 700] = 1.0986122886681098f; // log(3) - expected_values[row] = 675.0f; - } - - // custom_sdpa uses [batch, sequence, heads, head_dim]. - auto query = tfFloat.zeros({1, q_size, 1, 1}); - auto key = tfFloat.make({1, kv_size, 1, 1}, key_values); - auto value = tfFloat.make({1, kv_size, 1, 1}, value_values); - auto attn_mask = tfFloat.make({q_size, kv_size}, mask_values); - auto out = tfFloat.zeros({1, q_size, 1, 1}); - - auto result = op_custom_sdpa(query, key, value, attn_mask, out); - - auto expected = tfFloat.make({1, q_size, 1, 1}, expected_values); - EXPECT_TENSOR_CLOSE_WITH_TOL(result, expected, 1e-5, 1e-5); -} - -TEST( - OpScaledDotProductAttentionTest, - CausalSparseMaskRangesCoverGqaAndReducedPrecision) { - TensorFactory tfHalf; - TensorFactory tfFloat; - - constexpr int32_t q_size = 256; - constexpr int32_t kv_size = 1024; - constexpr int32_t num_query_heads = 2; - std::vector value_values(kv_size); - std::vector mask_values( - q_size * kv_size, -std::numeric_limits::infinity()); - std::vector expected_values(num_query_heads * q_size); - - for (int32_t col = 0; col < kv_size; ++col) { - value_values[col] = static_cast(col); - } - for (int32_t row = 0; row < q_size; ++row) { - mask_values[row * kv_size] = 0.0f; - for (int32_t col = 100; col < q_size; ++col) { - mask_values[row * kv_size + col] = 0.0f; - } - const float expected = - row < 100 ? 0.0f : (100.0f + row) * (row - 99) / (2.0f * (row - 98)); - for (int32_t head = 0; head < num_query_heads; ++head) { - expected_values[head * q_size + row] = - static_cast(expected); - } - } - - auto query = tfHalf.zeros({1, num_query_heads, q_size, 1}); - auto key = tfHalf.zeros({1, 1, kv_size, 1}); - auto value = tfHalf.make({1, 1, kv_size, 1}, value_values); - auto attn_mask = tfFloat.make({q_size, kv_size}, mask_values); - auto out = tfHalf.zeros({1, num_query_heads, q_size, 1}); - - auto result = op_scaled_dot_product_attention( - query, - key, - value, - attn_mask, - /*dropout_p=*/0.0, - /*is_causal=*/true, - /*scale=*/std::nullopt, - out); - - auto expected = tfHalf.make({1, num_query_heads, q_size, 1}, expected_values); - EXPECT_TENSOR_CLOSE_WITH_TOL(result, expected, 1e-2, 1e-2); -} - -TEST(OpScaledDotProductAttentionTest, QuantizedSparseMaskTrimmedBoundary) { - TensorFactory tfChar; - TensorFactory tfFloat; - - constexpr int32_t kv_size = 1024; - std::vector value_values(kv_size); - std::vector mask_values( - kv_size, -std::numeric_limits::infinity()); - for (int32_t col = 0; col < kv_size; ++col) { - value_values[col] = static_cast(col % 100); - } - for (int32_t col = 100; col < 200; ++col) { - mask_values[col] = 0.0f; - } - - auto query = tfChar.zeros({1, 1, 1, 1}); - auto key = tfChar.zeros({1, kv_size, 1, 1}); - auto value = tfChar.make({1, kv_size, 1, 1}, value_values); - auto query_zero_points = tfChar.zeros({1, 1, 1, 1}); - auto key_zero_points = tfChar.zeros({1, kv_size, 1, 1}); - auto value_zero_points = tfChar.zeros({1, kv_size, 1, 1}); - auto query_scales = tfFloat.ones({1, 1, 1, 1}); - auto key_scales = tfFloat.ones({1, kv_size, 1, 1}); - auto value_scales = tfFloat.ones({1, kv_size, 1, 1}); - auto attn_mask = tfFloat.make({1, kv_size}, mask_values); - auto out = tfFloat.zeros({1, 1, 1, 1}); - - executorch::runtime::KernelRuntimeContext context{}; - auto result = torch::executor::native::custom_quantized_sdpa_out( - context, - query, - key, - value, - /*start_pos=*/0, - attn_mask, - /*dropout_p=*/0.0, - /*is_causal=*/false, - /*scale=*/std::nullopt, - query_zero_points, - query_scales, - key_zero_points, - key_scales, - value_zero_points, - value_scales, - /*is_seq_at_dim_2=*/false, - out); - - auto expected = tfFloat.make({1, 1, 1, 1}, {49.5f}); - EXPECT_TENSOR_CLOSE_WITH_TOL(result, expected, 1e-4, 1e-4); -} - TEST(OpScaledDotProductAttentionTest, CorrectnessTest_11) { TensorFactory tfFloat; From 3413fdfae9e173973ae709bff29abef0aed4df13 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?M=C3=A5ns=20Nilsson?= Date: Thu, 10 Sep 2026 02:52:12 +0200 Subject: [PATCH 128/190] XNNPACK: Lift constant mul scalars for partitioning (#20515) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit XNNPACK supports the tensor overload for multiply, but aten.mul.Scalar is normally decomposed before the XNNPACK partitioner can select it. Add a MulScalarConfig that preserves eligible scalar multiplications during partitioning. During XNNPACK preprocessing, lift the scalar into an exported constant-tensor placeholder and rewrite the operation to aten.mul.Tensor, allowing the existing multiply serialization path to handle it. This avoids introducing an aten.full operation and works with the default to_edge_transform_and_lower flow without explicit XNNPACK transform passes. Run scalar lifting after SDPA conversion so ConvertToSDPAPass can recover attention scales from the original aten.mul.Scalar operations. Also allow SDPA conversion to match graphs processed by the default no-op-expand transform. For DeIT Tiny, this removes 24 portable aten.mul.Scalar nodes and reduces delegate count from 62 to 50. In current local timing checks the latency impact is modest: about 1% faster on both Android SME2 and the aarch64 XNNPACK/KleidiAI NEON-class host runner. These are modest uplifts but may introduce more opportunities for improvements. cc @GregoryComer @digantdesai @cbilgin @freddan80 @per @zingo @oscarandersson8218 @Sebastian-Larsson @robell @rascani --------- Signed-off-by: Måns Nilsson Co-authored-by: Jacob Stevens --- backends/xnnpack/_passes/__init__.py | 4 + backends/xnnpack/_passes/convert_to_sdpa.py | 56 ++++--- .../lift_constant_scalar_operands_pass.py | 119 +++++++++++++++ backends/xnnpack/partition/config/__init__.py | 2 + .../partition/config/generic_node_configs.py | 33 ++++ backends/xnnpack/test/ops/test_multiply.py | 22 +++ ...test_lift_constant_scalar_operands_pass.py | 142 ++++++++++++++++++ .../xnnpack/test/test_xnnpack_partitioner.py | 66 ++++++++ 8 files changed, 420 insertions(+), 24 deletions(-) create mode 100644 backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py create mode 100644 backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py diff --git a/backends/xnnpack/_passes/__init__.py b/backends/xnnpack/_passes/__init__.py index f55144395af..675336f6274 100644 --- a/backends/xnnpack/_passes/__init__.py +++ b/backends/xnnpack/_passes/__init__.py @@ -29,6 +29,9 @@ from executorch.backends.xnnpack._passes.fuse_activation_pass import FuseActivationPass from executorch.backends.xnnpack._passes.fuse_batch_norm import FuseBatchNormPass from executorch.backends.xnnpack._passes.insert_pad_qdq import InsertPadQDQPass +from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import ( + LiftConstantScalarOperandsPass, +) from executorch.backends.xnnpack._passes.prelu_reshape_pass import PReLUReshapePass from executorch.backends.xnnpack._passes.propagate_custom_meta_pass import ( PropagateCustomMetaPass, @@ -79,6 +82,7 @@ def __init__( ConvertToLinearPass, PropagateCustomMetaPass, ConvertToSDPAPass, + LiftConstantScalarOperandsPass, ConstPropPass, FuseBatchNormPass, DecomposeBatchNorm, diff --git a/backends/xnnpack/_passes/convert_to_sdpa.py b/backends/xnnpack/_passes/convert_to_sdpa.py index c7982db750f..b926d62df8f 100644 --- a/backends/xnnpack/_passes/convert_to_sdpa.py +++ b/backends/xnnpack/_passes/convert_to_sdpa.py @@ -1,19 +1,22 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import logging +from copy import deepcopy from typing import Optional import torch from executorch.backends.transforms import get_shape - +from executorch.backends.xnnpack._passes.remove_noop_expand_copy_pass import ( + RemoveNoopExpandCopyPass, +) from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.partition.graphs import sdpa from executorch.exir.dialects._ops import ops as exir_ops - from torch.fx.passes.infra.pass_base import PassResult from torch.fx.passes.utils.matcher_utils import InternalMatch, SubgraphMatcher @@ -24,31 +27,31 @@ class ConvertToSDPAPass(XNNPACKPass): def get_scale(self, match: InternalMatch) -> Optional[float]: """ - Returns the scale of the SDPA op. + Return the SDPA scale recovered from the matched pre-QK^T multiplications. - Scale: Optional[float] doesn't change the graph pattern. - The default value can be calulated however we need to extract - it for lowering when it is the user supplied value anyway. + The decomposition applies the square root of the attention scale before + QK^T, so the extracted multiplier is squared to recover the original value. """ for node in match.nodes_map.values(): if ( - node.op == "call_function" - and node.target == exir_ops.edge.aten.mul.Scalar + node.op != "call_function" + or node.target != exir_ops.edge.aten.mul.Scalar ): - scale = node.args[1] + continue - dtype = torch.float - mul_val = node.meta.get("val", None) - if mul_val is not None: - dtype = mul_val.dtype + scale = node.args[1] - if isinstance(scale, float): - # Convert scale value to fp16 (reducing precision) - scale = torch.tensor(scale, dtype=dtype).item() + dtype = torch.float + mul_val = node.meta.get("val", None) + if mul_val is not None: + dtype = mul_val.dtype - # since scale we extracted this before the QK^T. - return scale**2 - break + if isinstance(scale, float): + # Convert scale value to fp16 (reducing precision) + scale = torch.tensor(scale, dtype=dtype).item() + + # since scale we extracted this before the QK^T. + return scale**2 return None def assert_2d_mask(self, match: InternalMatch) -> None: @@ -99,11 +102,16 @@ def call(self, graph_module: torch.fx.GraphModule): logger.debug("ConvertToSDPA Begin: ") logger.debug(graph_module.print_readable(print_output=False)) - for pattern in sdpa.get_graphs(): - sm = SubgraphMatcher(pattern.graph, ignore_literals=True) - matches = list(sm.match(graph_module.graph)) - for partition_to_replace in matches: - self.create_sdpa(graph_module, partition_to_replace) + for scalar_pattern in sdpa.get_graphs(): + normalized_pattern = RemoveNoopExpandCopyPass()( + deepcopy(scalar_pattern) + ).graph_module + + for pattern in (scalar_pattern, normalized_pattern): + sm = SubgraphMatcher(pattern.graph, ignore_literals=True) + matches = list(sm.match(graph_module.graph)) + for partition_to_replace in matches: + self.create_sdpa(graph_module, partition_to_replace) graph_module.recompile() graph_module = super().call(graph_module).graph_module diff --git a/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py new file mode 100644 index 00000000000..01524a3268d --- /dev/null +++ b/backends/xnnpack/_passes/lift_constant_scalar_operands_pass.py @@ -0,0 +1,119 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +from numbers import Number +from typing import Dict, Optional, Union + +import torch +from executorch.backends.transforms.utils import create_constant_placeholder +from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload +from executorch.exir.pass_base import PassResult +from torch._ops import OpOverload +from torch.export import ExportedProgram +from torch.export.graph_signature import InputKind + +ScalarOp = Union[EdgeOpOverload, OpOverload] + + +class LiftConstantScalarOperandsPass(XNNPACKPass): + """ + Lift scalar operands into tensor constants for selected binary ops. + + XNNPACK already supports the tensor overloads for these binary operations. + This pass converts explicitly listed scalar overloads to their tensor + overloads by replacing constant scalar operands with small tensor constants. + The constants are registered as exported-program constant tensor inputs. + Keep the op map narrow until each new scalar overload is covered by tests. + """ + + default_scalar_to_tensor_ops: Dict[ScalarOp, ScalarOp] = { + exir_ops.edge.aten.mul.Scalar: exir_ops.edge.aten.mul.Tensor, + } + + def __init__( + self, + exported_program: ExportedProgram, + scalar_to_tensor_ops: Optional[Dict[ScalarOp, ScalarOp]] = None, + ) -> None: + super().__init__(exported_program) + self.scalar_to_tensor_ops = ( + scalar_to_tensor_ops + if scalar_to_tensor_ops is not None + else self.default_scalar_to_tensor_ops + ) + + def _create_constant_node( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + value: Number, + ) -> torch.fx.Node: + input_node = node.args[0] + if not isinstance(input_node, torch.fx.Node): + raise RuntimeError("Expected scalar op input to be an FX node.") + + input_value = input_node.meta["val"] + tensor = torch.tensor(value, dtype=input_value.dtype, device=input_value.device) + name = self._get_new_constant_name(graph_module) + first_placeholder = next( + graph_node + for graph_node in graph_module.graph.nodes + if graph_node.op == "placeholder" + ) + with graph_module.graph.inserting_before(first_placeholder): + return create_constant_placeholder( + self.exported_program, + graph_module.graph, + name, + InputKind.CONSTANT_TENSOR, + tensor, + ) + + def _get_new_constant_name(self, graph_module: torch.fx.GraphModule) -> str: + prefix = "_tensor_constant_" + existing_names = {node.name for node in graph_module.graph.nodes} + existing_names.update(self.exported_program.constants) + existing_names.update(self.exported_program.state_dict) + index = 0 + while f"{prefix}{index}" in existing_names: + index += 1 + return f"{prefix}{index}" + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + modified = False + + for node in list(graph_module.graph.nodes): + if ( + node.op != "call_function" + or node.target not in self.scalar_to_tensor_ops + or len(node.args) != 2 + or not isinstance(node.args[0], torch.fx.Node) + or not isinstance(node.args[1], Number) + ): + continue + + input_value = node.args[0].meta.get("val") + output_value = node.meta.get("val") + if ( + not isinstance(input_value, torch.Tensor) + or not isinstance(output_value, torch.Tensor) + or input_value.dtype != output_value.dtype + ): + continue + + tensor_arg = self._create_constant_node(graph_module, node, node.args[1]) + node.args = (node.args[0], tensor_arg) + node.target = self.scalar_to_tensor_ops[node.target] + modified = True + + graph_module.graph.eliminate_dead_code() + graph_module.graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/xnnpack/partition/config/__init__.py b/backends/xnnpack/partition/config/__init__.py index 7775a674510..02ddc64874d 100644 --- a/backends/xnnpack/partition/config/__init__.py +++ b/backends/xnnpack/partition/config/__init__.py @@ -43,6 +43,7 @@ MeanDimConfig, MinimumConfig, MulConfig, + MulScalarConfig, NegConfig, PermuteConfig, PowConfig, @@ -107,6 +108,7 @@ MinimumConfig, MMConfig, MulConfig, + MulScalarConfig, NegConfig, PermuteConfig, PowConfig, diff --git a/backends/xnnpack/partition/config/generic_node_configs.py b/backends/xnnpack/partition/config/generic_node_configs.py index 8741962db38..907e972cedf 100644 --- a/backends/xnnpack/partition/config/generic_node_configs.py +++ b/backends/xnnpack/partition/config/generic_node_configs.py @@ -8,6 +8,7 @@ # pyre-unsafe import logging +from numbers import Number from typing import cast, List, Optional import numpy as np @@ -439,6 +440,38 @@ def supported_precision_types(self) -> List[ConfigPrecisionType]: return [ConfigPrecisionType.FP32, ConfigPrecisionType.STATIC_QUANT] +class MulScalarConfig(GenericNodePartitionerConfig): + target_name = "mul.Scalar" + + def check_constraints(self, node: torch.fx.Node, ep: ExportedProgram) -> bool: + if node.graph is not ep.graph: + return False + + if not self.check_common_constraints(node, ep): + return False + + if ( + len(node.args) != 2 + or not isinstance(node.args[0], torch.fx.Node) + or not isinstance(node.args[1], Number) + ): + return False + + input_value = node.args[0].meta.get("val") + output_value = node.meta.get("val") + return ( + isinstance(input_value, torch.Tensor) + and isinstance(output_value, torch.Tensor) + and input_value.dtype == output_value.dtype + ) + + def supported_precision_types(self) -> List[ConfigPrecisionType]: + return [ConfigPrecisionType.FP32] + + def get_original_aten(self) -> Optional[torch._ops.OpOverload]: + return torch.ops.aten.mul.Scalar + + class MaximumConfig(GenericNodePartitionerConfig): target_name = "maximum.default" diff --git a/backends/xnnpack/test/ops/test_multiply.py b/backends/xnnpack/test/ops/test_multiply.py index 3315200005d..a44b5e6b406 100644 --- a/backends/xnnpack/test/ops/test_multiply.py +++ b/backends/xnnpack/test/ops/test_multiply.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -29,6 +30,10 @@ def forward(self, x, y): z = torch.mul(x, y) * torch.functional.torch.mul(x, y) return z + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + class MulRelu(torch.nn.Module): def forward(self, x, y): z = x * y @@ -58,6 +63,23 @@ def test_fp32_mul(self): inputs = (torch.randn((1, 3)), torch.randn((4, 3))) self._test_mul(inputs) + def test_fp32_mul_scalar(self): + ( + Tester(self.MulScalar(), (torch.randn(2, 3),)) + .export() + .to_edge_transform_and_lower() + .check_count({"torch.ops.higher_order.executorch_call_delegate": 1}) + .check_not( + [ + "executorch_exir_dialects_edge__ops_aten_mul_Tensor", + "executorch_exir_dialects_edge__ops_aten_mul_Scalar", + ] + ) + .to_executorch() + .serialize() + .run_method_and_compare_outputs() + ) + def test_qs8_mul(self): inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 1)) ( diff --git a/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py new file mode 100644 index 00000000000..dd59aab5204 --- /dev/null +++ b/backends/xnnpack/test/passes/test_lift_constant_scalar_operands_pass.py @@ -0,0 +1,142 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch +from executorch.backends.xnnpack._passes import XNNPACKPassManager +from executorch.backends.xnnpack._passes.convert_to_sdpa import ConvertToSDPAPass +from executorch.backends.xnnpack._passes.lift_constant_scalar_operands_pass import ( + LiftConstantScalarOperandsPass, +) +from executorch.backends.xnnpack.utils.configs import ( + get_transform_passes, + get_xnnpack_edge_compile_config, +) +from executorch.exir import to_edge +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export.graph_signature import InputKind + + +class TestLiftConstantScalarOperandsPass(unittest.TestCase): + def setUp(self): + torch._dynamo.reset() + + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + + class AddScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.add.Scalar(x, 0.5) + + class SDPA(torch.nn.Module): + def forward(self, q, k, v, mask): + return torch.nn.functional.scaled_dot_product_attention(q, k, v, mask) + + def _to_edge_program_manager(self, module): + return to_edge( + torch.export.export(module, (torch.randn(2, 3),), strict=True), + compile_config=get_xnnpack_edge_compile_config(skip_dim_order=True), + ) + + def _lift(self, exported_program): + return XNNPACKPassManager( + exported_program, passes=[LiftConstantScalarOperandsPass] + ).transform() + + def test_lifts_mul_scalar_operand(self): + exported_program = self._lift( + self._to_edge_program_manager(self.MulScalar()).exported_program() + ) + graph = exported_program.graph_module.graph + + self.assertFalse( + any(node.target == exir_ops.edge.aten.mul.Scalar for node in graph.nodes) + ) + self.assertTrue( + any(node.target == exir_ops.edge.aten.mul.Tensor for node in graph.nodes) + ) + self.assertFalse(any(node.op == "get_attr" for node in graph.nodes)) + + constant_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CONSTANT_TENSOR + ] + self.assertEqual(len(constant_specs), 1) + constant_spec = constant_specs[0] + self.assertIn(constant_spec.target, exported_program.constants) + + placeholders = [node for node in graph.nodes if node.op == "placeholder"] + self.assertEqual(placeholders[0].name, constant_spec.arg.name) + mul_node = next( + node for node in graph.nodes if node.target == exir_ops.edge.aten.mul.Tensor + ) + self.assertIs(mul_node.args[1], placeholders[0]) + + def test_is_idempotent(self): + exported_program = self._lift( + self._to_edge_program_manager(self.MulScalar()).exported_program() + ) + exported_program = self._lift(exported_program) + + constant_specs = [ + spec + for spec in exported_program.graph_signature.input_specs + if spec.kind == InputKind.CONSTANT_TENSOR + ] + self.assertEqual(len(constant_specs), 1) + self.assertEqual(len(exported_program.constants), 1) + + def test_keeps_unmapped_scalar_op(self): + exported_program = self._lift( + self._to_edge_program_manager(self.AddScalar()).exported_program() + ) + graph = exported_program.graph_module.graph + + self.assertTrue( + any(node.target == exir_ops.edge.aten.add.Scalar for node in graph.nodes) + ) + self.assertFalse(exported_program.constants) + + def test_converts_sdpa_after_default_transform_passes(self): + q = torch.randn(2, 4, 8, 16) + k = torch.randn(2, 4, 8, 16) + v = torch.randn(2, 4, 8, 16) + mask = torch.randn(8, 8) + for use_default_transforms in (False, True): + with self.subTest(use_default_transforms=use_default_transforms): + edge = to_edge( + torch.export.export(self.SDPA(), (q, k, v, mask), strict=True), + compile_config=get_xnnpack_edge_compile_config(), + ) + if use_default_transforms: + edge = edge.transform(get_transform_passes()) + exported_program = XNNPACKPassManager( + edge.exported_program(), + passes=[ConvertToSDPAPass, LiftConstantScalarOperandsPass], + ).transform() + + graph = exported_program.graph_module.graph + self.assertTrue( + any( + node.target + == exir_ops.edge.aten.scaled_dot_product_attention.default + for node in graph.nodes + ) + ) + self.assertFalse( + any( + node.target == exir_ops.edge.aten.bmm.default + for node in graph.nodes + ) + ) + self.assertFalse( + any( + node.target == exir_ops.edge.aten.mul.Scalar + for node in graph.nodes + ) + ) diff --git a/backends/xnnpack/test/test_xnnpack_partitioner.py b/backends/xnnpack/test/test_xnnpack_partitioner.py index 894fab4098f..ff25c69655e 100644 --- a/backends/xnnpack/test/test_xnnpack_partitioner.py +++ b/backends/xnnpack/test/test_xnnpack_partitioner.py @@ -1,5 +1,6 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. +# Copyright 2026 Arm Limited and/or its affiliates. # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. @@ -30,6 +31,71 @@ def __init__(self): def forward(self, x): return self.linear(x) + class MulScalar(torch.nn.Module): + def forward(self, x): + return torch.ops.aten.mul.Scalar(x, 0.5) + + class MulScalarCond(torch.nn.Module): + def forward(self, pred, x): + def true_fn(value): + return torch.ops.aten.mul.Scalar(value, 0.5) + + def false_fn(value): + return torch.ops.aten.mul.Scalar(value, 2.0) + + return torch.cond(pred, true_fn, false_fn, (x,)) + + def test_mul_scalar_ops_to_not_decompose(self): + partitioner = XnnpackPartitioner() + exported_program = export(self.MulScalar(), (torch.randn(2, 3),)) + ops, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIn(torch.ops.aten.mul.Scalar, ops) + self.assertIsNotNone(filter_fn) + mul_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertTrue(filter_fn(mul_node)) + + def test_mul_scalar_ops_to_not_decompose_rejects_unsupported_dtype(self): + partitioner = XnnpackPartitioner() + exported_program = export( + self.MulScalar(), (torch.ones(2, 3, dtype=torch.int32),) + ) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIsNotNone(filter_fn) + mul_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertFalse(filter_fn(mul_node)) + + def test_mul_scalar_ops_to_not_decompose_rejects_cond_branches(self): + partitioner = XnnpackPartitioner() + exported_program = export( + self.MulScalarCond(), (torch.tensor(True), torch.randn(2, 3)) + ) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + + self.assertIsNotNone(filter_fn) + cond_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.higher_order.cond + ) + for branch_node in cond_node.args[1:3]: + branch = exported_program.graph_module.get_submodule(branch_node.target) + mul_node = next( + node + for node in branch.graph.nodes + if node.target == torch.ops.aten.mul.Scalar + ) + self.assertFalse(filter_fn(mul_node)) + def test_deprecation_warning_for_to_backend_workflow(self): """ Test that the deprecated to_edge + to_backend workflow shows a deprecation warning. From 897fc8e3dc63e542e50ce4570f6f782d3ee46d22 Mon Sep 17 00:00:00 2001 From: Ethan Ng Date: Wed, 9 Sep 2026 17:56:00 -0700 Subject: [PATCH 129/190] Move split/concat chain folding to OSS (#22628) Summary: Move `MergeSplitConcatChainPass` into `executorch.backends.transforms` package and implement it as a standard `PassBase`. The pass recognizes complete, ordered outputs from ATen `chunk`, `split`, and `split_with_sizes`, as well as edge `split_with_sizes_copy`, followed by `cat`. Replace the chain with `view_copy` only when the input is contiguous, element counts match, and every axis crossed between the split and concat dimensions is singleton. This preserves flattened storage order for both pre-edge and post-edge consumers. Reviewed By: DrJessop Differential Revision: D119242414 Pull Request resolved: https://github.com/pytorch/executorch/pull/22628 --- .../transforms/merge_split_concat_chain.py | 153 ++++++++++++++ backends/transforms/targets.bzl | 24 +++ .../test/test_merge_split_concat_chain.py | 197 ++++++++++++++++++ 3 files changed, 374 insertions(+) create mode 100644 backends/transforms/merge_split_concat_chain.py create mode 100644 backends/transforms/test/test_merge_split_concat_chain.py diff --git a/backends/transforms/merge_split_concat_chain.py b/backends/transforms/merge_split_concat_chain.py new file mode 100644 index 00000000000..b29ca68efab --- /dev/null +++ b/backends/transforms/merge_split_concat_chain.py @@ -0,0 +1,153 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import operator +from collections.abc import Callable, Sequence, Set + +import torch +from executorch.backends.transforms.permute_pass_utils import get_arg +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from torch.fx import GraphModule, Node +from torch.fx.passes.infra.pass_base import PassBase + + +_RAW_SPLIT_TARGETS: frozenset[Callable[..., object]] = frozenset( + { + torch.ops.aten.chunk.default, + torch.ops.aten.split.Tensor, + torch.ops.aten.split_with_sizes.default, + } +) +_EDGE_SPLIT_TARGETS: frozenset[Callable[..., object]] = frozenset( + {exir_ops.edge.aten.split_with_sizes_copy.default} +) + + +def _normalize_dim(dim: int, rank: int) -> int: + normalized = dim + rank if dim < 0 else dim + assert 0 <= normalized < rank + return normalized + + +def _ordered_split_inputs( + cat_node: Node, + split_targets: Set[Callable[..., object]], +) -> tuple[Node, list[Node]] | None: + cat_inputs = get_arg(cat_node, "tensors") + if not isinstance(cat_inputs, Sequence) or not cat_inputs: + return None + + getitem_nodes: list[Node] = [] + for inp in cat_inputs: + if not isinstance(inp, Node) or inp.target != operator.getitem: + return None + getitem_nodes.append(inp) + + split_node = getitem_nodes[0].args[0] + if not isinstance(split_node, Node) or split_node.target not in split_targets: + return None + + split_outputs = split_node.meta["val"] + if not isinstance(split_outputs, (tuple, list)) or len(getitem_nodes) != len( + split_outputs + ): + return None + for index, getitem_node in enumerate(getitem_nodes): + if getitem_node.args[0] != split_node or getitem_node.args[1] != index: + return None + return split_node, getitem_nodes + + +def _is_view_equivalent( + cat_node: Node, + split_node: Node, + getitem_nodes: Sequence[Node], +) -> bool: + split_input = get_arg(split_node, "input", Node) + input_val = split_input.meta["val"] + output_val = cat_node.meta["val"] + if not isinstance(input_val, torch.Tensor) or not isinstance( + output_val, torch.Tensor + ): + return False + if not input_val.is_contiguous() or input_val.numel() != output_val.numel(): + return False + + rank = input_val.ndim + split_dim = _normalize_dim(get_arg(split_node, "dim", int), rank) + cat_dim = _normalize_dim(get_arg(cat_node, "dim", int), rank) + + first_moved_dim, last_moved_dim = sorted((split_dim, cat_dim)) + # Moving split parts across axes preserves flattened storage order only + # when every crossed axis is singleton. + for getitem_node in getitem_nodes: + value = getitem_node.meta["val"] + if not isinstance(value, torch.Tensor) or any( + size != 1 for size in value.shape[first_moved_dim:last_moved_dim] + ): + return False + return True + + +class MergeSplitConcatChainPass(PassBase): + """Replace view-equivalent split/getitem/cat chains with a view. + + For example, splitting a contiguous ``[2, 1, 6, 4, 4]`` tensor into + three parts along dimension 2 and concatenating them along dimension 1 + is equivalent to a view with shape ``[2, 3, 2, 4, 4]``. + """ + + def _replace_cats( + self, + graph_module: GraphModule, + cat_target: Callable[..., object], + split_targets: Set[Callable[..., object]], + view_target: Callable[..., object], + ) -> bool: + modified = False + for cat_node in graph_module.graph.find_nodes( + op="call_function", target=cat_target + ): + match = _ordered_split_inputs(cat_node, split_targets) + if match is None: + continue + split_node, getitem_nodes = match + if not _is_view_equivalent(cat_node, split_node, getitem_nodes): + continue + + output_val = cat_node.meta["val"] + assert isinstance(output_val, torch.Tensor) + split_input = get_arg(split_node, "input", Node) + with graph_module.graph.inserting_before(cat_node): + replacement_view = graph_module.graph.call_function( + view_target, + (split_input, list(output_val.shape)), + ) + replacement_view.meta = cat_node.meta.copy() + cat_node.replace_all_uses_with(replacement_view) + modified = True + return modified + + def call(self, graph_module: GraphModule) -> PassResult: + modified = self._replace_cats( + graph_module, + torch.ops.aten.cat.default, + _RAW_SPLIT_TARGETS, + torch.ops.aten.view_copy.default, + ) + modified |= self._replace_cats( + graph_module, + exir_ops.edge.aten.cat.default, + _EDGE_SPLIT_TARGETS, + exir_ops.edge.aten.view_copy.default, + ) + if modified: + graph_module.graph.eliminate_dead_code() + graph_module.recompile() + return PassResult(graph_module, modified) diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index d34767b84ae..eb282ce23e4 100644 --- a/backends/transforms/targets.bzl +++ b/backends/transforms/targets.bzl @@ -521,6 +521,30 @@ def define_common_targets(): ], ) + runtime.python_library( + name = "merge_split_concat_chain", + srcs = ["merge_split_concat_chain.py"], + visibility = ["PUBLIC"], + deps = [ + ":permute_pass_utils", + "//caffe2:torch", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + + runtime.python_test( + name = "test_merge_split_concat_chain", + srcs = ["test/test_merge_split_concat_chain.py"], + deps = [ + ":merge_split_concat_chain", + "//caffe2:torch", + "//executorch/backends/test:graph_builder", + "//executorch/exir:pass_base", + "//executorch/exir/dialects:lib", + ], + ) + runtime.python_library( name = "fuse_cascaded_transpose_or_permute_ops", srcs = ["fuse_cascaded_transpose_or_permute_ops.py"], diff --git a/backends/transforms/test/test_merge_split_concat_chain.py b/backends/transforms/test/test_merge_split_concat_chain.py new file mode 100644 index 00000000000..153bed74fa6 --- /dev/null +++ b/backends/transforms/test/test_merge_split_concat_chain.py @@ -0,0 +1,197 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import copy +import operator +import unittest +from collections.abc import Callable, Sequence +from typing import cast + +import torch +from executorch.backends.test.graph_builder import GraphBuilder +from executorch.backends.transforms.merge_split_concat_chain import ( + MergeSplitConcatChainPass, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import PassResult +from torch.fx import GraphModule + + +def _build_split_cat_graph( + input_value: torch.Tensor, + split_target: Callable[..., object], + split_spec: int | list[int], + split_dim: int, + output_count: int, + cat_target: Callable[..., object], + cat_dim: int, + order: Sequence[int] | None = None, +) -> GraphModule: + builder = GraphBuilder() + input_node = builder.placeholder("input", input_value) + split = builder.call_operator( + split_target, + (input_node, split_spec, split_dim), + ) + output_order = range(output_count) if order is None else order + split_outputs = [ + builder.call_operator(operator.getitem, (split, index)) + for index in output_order + ] + cat = builder.call_operator(cat_target, (split_outputs, cat_dim)) + builder.output([cat]) + return builder.get_graph_module() + + +_VIEW_EQUIVALENT_CASES = ( + ( + "split", + torch.ops.aten.split.Tensor, + 2, + (2, 1, 6, 4, 4), + 2, + 3, + torch.ops.aten.cat.default, + 1, + torch.ops.aten.view_copy.default, + ), + ( + "chunk", + torch.ops.aten.chunk.default, + 3, + (2, 1, 5, 4, 4), + 2, + 3, + torch.ops.aten.cat.default, + 2, + torch.ops.aten.view_copy.default, + ), + ( + "split_with_sizes", + torch.ops.aten.split_with_sizes.default, + [1, 1, 1, 1], + (1, 4, 3, 2), + 1, + 4, + torch.ops.aten.cat.default, + 0, + torch.ops.aten.view_copy.default, + ), + ( + "edge_split_with_sizes", + exir_ops.edge.aten.split_with_sizes_copy.default, + [1, 1, 1, 1], + (4, 1, 3, 2), + 0, + 4, + exir_ops.edge.aten.cat.default, + 1, + exir_ops.edge.aten.view_copy.default, + ), +) + +_UNSAFE_CASES = ( + ("reordered_outputs", [1, 0, 2, 3], 0, False), + ("non_singleton_crossing", [0, 1, 2, 3], 3, False), + ("non_contiguous_input", [0, 1, 2, 3], 0, True), +) + +_SPLIT_CAT_DIALECTS = ( + ( + "aten", + torch.ops.aten.split_with_sizes.default, + torch.ops.aten.cat.default, + ), + ( + "edge", + exir_ops.edge.aten.split_with_sizes_copy.default, + exir_ops.edge.aten.cat.default, + ), +) + + +class MergeSplitConcatChainPassTest(unittest.TestCase): + def test_merges_view_equivalent_chains(self) -> None: + for ( + name, + split_target, + split_spec, + input_shape, + split_dim, + output_count, + cat_target, + cat_dim, + view_target, + ) in _VIEW_EQUIVALENT_CASES: + with self.subTest(name=name): + input_value = torch.randn(input_shape) + graph_module = _build_split_cat_graph( + input_value, + split_target, + split_spec, + split_dim, + output_count, + cat_target, + cat_dim, + ) + reference = copy.deepcopy(graph_module) + + result = cast(PassResult, MergeSplitConcatChainPass()(graph_module)) + + self.assertTrue(result.modified) + torch.testing.assert_close( + reference(input_value), result.graph_module(input_value) + ) + for target, expected_count in ( + (split_target, 0), + (cat_target, 0), + (view_target, 1), + ): + self.assertEqual( + expected_count, + len( + result.graph_module.graph.find_nodes( + op="call_function", target=target + ) + ), + ) + + def test_does_not_merge_unsafe_chains(self) -> None: + for dialect, split_target, cat_target in _SPLIT_CAT_DIALECTS: + for name, order, cat_dim, non_contiguous in _UNSAFE_CASES: + with self.subTest(dialect=dialect, name=name): + input_value = torch.randn(1, 4, 3, 2) + if non_contiguous: + input_value = input_value.transpose(2, 3) + graph_module = _build_split_cat_graph( + input_value, + split_target, + [1, 1, 1, 1], + 1, + 4, + cat_target, + cat_dim, + order, + ) + reference = copy.deepcopy(graph_module) + + result = cast(PassResult, MergeSplitConcatChainPass()(graph_module)) + + self.assertFalse(result.modified) + torch.testing.assert_close( + reference(input_value), result.graph_module(input_value) + ) + self.assertEqual( + 1, + len( + result.graph_module.graph.find_nodes( + op="call_function", + target=cat_target, + ) + ), + ) From 18c08031eb335bfc2b72ae8e08b09e98bfd46c8f Mon Sep 17 00:00:00 2001 From: Jiseong-oh Date: Thu, 10 Sep 2026 13:36:16 +0900 Subject: [PATCH 130/190] Samsung AI LiteCore - buildsystem: separate backend and example builds. (#22582) ### Summary To keep backend and example artifacts distinct. AS-IS: - The backend library and 'enn_executor_runner' were built together. - The example artifact 'enn_executor_runner' was placed under 'backend/samsung', it is not proper location for example artifacts. TO-BE: - Build the backend library first, then let the example projects refer to already-built backend library. - Adjust CMakeLists so that 'enn_executor_runner' is compiled in the 'examples/samsung/' directory and its output is placed in there. ### Test plan ${EXECUTORCH_ROOT}/backends/samsung/build.sh -b all cc @SS-JIA @digantdesai @kimishpatel --------- Signed-off-by: Jiseong.oh Signed-off-by: jiseong.oh --- .ci/scripts/setup-samsung-linux-deps.sh | 1 + backends/samsung/CMakeLists.txt | 18 ------ backends/samsung/README.md | 9 ++- backends/samsung/build.sh | 10 +++ .../samsung/test/utils/runtime_executor.py | 2 +- examples/samsung/CMakeLists.txt | 61 +++++++++++++++++++ examples/samsung/README.md | 19 ++++-- examples/samsung/build.sh | 45 ++++++++++++++ 8 files changed, 139 insertions(+), 26 deletions(-) create mode 100644 examples/samsung/CMakeLists.txt create mode 100755 examples/samsung/build.sh diff --git a/.ci/scripts/setup-samsung-linux-deps.sh b/.ci/scripts/setup-samsung-linux-deps.sh index a8339b16c7d..548f30bf161 100644 --- a/.ci/scripts/setup-samsung-linux-deps.sh +++ b/.ci/scripts/setup-samsung-linux-deps.sh @@ -165,6 +165,7 @@ install_enn_backend() { echo "NDK will be installed/used at: ${ANDROID_NDK_ROOT}" bash backends/samsung/build.sh --build all + bash examples/samsung/build.sh export EXECUTORCH_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" export PYTHONPATH="${PYTHONPATH:-}:${EXECUTORCH_ROOT}/.." diff --git a/backends/samsung/CMakeLists.txt b/backends/samsung/CMakeLists.txt index 1f647c5bbe2..c8b4bedbc4a 100644 --- a/backends/samsung/CMakeLists.txt +++ b/backends/samsung/CMakeLists.txt @@ -146,24 +146,6 @@ if(${ANDROID}) executorch_target_link_options_shared_lib(enn_backend) target_compile_options(enn_backend PRIVATE -Wno-deprecated-declarations) - set(__enn_executor_runner_srcs - ${EXECUTORCH_SOURCE_DIR}/examples/samsung/executor_runner/enn_executor_runner.cpp - ) - add_executable(enn_executor_runner ${__enn_executor_runner_srcs}) - add_dependencies(enn_executor_runner enn_backend) - target_link_libraries( - enn_executor_runner - PRIVATE enn_logging - enn_backend - gflags - executorch - extension_data_loader - portable_ops_lib - android - ) - set_target_properties( - enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden - ) install( TARGETS enn_backend enn_logging EXPORT ExecuTorchTargets diff --git a/backends/samsung/README.md b/backends/samsung/README.md index f498bd0a098..30c855fe365 100644 --- a/backends/samsung/README.md +++ b/backends/samsung/README.md @@ -49,13 +49,16 @@ Generates python artifacts that allow user call `Compile` interface to lower a m ./backends/samsung/build.sh -b x86_64 ``` -### Build ENN Executor Runner +### Build Backend Delegate ```bash ./backends/samsung/build.sh -b android --ndk ${ANDROID_NDK} ``` -ANDROID_ABI=arm64-v8a is default, necessary runtime executable generated in `build_exynos_android` directory. +ANDROID_ABI=arm64-v8a is default, necessary runtime backend library generated in `build_samsung_android` directory. -### Build Anroid Extension +### Build Executable +Please see the [README.md](../../examples/samsung/README.md). + +### Build Android Extension This is later exposed Java app. Please turn on CMake option `EXECUTORCH_BUILD_ENN`, and ENN runtime will be added. ```bash cmake extension/android \ diff --git a/backends/samsung/build.sh b/backends/samsung/build.sh index a4871feb50c..e6258152785 100755 --- a/backends/samsung/build.sh +++ b/backends/samsung/build.sh @@ -66,7 +66,17 @@ function build_android() { ANDROID_ABI=arm64-v8a ANDROID_PLATFORM=android-28 # Trace requires over android-23 + local host_flatcc=${X86_64_BUILD_DIR}/third-party/flatcc_ep/bin/flatcc + local flatcc_args=() + if [[ -x ${host_flatcc} ]]; then + flatcc_args=(-DFLATCC_EXECUTABLE=${host_flatcc}) + else + echo "Warning: ${host_flatcc} not found. Build the x86_64 target first" \ + "('-b x86_64' or '-b all') if executor_runner fails to link libflatccrt.a." + fi + cmake \ + "${flatcc_args[@]}" \ -DCMAKE_INSTALL_PREFIX=${ANDROID_BUILD_DIR} \ -DCMAKE_TOOLCHAIN_FILE="${ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake" \ -DANDROID_NDK=${ANDROID_NDK} \ diff --git a/backends/samsung/test/utils/runtime_executor.py b/backends/samsung/test/utils/runtime_executor.py index 9bc274799d7..4d642657d23 100644 --- a/backends/samsung/test/utils/runtime_executor.py +++ b/backends/samsung/test/utils/runtime_executor.py @@ -26,7 +26,7 @@ def get_runner_path() -> Path: cwd=os.path.dirname(os.path.realpath(__file__)), text=True, ).strip() - return Path(git_root) / "build_samsung_android/backends/samsung/enn_executor_runner" + return Path(git_root) / "build_samsung_android/examples/samsung/enn_executor_runner" class EDBTestManager: diff --git a/examples/samsung/CMakeLists.txt b/examples/samsung/CMakeLists.txt new file mode 100644 index 00000000000..28a6dcc6660 --- /dev/null +++ b/examples/samsung/CMakeLists.txt @@ -0,0 +1,61 @@ +# Copyright (c) 2025 Samsung Electronics Co. LTD +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +cmake_minimum_required(VERSION 3.15) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +project(samsung_example) + +# Source root directory for executorch. +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../..) +endif() + +include(${EXECUTORCH_ROOT}/tools/cmake/Utils.cmake) +include(${EXECUTORCH_ROOT}/tools/cmake/Codegen.cmake) + +if(NOT PYTHON_EXECUTABLE) + resolve_python_executable() +endif() + +if(NOT CMAKE_BUILD_TYPE) + set(CMAKE_BUILD_TYPE RelWithDebInfo) +endif() + +if(CMAKE_TOOLCHAIN_FILE MATCHES ".*(iOS|ios\.toolchain)\.cmake$") + message(FATAL_ERROR "ios is not supported by Samsung AI system.") +endif() + +# prebuilt libraries. executorch package should contain portable_ops_lib, +# etdump, bundled_program. +find_package(executorch CONFIG REQUIRED) +target_compile_options(executorch INTERFACE -DET_EVENT_TRACER_ENABLED) +find_package(gflags REQUIRED) + +set(_common_compile_options -Wno-deprecated-declarations -fPIC) + +add_compile_options(-Wall -Werror -fPIC) + +message("Build Samsung Android Examples") + +set(__enn_executor_runner__srcs + ${CMAKE_CURRENT_LIST_DIR}/executor_runner/enn_executor_runner.cpp +) + +add_executable(enn_executor_runner ${__enn_executor_runner__srcs}) + +target_include_directories(enn_executor_runner PRIVATE ${EXECUTORCH_ROOT}/..) + +target_compile_options(enn_executor_runner PRIVATE ${_common_compile_options}) + +target_link_libraries( + enn_executor_runner PRIVATE enn_logging enn_backend gflags executorch + extension_data_loader portable_ops_lib +) + +set_target_properties( + enn_executor_runner PROPERTIES CXX_VISIBILITY_PRESET hidden +) diff --git a/examples/samsung/README.md b/examples/samsung/README.md index 8b21c48a34f..a64169bd9a1 100644 --- a/examples/samsung/README.md +++ b/examples/samsung/README.md @@ -1,9 +1,9 @@ -# Exynos backend Examples +# Exynos Backend examples This directory contains examples for some AI models. -Please make sure you have built the library and executable before -you start, if you have no idea how to build, please refer to [backend README](../../backends/samsung/README.md). +Please make sure you have built the library before you start, +if you have no idea how to build, please refer to [backend README](../../backends/samsung/README.md). ## Environment We set up `PYTHONPATH` because it's easier to develop and import executorch Python APIs. @@ -42,15 +42,26 @@ Examples use "PerformanceMode.HIGH_PERFORMANCE" mode, this mode is experimental. If you want to use this mode on your model, verify your model on devicefarm which can use samsung developer society site firstly for checking stability. (https://soc-developer.semiconductor.samsung.com/) +## Building Executable +### Prerequisites +Please set up the backend before building the executable, See the [backend README](../../backends/samsung/README.md) for details. +### Building 'enn_executor_runner' for Android +```bash +export EXYNOS_AI_LITECORE_ROOT=/path/to/enn_sdk +export ANDROID_NDK_ROOT=/path/to/android_ndk +${EXECUTORCH_ROOT}/examples/samsung/build.sh +``` +After the build completes, `enn_executor_runner` can be found at `${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/` ## Execution + After lowering, we could get a pte model and then run it on mobile phone. #### Step 1: Push required ENN libraries and executor runner to device ```bash DEVICE_DIR=/data/local/tmp/executorch adb shell mkdir ${DEVICE_DIR} -adb push ${EXECUTORCH_ROOT}/cmake-android-out/backends/samsung/enn_executor_runner ${DEVICE_DIR} +adb push ${EXECUTORCH_ROOT}/build_samsung_android/examples/samsung/enn_executor_runner ${DEVICE_DIR} ``` #### Step 2: Indicate dynamic linkers and execute model diff --git a/examples/samsung/build.sh b/examples/samsung/build.sh new file mode 100755 index 00000000000..4568180777a --- /dev/null +++ b/examples/samsung/build.sh @@ -0,0 +1,45 @@ +#!/bin/bash + +set -e + +BASE_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" >/dev/null 2>&1 && pwd )" +PROJECT_DIR=$(realpath ${BASE_DIR}/../../) + +echo PROJECT_DIR=${PROJECT_DIR} + +if [[ -z ${ANDROID_NDK_ROOT} ]]; then + echo "Please export ANDROID_NDK_ROOT" + exit 1 +fi + +ANDROID_ABI=arm64-v8a +ANDROID_PLATFORM=android-28 # Trace requires over android-23 + +echo ANDROID_NDK_ROOT=${ANDROID_NDK_ROOT} +echo ANDROID_ABI=${ANDROID_ABI} +echo ANDROID_PLATFORM=${ANDROID_PLATFORM} + +main() { + cd "$PROJECT_DIR" + local build_dir_root="build_samsung_android" + local example_root="examples/samsung" + local build_dir_example="$PROJECT_DIR/${build_dir_root}/${example_root}" + local cmake_prefix_path="$PROJECT_DIR/${build_dir_root}/lib/cmake/ExecuTorch;$PROJECT_DIR/${build_dir_root}/third-party/gflags;$PROJECT_DIR/${build_dir_root}/lib/cmake/tokenizers;$PROJECT_DIR/${build_dir_root}/lib/cmake/re2;$PROJECT_DIR/${build_dir_root}/lib/cmake/absl;" + + echo build_dir=${build_dir_root} + echo build_dir_example=${build_dir_example} + echo cmake_prefix_path=${cmake_prefix_path} + + cmake -DCMAKE_PREFIX_PATH=${cmake_prefix_path} \ + -DCMAKE_TOOLCHAIN_FILE="$ANDROID_NDK_ROOT/build/cmake/android.toolchain.cmake" \ + -DANDROID_NDK=$ANDROID_NDK \ + -DANDROID_ABI="$ANDROID_ABI" \ + -DANDROID_PLATFORM=$ANDROID_PLATFORM \ + -DCMAKE_FIND_ROOT_PATH_MODE_PACKAGE=BOTH \ + -DCMAKE_BUILD_TYPE=Release \ + -B"${build_dir_example}" \ + "$PROJECT_DIR/${example_root}" + cmake --build build_samsung_android/examples/samsung/ --config Release +} + +main "$@" From d1d5f406e7316c33d6252b6c5f470e6955471c2c Mon Sep 17 00:00:00 2001 From: Jiseong-oh Date: Thu, 10 Sep 2026 13:37:48 +0900 Subject: [PATCH 131/190] Samsung Exynos AI LiteCore - Samsung Backend in test suite of ExecuTorch (#22583) ### Summary Release notes: TestSuite The Samsung Exynos backend in TestSuite of ExecuTorch can be verfied and it was passed Testsuite for E9955/E9965 SoCs. ### Test plan pytest -c /dev/nul backends/test/suite/operators/test_add.py -m backend_samsung -vs --html=./report/add_report.html pytest -c /dev/nul backends/test/suite/models/ -m backend_samsung --html=./report/samsung_models.html cc @SS-JIA @digantdesai @kimishpatel --------- Signed-off-by: Jiseong Oh Signed-off-by: jiseong.oh Co-authored-by: Guochao Yang Co-authored-by: Jingya Zhang --- .../samsung/_passes/remove_useless_ops.py | 1 - .../samsung/_passes/replace_scalar_ops.py | 9 +++- backends/samsung/builders/__init__.py | 4 ++ backends/samsung/builders/node_visitor.py | 6 ++- backends/samsung/builders/op_add.py | 9 +++- backends/samsung/builders/op_avg_pool2d.py | 8 ++- backends/samsung/builders/op_batch_norm.py | 16 +++++- backends/samsung/builders/op_bmm.py | 4 +- backends/samsung/builders/op_cat.py | 4 +- backends/samsung/builders/op_clamp.py | 10 +++- .../samsung/builders/op_constant_pad_nd.py | 4 +- backends/samsung/builders/op_conv2d.py | 22 +++++++- backends/samsung/builders/op_cos.py | 4 +- backends/samsung/builders/op_div.py | 4 +- backends/samsung/builders/op_embedding.py | 4 +- backends/samsung/builders/op_exp.py | 33 ++++++++++++ backends/samsung/builders/op_expand_copy.py | 14 ++++-- backends/samsung/builders/op_gelu.py | 4 +- backends/samsung/builders/op_getitem.py | 4 +- backends/samsung/builders/op_group_norm.py | 4 +- backends/samsung/builders/op_hardsigmoid.py | 4 +- backends/samsung/builders/op_hardswish.py | 4 +- backends/samsung/builders/op_hardtanh.py | 4 +- backends/samsung/builders/op_index.py | 4 +- backends/samsung/builders/op_layer_norm.py | 4 +- backends/samsung/builders/op_leaky_relu.py | 4 +- backends/samsung/builders/op_linear.py | 4 +- backends/samsung/builders/op_log.py | 4 +- backends/samsung/builders/op_log_softmax.py | 4 +- backends/samsung/builders/op_max_pool2d.py | 8 ++- backends/samsung/builders/op_maximum.py | 4 +- backends/samsung/builders/op_mean_dim.py | 13 ++++- backends/samsung/builders/op_minimum.py | 4 +- backends/samsung/builders/op_mul.py | 4 +- backends/samsung/builders/op_permute.py | 6 ++- backends/samsung/builders/op_pixel_shuffle.py | 4 +- backends/samsung/builders/op_placeholder.py | 4 +- backends/samsung/builders/op_pow.py | 15 ++++-- backends/samsung/builders/op_quantize.py | 41 ++++++++------- backends/samsung/builders/op_relu.py | 4 +- backends/samsung/builders/op_reshape.py | 12 +++-- backends/samsung/builders/op_rms_norm.py | 4 +- backends/samsung/builders/op_rsqrt.py | 4 +- backends/samsung/builders/op_select.py | 4 +- backends/samsung/builders/op_sigmoid.py | 4 +- backends/samsung/builders/op_sin.py | 4 +- backends/samsung/builders/op_skip.py | 32 ++++++++++++ backends/samsung/builders/op_slice_copy.py | 4 +- backends/samsung/builders/op_softmax.py | 4 +- .../builders/op_split_with_sizes_copy.py | 8 ++- backends/samsung/builders/op_sqrt.py | 4 +- backends/samsung/builders/op_squeeze.py | 4 +- backends/samsung/builders/op_sub.py | 9 +++- backends/samsung/builders/op_sum_int_list.py | 4 +- backends/samsung/builders/op_tanh.py | 4 +- backends/samsung/builders/op_to_copy.py | 4 +- backends/samsung/builders/op_topk.py | 4 +- backends/samsung/builders/op_unsqueeze.py | 4 +- .../builders/op_upsample_bilinear2d.py | 10 +++- .../samsung/builders/op_upsample_nearest2d.py | 8 ++- backends/samsung/partition/enn_partitioner.py | 23 +++++++-- .../samsung/test/tester/samsung_tester.py | 12 ++++- .../samsung/test/utils/runtime_executor.py | 2 +- backends/samsung/utils/export_utils.py | 5 ++ backends/test/suite/flow.py | 13 +++++ backends/test/suite/flows/samsung.py | 43 ++++++++++++++++ backends/test/suite/models/test_torchaudio.py | 2 +- .../executor_runner/enn_executor_runner.cpp | 50 +++++++++++-------- 68 files changed, 475 insertions(+), 121 deletions(-) create mode 100644 backends/samsung/builders/op_exp.py create mode 100644 backends/samsung/builders/op_skip.py create mode 100644 backends/test/suite/flows/samsung.py diff --git a/backends/samsung/_passes/remove_useless_ops.py b/backends/samsung/_passes/remove_useless_ops.py index c88a2d4a5d8..9c965844749 100644 --- a/backends/samsung/_passes/remove_useless_ops.py +++ b/backends/samsung/_passes/remove_useless_ops.py @@ -15,7 +15,6 @@ class RemoveUselessOpPass(ExportPass): USELESS_OP_SET = { exir_ops.edge.aten._to_copy.default, exir_ops.edge.aten.clone.default, - exir_ops.edge.aten.clone.default, exir_ops.edge.aten.alias.default, exir_ops.edge.aten.lift_fresh_copy.default, exir_ops.edge.dim_order_ops._to_dim_order_copy.default, diff --git a/backends/samsung/_passes/replace_scalar_ops.py b/backends/samsung/_passes/replace_scalar_ops.py index 8ae54b0dc98..22a74c15f61 100644 --- a/backends/samsung/_passes/replace_scalar_ops.py +++ b/backends/samsung/_passes/replace_scalar_ops.py @@ -38,9 +38,16 @@ def call_operator( if op not in self._ops_with_scalar: return super().call_operator(op, args, kwargs, meta) + # For pow operation, convert int scalar to float32 tensor + # because the PowVisitor requires both inputs to be float32 + if op == exir_ops.edge.aten.pow.Tensor_Scalar and isinstance(args[1], int): + args1 = torch.tensor(float(args[1]), dtype=torch.float32) + else: + args1 = torch.tensor(args[1]) + return super().call_operator( op=self._ops_with_scalar.get(op, op), - args=(args[0], torch.tensor(args[1])), + args=(args[0], args1), kwargs=kwargs, meta=meta, ) diff --git a/backends/samsung/builders/__init__.py b/backends/samsung/builders/__init__.py index 14b9a17a6c9..60f97cf60f7 100644 --- a/backends/samsung/builders/__init__.py +++ b/backends/samsung/builders/__init__.py @@ -18,6 +18,7 @@ op_dequantize, op_div, op_embedding, + op_exp, op_expand_copy, op_gelu, op_getitem, @@ -48,6 +49,7 @@ op_select, op_sigmoid, op_sin, + op_skip, op_slice_copy, op_softmax, op_split_with_sizes_copy, @@ -77,6 +79,7 @@ op_dequantize, op_div, op_embedding, + op_exp, op_expand_copy, op_gelu, op_getitem, @@ -107,6 +110,7 @@ op_select, op_sigmoid, op_sin, + op_skip, op_slice_copy, op_softmax, op_split_with_sizes_copy, diff --git a/backends/samsung/builders/node_visitor.py b/backends/samsung/builders/node_visitor.py index 0d2707da8f5..cb7d4db690a 100644 --- a/backends/samsung/builders/node_visitor.py +++ b/backends/samsung/builders/node_visitor.py @@ -31,7 +31,7 @@ def __init__(self, exported_program: ExportedProgram) -> None: def exported_program(self) -> ExportedProgram: return self._exported_program - def define_node(self, node: torch.fx.Node, enn_graph: EnnGraph): + def define_node(self, node: torch.fx.Node, enn_graph: EnnGraph) -> bool: raise NotImplementedError("NodeVisitor must be extended!") def define_tensor( @@ -58,7 +58,9 @@ def define_tensor( if is_param_node(self.exported_program, node): if swap_nc_for_weights: tensor = torch.swapdims(tensor, 0, 1) - const_data = tensor.contiguous().detach().numpy() + if not isinstance(tensor, torch._subclasses.fake_tensor.FakeTensor): + # .numpy() is not supported for tensor subclasses if the tensor is a fake tensor. + const_data = tensor.contiguous().detach().numpy() dims = [1] if len(tensor.size()) == 0 else list(tensor.size()) diff --git a/backends/samsung/builders/op_add.py b/backends/samsung/builders/op_add.py index a6eb79897dd..177f700f7e1 100644 --- a/backends/samsung/builders/op_add.py +++ b/backends/samsung/builders/op_add.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import Dict import torch @@ -26,16 +27,22 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) params = {} self._update_params_qdtype(node, params) input2 = node.args[1] input_id_2 = self.define_tensor(input2, enn_graph, vals_to_ids) + alpha = node.kwargs.get("alpha", 1.0) + if alpha != 1.0: + logging.warning("Currently, only alpha 1 for add is supported.") + return False output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op( node.name, "ELTSUM", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_avg_pool2d.py b/backends/samsung/builders/op_avg_pool2d.py index bfca8b89b22..529a3156030 100644 --- a/backends/samsung/builders/op_avg_pool2d.py +++ b/backends/samsung/builders/op_avg_pool2d.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -52,10 +52,6 @@ def define_node( params["explicit_padding"] = explicit_padding self._update_params_qdtype(node, params) - if len(node.args) > 4: - ceil_mode = cast(bool, node.args[4]) - assert not ceil_mode, "Not support ceil_mode = True." - if len(node.args) > 5: params["count_include_pad"] = cast(bool, node.args[5]) else: @@ -68,3 +64,5 @@ def define_node( ), "Not supported divisor_override which is not equal to pooling region." output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "AVGPOOL2D", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_batch_norm.py b/backends/samsung/builders/op_batch_norm.py index e5373a8223a..990b8e0ca20 100644 --- a/backends/samsung/builders/op_batch_norm.py +++ b/backends/samsung/builders/op_batch_norm.py @@ -25,7 +25,16 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: + index_zero_getitems = [] + for user in node.users.keys(): + if user.target.__name__ != "getitem": + continue + if user.args[1] == 0: + index_zero_getitems.append(user) + elif len(user.users) > 0: + return False + all_input_tensors = [] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -51,6 +60,11 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids, output_idx=0) + for getitem in index_zero_getitems: + vals_to_ids[getitem] = output_id + enn_graph.define_op( node.name, "BatchNormalization", all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_bmm.py b/backends/samsung/builders/op_bmm.py index 13e0d19cb14..2ac96a4a1b7 100644 --- a/backends/samsung/builders/op_bmm.py +++ b/backends/samsung/builders/op_bmm.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -41,3 +41,5 @@ def define_node( enn_graph.define_op( node.name, "BATCH_MATMUL", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_cat.py b/backends/samsung/builders/op_cat.py index 09387f2e361..82762cfcdbf 100644 --- a/backends/samsung/builders/op_cat.py +++ b/backends/samsung/builders/op_cat.py @@ -28,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: tensors = cast(List[torch.fx.Node], node.args[0]) input_tensor_ids = [] constant_idx = None @@ -48,3 +48,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "CONCAT", input_tensor_ids, [output_id], params) + + return True diff --git a/backends/samsung/builders/op_clamp.py b/backends/samsung/builders/op_clamp.py index 74af83212a5..b69066c3337 100644 --- a/backends/samsung/builders/op_clamp.py +++ b/backends/samsung/builders/op_clamp.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import cast, Dict import torch @@ -11,6 +12,7 @@ NodeVisitor, register_node_visitor, ) +from executorch.backends.samsung.builders.utils import get_tensor from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph @@ -26,9 +28,13 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) + input_tensor = get_tensor(self.exported_program, input) + if input_tensor.dtype == torch.int64: + logging.warning("Currently, int64 clip is unsupported!") + return False # The default value of lower bound and upper bound output_min = torch.finfo(torch.float32).min @@ -45,3 +51,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "CLIP", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_constant_pad_nd.py b/backends/samsung/builders/op_constant_pad_nd.py index 006f52619ff..bc72305bc13 100644 --- a/backends/samsung/builders/op_constant_pad_nd.py +++ b/backends/samsung/builders/op_constant_pad_nd.py @@ -29,7 +29,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -54,3 +54,5 @@ def define_node( } self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "PAD", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_conv2d.py b/backends/samsung/builders/op_conv2d.py index ab77d8df626..87f76f610bf 100644 --- a/backends/samsung/builders/op_conv2d.py +++ b/backends/samsung/builders/op_conv2d.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import cast, Dict, List import torch @@ -27,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -52,9 +53,24 @@ def define_node( padding = cast(List[int], node.args[4]) dilation = cast(List[int], node.args[5]) groups = cast(int, node.args[8]) + if is_transpose_conv and groups != 1: + logging.warning("Don't support groups for transpose conv.") + return False + output_padding = cast(List[int], node.args[7]) + if is_transpose_conv and output_padding != [0, 0]: + logging.warning("Don't support output padding for transpose conv.") + return False + if len(padding) < 2: + logging.warning( + "For conv1d decomposed to conv2d(with conv1d params), Conv1dToConv2d pass will update the params." + ) + return True explicit_padding = [padding[0], padding[1], padding[0], padding[1]] input_shape = get_shape(input) + if len(input_shape) > 4: + logging.warning("Currently, only conv2d is supported.") + return False kernel_shape = get_shape(weight_node) params = {} self._update_params_qdtype(node, params) @@ -72,7 +88,7 @@ def define_node( params["explicit_padding"] = explicit_padding params["in_channels"] = input_shape[1] params["out_channels"] = kernel_shape[0] * kernel_shape[1] * groups - params["out_channels"] //= input_shape[1] * input_shape[0] + params["out_channels"] //= input_shape[1] output_id = self.define_tensor(node, enn_graph, vals_to_ids) @@ -87,3 +103,5 @@ def define_node( enn_graph.define_op( node.name, conv_type, all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_cos.py b/backends/samsung/builders/op_cos.py index bd746db91cd..f1d4d917b7d 100644 --- a/backends/samsung/builders/op_cos.py +++ b/backends/samsung/builders/op_cos.py @@ -23,9 +23,11 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "Cos", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_div.py b/backends/samsung/builders/op_div.py index 8b0e7cdd5af..7afc23220fe 100644 --- a/backends/samsung/builders/op_div.py +++ b/backends/samsung/builders/op_div.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -40,3 +40,5 @@ def define_node( enn_graph.define_op( node.name, "ELTDIV", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_embedding.py b/backends/samsung/builders/op_embedding.py index a500ea051fd..bc28ae4e55e 100644 --- a/backends/samsung/builders/op_embedding.py +++ b/backends/samsung/builders/op_embedding.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: weight_node = node.args[0] weight_id = self.define_tensor(weight_node, enn_graph, vals_to_ids) @@ -40,3 +40,5 @@ def define_node( enn_graph.define_op( node.name, "GATHER", [weight_id, input_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_exp.py b/backends/samsung/builders/op_exp.py new file mode 100644 index 00000000000..63dd9ca3219 --- /dev/null +++ b/backends/samsung/builders/op_exp.py @@ -0,0 +1,33 @@ +# Copyright (c) 2026 Samsung Electronics Co. LTD +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Dict + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph + + +@register_node_visitor +class ExpVisitor(NodeVisitor): + target = "aten.exp.default" + + def define_node( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.Tensor, int], + ) -> bool: + input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) + + output_id = self.define_tensor(node, enn_graph, vals_to_ids) + + enn_graph.define_op(node.name, "Exp", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_expand_copy.py b/backends/samsung/builders/op_expand_copy.py index f4c707b8e62..5fb6b6c5166 100644 --- a/backends/samsung/builders/op_expand_copy.py +++ b/backends/samsung/builders/op_expand_copy.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import cast, Dict, List import torch @@ -27,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: # inputs input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -35,6 +36,8 @@ def define_node( in_shape = get_shape(input) sizes = cast(List[int], node.args[1]) expand_dims = self.check_expand_dims(sizes, in_shape) + if expand_dims is None: + return False # output output_id = self.define_tensor(node, enn_graph, vals_to_ids) @@ -53,7 +56,10 @@ def define_node( params, ) else: - raise NotImplementedError("Don't support expanding at more than one axes.") + logging.warning("Don't support expanding at more than one axes.") + return False + + return True def check_expand_dims(self, sizes, in_shape): expand_dims = [] @@ -72,6 +78,8 @@ def check_expand_dims(self, sizes, in_shape): while new_size_index > 0: new_size_index -= 1 - assert sizes[new_size_index] == 1, "Current expand is unsupported!" + if sizes[new_size_index] != 1: + logging.warning("Current expand is unsupported!") + return None return expand_dims diff --git a/backends/samsung/builders/op_gelu.py b/backends/samsung/builders/op_gelu.py index 88417f688f9..6a064194561 100644 --- a/backends/samsung/builders/op_gelu.py +++ b/backends/samsung/builders/op_gelu.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # input1 input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -38,3 +38,5 @@ def define_node( self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "GELU", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_getitem.py b/backends/samsung/builders/op_getitem.py index 901ec73cf7d..bc7c0441e3c 100644 --- a/backends/samsung/builders/op_getitem.py +++ b/backends/samsung/builders/op_getitem.py @@ -28,5 +28,5 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: - return + ) -> bool: + return True diff --git a/backends/samsung/builders/op_group_norm.py b/backends/samsung/builders/op_group_norm.py index 55c7bb6732a..b0509005053 100644 --- a/backends/samsung/builders/op_group_norm.py +++ b/backends/samsung/builders/op_group_norm.py @@ -23,7 +23,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) all_input_tensors.append(input_id) @@ -44,3 +44,5 @@ def define_node( enn_graph.define_op( node.name, "GROUPNORM", all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_hardsigmoid.py b/backends/samsung/builders/op_hardsigmoid.py index 3a50d65da41..58cc0a12d5b 100644 --- a/backends/samsung/builders/op_hardsigmoid.py +++ b/backends/samsung/builders/op_hardsigmoid.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) params = {} self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "HardSigmoid", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_hardswish.py b/backends/samsung/builders/op_hardswish.py index 8c30125e8a4..fc9ec134418 100644 --- a/backends/samsung/builders/op_hardswish.py +++ b/backends/samsung/builders/op_hardswish.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) params = {} self._update_params_qdtype(node, params) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "HARDSWISH", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_hardtanh.py b/backends/samsung/builders/op_hardtanh.py index 7d65e97a566..4c60d7227dc 100644 --- a/backends/samsung/builders/op_hardtanh.py +++ b/backends/samsung/builders/op_hardtanh.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -40,3 +40,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "CLIP", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_index.py b/backends/samsung/builders/op_index.py index b7765e35b3f..0145616de56 100644 --- a/backends/samsung/builders/op_index.py +++ b/backends/samsung/builders/op_index.py @@ -23,7 +23,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -47,3 +47,5 @@ def define_node( enn_graph.define_op( node.name, "GATHER", [input_id, indices_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_layer_norm.py b/backends/samsung/builders/op_layer_norm.py index 098bc92dc84..937168c36e9 100644 --- a/backends/samsung/builders/op_layer_norm.py +++ b/backends/samsung/builders/op_layer_norm.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input_node = node.args[0] input_id = self.define_tensor(input_node, enn_graph, vals_to_ids) @@ -51,3 +51,5 @@ def define_node( enn_graph.define_op( node.name, "LAYERNORM", all_input_tensors, [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_leaky_relu.py b/backends/samsung/builders/op_leaky_relu.py index c7ed37d12e5..28f4a7851d4 100644 --- a/backends/samsung/builders/op_leaky_relu.py +++ b/backends/samsung/builders/op_leaky_relu.py @@ -27,7 +27,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) all_input_tensors.append(input_id) @@ -56,3 +56,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "PRELU", all_input_tensors, [output_id]) + + return True diff --git a/backends/samsung/builders/op_linear.py b/backends/samsung/builders/op_linear.py index 720439de976..dffb6b0108c 100644 --- a/backends/samsung/builders/op_linear.py +++ b/backends/samsung/builders/op_linear.py @@ -27,7 +27,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: all_input_tensors = [] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -49,3 +49,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "FC", all_input_tensors, [output_id], params) + + return True diff --git a/backends/samsung/builders/op_log.py b/backends/samsung/builders/op_log.py index 97127dd94ba..a20de9bd95d 100644 --- a/backends/samsung/builders/op_log.py +++ b/backends/samsung/builders/op_log.py @@ -23,10 +23,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "LOG", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_log_softmax.py b/backends/samsung/builders/op_log_softmax.py index f2d87601cbb..f26a36e2af7 100644 --- a/backends/samsung/builders/op_log_softmax.py +++ b/backends/samsung/builders/op_log_softmax.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( meta_data = {"axis": axis} enn_graph.define_op(node.name, "LOGSOFTMAX", [input_id], [output_id], meta_data) + + return True diff --git a/backends/samsung/builders/op_max_pool2d.py b/backends/samsung/builders/op_max_pool2d.py index 57b716fcb34..bec8da5f683 100644 --- a/backends/samsung/builders/op_max_pool2d.py +++ b/backends/samsung/builders/op_max_pool2d.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -75,10 +75,6 @@ def define_node( params["dilation_w"] = dilation[1] self._update_params_qdtype(node, params) - if len(node.args) > 5: - ceil_mode = cast(bool, node.args[5]) - assert not ceil_mode, "Not support ceil_mode = True." - if not is_indices: output_id = self.define_tensor( node, @@ -94,3 +90,5 @@ def define_node( ) enn_graph.define_op(node.name, "MAXPOOL2D", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_maximum.py b/backends/samsung/builders/op_maximum.py index d3358d736f3..d81dfdd4a35 100644 --- a/backends/samsung/builders/op_maximum.py +++ b/backends/samsung/builders/op_maximum.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # inputs input_id_1 = self.define_tensor(node.args[0], enn_graph, vals_to_ids) input_id_2 = self.define_tensor(node.args[1], enn_graph, vals_to_ids) @@ -35,3 +35,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "MAXIMUM", [input_id_1, input_id_2], [output_id]) + + return True diff --git a/backends/samsung/builders/op_mean_dim.py b/backends/samsung/builders/op_mean_dim.py index 3d0377703a7..3396102c045 100644 --- a/backends/samsung/builders/op_mean_dim.py +++ b/backends/samsung/builders/op_mean_dim.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import cast, Dict, List import torch @@ -11,6 +12,7 @@ NodeVisitor, register_node_visitor, ) +from executorch.backends.samsung.builders.utils import get_tensor from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph from executorch.backends.transforms import get_shape @@ -27,7 +29,11 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: + output_tensor = get_tensor(self.exported_program, node) + if output_tensor.dtype == torch.float64: + logging.warning("float64 for mean has not supported yet.") + return False # input input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,6 +43,9 @@ def define_node( dims = cast(List[int], node.args[1]) reduce_axes = [] in_shape = get_shape(input) + if dims is None: + logging.warning("dims is None for this case.") + return False for dim in dims: reduce_axes.append(dim % len(in_shape)) @@ -47,3 +56,5 @@ def define_node( params = {"keep_dims": keep_dim, "axis": reduce_axes} self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "REDUCEMEAN", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_minimum.py b/backends/samsung/builders/op_minimum.py index a32b462d45f..4c612df34c7 100644 --- a/backends/samsung/builders/op_minimum.py +++ b/backends/samsung/builders/op_minimum.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # inputs input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "MIN", [input_id_1, input_id_2], [output_id]) + + return True diff --git a/backends/samsung/builders/op_mul.py b/backends/samsung/builders/op_mul.py index 6dd7c0dd9f0..e703ddff253 100644 --- a/backends/samsung/builders/op_mul.py +++ b/backends/samsung/builders/op_mul.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) @@ -41,3 +41,5 @@ def define_node( enn_graph.define_op( node.name, "ELTMUL", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_permute.py b/backends/samsung/builders/op_permute.py index 646eac4c06a..42286dddfef 100644 --- a/backends/samsung/builders/op_permute.py +++ b/backends/samsung/builders/op_permute.py @@ -25,13 +25,17 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) # permutation permute_order = cast(List[int], node.args[1]) + # to prevent negative values + permute_order = [x % len(permute_order) for x in permute_order] params = {"perm": permute_order} output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "TRANSPOSE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_pixel_shuffle.py b/backends/samsung/builders/op_pixel_shuffle.py index 28259299c81..db0aaaaef08 100644 --- a/backends/samsung/builders/op_pixel_shuffle.py +++ b/backends/samsung/builders/op_pixel_shuffle.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) scale_factor = cast(int, node.args[1]) @@ -36,3 +36,5 @@ def define_node( enn_graph.define_op( node.name, "DEPTH_TO_SPACE", [input_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_placeholder.py b/backends/samsung/builders/op_placeholder.py index b4b606f56ea..8c6a89a5eb5 100644 --- a/backends/samsung/builders/op_placeholder.py +++ b/backends/samsung/builders/op_placeholder.py @@ -31,7 +31,9 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: if is_param_node(self.exported_program, node): return self.define_tensor(node, enn_graph, vals_to_ids) + + return True diff --git a/backends/samsung/builders/op_pow.py b/backends/samsung/builders/op_pow.py index cd6ec7f81ef..d417685126e 100644 --- a/backends/samsung/builders/op_pow.py +++ b/backends/samsung/builders/op_pow.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import Dict import torch @@ -24,15 +25,17 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input1 = node.args[0] input2 = node.args[1] input_tensor_1 = get_tensor(self.exported_program, input1) input_tensor_2 = get_tensor(self.exported_program, input2) - assert ( - input_tensor_1.dtype == torch.float32 - and input_tensor_2.dtype == torch.float32 - ), "Requires the two inputs are all float type" + if ( + input_tensor_1.dtype != torch.float32 + or input_tensor_2.dtype != torch.float32 + ): + logging.warning("Requires the two inputs are all float type.") + return False input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) input_id_2 = self.define_tensor(input2, enn_graph, vals_to_ids) @@ -40,3 +43,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "POW", [input_id_1, input_id_2], [output_id]) + + return True diff --git a/backends/samsung/builders/op_quantize.py b/backends/samsung/builders/op_quantize.py index dcf30e291f9..771ab419888 100644 --- a/backends/samsung/builders/op_quantize.py +++ b/backends/samsung/builders/op_quantize.py @@ -24,32 +24,39 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # input input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) scales = node.args[1] - if isinstance(scales, torch.Tensor): - scales = scales.tolist() - elif not isinstance(scales, list): - scales = torch.tensor(scales).reshape([1]).tolist() zero_points = node.args[2] - if isinstance(zero_points, torch.Tensor): - zero_points = zero_points.tolist() - elif not isinstance(zero_points, list): - zero_points = torch.tensor(zero_points).reshape([1]).tolist() + if not isinstance(scales, torch.fx.Node) and not isinstance( + zero_points, torch.fx.Node + ): + if isinstance(scales, torch.Tensor): + scales = scales.tolist() + elif not isinstance(scales, list): + scales = torch.tensor(scales).reshape([1]).tolist() + if isinstance(zero_points, torch.Tensor): + zero_points = zero_points.tolist() + elif not isinstance(zero_points, list): + zero_points = torch.tensor(zero_points).reshape([1]).tolist() - output_id = self.define_tensor(node, enn_graph, vals_to_ids) + output_id = self.define_tensor(node, enn_graph, vals_to_ids) - params = {"scales": scales, "zero_points": zero_points} + params = {"scales": scales, "zero_points": zero_points} - if node.target in QuantConstants.QUANT_OPS_KEY_MAP: - enn_graph.define_op(node.name, "QUANTIZE", [input_id], [output_id], params) - else: - enn_graph.define_op( - node.name, "DEQUANTIZE", [input_id], [output_id], params - ) + if node.target in QuantConstants.QUANT_OPS_KEY_MAP: + enn_graph.define_op( + node.name, "QUANTIZE", [input_id], [output_id], params + ) + else: + enn_graph.define_op( + node.name, "DEQUANTIZE", [input_id], [output_id], params + ) + + return True @register_node_visitor diff --git a/backends/samsung/builders/op_relu.py b/backends/samsung/builders/op_relu.py index a4a2b6bc4f0..fb2668e46fc 100644 --- a/backends/samsung/builders/op_relu.py +++ b/backends/samsung/builders/op_relu.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -35,3 +35,5 @@ def define_node( self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "RELU", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_reshape.py b/backends/samsung/builders/op_reshape.py index 1f4e85ac059..bb413ed793d 100644 --- a/backends/samsung/builders/op_reshape.py +++ b/backends/samsung/builders/op_reshape.py @@ -7,6 +7,7 @@ NodeVisitor, register_node_visitor, ) +from executorch.backends.samsung.builders.utils import get_tensor from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph @@ -22,13 +23,18 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) - new_shape = node.args[1] + # node.args[1] may contain "sym_size" + tensor = get_tensor(self.exported_program, node) + shape = [1] if len(tensor.size()) == 0 else list(tensor.size()) + enn_graph.define_op( - node.name, "RESHAPE", [input_id], [output_id], {"new_shape": new_shape} + node.name, "RESHAPE", [input_id], [output_id], {"new_shape": shape} ) + + return True diff --git a/backends/samsung/builders/op_rms_norm.py b/backends/samsung/builders/op_rms_norm.py index 6a58d62a5ce..0ff01701d1b 100644 --- a/backends/samsung/builders/op_rms_norm.py +++ b/backends/samsung/builders/op_rms_norm.py @@ -24,7 +24,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # args of node : ['input', 'normalized_shape', 'weight', 'eps'] input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -50,3 +50,5 @@ def define_node( enn_graph.define_op( node.name, "RMSNORM", [input_id, gamma_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_rsqrt.py b/backends/samsung/builders/op_rsqrt.py index b3600d41ee2..55e9ddf54f9 100644 --- a/backends/samsung/builders/op_rsqrt.py +++ b/backends/samsung/builders/op_rsqrt.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "RSQRT", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_select.py b/backends/samsung/builders/op_select.py index 26f455b2548..3f3550d3cf8 100644 --- a/backends/samsung/builders/op_select.py +++ b/backends/samsung/builders/op_select.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -50,3 +50,5 @@ def define_node( } enn_graph.define_op(node.name, "STRIDEDSLICE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_sigmoid.py b/backends/samsung/builders/op_sigmoid.py index e87973f9a85..aef9d90ec52 100644 --- a/backends/samsung/builders/op_sigmoid.py +++ b/backends/samsung/builders/op_sigmoid.py @@ -23,10 +23,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "SIGMOID", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_sin.py b/backends/samsung/builders/op_sin.py index 5fd22e8275e..11dc6bccb46 100644 --- a/backends/samsung/builders/op_sin.py +++ b/backends/samsung/builders/op_sin.py @@ -23,9 +23,11 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input_id = self.define_tensor(node.args[0], enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "Sin", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_skip.py b/backends/samsung/builders/op_skip.py new file mode 100644 index 00000000000..3413e603983 --- /dev/null +++ b/backends/samsung/builders/op_skip.py @@ -0,0 +1,32 @@ +# Copyright (c) 2026 Samsung Electronics Co. LTD +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Dict + +import torch +from executorch.backends.samsung.builders.node_visitor import ( + NodeVisitor, + register_node_visitor, +) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph + + +@register_node_visitor +class OpSkipVisitor(NodeVisitor): + target = ["sym_size.int", "add", "floordiv"] + """ + do nothing + """ + + def __init__(self, *args) -> None: + super().__init__(*args) + + def define_node( + self, + node: torch.fx.Node, + enn_graph: EnnGraph, + vals_to_ids: Dict[torch.Tensor, int], + ) -> bool: + return True diff --git a/backends/samsung/builders/op_slice_copy.py b/backends/samsung/builders/op_slice_copy.py index e85b6bf60c3..4b837db1b12 100644 --- a/backends/samsung/builders/op_slice_copy.py +++ b/backends/samsung/builders/op_slice_copy.py @@ -27,7 +27,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -61,3 +61,5 @@ def define_node( params = {"begin": begin, "end": end, "strides": strides} enn_graph.define_op(node.name, "STRIDEDSLICE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_softmax.py b/backends/samsung/builders/op_softmax.py index 7f569cea6fc..e96b5638db3 100644 --- a/backends/samsung/builders/op_softmax.py +++ b/backends/samsung/builders/op_softmax.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ): + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( params = {"axis": axis} self._update_params_qdtype(node, params) enn_graph.define_op(node.name, "SOFTMAX", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_split_with_sizes_copy.py b/backends/samsung/builders/op_split_with_sizes_copy.py index b67b5331627..48612ba9a6d 100644 --- a/backends/samsung/builders/op_split_with_sizes_copy.py +++ b/backends/samsung/builders/op_split_with_sizes_copy.py @@ -23,7 +23,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -39,6 +39,10 @@ def define_node( ) all_output_tensors.append(output_id) + for user in node.users.keys(): + if user.target.__name__ == "getitem" and len(user.args) > 1: + vals_to_ids[user] = all_output_tensors[user.args[1]] + axis = node.args[2] if len(node.args) > 2 else 0 params = {} @@ -46,3 +50,5 @@ def define_node( params["point"] = node.args[1] enn_graph.define_op(node.name, "SPLIT", [input_id], all_output_tensors, params) + + return True diff --git a/backends/samsung/builders/op_sqrt.py b/backends/samsung/builders/op_sqrt.py index 3560542a0bc..77793e4589d 100644 --- a/backends/samsung/builders/op_sqrt.py +++ b/backends/samsung/builders/op_sqrt.py @@ -26,10 +26,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "SQRT", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_squeeze.py b/backends/samsung/builders/op_squeeze.py index 82fa17fbc95..ff5286fee28 100644 --- a/backends/samsung/builders/op_squeeze.py +++ b/backends/samsung/builders/op_squeeze.py @@ -26,7 +26,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -35,3 +35,5 @@ def define_node( params = {"new_shape": [*node.meta["val"].shape]} enn_graph.define_op(node.name, "RESHAPE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_sub.py b/backends/samsung/builders/op_sub.py index 7dc97bfa7ca..c12f4c85f4d 100644 --- a/backends/samsung/builders/op_sub.py +++ b/backends/samsung/builders/op_sub.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import Dict import torch @@ -26,12 +27,16 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: # inputs input1 = node.args[0] input_id_1 = self.define_tensor(input1, enn_graph, vals_to_ids) input2 = node.args[1] input_id_2 = self.define_tensor(input2, enn_graph, vals_to_ids) + alpha = node.kwargs.get("alpha", 1.0) + if alpha != 1.0: + logging.warning("Currently, only alpha 1 for sub is supported.") + return False # output output_id = self.define_tensor(node, enn_graph, vals_to_ids) @@ -41,3 +46,5 @@ def define_node( enn_graph.define_op( node.name, "SUB", [input_id_1, input_id_2], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_sum_int_list.py b/backends/samsung/builders/op_sum_int_list.py index 7743e6632dd..a8c5367371c 100644 --- a/backends/samsung/builders/op_sum_int_list.py +++ b/backends/samsung/builders/op_sum_int_list.py @@ -24,7 +24,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -37,3 +37,5 @@ def define_node( output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "REDUCESUM", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_tanh.py b/backends/samsung/builders/op_tanh.py index 5b002890075..106256c791b 100644 --- a/backends/samsung/builders/op_tanh.py +++ b/backends/samsung/builders/op_tanh.py @@ -23,10 +23,12 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op(node.name, "TANH", [input_id], [output_id]) + + return True diff --git a/backends/samsung/builders/op_to_copy.py b/backends/samsung/builders/op_to_copy.py index c770602bb5f..143ab007cba 100644 --- a/backends/samsung/builders/op_to_copy.py +++ b/backends/samsung/builders/op_to_copy.py @@ -28,7 +28,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: memory_format_target = node.kwargs.get("memory_format", torch.contiguous_format) to_contiguous = bool(memory_format_target == torch.contiguous_format) assert to_contiguous, "Don't support other param in _to_copy" @@ -42,3 +42,5 @@ def define_node( params["out_dtype"] = get_map_dtype(out_tensor.dtype) enn_graph.define_op(node.name, "CAST", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_topk.py b/backends/samsung/builders/op_topk.py index e4cda0ef148..6a4de9ccc91 100644 --- a/backends/samsung/builders/op_topk.py +++ b/backends/samsung/builders/op_topk.py @@ -24,7 +24,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -72,3 +72,5 @@ def define_node( raise AssertionError("Not supported sorted = False.") enn_graph.define_op(node.name, "TopK", [input_id], all_output_tensors, params) + + return True diff --git a/backends/samsung/builders/op_unsqueeze.py b/backends/samsung/builders/op_unsqueeze.py index 61fa06e6310..18f7c2d33b2 100644 --- a/backends/samsung/builders/op_unsqueeze.py +++ b/backends/samsung/builders/op_unsqueeze.py @@ -25,7 +25,7 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) @@ -33,3 +33,5 @@ def define_node( params = {"new_shape": [*node.meta["val"].shape]} enn_graph.define_op(node.name, "RESHAPE", [input_id], [output_id], params) + + return True diff --git a/backends/samsung/builders/op_upsample_bilinear2d.py b/backends/samsung/builders/op_upsample_bilinear2d.py index d4b040460e3..7374e687101 100644 --- a/backends/samsung/builders/op_upsample_bilinear2d.py +++ b/backends/samsung/builders/op_upsample_bilinear2d.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import cast, Dict, List import torch @@ -27,11 +28,14 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) in_shape = get_shape(input) output_size = cast(List[int], node.args[1]) + if output_size is None: + logging.warning("output is None for this case.") + return False scale_factor = [ output_size[0] * 1.0 / in_shape[-2], output_size[1] * 1.0 / in_shape[-1], @@ -44,10 +48,12 @@ def define_node( params = { "align_corners": align_corners, "upsampling_factor": scale_factor, - "half_pixel_centers": True, + "half_pixel_centers": not align_corners, } self._update_params_qdtype(node, params) output_id = self.define_tensor(node, enn_graph, vals_to_ids) enn_graph.define_op( node.name, "RESIZE_BILINEAR", [input_id], [output_id], params ) + + return True diff --git a/backends/samsung/builders/op_upsample_nearest2d.py b/backends/samsung/builders/op_upsample_nearest2d.py index 9859cd8f07e..6af5402d56c 100644 --- a/backends/samsung/builders/op_upsample_nearest2d.py +++ b/backends/samsung/builders/op_upsample_nearest2d.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +import logging from typing import cast, Dict, List import torch @@ -27,11 +28,14 @@ def define_node( node: torch.fx.Node, enn_graph: EnnGraph, vals_to_ids: Dict[torch.Tensor, int], - ) -> None: + ) -> bool: input = node.args[0] input_id = self.define_tensor(input, enn_graph, vals_to_ids) in_shape = get_shape(input) output_size = cast(List[int], node.args[1]) + if output_size is None: + logging.warning("output is None for this case.") + return False scale_factor = [ output_size[0] * 1.0 / in_shape[-2], output_size[1] * 1.0 / in_shape[-1], @@ -50,3 +54,5 @@ def define_node( enn_graph.define_op( node.name, "RESIZE_NEAREST_NEIGHBOR", [input_id], [output_id], params ) + + return True diff --git a/backends/samsung/partition/enn_partitioner.py b/backends/samsung/partition/enn_partitioner.py index 91f496e7a5c..3e450cdf1bd 100644 --- a/backends/samsung/partition/enn_partitioner.py +++ b/backends/samsung/partition/enn_partitioner.py @@ -15,6 +15,7 @@ from executorch.backends.samsung.serialization.compile_options import ( ENN_COMPILE_OPTION_TITLE, ) +from executorch.backends.samsung.serialization.enn_graph_schema import EnnGraph from executorch.backends.samsung.utils.utils import get_compile_spec from executorch.exir.backend.backend_details import CompileSpec from executorch.exir.backend.canonical_partitioners.pattern_op_partitioner import ( @@ -70,9 +71,16 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool: ]: return False - if node.target in SUPPORTED_OPS or node.target.__name__ in self.node_visitors: + if node.target in SUPPORTED_OPS: return True + if node.target.__name__ in self.node_visitors: + enn_graph = EnnGraph() + vals_to_ids: Dict[torch.fx.Node, int] = {} + return self.node_visitors[node.target.__name__].define_node( + node, enn_graph, vals_to_ids + ) + supported = self.enn_wrapper.IsNodeSupportedByBackend() return supported @@ -91,10 +99,19 @@ def generate_partitions( self, edge_program: torch.export.ExportedProgram ) -> List[Any]: self.op_support_checker = EnnOperatorSupport(edge_program, self.compile_specs) - return generate_partitions_from_list_of_nodes( + partition_list = generate_partitions_from_list_of_nodes( edge_program.graph_module, op_support=self.op_support_checker, ) + if len(partition_list) == 1 and partition_list[0].size() == 1: + first_node = list(partition_list[0].nodes.keys())[0] + # If there is only one partition graph containing a single "aten.clone.default" that is a useless operation, + # the RemoveUselessOpPass will remove this operation and cause a graph error. + # Therefore, we delete this node to prevent this graph error. + # For example, in the test_index_put_in_place_dtype case partition_list is [{aten_clone_default: 2}] + if first_node.target == exir_ops.edge.aten.clone.default: + del partition_list[0] + return partition_list def tag_nodes(self, partitions: List[Partition]) -> None: partition_tags: Dict[str, DelegationSpec] = {} @@ -127,8 +144,6 @@ def ops_to_not_decompose( torch.ops.aten.max_pool2d.default, torch.ops.aten.linear.default, torch.ops.aten._safe_softmax.default, - torch.ops.aten.upsample_bilinear2d.vec, - torch.ops.aten.upsample_nearest2d.vec, torch.ops.aten.prelu.default, torch.ops.aten.layer_norm.default, torch.ops.aten.pixel_shuffle.default, diff --git a/backends/samsung/test/tester/samsung_tester.py b/backends/samsung/test/tester/samsung_tester.py index 258aef191d0..001b7e3bb0d 100644 --- a/backends/samsung/test/tester/samsung_tester.py +++ b/backends/samsung/test/tester/samsung_tester.py @@ -11,8 +11,12 @@ import torch from executorch.backends.samsung.partition.enn_partitioner import EnnPartitioner from executorch.backends.samsung.quantizer.quantizer import EnnQuantizer, Precision +from executorch.backends.samsung.serialization.compile_options import ( + gen_samsung_backend_compile_spec, +) from executorch.backends.samsung.test.utils import RuntimeExecutor from executorch.backends.samsung.test.utils.quant_checkers import get_checker +from executorch.backends.samsung.test.utils.utils import TestConfig from executorch.backends.samsung.utils.export_utils import get_edge_compile_config from executorch.backends.test.harness import Tester as TesterBase from executorch.backends.test.harness.stages import StageType @@ -112,6 +116,7 @@ def run( transform_passes=self.transform_passes, partitioner=self.partitioners, compile_config=self.edge_compile_config, + generate_etrecord=generate_etrecord, ) @@ -145,6 +150,8 @@ def __init__( self.original_module = module self.exported_module = module self.example_inputs = example_inputs + if compile_specs is None: + compile_specs = [gen_samsung_backend_compile_spec(TestConfig.chipset)] self.compile_specs = compile_specs def quantize( @@ -167,9 +174,12 @@ def quantize( def to_edge_transform_and_lower( self, edge_compile_config: Optional[EdgeCompileConfig] = None, + generate_etrecord: bool = False, ): to_edge_transform_and_lower_stage = ToEdgeTransformAndLower( self.compile_specs, edge_compile_config ) - return super().to_edge_transform_and_lower(to_edge_transform_and_lower_stage) + return super().to_edge_transform_and_lower( + to_edge_transform_and_lower_stage, generate_etrecord + ) diff --git a/backends/samsung/test/utils/runtime_executor.py b/backends/samsung/test/utils/runtime_executor.py index 4d642657d23..1117ffa3b84 100644 --- a/backends/samsung/test/utils/runtime_executor.py +++ b/backends/samsung/test/utils/runtime_executor.py @@ -157,7 +157,7 @@ def run_on_device(self) -> Tuple[torch.Tensor]: output_tensor = ( torch.from_numpy(output_array) .view(dtype=model_outputs[idx].dtype) - .view(*model_outputs[idx].shape) + .reshape(model_outputs[idx].shape) ) result.append(output_tensor) diff --git a/backends/samsung/utils/export_utils.py b/backends/samsung/utils/export_utils.py index 22f1833bd18..86512b75dec 100644 --- a/backends/samsung/utils/export_utils.py +++ b/backends/samsung/utils/export_utils.py @@ -37,6 +37,11 @@ def get_edge_compile_config(): exir_ops.edge.aten.layer_norm.default, exir_ops.edge.aten.matmul.default, exir_ops.edge.aten.hardsigmoid.default, + exir_ops.edge.aten.round.decimals, + exir_ops.edge.aten.median.dim, + exir_ops.edge.aten.median.default, + exir_ops.edge.aten.adaptive_max_pool2d.default, + exir_ops.edge.aten.adaptive_max_pool3d.default, ], ) diff --git a/backends/test/suite/flow.py b/backends/test/suite/flow.py index 547c79326e6..7331adc4a49 100644 --- a/backends/test/suite/flow.py +++ b/backends/test/suite/flow.py @@ -227,4 +227,17 @@ def all_flows() -> dict[str, TestFlow]: except Exception as e: logger.info(f"Skipping MLX flow registration: {e}") + try: + from executorch.backends.test.suite.flows.samsung import ( + SAMSUNG_A8W8_TEST_FLOW, + SAMSUNG_TEST_FLOW, + ) + + flows += [ + SAMSUNG_TEST_FLOW, + SAMSUNG_A8W8_TEST_FLOW, + ] + except Exception as e: + logger.info(f"Skipping SAMSUNG flow registration: {e}") + return {f.name: f for f in flows if f is not None} diff --git a/backends/test/suite/flows/samsung.py b/backends/test/suite/flows/samsung.py new file mode 100644 index 00000000000..577fde77083 --- /dev/null +++ b/backends/test/suite/flows/samsung.py @@ -0,0 +1,43 @@ +import logging + +from executorch.backends.samsung.quantizer.quantizer import EnnQuantizer, Precision +from executorch.backends.samsung.test.tester.samsung_tester import ( + Quantize, + SamsungTester, +) +from executorch.backends.test.suite.flow import TestFlow + +logger = logging.getLogger(__name__) +logger.setLevel(logging.INFO) + + +def _create_samsung_flow( + name: str, + quantize: bool = False, + quant_dtype: Precision | None = None, + is_per_channel: bool = True, + is_qat: bool = False, +) -> TestFlow: + if quantize and quant_dtype is None: + raise RuntimeError("Quant dtype must be provided when quantize is true.") + + def create_quantize_stage() -> Quantize: + quantizer = EnnQuantizer() + quantizer.setup_quant_params(quant_dtype, is_per_channel, is_qat) + return Quantize(quantizer=quantizer) + + return TestFlow( + name, + backend="samsung", + tester_factory=SamsungTester, + quantize=quantize, + quantize_stage_factory=create_quantize_stage if quantize else None, + supports_serialize=False, + ) + + +SAMSUNG_TEST_FLOW = _create_samsung_flow("samsung") + +SAMSUNG_A8W8_TEST_FLOW = _create_samsung_flow( + "samsung_a8w8", quantize=True, quant_dtype=Precision.A8W8 +) diff --git a/backends/test/suite/models/test_torchaudio.py b/backends/test/suite/models/test_torchaudio.py index 2287b226c37..b6879162297 100644 --- a/backends/test/suite/models/test_torchaudio.py +++ b/backends/test/suite/models/test_torchaudio.py @@ -62,7 +62,7 @@ def test_conformer(test_runner, dtype: torch.dtype, use_dynamic_shapes: bool): encoder_padding_mask, ) - test_runner.lower_and_run_model(model, inputs) + test_runner.lower_and_run_model(model, inputs, generate_random_test_inputs=False) @pytest.mark.parametrize("dtype", [torch.float32], ids=dtype_to_str) diff --git a/examples/samsung/executor_runner/enn_executor_runner.cpp b/examples/samsung/executor_runner/enn_executor_runner.cpp index bdbff74088c..14dfbf58ee2 100644 --- a/examples/samsung/executor_runner/enn_executor_runner.cpp +++ b/examples/samsung/executor_runner/enn_executor_runner.cpp @@ -27,6 +27,7 @@ #include #include +#include #include #include #include @@ -96,10 +97,26 @@ class DataReader { return data_set_[index].size(); } + // Some ops rewrite their input buffer in place during execute(), so a + // repeated execution would otherwise run on whatever the previous + // execution left behind. Snapshot the pristine bytes once inputs are set, + // then restore() before every execution. + void snapshot() { + pristine_ = data_set_; + } + + void restore() { + for (size_t i = 0; i < data_set_.size(); ++i) { + std::copy( + pristine_[i].cbegin(), pristine_[i].cend(), data_set_[i].begin()); + } + } + ~DataReader() = default; private: std::vector data_set_; + std::vector pristine_; int32_t index_ = 0; }; @@ -296,46 +313,37 @@ int main(int argc, char** argv) { ET_CHECK_MSG(ret == Error::Ok, "Failed to set input tensor: %d", ret); } EXYNOS_ATRACE_END(); + input_data_reader.snapshot(); // Warm up ET_LOG(Info, "Perform %d inference for warming up", FLAGS_warm_up); Error status; for (int i = 0; i < FLAGS_warm_up; ++i) { + input_data_reader.restore(); status = method->execute(); } - // Run the model. - ET_LOG(Info, "Start 1st inference."); - auto before_exec = std::chrono::high_resolution_clock::now(); - status = method->execute(); - auto after_exec = std::chrono::high_resolution_clock::now(); - double interval_1st_infs = - std::chrono::duration_cast( - after_exec - before_exec) - .count() / - 1000.0; - ET_LOG(Info, "Start inference."); - before_exec = std::chrono::high_resolution_clock::now(); + std::chrono::microseconds infs_duration{0}; for (int i = 0; i < FLAGS_num_executions; ++i) { + // Restored outside the timed section so it measures execute() alone, + // not the cost of undoing the previous iteration's input mutation. + input_data_reader.restore(); + auto before_exec = std::chrono::high_resolution_clock::now(); status = method->execute(); + auto after_exec = std::chrono::high_resolution_clock::now(); + infs_duration += std::chrono::duration_cast( + after_exec - before_exec); } - after_exec = std::chrono::high_resolution_clock::now(); - double interval_infs = std::chrono::duration_cast( - after_exec - before_exec) - .count() / - 1000.0; + double interval_infs = infs_duration.count() / 1000.0; if (FLAGS_dump_statistics) { auto output_file_name = "statistics.txt"; std::ofstream fout(output_file_name); fout << "init: " + std::to_string(interval_init) << "\nload: " + std::to_string(interval_load) - << "\n1st: " + std::to_string(interval_1st_infs) << "\navg: " + - std::to_string( - (interval_infs + interval_1st_infs) / - ((float)FLAGS_num_executions + 1.f)) + std::to_string(interval_infs / (float)FLAGS_num_executions) << std::endl; fout.close(); } From d0da5f8c2e32206047dcdacef1a172d95bbd8387 Mon Sep 17 00:00:00 2001 From: Andrew Grebenisan <33402477+DrJessop@users.noreply.github.com> Date: Wed, 9 Sep 2026 23:19:58 -0700 Subject: [PATCH 132/190] Add support to PropagateSlice for custom unary/binary ops Summary: As titled. Differential Revision: D119399660 Pull Request resolved: https://github.com/pytorch/executorch/pull/22657 --- backends/cadence/aot/reorder_ops.py | 15 +++-- .../aot/tests/test_reorder_ops_passes.py | 67 +++++++++++++++++++ 2 files changed, 78 insertions(+), 4 deletions(-) diff --git a/backends/cadence/aot/reorder_ops.py b/backends/cadence/aot/reorder_ops.py index 3d7ae7ac3d1..3df555e0fba 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/backends/cadence/aot/reorder_ops.py @@ -1179,17 +1179,23 @@ class PropagateSlice(RemoveOrReplacePassInterface): Handles any slice dim and any step size. """ - def __init__(self) -> None: + def __init__( + self, + additional_unary_targets: Optional[list[EdgeOpOverload]] = None, + additional_binary_targets: Optional[list[EdgeOpOverload]] = None, + ) -> None: super().__init__() - elementwise_targets = [ + unary_targets = [ exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, exir_ops.edge.cadence.quantize_per_tensor.default, exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, exir_ops.edge.cadence.dequantize_per_tensor.default, + *(additional_unary_targets or []), ] binary_targets = [ exir_ops.edge.aten.add.Tensor, exir_ops.edge.aten.mul.Tensor, + *(additional_binary_targets or []), ] self._dispatch: dict[ EdgeOpOverload, @@ -1198,7 +1204,7 @@ def __init__(self) -> None: Callable[[torch.fx.Node, torch.fx.Node], bool], ], ] = {} - for t in elementwise_targets: + for t in unary_targets: self._dispatch[t] = ( self._should_swap_elementwise, self._swap_elementwise_slice, @@ -1224,7 +1230,8 @@ def _should_swap_elementwise( def _swap_elementwise_slice( self, op_node: torch.fx.Node, slice_node: torch.fx.Node ) -> bool: - op_input = get_arg(op_node, "input", torch.fx.Node) + op_input = op_node.args[0] + assert isinstance(op_input, torch.fx.Node) graph = slice_node.graph slice_dim = get_arg(slice_node, "dim", int) diff --git a/backends/cadence/aot/tests/test_reorder_ops_passes.py b/backends/cadence/aot/tests/test_reorder_ops_passes.py index fee34dabaa5..3c4443aafa1 100644 --- a/backends/cadence/aot/tests/test_reorder_ops_passes.py +++ b/backends/cadence/aot/tests/test_reorder_ops_passes.py @@ -1358,6 +1358,73 @@ def test_unsupported_parent_not_swapped(self) -> None: self.assertFalse(result.modified) + def test_swap_additional_unary_target(self) -> None: + x_data = torch.randn(4, 60, 1, 1) + builder = GraphBuilder() + x = builder.placeholder("x", x_data) + relu = builder.call_operator(exir_ops.edge.aten.relu.default, args=(x,)) + sliced = builder.call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + args=(relu, 0, 0, 4, 2), + ) + builder.output([sliced]) + gm = builder.get_graph_module() + + result = transform_and_check_numerics( + gm, + (x_data,), + PropagateSlice(additional_unary_targets=[exir_ops.edge.aten.relu.default]), + ) + + self.assertTrue(result.modified) + slice_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + ) + self.assertEqual(len(slice_nodes), 1) + relu_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.relu.default + ) + self.assertEqual(len(relu_nodes), 1) + self.assertIs(relu_nodes[0].args[0], slice_nodes[0]) + self.assertEqual(list(relu_nodes[0].meta["val"].shape), [2, 60, 1, 1]) + + def test_swap_additional_binary_target(self) -> None: + lhs_data = torch.randn(1, 60, 1, 1) + rhs_data = torch.randn(4, 60, 1, 1) + builder = GraphBuilder() + lhs = builder.placeholder("lhs", lhs_data) + rhs = builder.placeholder("rhs", rhs_data) + sub = builder.call_operator( + exir_ops.edge.aten.sub.Tensor, + args=(lhs, rhs), + ) + sliced = builder.call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + args=(sub, 0, 0, 4, 2), + ) + builder.output([sliced]) + gm = builder.get_graph_module() + + result = transform_and_check_numerics( + gm, + (lhs_data, rhs_data), + PropagateSlice(additional_binary_targets=[exir_ops.edge.aten.sub.Tensor]), + ) + + self.assertTrue(result.modified) + slice_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + ) + self.assertEqual(len(slice_nodes), 1) + self.assertEqual(slice_nodes[0].args[0].name, "rhs") + sub_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.sub.Tensor + ) + self.assertEqual(len(sub_nodes), 1) + self.assertIs(sub_nodes[0].args[0], lhs.node) + self.assertIs(sub_nodes[0].args[1], slice_nodes[0]) + self.assertEqual(list(sub_nodes[0].meta["val"].shape), [2, 60, 1, 1]) + def test_swap_broadcast_mul_slice_on_broadcast_dim(self) -> None: """[1,60,1,1] * [4,1,1,1] → [4,60,1,1] → slice(dim=0, step=2) Only the [4,1,1,1] input should be sliced.""" From d3a225e6816d8f3ef919f3882bf7244702a56696 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Thu, 10 Sep 2026 00:23:59 -0700 Subject: [PATCH 133/190] [Reland] Move pull.yml to linux_job_v3 (#22632) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit **Draft.** Reverts #22626, restoring `pull.yml` to its state as merged in a519efa64e, and folds in the teardown fix that #22626's failures also needed. **Why it was reverted.** Ten migrated Linux jobs failed writing to the read-only `/mnt/hf_cache`. The original PR passed only because it carried `ci-refresh-hf-cache`, which sets `HF_CACHE_REFRESH=1` and repoints `HF_HOME`/`HF_DATASETS_CACHE` at writable job-local paths, so the ordinary non-refresh path was never exercised. This PR is deliberately **unlabelled**. **Two fixes, in two places.** pytorch/test-infra#8756 redirects the writes that a read-only cache can't take — dataset `.lock` files, `.cache/meta_checkpoints` (which `hf_download.py` derives from `HF_HOME`), and `stored_tokens` — while keeping models shared via `HF_HUB_CACHE`. That's pytorch's `_linux-test.yml` shape, which this workflow's port had only applied on the refresh path. The second commit here fixes `test_lora.sh` and `test_lora_multimethod.sh`, which resolved model paths through `snapshot_download` and then `rm -rf`'d them on the way out. Those paths are *inside* the hub cache, so on OSDC the `rm` failed and, under `set -e`, failed the job after the tests had already printed `Multimethod tests passed!`. Off OSDC it succeeded and discarded a cache entry the next job re-downloaded. `refs/main` writes on the read-only mount are logged by `huggingface_hub` as `Ignored error while writing commit hash` and are non-fatal, so seeded models need no further work. The third commit points the 45 call sites at #8756 for validation and must come off before landing. --- Third of six splitting up #22107. Stacked on #22246. 45 call sites in one file, all mechanical: the v2 -> v3 rename, EC2 runner labels swapped for their OSDC equivalents, `use-custom-docker-registry` dropped since v3 ignores it, and the ECR image spelled out against the hash `_docker-image.yml` resolves. | EC2 | OSDC | |---|---| | `linux.2xlarge`, `linux.2xlarge.memory` | `mt-l-x86iavx512-8-64` | | `linux.4xlarge.memory` | `mt-l-x86iavx512-16-128` | | `linux.24xlarge` | `mt-l-x86iavx512-94-192` | | `linux.24xlarge.memory` | `mt-l-x86iavx512-94-768` | | `linux.arm64.2xlarge` | `mt-l-arm64g4-16-62` | | `linux.g5.4xlarge.nvidia.gpu` | `mt-l-x86aavx2-29-113-a10g` | Two jobs needed more than the rename: - `test-qnn-delegate-linux` goes to `mt-l-x86iamx-8-64` instead. The vendor letter names the ISA, not the silicon, so `x86iavx512` is `r7a`, i.e. AMD EPYC. The QNN backend's `disable_mkldnn_on_amd()` only runs on an AMD host, where it sets `torch.backends.mkldnn.enabled` outside a `flags()` context and raises once the tests have frozen the flags, failing all 1040. `mt-l-x86iamx-8-64` is `r7i` at the same 8 vCPU / 64Gi, so the host stays Intel as it was on `c5`. Worth fixing in the backend separately for anyone who does run it on AMD. - The pytest-xdist worker cap, since `auto` and `logical` size themselves from the machine's cores rather than the pod's limit and get OOM-killed. Authored with Claude Code. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- .ci/scripts/test_lora.sh | 13 +- .ci/scripts/test_lora_multimethod.sh | 13 +- .github/workflows/pull.yml | 361 +++++++++++++++------------ 3 files changed, 222 insertions(+), 165 deletions(-) diff --git a/.ci/scripts/test_lora.sh b/.ci/scripts/test_lora.sh index 102347a08fd..91c78ae5f3e 100644 --- a/.ci/scripts/test_lora.sh +++ b/.ci/scripts/test_lora.sh @@ -28,11 +28,14 @@ cmake_build_llama_runner() { } cleanup_files() { - echo "Deleting downloaded and generated files" - rm -rf "${HF_QWEN_PATH}/" - rm -rf "${HF_ADAPTER_PATH}/" - rm -rf *.pte *.ptd - rm result*.txt + # Only what this script generated. HF_QWEN_PATH and HF_ADAPTER_PATH point + # inside the huggingface_hub cache: on OSDC that is a shared read-only mount, + # so removing them failed the job after the tests had already passed, and + # anywhere else it discards a cache entry the next job wants. A teardown also + # must not fail a run whose tests passed. + echo "Deleting generated files" + rm -rf ./*.pte ./*.ptd || true + rm -f result*.txt || true } matches_base_response_prefix() { diff --git a/.ci/scripts/test_lora_multimethod.sh b/.ci/scripts/test_lora_multimethod.sh index f0b30bd4be1..afcde517cd6 100755 --- a/.ci/scripts/test_lora_multimethod.sh +++ b/.ci/scripts/test_lora_multimethod.sh @@ -28,11 +28,14 @@ cmake_build_llama_runner() { } cleanup_files() { - echo "Deleting downloaded and generated files" - rm -rf "${HF_QWEN_PATH}/" - rm -rf "${HF_ADAPTER_PATH}/" - rm -rf *.pte - rm -f result*.txt + # Only what this script generated. HF_QWEN_PATH and HF_ADAPTER_PATH point + # inside the huggingface_hub cache: on OSDC that is a shared read-only mount, + # so removing them failed the job after the tests had already passed, and + # anywhere else it discards a cache entry the next job wants. A teardown also + # must not fail a run whose tests passed. + echo "Deleting generated files" + rm -rf ./*.pte || true + rm -f result*.txt || true } matches_base_response_prefix() { diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 2522cca3113..d17be95b3aa 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -13,6 +13,10 @@ concurrency: cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + # Emits the list of changed files for the current PR or push commit. # On PR: PR diff. On push: diff against `github.event.before`. # On events without a diff base (workflow_dispatch, tag creation, @@ -33,8 +37,9 @@ jobs: uses: ./.github/workflows/_ci-run-decision.yml test-qnn-wheel-packages-linux: + needs: docker-image name: test-qnn-wheel-packages-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -43,8 +48,8 @@ jobs: matrix: python-version: [ "3.10", "3.11", "3.12", "3.13" ] with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -77,7 +82,7 @@ jobs: contents: read test-minimal-wheel-linux: - needs: changed-files + needs: [docker-image, changed-files] if: | github.event_name != 'pull_request' || contains(needs.changed-files.outputs.changed-files, '.ci/scripts/test_minimal_wheel.sh') || @@ -90,13 +95,13 @@ jobs: contains(needs.changed-files.outputs.changed-files, 'setup.py') || contains(needs.changed-files.outputs.changed-files, 'tools/cmake/') name: test-minimal-wheel-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -107,16 +112,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_minimal_wheel.sh test-setup-linux-gcc: + needs: docker-image name: test-setup-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -132,8 +138,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "add" "${BUILD_TOOL}" "portable" test-models-linux-basic: + needs: docker-image name: test-models-linux-basic - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -142,23 +149,23 @@ jobs: model: [mv3, vit] backend: [portable, xnnpack-quantization-delegation] build-tool: [cmake, buck2] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 # TODO: Need to figure out why buck2 doesnt work on Graviton instances. - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 build-tool: buck2 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -176,8 +183,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-models-linux: + needs: docker-image name: test-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -185,33 +193,33 @@ jobs: matrix: model: [linear, add, add_mul, ic3, mv2, resnet18, resnet50, mobilebert, emformer_transcribe] backend: [portable, xnnpack-quantization-delegation] - runner: [linux.2xlarge] + runner: [mt-l-x86iavx512-8-64] include: - model: ic4 backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: ic4 backend: xnnpack-quantization-delegation - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: emformer_join backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: emformer_join backend: xnnpack-quantization-delegation - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: phi_4_mini backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: llama3_2_vision_encoder backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 - model: w2l backend: portable - runner: linux.4xlarge.memory + runner: mt-l-x86iavx512-16-128 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -229,16 +237,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-parakeet-xnnpack-linux: + needs: docker-image name: test-parakeet-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.4xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -262,16 +271,17 @@ jobs: echo "::endgroup::" test-voxtral-realtime-xnnpack-linux: + needs: docker-image name: test-voxtral-realtime-xnnpack-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.4xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-16-128 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -298,9 +308,10 @@ jobs: echo "::endgroup::" test-llama-runner-linux: + needs: docker-image # Test Both linux x86 and linux aarch64 name: test-llama-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -308,25 +319,25 @@ jobs: matrix: dtype: [fp32] mode: [xnnpack+custom+qe,xnnpack+custom+quantize_kv,xnnpack+quantize_kv] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] include: - dtype: bf16 mode: custom - runner: linux.2xlarge + runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-clang12 # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -349,16 +360,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -dtype "${DTYPE}" -mode "${MODE}" -upload "${ARTIFACTS_DIR_NAME}" test-llama-runner-linux-android: + needs: docker-image name: test-llama-runner-linux-android - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -374,16 +386,17 @@ jobs: bash .ci/scripts/build_llama_android.sh "${BUILD_TOOL}" test-custom-ops-linux: + needs: docker-image name: test-custom-ops-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -398,16 +411,17 @@ jobs: PYTHON_EXECUTABLE=python bash examples/portable/custom_ops/test_custom_ops.sh "${BUILD_TOOL}" test-selective-build-linux: + needs: docker-image name: test-selective-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -422,9 +436,10 @@ jobs: PYTHON_EXECUTABLE=python bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" test-multimodal-linux: + needs: docker-image if: ${{ !github.event.pull_request.head.repo.fork }} name: test-multimodal-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -435,8 +450,8 @@ jobs: model: ["gemma3-4b"] # llava gives segfault so not covering. with: secrets-env: EXECUTORCH_HF_TOKEN - runner: linux.24xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-768 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -461,16 +476,17 @@ jobs: echo "::endgroup::" test-moshi-linux: + needs: docker-image name: test-moshi-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -493,16 +509,17 @@ jobs: python -m unittest examples.models.moshi.mimi.test_mimi test-quantized-aot-lib-linux: + needs: docker-image name: test-quantized-aot-lib-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -516,16 +533,17 @@ jobs: PYTHON_EXECUTABLE=python bash examples/xnnpack/quantization/test_quantize.sh "${BUILD_TOOL}" mv2 test-binary-size-linux-gcc: + needs: docker-image name: test-binary-size-linux-gcc - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc9-nopytorch + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc9-nopytorch-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -559,16 +577,17 @@ jobs: fi test-binary-size-linux: + needs: docker-image name: test-binary-size-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -603,8 +622,9 @@ jobs: fi test-arm-cortex-m-size-test: + needs: docker-image name: test-arm-cortex-m-size-test - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -613,8 +633,8 @@ jobs: os: [bare_metal, zephyr-preset] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -696,14 +716,15 @@ jobs: fi test-mcu-cortex-m-backend: + needs: docker-image name: test-mcu-cortex-m-backend - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -763,16 +784,17 @@ jobs: docker-image: ci-image:executorch-ubuntu-22.04-clang12 test-qnn-buck-build-linux: + needs: docker-image name: test-qnn-buck-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -797,8 +819,9 @@ jobs: buck2 build //backends/qualcomm/... test-arm-backend-no-driver: + needs: docker-image name: test-arm-backend-no-driver - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -811,8 +834,8 @@ jobs: - test_arm_backend: test_run_tosa fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -832,14 +855,15 @@ jobs: backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-public-api-backward-compatibility: + needs: docker-image name: test-arm-backend-public-api-backward-compatibility - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-24.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -859,8 +883,9 @@ jobs: python backends/arm/test/public_api_bc/run_public_api_bc_scenarios.py test-llama-runner-qnn-linux: + needs: docker-image name: test-llama-runner-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -871,8 +896,8 @@ jobs: mode: [qnn] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -898,8 +923,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama.sh -model stories110M -build_tool "${BUILD_TOOL}" -mode "${MODE}" -dtype "${DTYPE}" -pt2e_quantize "${PT2E_QUANTIZE}" test-static-llama-qnn-linux: + needs: docker-image name: test-static-llama-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -908,8 +934,8 @@ jobs: task: [stories_110m, stories_260k_bc] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -932,8 +958,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} test-sqnr-static-llm-qnn-linux: + needs: docker-image name: test-sqnr-static-llm-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -942,8 +969,8 @@ jobs: task: [smollm2_135m] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -966,8 +993,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_qnn_static_llm.sh ${{ matrix.task }} sqnr test-qnn-models-linux: + needs: docker-image name: test-qnn-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -976,8 +1004,8 @@ jobs: model: [mv2, mv3, dl3] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -991,14 +1019,15 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn" test-qnn-direct-build-linux: + needs: docker-image name: test-qnn-direct-build-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1021,22 +1050,23 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true - # No runner-linux, so this takes the memory-optimized default. The suite - # runs one export worker per core, so what matters is memory per core, not - # core count: the previous instance gave each worker about 4 GiB and the - # job was killed. The memory-optimized default gives each worker more. + # No runner-linux, so this takes _test_backend.yml's default. The suite + # runs one export worker per core, so what matters is memory per core: + # the instance this used to run on gave each worker about 4 GiB and the + # job was killed. The default label is a little under 8 GiB per core. test-qnn-passes-linux: + needs: docker-image name: test-qnn-passes-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 30 @@ -1065,16 +1095,21 @@ jobs: pytest -xvs backends/qualcomm/tests/test_import_side_effects.py test-qnn-delegate-linux: + needs: docker-image name: test-qnn-delegate-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + # Intel, unlike the avx512 labels: those are r7a, i.e. AMD EPYC, and the + # QNN backend disables MKLDNN on an AMD host through a non-bracketed + # torch.backends mutation that raises once the tests have frozen the + # flags. Same 8 vCPU / 64Gi, on r7i. + runner: mt-l-x86iamx-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1103,16 +1138,17 @@ jobs: -k "TestQNNFloatingPointOperator or TestQNNQuantizedOperator" test-phi-3-mini-runner-linux: + needs: docker-image name: test-phi-3-mini-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1133,16 +1169,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_phi_3_mini.sh Release test-qnn-python-imports-linux: + needs: docker-image name: test-qnn-python-imports-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 15 @@ -1181,16 +1218,17 @@ jobs: --module-prefix executorch.examples.qualcomm test-eval_llama-wikitext-linux: + needs: docker-image name: test-eval_llama-wikitext-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1210,7 +1248,7 @@ jobs: # TODO(larryliu0820): Fix this issue before reenabling it: https://gist.github.com/larryliu0820/7377ecd0d79dbc06076cec8d9f2b85d2 # test-eval_llama-mmlu-linux: # name: test-eval_llama-mmlu-linux - # uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + # uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main # permissions: # id-token: write # contents: read @@ -1236,16 +1274,17 @@ jobs: # PYTHON_EXECUTABLE=python bash .ci/scripts/test_eval_llama_mmlu.sh test-llama_runner_eager-linux: + needs: docker-image name: test-llama_runner_eager-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1263,16 +1302,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_llama_runner_eager.sh test-lora-linux: + needs: docker-image name: test-lora-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1290,16 +1330,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora.sh test-lora-multimethod-linux: + needs: docker-image name: test-lora-multimethod-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1317,16 +1358,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_lora_multimethod.sh test-mediatek-models-linux: + needs: docker-image name: test-mediatek-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-mediatek-sdk + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-mediatek-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1344,16 +1386,17 @@ jobs: # placeholder for mediatek to add more tests test-openvino-linux: + needs: docker-image name: test-openvino-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-gcc11 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1366,16 +1409,17 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_openvino.sh test-build-wasm-linux: + needs: docker-image name: test-build-wasm-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1394,8 +1438,9 @@ jobs: PYTHON_EXECUTABLE=python bash examples/wasm/test_build_wasm.sh unittest-wasm-bindings: + needs: docker-image name: unittest-wasm-bindings - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -1404,8 +1449,8 @@ jobs: enable-etdump: ['', '--enable-etdump'] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1440,13 +1485,14 @@ jobs: pnpm test unittest-nxp-neutron: - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + needs: docker-image + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 150 @@ -1484,18 +1530,19 @@ jobs: bash backends/nxp/run_unittests.sh test-samsung-quantmodels-linux: + needs: docker-image name: test-samsung-quantmodels-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 180 @@ -1522,18 +1569,19 @@ jobs: done test-samsung-models-linux: + needs: docker-image name: test-samsung-models-linux # Skip this job if the pull request is from a fork (secrets are not available) if: github.event.pull_request.head.repo.full_name == github.repository || github.event_name != 'pull_request' - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read secrets: inherit with: secrets-env: SAMSUNG_AI_LITECORE_KEY - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12-android + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-android-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 360 @@ -1564,14 +1612,15 @@ jobs: python -m unittest discover -s backends/samsung/test/models -p "test_*.py" test-vulkan-models-linux: + needs: docker-image name: test-vulkan-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1605,14 +1654,15 @@ jobs: done test-vulkan-operators-linux: + needs: docker-image name: test-vulkan-operators-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -1697,14 +1747,15 @@ jobs: echo "::endgroup::" nxp-build-test: + needs: docker-image name: nxp-build-test - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 From 095bf6bd0811c450e27dea6060173bffb07dc670 Mon Sep 17 00:00:00 2001 From: Martin Pavella Date: Thu, 10 Sep 2026 09:37:24 +0200 Subject: [PATCH 134/190] NXP backend: Add QAT support to Neutron lowering recipes. (#22642) ### Summary Add 2 QAT export recipes for the Neutron backend. The `INT8_QAT_NEUTRON` recipe is a direct equivalent to the imperative QAT lowering that was used until now. ### Test plan `pytest backends/nxp/tests/generic_tests/test_recipe_export.py` cc @robert-kalmar @JakeStevens @digantdesai @rascani --- backends/nxp/recipes/nxp_recipe_provider.py | 114 +++++- backends/nxp/recipes/nxp_recipe_types.py | 10 + .../tests/generic_tests/test_recipe_export.py | 336 +++++++++++++++++- 3 files changed, 438 insertions(+), 22 deletions(-) diff --git a/backends/nxp/recipes/nxp_recipe_provider.py b/backends/nxp/recipes/nxp_recipe_provider.py index 5617bcdb6f2..6592f1b0dd9 100644 --- a/backends/nxp/recipes/nxp_recipe_provider.py +++ b/backends/nxp/recipes/nxp_recipe_provider.py @@ -9,6 +9,15 @@ from functools import partial from typing import Any, Callable, cast, Iterable, Optional, Sequence +import torch + +from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import ( + FuseBatchNormWithLinearPass, +) +from executorch.backends.nxp.aten_passes.simulated_linear_bn_fusion_passes import ( + AddSimulatedLinearBatchNormFusionQATPass, + RemoveSimulatedLinearBatchNormFusionQATPass, +) from executorch.backends.nxp.backend.custom_delegation_options import ( CustomDelegationOptions, ) @@ -29,6 +38,9 @@ default_preserve_ops, generate_neutron_compile_spec, ) +from executorch.backends.nxp.quantizer.utils import ( + _replace_histogram_observers_for_integer_inputs, +) from executorch.backends.nxp.recipes.nxp_recipe_types import NXP_BACKEND, NXPRecipeType from executorch.backends.nxp.tests.executorch_pipeline import ( get_default_quantizer, @@ -36,6 +48,9 @@ ModelInputSpec, to_model_input_spec, ) +from executorch.backends.transforms.quantize_fused_convbn_bias_pass import ( + QuantizeFusedConvBnBiasAtenPass, +) from executorch.exir import ( EdgeCompileConfig, EdgeProgramManager, @@ -72,7 +87,7 @@ def __call__( class NeutronRecipeConfig: """Configuration shared by all NXP recipe types. - Parameters that vary the *type* of export (delegate vs no-delegate) + Parameters that vary the *type* of export (delegate vs no-delegate, PTQ vs QAT) are expressed by choosing a different NXPRecipeType rather than by flags here. Attributes: @@ -95,6 +110,9 @@ class NeutronRecipeConfig: dump_kernel_selection_code: Generate kernel-selection files after compilation. use_profiling: Enable Neutron execution profiling. IMPORTANT: To also generate an ETRecord, pass generate_etrecord=True to export() separately. + train_fn: Training function required for QAT recipe types (INT8_QAT_NEUTRON and + INT8_QAT_NO_DELEGATE). Receives the prepared GraphModule and must + perform the training loop. Ignored for PTQ recipe types. """ input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]] @@ -109,6 +127,7 @@ class NeutronRecipeConfig: fetch_constants_to_sram: bool = False dump_kernel_selection_code: bool = False use_profiling: bool = False + train_fn: Callable[["torch.fx.GraphModule"], None] | None = None class NXPRecipeProvider(BackendRecipeProvider): @@ -149,6 +168,10 @@ def create_recipe( return self._build_recipe(recipe_type, rc, is_qat=False, delegate=True) case NXPRecipeType.INT8_PTQ_NO_DELEGATE: return self._build_recipe(recipe_type, rc, is_qat=False, delegate=False) + case NXPRecipeType.INT8_QAT_NEUTRON: + return self._build_recipe(recipe_type, rc, is_qat=True, delegate=True) + case NXPRecipeType.INT8_QAT_NO_DELEGATE: + return self._build_recipe(recipe_type, rc, is_qat=True, delegate=False) case _: raise NotImplementedError( f"NXP backend: Recipe `{recipe_type}` is not supported." @@ -162,10 +185,10 @@ def _build_recipe( is_qat: bool, delegate: bool, ) -> ExportRecipe: - # is_qat=True is reserved for future QAT support; always False for now. - if is_qat: - raise NotImplementedError( - "NXP recipe with QAT (quantization aware training) is not yet supported." + if is_qat and rc.train_fn is None: + raise ValueError( + f"NXP backend: Recipe `{recipe_type}` requires `train_fn` to be set in " + f"NeutronRecipeConfig. Provide a callable that trains the prepared model." ) neutron_target_spec = NeutronTargetSpec(rc.target) @@ -175,7 +198,7 @@ def _build_recipe( get_default_quantizer, neutron_target_spec, is_qat ) - quantization_recipe = _build_quantization_recipe(rc) + quantization_recipe = _build_quantization_recipe(rc, is_qat) compile_spec = generate_neutron_compile_spec( rc.target, intermediates_dir=rc.intermediates_dir, @@ -199,23 +222,88 @@ def _build_recipe( ) +# --------------------------------------------------------------------------- +# Pass wrappers +# --------------------------------------------------------------------------- +# ExirPassBase subclasses return a PassResult with a .graph_module attribute, +# but QuantizeStage._apply_passes expects callable(GraphModule) -> GraphModule. +# These thin wrappers bridge the two conventions. + + +def _wrap_exir_pass(pass_cls, *args, **kwargs): + """Return a callable(GraphModule) -> GraphModule wrapping an ExirPass instance.""" + _pass_instance = pass_cls(*args, **kwargs) + + def _wrapped(m): + return _pass_instance(m).graph_module + + _wrapped.__qualname__ = f"_wrap_exir_pass({pass_cls.__name__})" + return _wrapped + + +def _histogram_observer_fix_pass(m): + """Callable(GraphModule) -> GraphModule that replaces HistogramObserver for integer inputs.""" + _replace_histogram_observers_for_integer_inputs(m) + return m + + # --------------------------------------------------------------------------- # Module-level builder helpers # --------------------------------------------------------------------------- -def _build_quantization_recipe(rc: NeutronRecipeConfig) -> QuantizationRecipe: - """Build the QuantizationRecipe for PTQ. +def _build_quantization_recipe( + rc: NeutronRecipeConfig, is_qat: bool +) -> QuantizationRecipe: + """Build the QuantizationRecipe for PTQ or QAT. PTQ uses the standard QuantizeStage flow (prepare_pt2e -> calibrate -> convert_pt2e). - The example_inputs passed to the export session are used directly for calibration. - Multiple PTQ recipes can be combined with ExportRecipe.combine(). + QAT uses the QAT flow (prepare_qat_pt2e -> BN-fusion passes -> train_fn -> convert_pt2e). + + The NXP-specific passes are injected via the QuantizationRecipe hook lists so + that QuantizeStage executes them in the correct order. """ _quantizer = rc.get_quantizer_fn() - return QuantizationRecipe( - quantizers=[_quantizer], - ) + # post_prepare_passes: always fix HistogramObserver for non-float inputs. + # For QAT, also insert the simulated linear-BN fusion before training so + # fake-quantize nodes see fused weights during the training loop. + post_prepare: list[Callable] = [] + if is_qat: + post_prepare.append(_wrap_exir_pass(AddSimulatedLinearBatchNormFusionQATPass)) + post_prepare.append(_histogram_observer_fix_pass) + + if is_qat: + # pre_convert_passes: tear down the simulated fusion and fold BN into + # the linear weights before convert_pt2e. + pre_convert: list[Callable] = [ + _wrap_exir_pass(RemoveSimulatedLinearBatchNormFusionQATPass), + _wrap_exir_pass(FuseBatchNormWithLinearPass), + ] + + # post_convert_passes: fix up quantization parameters for fused conv+BN + # bias nodes after convert_pt2e has inserted the quantize/dequantize ops. + post_convert: list[Callable] = [ + _wrap_exir_pass( + QuantizeFusedConvBnBiasAtenPass, + default_zero_bias=False, + symmetric_quant=True, + ) + ] + + return QuantizationRecipe( + quantizers=[_quantizer], + is_qat=True, + train_fn=rc.train_fn, + post_prepare_passes=post_prepare, + pre_convert_passes=pre_convert, + post_convert_passes=post_convert, + ) + else: + return QuantizationRecipe( + quantizers=[_quantizer], + post_prepare_passes=post_prepare, + ) def _build_lowering_recipe( diff --git a/backends/nxp/recipes/nxp_recipe_types.py b/backends/nxp/recipes/nxp_recipe_types.py index a5c655c2fce..b816a064c67 100644 --- a/backends/nxp/recipes/nxp_recipe_types.py +++ b/backends/nxp/recipes/nxp_recipe_types.py @@ -15,6 +15,8 @@ class NXPRecipeType(RecipeType): Choose the recipe that matches your intended export configuration: - INT8_PTQ_NEUTRON: standard post-training quantization, delegates to Neutron NPU. - INT8_PTQ_NO_DELEGATE: PTQ without NPU delegation (useful for debugging or CPU-only deployment). + - INT8_QAT_NEUTRON: quantization-aware training, delegates to Neutron NPU. + - INT8_QAT_NO_DELEGATE: QAT without NPU delegation (useful for accuracy evaluation). """ # INT8 static PTQ (weights + activations). Calibration dataset required. @@ -25,6 +27,14 @@ class NXPRecipeType(RecipeType): # Useful for accuracy evaluation or debugging before enabling delegation. INT8_PTQ_NO_DELEGATE = "nxp_int8_ptq_no_delegate" + # INT8 QAT (weights + activations). A train_fn must be provided in NeutronRecipeConfig. + # Applicable operators are delegated to the Neutron NPU. + INT8_QAT_NEUTRON = "nxp_int8_qat_neutron" + + # INT8 QAT without NPU delegation. Produces a quantized graph that runs on CPU. + # Useful for accuracy evaluation or debugging before enabling delegation. + INT8_QAT_NO_DELEGATE = "nxp_int8_qat_no_delegate" + @classmethod def get_backend_name(cls) -> str: return NXP_BACKEND diff --git a/backends/nxp/tests/generic_tests/test_recipe_export.py b/backends/nxp/tests/generic_tests/test_recipe_export.py index e820b71d69c..2403f2e62a7 100644 --- a/backends/nxp/tests/generic_tests/test_recipe_export.py +++ b/backends/nxp/tests/generic_tests/test_recipe_export.py @@ -7,15 +7,24 @@ import torch import torch.nn +from executorch.backends.nxp.aten_passes.fuse_batch_norm_with_linear_pass import ( + FuseBatchNormWithLinearPass, +) +from executorch.backends.nxp.aten_passes.simulated_linear_bn_fusion_passes import ( + AddSimulatedLinearBatchNormFusionQATPass, + RemoveSimulatedLinearBatchNormFusionQATPass, +) from executorch.backends.nxp.backend.custom_delegation_options import ( CustomDelegationOptions, ) +from executorch.backends.nxp.backend.graph_utils import batch_norm_target_ops from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( NeutronEdgePassManager, ) from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.recipes.nxp_recipe_provider import ( + _histogram_observer_fix_pass, NEUTRON_RECIPE_CONFIG_KEY, NeutronRecipeConfig, NXPRecipeProvider, @@ -26,6 +35,10 @@ graph_contains_any, graph_contains_any_of_ops, ) +from executorch.backends.nxp.tests.models import ConvBatchNormModule +from executorch.backends.transforms.quantize_fused_convbn_bias_pass import ( + QuantizeFusedConvBnBiasAtenPass, +) from executorch.export import export from executorch.export.recipe import ExportRecipe from torch._inductor.lowering import quantized_decomposed @@ -74,7 +87,10 @@ def is_cnn_op(n): assert not graph_contains_any(graph, is_cnn_op) nodes = list(graph.nodes) - first_call = next(n for n in nodes if n.op == "call_function" and n.name != "alloc") + # Skip alloc nodes (e.g. "alloc", "alloc_1") which also have op == "call_function". + first_call = next( + n for n in nodes if n.op == "call_function" and not n.name.startswith("alloc") + ) last_call = next(n for n in reversed(nodes) if n.op == "call_function") assert first_call.target == quantized_decomposed.quantize_per_tensor.out assert last_call.target == quantized_decomposed.dequantize_per_tensor.out @@ -97,14 +113,6 @@ def is_cnn_op(n): # With no delegation, original ops should be visible in the graph. assert graph_contains_any(graph, is_cnn_op) - def test__recipe_has_empty_partitioners(self): - """INT8_PTQ_NO_DELEGATE recipe has an empty partitioner list.""" - rc = NeutronRecipeConfig(INPUT_SHAPE) - recipe = NXPRecipeProvider().create_recipe( - NXPRecipeType.INT8_PTQ_NO_DELEGATE, neutron_recipe_config=rc - ) - assert recipe.lowering_recipe.partitioners == [] - class TestNeutronRecipeConfigFlags: def test_operators_not_to_delegate(self): @@ -413,3 +421,313 @@ def test__qdq_pass_callable_returns_pass_manager(self, mocker): qdq_callable = recipe.lowering_recipe.edge_manager_transform_passes[0] result = qdq_callable(mocker.MagicMock()) assert isinstance(result, NeutronEdgePassManager) + + +# --------------------------------------------------------------------------- +# Helpers shared by QAT tests +# --------------------------------------------------------------------------- + + +def _noop_train_fn(model: torch.fx.GraphModule) -> None: + """A no-op train_fn used by structural/unit tests that only inspect pass shape.""" + pass + + +def _minimal_train_fn(model: torch.fx.GraphModule, shape=(1, 3, 5, 5)) -> None: + """Run a few SGD steps on random data so fake-quant observer statistics are populated. + Used by end-to-end tests. + """ + optimizer = torch.optim.SGD(model.parameters(), lr=1e-4) + for _ in range(3): + optimizer.zero_grad() + out = model(torch.randn(shape)) + loss = out.sum() + loss.backward() + optimizer.step() + + +def _run_qat_export(model, train_fn=None, recipe_type=NXPRecipeType.INT8_QAT_NEUTRON): + if train_fn is None: + train_fn = _noop_train_fn + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=train_fn) + return _run_export(model, rc, recipe_type=recipe_type) + + +# --------------------------------------------------------------------------- +# QAT recipe: end-to-end tests +# --------------------------------------------------------------------------- + + +# Both QAT recipe types must reject a missing train_fn. +@pytest.mark.parametrize( + "recipe_type", + [NXPRecipeType.INT8_QAT_NEUTRON, NXPRecipeType.INT8_QAT_NO_DELEGATE], + ids=lambda r: r.value, +) +def test__qat_requires_train_fn(recipe_type): + """Any QAT recipe raises ValueError when train_fn is absent from NeutronRecipeConfig.""" + rc = NeutronRecipeConfig(INPUT_SHAPE) # train_fn=None (default) + with pytest.raises(ValueError, match="train_fn"): + NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + + +# Both _NO_DELEGATE recipe types (PTQ and QAT) must produce an empty partitioner list. +@pytest.mark.parametrize( + "recipe_type", + [NXPRecipeType.INT8_PTQ_NO_DELEGATE, NXPRecipeType.INT8_QAT_NO_DELEGATE], + ids=lambda r: r.value, +) +def test__no_delegate_recipe_has_empty_partitioners(recipe_type): + """Both NO_DELEGATE recipe types produce an empty partitioner list.""" + # train_fn is required by QAT recipes; PTQ ignores it, so always pass it. + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + recipe = NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + assert recipe.lowering_recipe.partitioners == [] + + +class TestInt8QATNeutron: + + def test__basic(self): + """INT8_QAT_NEUTRON: full export succeeds and the graph contains a delegate call.""" + model = SimpleCNN() + sess = _run_qat_export(model) + graph = _get_graph(sess) + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def test__train_fn_is_called(self): + """train_fn is invoked exactly once during the QAT export pipeline.""" + model = SimpleCNN() + call_count = [] + + def counting_train_fn(m): + call_count.append(1) + + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=counting_train_fn) + _run_export(model, rc, recipe_type=NXPRecipeType.INT8_QAT_NEUTRON) + assert ( + len(call_count) == 1 + ), f"Expected train_fn called once, got {len(call_count)}" + + def test__recipe_name(self): + """INT8_QAT_NEUTRON recipe has the expected name.""" + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc + ) + assert recipe.name == NXPRecipeType.INT8_QAT_NEUTRON.value + + def test__recipe_structure(self): + """INT8_QAT_NEUTRON recipe has is_qat=True, one quantizer, one partitioner.""" + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + recipe = NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_QAT_NEUTRON, neutron_recipe_config=rc + ) + qr = recipe.quantization_recipe + assert qr is not None + assert qr.is_qat is True + assert qr.train_fn is _noop_train_fn + assert len(qr.quantizers) == 1 + assert recipe.lowering_recipe.partitioners is not None + assert len(recipe.lowering_recipe.partitioners) == 1 + + def test__io_is_quantized_by_default(self): + """QAT export with default settings: IO boundary has quantize/dequantize ops.""" + model = SimpleCNN() + sess = _run_qat_export(model) + graph = _get_graph(sess) + nodes = list(graph.nodes) + # Skip alloc nodes (e.g. "alloc", "alloc_1") which also have op == "call_function". + first_call = next( + n + for n in nodes + if n.op == "call_function" and not n.name.startswith("alloc") + ) + last_call = next(n for n in reversed(nodes) if n.op == "call_function") + assert first_call.target == quantized_decomposed.quantize_per_tensor.out + assert last_call.target == quantized_decomposed.dequantize_per_tensor.out + + +class TestInt8QATNoDelegate: + + def test__basic(self): + """INT8_QAT_NO_DELEGATE: export succeeds without any delegate call.""" + model = SimpleCNN() + sess = _run_qat_export(model, recipe_type=NXPRecipeType.INT8_QAT_NO_DELEGATE) + graph = _get_graph(sess) + assert not graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + +# --------------------------------------------------------------------------- +# QAT recipe: NXP-specific pass structure tests +# --------------------------------------------------------------------------- + +# Both QAT recipe types are built from the same _build_quantization_recipe(is_qat=True) +# call, so their pass lists must be identical. The parametrization below makes this +# explicit and catches any accidental divergence. +_QAT_RECIPE_TYPES = [NXPRecipeType.INT8_QAT_NEUTRON, NXPRecipeType.INT8_QAT_NO_DELEGATE] + + +@pytest.mark.parametrize("recipe_type", _QAT_RECIPE_TYPES, ids=lambda r: r.value) +class TestQATNXPPasses: + + def _get_qat_recipe(self, recipe_type: NXPRecipeType) -> "ExportRecipe": + rc = NeutronRecipeConfig(INPUT_SHAPE, train_fn=_noop_train_fn) + return NXPRecipeProvider().create_recipe(recipe_type, neutron_recipe_config=rc) + + def test__post_prepare_passes_start_with_add_bn_fusion(self, recipe_type): + """QAT post_prepare_passes: first pass is AddSimulatedLinearBatchNormFusionQATPass wrapper.""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.post_prepare_passes is not None + # The first post-prepare pass must wrap AddSimulatedLinearBatchNormFusionQATPass. + # We verify by inspecting the __qualname__ set by _wrap_exir_pass. + first_pass = qr.post_prepare_passes[0] + assert ( + AddSimulatedLinearBatchNormFusionQATPass.__name__ in first_pass.__qualname__ + ) + + def test__post_prepare_passes_end_with_histogram_observer_fix(self, recipe_type): + """QAT post_prepare_passes: last pass is _histogram_observer_fix_pass.""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.post_prepare_passes is not None + last_pass = qr.post_prepare_passes[-1] + assert last_pass is _histogram_observer_fix_pass + + def test__pre_convert_passes_include_remove_bn_fusion_and_fold(self, recipe_type): + """QAT pre_convert_passes: contains RemoveSimulatedLinearBatchNormFusionQATPass + followed by FuseBatchNormWithLinearPass (each applied once).""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.pre_convert_passes is not None + assert len(qr.pre_convert_passes) == 2, ( + "Expected 2 pre_convert passes (remove + fuse), " + f"got {len(qr.pre_convert_passes)}" + ) + qualnames = [p.__qualname__ for p in qr.pre_convert_passes] + assert ( + qualnames[0] + == f"_wrap_exir_pass({RemoveSimulatedLinearBatchNormFusionQATPass.__name__})" + ) + assert ( + qualnames[1] == f"_wrap_exir_pass({FuseBatchNormWithLinearPass.__name__})" + ) + + def test__post_convert_passes_include_quant_fused_conv_bn_bias(self, recipe_type): + """QAT post_convert_passes: contains QuantizeFusedConvBnBiasAtenPass wrapper.""" + recipe = self._get_qat_recipe(recipe_type) + qr = recipe.quantization_recipe + assert qr.post_convert_passes is not None + assert len(qr.post_convert_passes) == 1 + assert f"_wrap_exir_pass({QuantizeFusedConvBnBiasAtenPass.__name__})" in ( + qr.post_convert_passes[0].__qualname__ + ) + + +# --------------------------------------------------------------------------- +# PTQ recipe: pass structure tests (complement to TestQATNXPPasses above) +# --------------------------------------------------------------------------- + + +class TestPTQNXPPasses: + """Verifies the pass structure of INT8_PTQ_NEUTRON recipes. + + These tests are separate from TestQATNXPPasses because the assertions are + PTQ-specific and independent of which QAT recipe type is being tested. + """ + + def _get_ptq_recipe(self) -> "ExportRecipe": + rc = NeutronRecipeConfig(INPUT_SHAPE) + return NXPRecipeProvider().create_recipe( + NXPRecipeType.INT8_PTQ_NEUTRON, neutron_recipe_config=rc + ) + + def test__no_pre_or_post_convert_passes(self): + """PTQ recipe does not set pre_convert_passes or post_convert_passes.""" + qr = self._get_ptq_recipe().quantization_recipe + assert qr.pre_convert_passes is None + assert qr.post_convert_passes is None + + def test__post_prepare_passes_include_histogram_observer_fix(self): + """PTQ recipe post_prepare_passes contains _histogram_observer_fix_pass.""" + qr = self._get_ptq_recipe().quantization_recipe + assert qr.post_prepare_passes is not None + assert _histogram_observer_fix_pass in qr.post_prepare_passes + + def test__post_prepare_passes_do_not_include_add_bn_fusion(self): + """PTQ recipe post_prepare_passes must NOT contain AddSimulatedLinearBatchNormFusionQATPass.""" + qr = self._get_ptq_recipe().quantization_recipe + for p in qr.post_prepare_passes or []: + assert AddSimulatedLinearBatchNormFusionQATPass.__name__ not in getattr( + p, "__qualname__", "" + ) + + +# --------------------------------------------------------------------------- +# QAT recipe: e2e test on a real model - equivalent to imperative QAT tests +# --------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "bias", [True, False], ids=lambda b: "bias" if b else "no_bias" +) +class TestQATEquivalentToImperative: + """Recipe-path QAT tests that mirror the imperative-path tests in test_batch_norm_fusion.py. + + The imperative reference is test_biasless_convbn_fusion_qat (and its bias=True variant). + The recipe path must produce an equivalent result: the graph is delegated and the + BN is fully fused away by the QAT passes. + """ + + # Use (1, 3, 5, 5) so that Conv2d(kernel_size=3) produces a (1, 3, 3, 3) feature map, + # giving BatchNorm > 1 value per channel in training mode (QAT requires train mode). + _CONVBN_INPUT_SHAPE = (1, 3, 5, 5) + + def test__convbn_qat_produces_delegate_call(self, bias): + """INT8_QAT_NEUTRON on ConvBatchNormModule produces a delegate call. + Equivalent imperative test: test_biasless_convbn_fusion_qat / test_batch_norm_conv_fusing + in backends/nxp/tests/generic_tests/test_batch_norm_fusion.py.""" + model = ConvBatchNormModule( + bias=bias, + input_rank=len(self._CONVBN_INPUT_SHAPE), + num_features=self._CONVBN_INPUT_SHAPE[1], + ) + rc = NeutronRecipeConfig( + self._CONVBN_INPUT_SHAPE, + train_fn=_minimal_train_fn, + use_neutron_for_format_conversion=False, + ) + sess = _run_export( + model, + rc, + recipe_type=NXPRecipeType.INT8_QAT_NEUTRON, + input_shape=self._CONVBN_INPUT_SHAPE, + ) + graph = _get_graph(sess) + + # Same assertion as the imperative path: the model is delegated. + assert graph_contains_any_of_ops(graph, [ExecutorchDelegateCall]) + + def test__convbn_qat_bn_is_fused_away(self, bias): + """INT8_QAT_NEUTRON on ConvBatchNormModule: BN is fused away by QAT passes. + Equivalent imperative test: test_batch_norm_conv_fusing__full_pipeline__2d + in backends/nxp/tests/generic_tests/test_batch_norm_fusion.py.""" + model = ConvBatchNormModule( + bias=bias, + input_rank=len(self._CONVBN_INPUT_SHAPE), + num_features=self._CONVBN_INPUT_SHAPE[1], + ) + rc = NeutronRecipeConfig( + self._CONVBN_INPUT_SHAPE, + train_fn=_minimal_train_fn, + use_neutron_for_format_conversion=False, + ) + sess = _run_export( + model, + rc, + recipe_type=NXPRecipeType.INT8_QAT_NEUTRON, + input_shape=self._CONVBN_INPUT_SHAPE, + ) + # The edge program (before delegation) must not contain any BN ops. + edge_graph = sess.get_edge_program_manager().exported_program().graph + assert not graph_contains_any_of_ops(edge_graph, batch_norm_target_ops) From 176fd9d2bd0f3417df1ff67d1a928211de5d19ec Mon Sep 17 00:00:00 2001 From: DannyYuyang-quic Date: Thu, 10 Sep 2026 15:53:58 +0800 Subject: [PATCH 135/190] Qualcomm AI Engine Direct - [LLM QAT] Embedding + LM head W4 PCQ support (#21879) ### Summary - With W4 PCQ applied to the embedding and LM head, we can further compress the embedding size. W4 PCQ. - Note: PCQ for embeddings is supported only in `QNN SDK` 2.48 or later. ### E2E script ``` bash python examples/qualcomm/oss_scripts/llama/llama.py --build_folder build-android --device ${SERIAL_NUM} --soc_model {SOC_MODEL} --decoder_model smollm2_135m --model_mode hybrid --max_seq_len 1024 --prompt "What dose it mean to edit written content?" --qat --train_hf_dataset "HuggingFaceTB/smol-smoltalk" --train_hf_limit 8000 --calib_hf_dataset "HuggingFaceTB/smol-smoltalk" --calib_hf_limit 8000 --batch_size 2 --grad_accum_steps 4 ``` ### Test plan ``` bash python backends/qualcomm/tests/test_qnn_delegate.py TestExampleLLMScript.test_static_llm_qat --device ${SERIAL_NUM} --soc_model ${SOC_MODEL} -r . -a . --build_folder build-android ``` --- backends/qualcomm/quantizer/quant_recipe.py | 14 ++++++++++++++ .../oss_scripts/llama/static_llm_quant_recipe.py | 15 ++------------- 2 files changed, 16 insertions(+), 13 deletions(-) diff --git a/backends/qualcomm/quantizer/quant_recipe.py b/backends/qualcomm/quantizer/quant_recipe.py index b2eb41841f0..347c82c1f82 100644 --- a/backends/qualcomm/quantizer/quant_recipe.py +++ b/backends/qualcomm/quantizer/quant_recipe.py @@ -18,6 +18,10 @@ QuantizationConfig, ) from executorch.backends.qualcomm.quantizer.rules import OpQuantRule +from executorch.backends.qualcomm.utils.check_qnn_version import ( + get_sdk_build_id, + is_qnn_sdk_version_less_than, +) from tabulate import tabulate from torch._ops import OpOverload from torchao.quantization.pt2e import UniformQuantizationObserverBase @@ -91,6 +95,7 @@ def __init__( is_qat=self.is_qat, is_conv_per_channel=True, is_linear_per_channel=True, + is_embedding_per_channel=True, act_observer=self.act_observer, act_symmetric=self.act_symmetric, ) @@ -108,6 +113,15 @@ def get_quant_config(self, node: torch.fx.Node) -> Optional[QuantizationConfig]: if self.granularity == QuantGranularity.PER_TENSOR: return self.quant_config.quant_config elif self.granularity == QuantGranularity.PER_CHANNEL: + if op == torch.ops.aten.embedding.default and is_qnn_sdk_version_less_than( + "2.48" + ): + raise RuntimeError( + "Per-channel embedding quantization requires QNN SDK version " + f">= 2.48. Current QNN SDK version: {get_sdk_build_id()}. Please " + "upgrade your QNN SDK, or use QuantGranularity.PER_TENSOR for the " + "embedding layer instead." + ) ch_axis = self.quant_config.use_per_channel_weight_quant_ops.get(op) assert ( ch_axis is not None diff --git a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py index 29eba63a07d..521dba38ef7 100644 --- a/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py +++ b/examples/qualcomm/oss_scripts/llama/static_llm_quant_recipe.py @@ -697,11 +697,7 @@ def __init__(self, verbose: bool = False): class Smollm2QATQuantRecipe(StaticLLMQATRecipe): - default_quant_dtype = QuantDtype.use_16a8w - frozen_param_patterns: List[str] = [ - r"tok_embedding", # Freeze token embeddings to prevent drift in the token space. - r"output\.conv", # Freeze lm head to prevent drift in the token space. - ] + default_quant_dtype = QuantDtype.use_16a4w def __init__(self, verbose: bool = False): super().__init__() @@ -725,14 +721,7 @@ def __init__(self, verbose: bool = False): ) .add_regex( {r"tok_embeddings"}, - QuantDtype.use_16a8w, - True, - act_observer=MovingAverageMinMaxObserver, - granularity=QuantGranularity.PER_TENSOR, - ) - .add_regex( - {r"output\.conv"}, - QuantDtype.use_16a8w, + QuantDtype.use_16a4w, True, act_observer=MovingAverageMinMaxObserver, granularity=QuantGranularity.PER_CHANNEL, From a457114d6b01415650de76ff2d6690ebd13329a5 Mon Sep 17 00:00:00 2001 From: Tom Allsop <72802373+tom-arm@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:22:48 +0100 Subject: [PATCH 136/190] Arm backend: Align static KV-cache tests to extensions/llm (#22641) * Add per tensor path to StaticQuantizedKVCache to enable delegation * Rearchitect StaticQuantizedCacheModule to use StaticQuantizedKVCache * Use StaticCacheModule for FP32 tests only * Use StaticQuantizedCacheModule for INT tests * Ensure cache buffers are initialized correctly Change-Id: I3f39f19ad75c056e10bd9f336b711ad7271d95ff cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Tom Allsop --- .../arm/test/modules/test_static_cache.py | 171 ++++-------------- .../source_transformation/custom_kv_cache.py | 77 ++++++-- 2 files changed, 92 insertions(+), 156 deletions(-) diff --git a/backends/arm/test/modules/test_static_cache.py b/backends/arm/test/modules/test_static_cache.py index 86649f1e589..20ddac33f79 100644 --- a/backends/arm/test/modules/test_static_cache.py +++ b/backends/arm/test/modules/test_static_cache.py @@ -12,6 +12,7 @@ InsertInt32CastsAfterInt64PlaceholdersPass, ) from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ToExecutorch from executorch.backends.arm.test.tester.test_pipeline import ( EthosU55PipelineINT, EthosU85PipelineINT, @@ -19,6 +20,11 @@ TosaPipelineINT, VgfPipeline, ) +from executorch.examples.models.llama.source_transformation.custom_kv_cache import ( + StaticQuantizedKVCache, +) +from executorch.exir import ExecutorchBackendConfig +from executorch.exir.passes.init_mutable_pass import InitializedMutableBufferPass from torch.export.graph_signature import InputKind, OutputKind from transformers import LlamaConfig @@ -41,39 +47,23 @@ InputKind.USER_INPUT: 3, } -EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS = { - InputKind.BUFFER: 4, - InputKind.USER_INPUT: 3, -} - EXPECTED_OUTPUT_COUNTS = { OutputKind.BUFFER_MUTATION: 2, OutputKind.USER_OUTPUT: 2, } -DYNAMIC_KVQ_OPS = [ - "torch.ops.quantized_decomposed.choose_qparams_per_token_asymmetric.default", - "torch.ops.quantized_decomposed.quantize_per_token.default", - "torch.ops.quantized_decomposed.dequantize_per_token.default", - "torch.ops.llama.update_cache.default", - "torch.ops.llama.update_cache_with_indices.default", -] - - -def _reject_dynamic_kvq_ops(pipeline): - pipeline.add_stage_after( - "export", pipeline.tester.check_not, DYNAMIC_KVQ_OPS, suffix="dynamic_kvq_ops" +def _initialize_cache_buffers(pipeline, pattern: list[str]) -> None: + pipeline.change_args( + "to_executorch", + ToExecutorch( + ExecutorchBackendConfig(passes=[InitializedMutableBufferPass(pattern)]) + ), ) @torch.no_grad() class StaticQuantizedCacheModule(torch.nn.Module): - key_cache: torch.Tensor - value_cache: torch.Tensor - key_scale: torch.Tensor - value_scale: torch.Tensor - def __init__( self, config: LlamaConfig, @@ -90,40 +80,23 @@ def __init__( self.hidden_size = hidden_size self.num_attention_heads = num_attention_heads self.head_dim = self.hidden_size // self.num_attention_heads - cache_shape = (1, self.num_attention_heads, max_cache_len, self.head_dim) - scale_shape = (1, 1, 1, self.head_dim) - - self.register_buffer("key_cache", torch.zeros(cache_shape, dtype=torch.int8)) - self.register_buffer("value_cache", torch.zeros(cache_shape, dtype=torch.int8)) - self.register_buffer( - "key_scale", torch.full(scale_shape, scale, dtype=torch.float32) - ) - self.register_buffer( - "value_scale", torch.full(scale_shape, scale, dtype=torch.float32) + self.cache = StaticQuantizedKVCache( + max_batch_size=1, + max_context_length=max_cache_len, + n_heads=self.num_attention_heads, + head_dim=self.head_dim, + scale=scale, + use_custom_update_cache_op=False, + use_per_channel=False, ) - # PT2E activation quantization does not create persistent int8 mutable buffers. - def _quantize(self, value: torch.Tensor, scale: torch.Tensor) -> torch.Tensor: - return torch.clamp(torch.round(value / scale), -128, 127).to(torch.int8) - def forward( self, key_states: torch.Tensor, value_states: torch.Tensor, cache_position: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: - key_q = self._quantize(key_states, self.key_scale) - value_q = self._quantize(value_states, self.value_scale) - - self.key_cache[:, :, cache_position] = key_q - self.value_cache[:, :, cache_position] = value_q - - key = self.key_cache.to(torch.float32) * self.key_scale - value = self.value_cache.to(torch.float32) * self.value_scale - key[:, :, cache_position] = key_states - value[:, :, cache_position] = value_states - - return key.clone(), value.clone() + return self.cache.update(cache_position, key_states, value_states) def get_inputs(self) -> input_t: key_states = torch.randn( @@ -236,17 +209,18 @@ def test_static_cache_tosa_FP(test_data): exir_op=[], transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], ) + _initialize_cache_buffers(pipeline, ["cache_layer_"]) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() -@pytest.mark.xfail(reason="BUFFER_MUTATION count mismatch: MLETORCH-1971") @common.parametrize("test_data", test_configs) def test_static_cache_tosa_INT(test_data): - module = StaticCacheModule(test_data).eval() + module = StaticQuantizedCacheModule(test_data).eval() pipeline = TosaPipelineINT[input_t]( - module, module.get_inputs(), aten_op=[], exir_op=[], fold_quantize=False + module, module.get_inputs(), aten_op=[], exir_op=[] ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() @@ -255,41 +229,26 @@ def test_static_cache_tosa_INT(test_data): @pytest.mark.xfail(reason="Scatter operator is not supported on U55.") @common.parametrize("test_data", test_configs) def test_static_cache_u55_INT(test_data): - module = StaticCacheModule(test_data).eval() + module = StaticQuantizedCacheModule(test_data).eval() pipeline = EthosU55PipelineINT[input_t]( module, module.get_inputs(), aten_ops=[], ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) pipeline.run() -@common.parametrize( - "test_data", - test_configs, - xfails={ - "multihead_attention": ( - "BUFFER_MUTATION count mismatch: MLETORCH-1971" - "Incorrect numerical behavior: MLBEDSW-11589" - ), - "grouped_query_attention": ( - "BUFFER_MUTATION count mismatch: MLETORCH-1971" - "Incorrect numerical behavior: MLBEDSW-11589" - ), - "multi_query_attention": ( - "BUFFER_MUTATION count mismatch: MLETORCH-1971" - "Incorrect numerical behavior: MLBEDSW-11589" - ), - }, -) +@common.XfailIfNoCorstone320 +@common.parametrize("test_data", test_configs) def test_static_cache_u85_INT(test_data): - module = StaticCacheModule(test_data).eval() + module = StaticQuantizedCacheModule(test_data).eval() pipeline = EthosU85PipelineINT[input_t]( module, module.get_inputs(), aten_ops=[], - fold_quantize=False, ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) # U85: keep _to_dim_order_copy portable for int64->int32 cast of cache_position (not delegatable). pipeline.tester.use_portable_ops = True pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) @@ -308,72 +267,14 @@ def test_static_cache_vgf_no_quant(test_data): transform_passes=[InsertInt32CastsAfterInt64PlaceholdersPass()], quantize=False, ) + _initialize_cache_buffers(pipeline, ["cache_layer_"]) pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() @common.SkipIfNoModelConverter -@pytest.mark.xfail(reason="BUFFER_MUTATION count mismatch: MLETORCH-1971") @common.parametrize("test_data", test_configs) def test_static_cache_vgf_quant(test_data): - module = StaticCacheModule(test_data).eval() - pipeline = VgfPipeline[input_t]( - module, - module.get_inputs(), - aten_op=[], - exir_op=[], - quantize=True, - fold_quantize=False, - tosa_spec="TOSA-1.0+INT", - ) - pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) - pipeline.run() - - -@common.parametrize("test_data", test_configs) -def test_static_quantized_cache_tosa_INT(test_data): - module = StaticQuantizedCacheModule(test_data).eval() - pipeline = TosaPipelineINT[input_t]( - module, module.get_inputs(), aten_op=[], exir_op=[], fold_quantize=False - ) - _reject_dynamic_kvq_ops(pipeline) - pipeline.change_args( - "check_count.exir", - {"torch.ops.higher_order.executorch_call_delegate": 2}, - ) - pipeline.count_program_io_kinds( - EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS - ) - pipeline.run() - - -@common.parametrize( - "test_data", - test_configs, - xfails={ - config: "Incorrect numerical behavior: MLBEDSW-11589" for config in test_configs - }, -) -def test_static_quantized_cache_u85_INT(test_data): - module = StaticQuantizedCacheModule(test_data).eval() - pipeline = EthosU85PipelineINT[input_t]( - module, module.get_inputs(), aten_ops=[], fold_quantize=False - ) - _reject_dynamic_kvq_ops(pipeline) - pipeline.change_args( - "check_count.exir", - {"torch.ops.higher_order.executorch_call_delegate": 2}, - ) - pipeline.tester.use_portable_ops = True - pipeline.count_program_io_kinds( - EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS - ) - pipeline.run() - - -@common.SkipIfNoModelConverter -@common.parametrize("test_data", test_configs) -def test_static_quantized_cache_vgf_quant(test_data): module = StaticQuantizedCacheModule(test_data).eval() pipeline = VgfPipeline[input_t]( module, @@ -381,12 +282,8 @@ def test_static_quantized_cache_vgf_quant(test_data): aten_op=[], exir_op=[], quantize=True, - fold_quantize=False, tosa_spec="TOSA-1.0+INT", - n_expected_delegates=2, - ) - _reject_dynamic_kvq_ops(pipeline) - pipeline.count_program_io_kinds( - EXPECTED_STATIC_QUANTIZED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS ) + _initialize_cache_buffers(pipeline, ["k_cache", "v_cache"]) + pipeline.count_program_io_kinds(EXPECTED_INPUT_COUNTS, EXPECTED_OUTPUT_COUNTS) pipeline.run() diff --git a/examples/models/llama/source_transformation/custom_kv_cache.py b/examples/models/llama/source_transformation/custom_kv_cache.py index dbaac9accf4..71cfe33d753 100644 --- a/examples/models/llama/source_transformation/custom_kv_cache.py +++ b/examples/models/llama/source_transformation/custom_kv_cache.py @@ -265,6 +265,7 @@ def __init__( scale: float = 1.0 / 127.0, use_custom_update_cache_op: bool = True, return_float_values: bool = True, + use_per_channel: bool = True, dtype: torch.dtype = torch.float32, ): super().__init__() @@ -276,9 +277,12 @@ def __init__( self.quantized_cache_dtype = torch.int8 self.return_float_values = return_float_values self.max_context_length = max_context_length + self.use_per_channel = use_per_channel + self.k_cache_scale = scale + self.v_cache_scale = scale self.calibration_enabled = False cache_shape = (max_batch_size, max_context_length, n_heads, head_dim) - scale_shape = (1, 1, 1, head_dim) + scale_shape = (1, 1, 1, head_dim) if use_per_channel else (1,) self.register_buffer( "k_cache", torch.zeros(cache_shape, dtype=self.quantized_cache_dtype), @@ -324,9 +328,11 @@ def finalize_calibration(self): k_scales = self.k_observed_max.to(self.k_cache_scales.dtype) / 127.0 v_scales = self.v_observed_max.to(self.v_cache_scales.dtype) / 127.0 if torch.any(k_scales == 0) or torch.any(v_scales == 0): + qparam_scope = "channel" if self.use_per_channel else "cache" logging.warning( - "Static KV cache calibration observed an all-zero K/V channel; " - "using the smallest positive scale for that channel." + "Static KV cache calibration observed an all-zero K/V %s; " + "using the smallest positive scale.", + qparam_scope, ) # This floor prevents division by zero; it is not an accuracy threshold. self.k_cache_scales.copy_( @@ -335,6 +341,9 @@ def finalize_calibration(self): self.v_cache_scales.copy_( v_scales.clamp_min(torch.finfo(self.v_cache_scales.dtype).tiny) ) + if not self.use_per_channel: + self.k_cache_scale = self.k_cache_scales.item() + self.v_cache_scale = self.v_cache_scales.item() self.calibration_enabled = False self.k_calibration_cache = None self.v_calibration_cache = None @@ -347,30 +356,54 @@ def _observe_and_update(self, input_pos, k_val, v_val): self.k_observed_max.copy_( torch.maximum( self.k_observed_max, - k_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True), + ( + k_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True) + if self.use_per_channel + else k_val.detach().abs().amax().reshape_as(self.k_observed_max) + ), ) ) self.v_observed_max.copy_( torch.maximum( self.v_observed_max, - v_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True), + ( + v_val.detach().abs().amax(dim=(0, 1, 2), keepdim=True) + if self.use_per_channel + else v_val.detach().abs().amax().reshape_as(self.v_observed_max) + ), ) ) self.k_calibration_cache[:, input_pos] = k_val self.v_calibration_cache[:, input_pos] = v_val return self.k_calibration_cache, self.v_calibration_cache - def _quantize(self, value, scales): - # torchao affine custom ops do not yet have the required Arm/TOSA - # lowering and ExecuTorch out-variant runtime support. - qmin = torch.iinfo(self.quantized_cache_dtype).min - qmax = torch.iinfo(self.quantized_cache_dtype).max - return torch.clamp(torch.round(value / scales), qmin, qmax).to( - self.quantized_cache_dtype + def _quantize(self, value, scale): + if self.use_per_channel: + qmin = torch.iinfo(self.quantized_cache_dtype).min + qmax = torch.iinfo(self.quantized_cache_dtype).max + return torch.clamp(torch.round(value / scale), qmin, qmax).to( + self.quantized_cache_dtype + ) + return torch.ops.quantized_decomposed.quantize_per_tensor.default( + value, + scale, + 0, + torch.iinfo(self.quantized_cache_dtype).min, + torch.iinfo(self.quantized_cache_dtype).max, + self.quantized_cache_dtype, ) - def _dequantize(self, value, scales, dtype): - return value.to(dtype) * scales.to(dtype) + def _dequantize(self, value, scale, dtype): + if self.use_per_channel: + return value.to(dtype) * scale.to(dtype) + return torch.ops.quantized_decomposed.dequantize_per_tensor.default( + value, + scale, + 0, + torch.iinfo(self.quantized_cache_dtype).min, + torch.iinfo(self.quantized_cache_dtype).max, + self.quantized_cache_dtype, + ).to(dtype) def _update_cache(self, value, cache, input_pos, indices=None): start_pos = input_pos[0].item() @@ -386,8 +419,10 @@ def _update_cache(self, value, cache, input_pos, indices=None): cache[:, input_pos] = value def _quantize_and_update(self, input_pos, k_val, v_val, indices=None): - quantized_k_val = self._quantize(k_val, self.k_cache_scales) - quantized_v_val = self._quantize(v_val, self.v_cache_scales) + k_scale = self.k_cache_scales if self.use_per_channel else self.k_cache_scale + v_scale = self.v_cache_scales if self.use_per_channel else self.v_cache_scale + quantized_k_val = self._quantize(k_val, k_scale) + quantized_v_val = self._quantize(v_val, v_scale) self._update_cache(quantized_k_val, self.k_cache, input_pos, indices) self._update_cache(quantized_v_val, self.v_cache, input_pos, indices) @@ -395,8 +430,10 @@ def _quantize_and_update(self, input_pos, k_val, v_val, indices=None): def _update_and_return_float_values(self, input_pos, k_val, v_val, indices=None): self._quantize_and_update(input_pos, k_val, v_val, indices) - k_out = self._dequantize(self.k_cache, self.k_cache_scales, k_val.dtype) - v_out = self._dequantize(self.v_cache, self.v_cache_scales, v_val.dtype) + k_scale = self.k_cache_scales if self.use_per_channel else self.k_cache_scale + v_scale = self.v_cache_scales if self.use_per_channel else self.v_cache_scale + k_out = self._dequantize(self.k_cache, k_scale, k_val.dtype) + v_out = self._dequantize(self.v_cache, v_scale, v_val.dtype) self._update_cache(k_val, k_out, input_pos, indices) self._update_cache(v_val, v_out, input_pos, indices) @@ -414,7 +451,7 @@ def update(self, input_pos, k_val, v_val, indices=None): """ k_val, v_val: [B, H, S, D] return: [B, H, S, D] - Storage is [B, S, H, D], with static per-head-dim qparams. + Storage is [B, S, H, D], with static per-head-dim or per-tensor qparams. """ k_val = k_val.transpose(1, 2) @@ -440,6 +477,7 @@ def from_float( kv_cache, scale: float = 1.0 / 127.0, use_custom_update_cache_op: bool = True, + use_per_channel: bool = True, ): if isinstance(kv_cache, CustomKVCache): max_batch_size, max_context_length, n_heads, head_dim = ( @@ -456,6 +494,7 @@ def from_float( head_dim, scale=scale, use_custom_update_cache_op=use_custom_update_cache_op, + use_per_channel=use_per_channel, dtype=kv_cache.k_cache.dtype, ) From dee51bb9025da12d6a76eeda8d91a72df2b94064 Mon Sep 17 00:00:00 2001 From: Emma Kujala <47500215+emmakujala@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:47:26 +0200 Subject: [PATCH 137/190] Arm backend: Add IO copy tracking to perf_monitor (#22672) Add an optional tracking of IO copying that reports how many input and/or output copy calls were made, as well as how many bytes of data was copied. Signed-off-by: Emma Kujala --- backends/arm/runtime/EthosUBackend.cpp | 10 +++ .../arm/runtime/EthosUBackend_Cortex_M.cpp | 3 + backends/arm/runtime/EthosUBackend_Internal.h | 4 + examples/arm/executor_runner/CMakeLists.txt | 12 +++ .../arm/executor_runner/arm_perf_monitor.cpp | 81 +++++++++++++++++++ 5 files changed, 110 insertions(+) diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index 2e752e99995..e7325f96a94 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -60,6 +60,10 @@ void __attribute__((weak)) EthosUBackend_execute_end() {} void __attribute__((weak)) EthosUBackend_delegate_begin(const void*) {} void __attribute__((weak)) EthosUBackend_delegate_end() {} #endif +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +void __attribute__((weak)) EthosUBackend_input_memcpy(size_t) {} +void __attribute__((weak)) EthosUBackend_output_memcpy(size_t) {} +#endif __attribute__((weak)) unsigned char* ethosu_fast_scratch = nullptr; __attribute__((weak)) size_t ethosu_fast_scratch_size = 0; } @@ -268,6 +272,9 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { event_tracer, "+EthosUBackend::execute()handles.input.memcpy()"); // Sizes match and elt size matches so memcpy. // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate. +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + EthosUBackend_input_memcpy(tensor_in.nbytes()); +#endif arm_ethos_io_memcpy( scratch_addr, tensor_in.mutable_data_ptr(), @@ -422,6 +429,9 @@ Error copy_with_layout_adjustment( const char* src_bytes = src; for (size_t chunk_idx = 0; chunk_idx < chunk_count; ++chunk_idx) { // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate. +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + EthosUBackend_output_memcpy(chunk_size); +#endif arm_ethos_io_memcpy(dest, src_bytes, chunk_size); src_bytes += vela_chunk_size; dest += chunk_size; diff --git a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp index 82cbe99afad..f33a523986d 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp @@ -152,6 +152,9 @@ Error platform_execute( io_bytes_total += tensor_bytes; } else { // Routed through arm_ethos_io_memcpy so firmware can DMA-accelerate. +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + EthosUBackend_output_memcpy(tensor_bytes); +#endif arm_ethos_io_memcpy( tensor_out.mutable_data_ptr(), static_cast(output_addr), diff --git a/backends/arm/runtime/EthosUBackend_Internal.h b/backends/arm/runtime/EthosUBackend_Internal.h index 01069a9af72..5b3cc58858f 100644 --- a/backends/arm/runtime/EthosUBackend_Internal.h +++ b/backends/arm/runtime/EthosUBackend_Internal.h @@ -78,6 +78,10 @@ void EthosUBackend_execute_end(); void EthosUBackend_delegate_begin(const void* handle); void EthosUBackend_delegate_end(); #endif +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +void EthosUBackend_input_memcpy(size_t size); +void EthosUBackend_output_memcpy(size_t size); +#endif extern unsigned char* ethosu_fast_scratch; extern size_t ethosu_fast_scratch_size; } diff --git a/examples/arm/executor_runner/CMakeLists.txt b/examples/arm/executor_runner/CMakeLists.txt index 4280dde15d8..6313ceedbbc 100644 --- a/examples/arm/executor_runner/CMakeLists.txt +++ b/examples/arm/executor_runner/CMakeLists.txt @@ -57,6 +57,9 @@ option(ET_LOG_DUMP_OUTPUT "Dump output in log" ON) option(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING "Report Ethos-U PMU statistics per delegate" OFF ) +option(ET_ARM_ETHOSU_PROFILE_IO_COPIES + "Count Ethos-U backend I/O memcpy calls and bytes" OFF +) set(ET_ARM_ETHOSU_MAX_PROFILED_DELEGATES "16" CACHE STRING "Maximum delegates tracked by per-delegate profiling" @@ -363,6 +366,15 @@ if(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) ) endif() +if(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + target_compile_definitions( + arm_executor_runner PRIVATE ET_ARM_ETHOSU_PROFILE_IO_COPIES + ) + target_compile_definitions( + executorch_delegate_ethos_u PRIVATE ET_ARM_ETHOSU_PROFILE_IO_COPIES + ) +endif() + if(ET_LOG_DUMP_INPUT) target_compile_definitions(arm_executor_runner PUBLIC ET_LOG_DUMP_INPUT) endif() diff --git a/examples/arm/executor_runner/arm_perf_monitor.cpp b/examples/arm/executor_runner/arm_perf_monitor.cpp index 31fa5937fd3..4f75d93d0c7 100644 --- a/examples/arm/executor_runner/arm_perf_monitor.cpp +++ b/examples/arm/executor_runner/arm_perf_monitor.cpp @@ -40,6 +40,14 @@ uint64_t ethosu_ArmBackendExecuteCycleCount = 0; uint64_t ethosu_ArmWhenNPURunCycleCountStart = 0; uint64_t ethosu_ArmWhenNPURunCycleCount = 0; uint64_t ethosu_pmuCycleCount = 0; +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +struct IOCopyStats { + uint64_t calls = 0; + uint64_t bytes = 0; +}; +IOCopyStats ethosu_inputCopyStats; +IOCopyStats ethosu_outputCopyStats; +#endif #if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) struct DelegateStats { const void* handle = nullptr; @@ -47,6 +55,10 @@ struct DelegateStats { uint64_t npu_invocations = 0; uint64_t pmu_cycles = 0; std::array pmu_events{}; +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + IOCopyStats input_copies; + IOCopyStats output_copies; +#endif }; std::array @@ -80,6 +92,30 @@ static_assert(ETHOSU_PMU_NCOUNTERS >= ethosu_pmuCountersUsed); extern "C" { +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) +void EthosUBackend_input_memcpy(size_t size) { + ethosu_inputCopyStats.calls++; + ethosu_inputCopyStats.bytes += size; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->input_copies.calls++; + ethosu_activeDelegate->input_copies.bytes += size; + } +#endif +} + +void EthosUBackend_output_memcpy(size_t size) { + ethosu_outputCopyStats.calls++; + ethosu_outputCopyStats.bytes += size; +#if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) + if (ethosu_activeDelegate != nullptr) { + ethosu_activeDelegate->output_copies.calls++; + ethosu_activeDelegate->output_copies.bytes += size; + } +#endif +} +#endif + #if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) void EthosUBackend_delegate_begin(const void* handle) { ethosu_activeDelegate = get_delegate_stats(handle); @@ -197,6 +233,10 @@ void StartMeasurements() { ethosu_ArmBackendExecuteCycleCount = 0; ethosu_ArmWhenNPURunCycleCount = 0; ethosu_pmuCycleCount = 0; +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + ethosu_inputCopyStats = {}; + ethosu_outputCopyStats = {}; +#endif #if defined(ET_ARM_ETHOSU_PER_DELEGATE_PROFILING) ethosu_delegateStats = {}; ethosu_delegateCount = 0; @@ -234,6 +274,32 @@ void StopMeasurements(int num_inferences) { "ethos-u : cycle_cnt : %" PRIu64 " cycles (%.2f per inference)", ethosu_ArmBackendExecuteCycleCount, (double)ethosu_ArmBackendExecuteCycleCount / num_inferences); +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + const uint64_t io_copy_calls = + ethosu_inputCopyStats.calls + ethosu_outputCopyStats.calls; + const uint64_t io_copy_bytes = + ethosu_inputCopyStats.bytes + ethosu_outputCopyStats.bytes; + ET_LOG( + Info, + "Ethos-U IO copy calls: %" PRIu64 " (%.2f per inference)", + io_copy_calls, + (double)io_copy_calls / num_inferences); + ET_LOG( + Info, + "Ethos-U IO copy bytes: %" PRIu64 " bytes (%.2f per inference)", + io_copy_bytes, + (double)io_copy_bytes / num_inferences); + ET_LOG( + Info, + "Ethos-U input copy: %" PRIu64 " calls, %" PRIu64 " bytes", + ethosu_inputCopyStats.calls, + ethosu_inputCopyStats.bytes); + ET_LOG( + Info, + "Ethos-U output copy: %" PRIu64 " calls, %" PRIu64 " bytes", + ethosu_outputCopyStats.calls, + ethosu_outputCopyStats.bytes); +#endif // We could print a list of the cycles used by the other delegates here in the // future but now we only print ethos-u: this means that "Operator(s) total: // ..." will be the same number as ethos-u : cycle_cnt and not the sum of all @@ -318,6 +384,21 @@ void StopMeasurements(int num_inferences) { event, stats.pmu_events[event]); } +#if defined(ET_ARM_ETHOSU_PROFILE_IO_COPIES) + ET_LOG( + Info, + "Ethos-U delegate %zu input copy: %" PRIu64 " calls, %" PRIu64 " bytes", + delegate_id, + stats.input_copies.calls, + stats.input_copies.bytes); + ET_LOG( + Info, + "Ethos-U delegate %zu output copy: %" PRIu64 " calls, %" PRIu64 + " bytes", + delegate_id, + stats.output_copies.calls, + stats.output_copies.bytes); +#endif } if (ethosu_delegateCapacityExceeded) { ET_LOG( From 79b149f24cdaa934bfd37b5ddd764184c01a6a25 Mon Sep 17 00:00:00 2001 From: Michiel Olieslagers <44864547+Michiel-Olieslagers@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:54:08 +0100 Subject: [PATCH 138/190] Arm backend: Added real-world dataset for NSS. (#22650) * Added a real-world dataset for the NSS model for more realistic calibration and evaluation of the model. * Updated model_gym and NSS model revision to "main". HF weights are now also properly loaded into the model. Signed-off-by: Michiel Olieslagers Change-Id: I69da4f0c93340f89fc6d1797113a46ba15ed87e4 Assisted-by: Codex cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- backends/arm/quantizer/arm_quantizer.py | 6 +- backends/arm/requirements-arm-models-test.txt | 1 + .../generate_neural_graphics_test_data.py | 435 ++++++++++++++++++ .../arm/scripts/install_models_for_test.sh | 11 +- .../arm/scripts/neural_graphics_test_data.py | 138 ++++++ backends/arm/test/models/test_nss.py | 120 +++-- 6 files changed, 680 insertions(+), 31 deletions(-) create mode 100644 backends/arm/scripts/generate_neural_graphics_test_data.py create mode 100644 backends/arm/scripts/neural_graphics_test_data.py diff --git a/backends/arm/quantizer/arm_quantizer.py b/backends/arm/quantizer/arm_quantizer.py index 121db2e0c90..8a1171421fc 100644 --- a/backends/arm/quantizer/arm_quantizer.py +++ b/backends/arm/quantizer/arm_quantizer.py @@ -14,7 +14,7 @@ import functools import logging from contextlib import contextmanager -from typing import Any, Callable, Dict, List, Optional +from typing import Any, Callable, Dict, Iterable, List, Optional import torch from executorch.backends.arm._passes import ArmPassManager @@ -959,7 +959,7 @@ def validate(self, model: GraphModule) -> None: def _quantize_with_submodules( self, model: GraphModule, - calibration_samples: list[tuple], + calibration_samples: Iterable[tuple], is_qat: bool = False, fold_quantize: bool = True, ): @@ -971,7 +971,7 @@ def _quantize_with_submodules( Args: model (GraphModule): The model to quantize. - calibration_samples (list[tuple]): A list of inputs to used to + calibration_samples (Iterable[tuple]): Inputs used to calibrate the model during quantization. To properly calibrate a model with submodules, at least one sample per code path is needed. diff --git a/backends/arm/requirements-arm-models-test.txt b/backends/arm/requirements-arm-models-test.txt index c6a1d94aef2..f59928bae53 100644 --- a/backends/arm/requirements-arm-models-test.txt +++ b/backends/arm/requirements-arm-models-test.txt @@ -8,4 +8,5 @@ diffusers[torch] @ git+https://github.com/huggingface/diffusers.git@a7cb14efbe4b pydantic slangtorch rich +torcheval==0.0.7 setuptools==80.10.2 diff --git a/backends/arm/scripts/generate_neural_graphics_test_data.py b/backends/arm/scripts/generate_neural_graphics_test_data.py new file mode 100644 index 00000000000..67bfa2a8754 --- /dev/null +++ b/backends/arm/scripts/generate_neural_graphics_test_data.py @@ -0,0 +1,435 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Generate NSS autoencoder calibration and verification data.""" + +from __future__ import annotations + +import json +import os +import shutil +from importlib.resources import files +from pathlib import Path + +import torch +import torch.nn.functional as F +from executorch.backends.arm.scripts.neural_graphics_test_data import ( + nss_input_shape, + nss_test_calibration_path, + nss_test_data_root, + nss_test_verification_path, +) +from safetensors.torch import save_file + +os.environ.setdefault("HF_HUB_DISABLE_XET", "1") + +from huggingface_hub import snapshot_download +from ng_model_gym.core.config.config_model import ( # type: ignore[import-not-found,import-untyped] + ConfigModel, +) +from ng_model_gym.core.data.data_utils import ( # type: ignore[import-not-found,import-untyped] + DataLoaderMode, + DatasetType, + tonemap_forward, + ToneMapperMode, +) +from ng_model_gym.usecases.nss.data.dataset import ( # type: ignore[import-not-found,import-untyped] + NSSDataset, +) + + +_DATASET_REPO_ID = "Arm/neural-graphics-dataset" +_CALIBRATION_SOURCE_ALLOW_PATTERNS = [ + "train/**/*.safetensors", + "nss/train/**/*.safetensors", +] +_EVALUATION_SOURCE_ALLOW_PATTERNS = [ + "test/test_full_resolution_sample.safetensors", + "nss/test/test_full_resolution_sample.safetensors", +] + +EPS = 1e-7 +NSS_V1_SPATIAL_MULTIPLE = 8 + + +def _luminance(rgb: torch.Tensor) -> torch.Tensor: + weights = torch.tensor( + [0.2126, 0.7152, 0.0722], + dtype=rgb.dtype, + device=rgb.device, + ).view(1, 3, 1, 1) + return torch.sum(rgb * weights, dim=1, keepdim=True) + + +def _motion_detector( + motion_lr: torch.Tensor, render_size: torch.Tensor +) -> torch.Tensor: + # render_size is stored as [height, width], matching the dataset writer. + size = render_size.to(dtype=torch.float32).view(-1, 2, 1, 1) + motion_norm = motion_lr.to(dtype=torch.float32) / torch.clamp(size, min=1.0) + motion_length = torch.linalg.vector_norm(motion_norm, dim=1, keepdim=True) + + pix_min = torch.linalg.vector_norm( + 1.0 / torch.clamp(size, min=1.0), dim=1 + ).unsqueeze(1) + pix_max = torch.linalg.vector_norm( + 200.0 / torch.clamp(size, min=1.0), dim=1 + ).unsqueeze(1) + detector = (torch.clamp(motion_length, pix_min, pix_max) - pix_min) / torch.clamp( + pix_max - pix_min, min=EPS + ) + return torch.sqrt(torch.clamp(detector, min=0.0)) + + +def _depth_edge(depth: torch.Tensor) -> torch.Tensor: + dx = F.pad(torch.abs(depth[..., :, 1:] - depth[..., :, :-1]), (0, 1, 0, 0)) + dy = F.pad(torch.abs(depth[..., 1:, :] - depth[..., :-1, :]), (0, 0, 0, 1)) + return torch.clamp((dx + dy) * 100.0, 0.0, 1.0) + + +def _reflect_pad_to_multiple( + tensor: torch.Tensor, + multiple: int = NSS_V1_SPATIAL_MULTIPLE, +) -> torch.Tensor: + h, w = tensor.shape[-2:] + pad_h = (multiple - (h % multiple)) % multiple + pad_w = (multiple - (w % multiple)) % multiple + if pad_h == 0 and pad_w == 0: + return tensor + return F.pad(tensor, (0, pad_w, 0, pad_h), mode="reflect") + + +def _model_gym_dataset(src: Path) -> NSSDataset: + config_path = files("ng_model_gym.usecases.nss.configs").joinpath( + "nss_v1_template.json" + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + for split in ("train", "validation", "test"): + config["dataset"]["path"][split] = str(src) + config["dataset"].update( + exposure=None, + tonemapper=ToneMapperMode.KARIS.value, + gt_augmentation=False, + ) + params = ConfigModel.model_validate(config) + return NSSDataset(params, DataLoaderMode.TEST, DatasetType.SAFETENSOR) + + +def _make_autoencoder_input( + current: dict[str, torch.Tensor], + previous: dict[str, torch.Tensor] | None = None, +) -> torch.Tensor: + colour_tm = current["colour"] + exposure = current["exposure"] + _, _, h, w = colour_tm.shape + same_sequence = previous is not None and torch.equal( + current["seq"], previous["seq"] + ) + + history_linear = torch.zeros_like(current["colour_linear"]) + if previous is not None and same_sequence: + history_linear = previous["ground_truth_linear"] + if history_linear.shape[-2:] != (h, w): + history_linear = F.interpolate( + history_linear, + size=(h, w), + mode="bilinear", + align_corners=False, + ) + history_tm = tonemap_forward(history_linear * exposure, mode=ToneMapperMode.KARIS) + + motion_signal = _motion_detector(current["motion_lr"], current["render_size"]) + luma = _luminance(colour_tm) + previous_luma = torch.zeros_like(luma) + feedback = torch.zeros((1, 4, h, w), dtype=torch.float32) + if previous is not None and same_sequence: + previous_luma = _luminance(previous["colour"]) + previous_luma_derivative = torch.clamp(previous_luma, 0.0, 1.0) + feedback[:, 0:1] = _motion_detector( + previous["motion_lr"], previous["render_size"] + ) + feedback[:, 1:2] = previous_luma_derivative + feedback[:, 2:3] = previous_luma + feedback[:, 3:4] = _depth_edge(previous["depth"]) + luma_derivative = torch.clamp(torch.abs(luma - previous_luma), 0.0, 1.0) + + autoencoder_input = torch.cat( + [ + _reflect_pad_to_multiple(history_tm), + _reflect_pad_to_multiple(colour_tm), + _reflect_pad_to_multiple(motion_signal), + _reflect_pad_to_multiple(feedback), + _reflect_pad_to_multiple(luma_derivative), + ], + dim=1, + ) + return autoencoder_input.to(torch.float16) + + +def _autoencoder_sample(dataset: NSSDataset, index: int) -> torch.Tensor: + current = dataset[index][0] + previous = dataset[index - 1][0] if index > 0 else None + return _make_autoencoder_input(current, previous) + + +def _metadata( + src: Path, tensor: torch.Tensor, shard: int | None = None +) -> dict[str, str]: + metadata = { + "format": "nss_v1_autoencoder_calibration", + "source": str(src), + "samples": str(tensor.shape[0]), + "shape": json.dumps(list(tensor.shape)), + "spatial_multiple": str(NSS_V1_SPATIAL_MULTIPLE), + "preprocess": "cpu_approximation_of_nss_v1_slang_pre_process", + "channels": json.dumps( + [ + "history.r", + "history.g", + "history.b", + "colour.r", + "colour.g", + "colour.b", + "motion_detector", + "feedback.r", + "feedback.g", + "feedback.b", + "feedback.a", + "luma_derivative", + ] + ), + } + if shard is not None: + metadata["shard"] = str(shard) + return metadata + + +def _write_tensor( + src: Path, dst: Path, tensor: torch.Tensor, shard: int | None = None +) -> None: + dst.parent.mkdir(parents=True, exist_ok=True) + save_file( + {"input": tensor.contiguous()}, dst, metadata=_metadata(src, tensor, shard) + ) + + +def _remove_stale_shards(dst: Path) -> None: + if not dst.exists(): + return + if not dst.is_dir(): + raise NotADirectoryError(f"Expected shard output directory, got {dst}") + for path in dst.glob("*.safetensors"): + path.unlink() + + +def _generate_verification_dataset( + src: Path, + dst: Path, + num_samples: int, + shard_size: int, +) -> None: + dataset = _model_gym_dataset(src) + sample_limit = min(num_samples, len(dataset)) + _remove_stale_shards(dst) + dst.mkdir(parents=True, exist_ok=True) + + shard_idx = 0 + tensors: list[torch.Tensor] = [] + sources: list[Path] = [] + for index in range(sample_limit): + tensors.append(_autoencoder_sample(dataset, index)) + sources.append(dataset.frame_indexes[index][0]) + if len(tensors) < shard_size and index + 1 < sample_limit: + continue + shard = torch.cat(tensors, dim=0) + _write_tensor( + sources[0], + dst / f"{shard_idx:04d}.safetensors", + shard, + shard_idx, + ) + shard_idx += 1 + tensors.clear() + sources.clear() + + +def _raw_source_path() -> Path: + return nss_test_data_root() / "source" + + +def _raw_dataset_root(snapshot_path: Path) -> Path: + if (snapshot_path / "nss" / "train").is_dir() or ( + snapshot_path / "nss" / "test" + ).is_dir(): + return snapshot_path / "nss" + return snapshot_path + + +def _env_int(name: str, default: int) -> int: + value = os.environ.get(name) + if value is None: + return default + parsed = int(value) + if parsed <= 0: + raise ValueError(f"{name} must be positive.") + return parsed + + +def _has_safetensors(path: Path) -> bool: + return path.is_file() or (path.is_dir() and any(path.glob("*.safetensors"))) + + +def _download_raw_sources( + allow_patterns: list[str], *, force_download: bool = False +) -> Path: + snapshot = Path( + snapshot_download( + repo_id=_DATASET_REPO_ID, + repo_type="dataset", + revision="5039ce015d7c877980fad44f87893fc5ac0927e2", + allow_patterns=allow_patterns, + local_dir=_raw_source_path(), + force_download=force_download, + ) + ) + return _raw_dataset_root(snapshot) + + +def _delete_raw_split(raw_root: Path, split: str, keep_env: str) -> None: + if not os.environ.get(keep_env): + shutil.rmtree(raw_root / split, ignore_errors=True) + + +def _delete_raw_sources() -> None: + if os.environ.get("NSS_KEEP_RAW_TRAIN_DATA") or os.environ.get( + "NSS_KEEP_RAW_EVALUATION_DATA" + ): + return + shutil.rmtree(_raw_source_path(), ignore_errors=True) + + +def _has_test_calibration_samples( + path: Path, num_samples: int, spatial_size: tuple[int, int] +) -> bool: + if not path.is_dir() or len(list(path.glob("*.safetensors"))) != num_samples: + return False + return all( + nss_input_shape(file_path)[1:] == (12, *spatial_size) + for file_path in path.glob("*.safetensors") + ) + + +def ensure_generated_test_calibration_dataset( + num_samples: int = 3663, + spatial_size: tuple[int, int] = (128, 128), + force_download: bool = False, +) -> Path: + """Generate evenly distributed, test-ready NSS calibration samples.""" + + if num_samples <= 0: + raise ValueError("num_samples must be positive.") + + calibration_path = nss_test_calibration_path(num_samples, spatial_size) + if _has_test_calibration_samples(calibration_path, num_samples, spatial_size): + return calibration_path + + raw_root = _download_raw_sources( + _CALIBRATION_SOURCE_ALLOW_PATTERNS, force_download=force_download + ) + dataset = _model_gym_dataset(raw_root / "train") + total_samples = len(dataset) + if num_samples > total_samples: + raise ValueError( + f"Requested {num_samples} calibration samples, but only found " + f"{total_samples}." + ) + + _remove_stale_shards(calibration_path) + calibration_path.mkdir(parents=True, exist_ok=True) + sample_indices = ( + [0] + if num_samples == 1 + else [ + index * (total_samples - 1) // (num_samples - 1) + for index in range(num_samples) + ] + ) + for output_index, sample_index in enumerate(sample_indices): + tensor = _autoencoder_sample(dataset, sample_index) + if tensor.shape[-2:] != spatial_size: + tensor = F.interpolate( + tensor.to(torch.float32), + size=spatial_size, + mode="bilinear", + align_corners=False, + ).to(torch.float16) + _write_tensor( + dataset.frame_indexes[sample_index][0], + calibration_path / f"{output_index:04d}.safetensors", + tensor, + output_index, + ) + + _delete_raw_split(raw_root, "train", "NSS_KEEP_RAW_TRAIN_DATA") + return calibration_path + + +def ensure_generated_verification_dataset( + force_download: bool = False, +) -> Path: + """Generate the held-out NSS verification input without calibration data.""" + + verification_path = nss_test_verification_path() + if _has_safetensors(verification_path): + return verification_path + + raw_root = _download_raw_sources( + _EVALUATION_SOURCE_ALLOW_PATTERNS, force_download=force_download + ) + _generate_verification_dataset( + raw_root / "test", + verification_path, + _env_int("NSS_GENERATED_EVALUATION_SAMPLES", 1), + _env_int("NSS_GENERATED_EVALUATION_SHARD_SIZE", 10), + ) + _delete_raw_split(raw_root, "test", "NSS_KEEP_RAW_EVALUATION_DATA") + return verification_path + + +def ensure_generated_test_datasets( + calibration_samples: int = 3663, + spatial_size: tuple[int, int] = (128, 128), +) -> tuple[Path, Path]: + """Generate the NSS artifacts consumed directly by ``test_nss.py``.""" + + calibration_path = ensure_generated_test_calibration_dataset( + calibration_samples, spatial_size + ) + verification_path = ensure_generated_verification_dataset() + _delete_raw_sources() + return calibration_path, verification_path + + +def generate_test_datasets_from_scratch( + calibration_samples: int = 3663, + spatial_size: tuple[int, int] = (128, 128), +) -> tuple[Path, Path]: + """Download and regenerate the NSS artifacts consumed by ``test_nss.py``.""" + + calibration_path = nss_test_calibration_path(calibration_samples, spatial_size) + verification_path = nss_test_verification_path() + for path in (calibration_path, verification_path, _raw_source_path()): + if path.is_dir(): + shutil.rmtree(path) + elif path.exists(): + path.unlink() + + calibration_path = ensure_generated_test_calibration_dataset( + calibration_samples, + spatial_size, + force_download=True, + ) + verification_path = ensure_generated_verification_dataset(force_download=True) + _delete_raw_sources() + return calibration_path, verification_path diff --git a/backends/arm/scripts/install_models_for_test.sh b/backends/arm/scripts/install_models_for_test.sh index 1e91cd9c08f..c7239ee2760 100644 --- a/backends/arm/scripts/install_models_for_test.sh +++ b/backends/arm/scripts/install_models_for_test.sh @@ -8,7 +8,8 @@ set -e pip install -r backends/arm/requirements-arm-models-test.txt # Install model gym repository -MODEL_GYM_REF="${MODEL_GYM_REF:-v0.3.0}" +MODEL_GYM_REF="${MODEL_GYM_REF:-main}" +rm -rf neural-graphics-model-gym git clone --depth 1 --branch "$MODEL_GYM_REF" https://github.com/arm/neural-graphics-model-gym.git cd neural-graphics-model-gym # Remove model-converter installation from model-gym repository (to prevent overwriting executorch version) @@ -20,3 +21,11 @@ fi pip install . --no-deps cd .. rm -rf neural-graphics-model-gym + +# Prepare the fixed NSS artifacts before pytest. The calibration data is +# generated from raw 128x128 training frames; evaluation retains the +# deployment-resolution input. +python3 -c ' +from executorch.backends.arm.scripts.generate_neural_graphics_test_data import generate_test_datasets_from_scratch +generate_test_datasets_from_scratch() +' diff --git a/backends/arm/scripts/neural_graphics_test_data.py b/backends/arm/scripts/neural_graphics_test_data.py new file mode 100644 index 00000000000..e404394dca2 --- /dev/null +++ b/backends/arm/scripts/neural_graphics_test_data.py @@ -0,0 +1,138 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +"""Load generated NSS autoencoder calibration and verification data.""" + +from __future__ import annotations + +import os +from pathlib import Path +from typing import Iterator + +import torch +from safetensors import safe_open + + +NSS_DATASET_REVISION = "main" +_TEST_CALIBRATION_DIR = "calibration/nss_v1_autoencoder_cpu_calibration" +_VERIFICATION_DIR = "evaluation/nss_v1_autoencoder_cpu_evaluation" + + +def nss_test_data_root() -> Path: + if "NSS_GENERATED_DATASET_ROOT" in os.environ: + return Path(os.environ["NSS_GENERATED_DATASET_ROOT"]) + else: + return ( + Path(__file__).resolve().parents[1] + / "test" + / "models" + / "nss_data" + / NSS_DATASET_REVISION + ) + + +def nss_test_calibration_path( + num_samples: int = 3663, spatial_size: tuple[int, int] = (128, 128) +) -> Path: + """Return the preprocessed calibration dataset path used by NSS tests.""" + + height, width = spatial_size + return nss_test_data_root() / ( + f"{_TEST_CALIBRATION_DIR}_{num_samples}_{height}x{width}" + ) + + +def nss_test_verification_path() -> Path: + """Return the preprocessed verification dataset path used by NSS tests.""" + + return nss_test_data_root() / _VERIFICATION_DIR + + +def nss_input_shape(path: Path) -> tuple[int, int, int, int]: + """Return the shape of the ``input`` tensor in a safetensors file.""" + + with safe_open(path, framework="pt", device="cpu") as handle: + keys = set(handle.keys()) + if "input" not in keys: + raise KeyError(f"{path} does not contain an `input` tensor. Found {keys}.") + + shape = tuple(handle.get_slice("input").get_shape()) + + if len(shape) != 4: + raise ValueError(f"Expected NCHW `input`, got shape {shape} in {path}.") + if shape[1] != 12: + raise ValueError(f"Expected 12 NSS input channels, got shape {shape}.") + return shape # type: ignore[return-value] + + +def _load_input_slice(path: Path, start: int, stop: int) -> torch.Tensor: + if not path.is_file(): + raise FileNotFoundError(path) + + shape = nss_input_shape(path) + if start < 0 or stop <= start or stop > shape[0]: + raise ValueError(f"Invalid slice [{start}:{stop}] for shape {shape}.") + + with safe_open(path, framework="pt", device="cpu") as handle: + tensor = ( + handle.get_slice("input")[start:stop].to(dtype=torch.float32).contiguous() + ) + + return tensor.to(memory_format=torch.channels_last) + + +def _safetensor_files(path: Path) -> list[Path]: + if path.is_file(): + return [path] + if path.is_dir(): + files = sorted(path.glob("*.safetensors")) + if files: + return files + raise FileNotFoundError(f"No safetensors found at {path}") + + +def _load_sample(path: Path, start: int = 0) -> tuple[torch.Tensor]: + if start < 0: + raise ValueError("start must be non-negative.") + + skipped = 0 + for file_path in _safetensor_files(path): + file_samples = nss_input_shape(file_path)[0] + if start < skipped + file_samples: + local_index = start - skipped + return (_load_input_slice(file_path, local_index, local_index + 1),) + skipped += file_samples + + raise ValueError(f"Sample {start} is outside the {skipped} samples in {path}.") + + +def iter_calibration_samples( + path: Path, + *, + num_samples: int = 8, +) -> Iterator[tuple[torch.Tensor]]: + """Stream NSS calibration samples without retaining them in memory.""" + + if num_samples <= 0: + raise ValueError("num_samples must be positive.") + + files = _safetensor_files(path) + if num_samples > len(files): + raise ValueError( + f"Requested {num_samples} samples from {path}, but only found {len(files)}." + ) + + for file_path in files[:num_samples]: + yield (_load_input_slice(file_path, 0, 1),) + + +def load_verification_inputs( + path: Path | None = None, + *, + start: int = 0, +) -> tuple[torch.Tensor]: + """Load one held-out verification sample in ``example_inputs`` format.""" + + path = nss_test_verification_path() if path is None else path + return _load_sample(path, start) diff --git a/backends/arm/test/models/test_nss.py b/backends/arm/test/models/test_nss.py index 61523964021..f3f1e0d54eb 100644 --- a/backends/arm/test/models/test_nss.py +++ b/backends/arm/test/models/test_nss.py @@ -3,10 +3,17 @@ # 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 +from pathlib import Path from typing import Tuple import pytest import torch +from executorch.backends.arm.scripts.neural_graphics_test_data import ( + iter_calibration_samples, + load_verification_inputs, + nss_test_calibration_path, +) from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -19,12 +26,30 @@ from huggingface_hub import hf_hub_download -from ng_model_gym.usecases.nss.model.model_blocks import ( # type: ignore[import-not-found,import-untyped] +from ng_model_gym.usecases.nss.model.model_blocks_v1 import ( # type: ignore[import-not-found,import-untyped] AutoEncoderV1, ) +from torch.export import Dim input_t = Tuple[torch.Tensor] # Input x +_RELEASE_REFS = ( + os.environ.get("GITHUB_REF", ""), + os.environ.get("GITHUB_REF_NAME", ""), + os.environ.get("GITHUB_BASE_REF", ""), +) +_IS_FROZEN_RELEASE = any( + ref.removeprefix("refs/heads/").startswith("release/") for ref in _RELEASE_REFS +) +pytestmark = pytest.mark.skipif( + _IS_FROZEN_RELEASE, + reason="NSS tests depend on resources fetched from main.", +) + +_NSS_HEIGHT = 8 * Dim("_nss_height", min=16, max=68) +_NSS_WIDTH = 8 * Dim("_nss_width", min=16, max=120) +_NSS_QUANTIZATION_DYNAMIC_SHAPES = ({2: _NSS_HEIGHT, 3: _NSS_WIDTH},) + class NSS(torch.nn.Module): def __init__(self, *args, **kwargs): @@ -35,46 +60,87 @@ def __init__(self, *args, **kwargs): def nss() -> AutoEncoderV1: """Get an instance of NSS with weights loaded.""" - weights = hf_hub_download( + weights = hf_hub_download( # nosec B615 repo_id="Arm/neural-super-sampling", - filename="nss_v0.1.0_fp32.pt", - revision="2e9b606acd9fa25071825a12f0764f1c3bef9480", + filename="nss_v1_0_1_high_fp32.pt", + revision="main", ) - nss_model = NSS() - nss_model.load_state_dict( - torch.load(weights, map_location=torch.device("cpu"), weights_only=True), - strict=False, + checkpoint = torch.load( + weights, map_location=torch.device("cpu"), weights_only=True ) + state_dict = { + f"auto_encoder.{key.removeprefix('autoencoder.')}": value + for key, value in checkpoint["model_state_dict"].items() + } + + nss_model = NSS() + nss_model.load_state_dict(state_dict, strict=True) return nss_model.auto_encoder def example_inputs(): - return (torch.randn((1, 12, 544, 960)),) + return load_verification_inputs() + + +def random_inputs(): + return (torch.rand((1, 12, 544, 960)),) + + +input_test_data = { + "real_data": True, + "random_data": False, +} -def test_nss_tosa_FP(): +def _nss_calibration_path() -> Path: + path = nss_test_calibration_path() + if not path.exists(): + raise RuntimeError( + "NSS calibration data is prepared by " + "backends/arm/scripts/install_models_for_test.sh." + ) + return path + + +def _set_nss_calibration_samples(pipeline): + quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0] + quantize_stage.dynamic_shapes = _NSS_QUANTIZATION_DYNAMIC_SHAPES + quantize_stage.calibration_samples = iter_calibration_samples( + _nss_calibration_path(), num_samples=3663 + ) + return pipeline + + +@common.parametrize("use_real_data", input_test_data) +def test_nss_tosa_FP(use_real_data): pipeline = TosaPipelineFP[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, ) - pipeline.add_stage_after("export", pipeline.tester.dump_operator_distribution) + if use_real_data: + pipeline.add_stage_after("export", pipeline.tester.dump_operator_distribution) pipeline.run() -def test_nss_tosa_INT(): +@common.parametrize("use_real_data", input_test_data) +def test_nss_tosa_INT(use_real_data): + pipeline_kwargs = ( + {"frobenius_threshold": 0.32, "qtol": 12} if use_real_data else {"qtol": 7} + ) pipeline = TosaPipelineINT[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, - frobenius_threshold=None, - cosine_threshold=None, + **pipeline_kwargs, ) + if use_real_data: + _set_nss_calibration_samples(pipeline) pipeline.run() @@ -89,6 +155,7 @@ def test_nss_u55_INT(): run_on_fvp=True, use_to_edge_transform_and_lower=True, ) + _set_nss_calibration_samples(pipeline) pipeline.run() @@ -105,17 +172,16 @@ def test_nss_u85_INT(): run_on_fvp=True, use_to_edge_transform_and_lower=True, ) + _set_nss_calibration_samples(pipeline) pipeline.run() -@pytest.mark.xfail( - reason="[MLETORCH-1430]: Double types are not supported in buffers in MSL" -) @common.SkipIfNoModelConverter -def test_nss_vgf_FP(): +@common.parametrize("use_real_data", input_test_data) +def test_nss_vgf_FP(use_real_data): pipeline = VgfPipeline[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], use_to_edge_transform_and_lower=True, @@ -128,10 +194,11 @@ def test_nss_vgf_FP(): @common.SkipIfNoModelConverter -def test_nss_vgf_INT(): +@common.parametrize("use_real_data", input_test_data) +def test_nss_vgf_INT(use_real_data): pipeline = VgfPipeline[input_t]( nss().eval(), - example_inputs(), + example_inputs() if use_real_data else random_inputs(), aten_op=[], exir_op=[], symmetric_io_quantization=True, @@ -140,9 +207,8 @@ def test_nss_vgf_INT(): quantize=True, # Override tosa version to test INT-only path tosa_version="TOSA-1.0+INT", + qtol=12 if use_real_data else 7, ) + if use_real_data: + _set_nss_calibration_samples(pipeline) pipeline.run() - - -ModelUnderTest = nss().eval() -ModelInputs = example_inputs() From 88958119c41033e587e0c97c18890b26803e4d9a Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Thu, 10 Sep 2026 13:59:24 +0100 Subject: [PATCH 139/190] Arm backend: Rewrite static roll with slice and concatenate patterns (#22639) Preserve supported static aten.roll nodes until the Arm backend rewrites them from: result = torch.roll(input, shifts=shifts, dims=dims) to this pattern for each shift and dimension: shift = shift % size result = torch.cat( ( result.slice(dim, size - shift, size), result.slice(dim, 0, size - shift), ), dim=dim, ) Apply the pattern once per nonzero shift so multi-dimensional and repeated-dimension rolls retain PyTorch semantics. Treat roll as shared-qparam data movement so quantized inputs and outputs reuse their scale and zero point. Leave dynamic, flattened, zero-sized, no-op, and unsupported-dtype forms on the existing fallback path. Change-Id: I70ddb6cf77611d08c4cd1102ba59ea41f63784dc Signed-off-by: Yufeng Shi --- backends/arm/_passes/__init__.py | 1 + backends/arm/_passes/arm_pass_manager.py | 2 + backends/arm/_passes/decompose_roll_pass.py | 131 +++++++++++++ backends/arm/quantizer/arm_quantizer_utils.py | 1 + .../arm/quantizer/quantization_annotator.py | 1 + backends/arm/test/ops/test_roll.py | 119 ++++++++++++ .../test/passes/test_decompose_roll_pass.py | 181 ++++++++++++++++++ .../test/quantizer/test_generic_annotater.py | 11 ++ backends/arm/tosa/partitioner.py | 56 +++++- .../cortex_m/test/misc/test_portable_int8.py | 7 + .../source/backends/arm-vgf/VGF_op_support.md | 3 +- 11 files changed, 511 insertions(+), 2 deletions(-) create mode 100644 backends/arm/_passes/decompose_roll_pass.py create mode 100644 backends/arm/test/ops/test_roll.py create mode 100644 backends/arm/test/passes/test_decompose_roll_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 57fe9ff414f..925faa35ca8 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -95,6 +95,7 @@ from .decompose_quant_nodes import DecomposeQuantNodesPass # noqa from .decompose_remainder_pass import DecomposeRemainderPass # noqa from .decompose_rnn_pass import DecomposeRnnPass # noqa +from .decompose_roll_pass import DecomposeRollPass # noqa from .decompose_round_pass import DecomposeRoundPass # noqa from .decompose_sdpa_pass import DecomposeScaledDotProductAttentionPass # noqa from .decompose_sdpa_with_regular_softmax_pass import ( # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 3873a92a442..8d10cbbcf33 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -88,6 +88,7 @@ DecomposeQuantNodesPass, DecomposeRemainderPass, DecomposeRnnPass, + DecomposeRollPass, DecomposeRoundPass, DecomposeScaledDotProductAttentionPass, DecomposeSDPAWithRegularSoftmaxPass, @@ -595,6 +596,7 @@ def _tosa_pipeline( DecomposeSinhPass(), DecomposeSignPass(), DecomposeFlipPass(), + DecomposeRollPass(), DecomposeFloorDividePass(), DecomposeGeluPass(), DecomposeAddSubAlphaPass(), diff --git a/backends/arm/_passes/decompose_roll_pass.py b/backends/arm/_passes/decompose_roll_pass.py new file mode 100644 index 00000000000..060732d153e --- /dev/null +++ b/backends/arm/_passes/decompose_roll_pass.py @@ -0,0 +1,131 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from collections.abc import Sequence +from typing import Set, Type + +from executorch.backends.arm._passes import ArmOpTargetedPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +RollParameters = tuple[tuple[int, int, int], ...] + + +def get_static_roll_parameters( + input_shape: Sequence[object], shifts: object, dims: object +) -> RollParameters | None: + """Normalize a statically-decomposable roll. + + Args: + input_shape (Sequence[object]): Shape of the roll input. + shifts (object): Roll shifts from the operator arguments. + dims (object): Roll dimensions from the operator arguments. + + Returns: + RollParameters | None: Normalized ``(shift, dim, size)`` tuples, or + ``None`` when the roll cannot be decomposed statically. + + """ + if not input_shape or any( + type(size) is not int or size <= 0 for size in input_shape + ): + return None + if not isinstance(shifts, (list, tuple)) or not isinstance(dims, (list, tuple)): + return None + if not shifts or len(shifts) != len(dims): + return None + if any(type(value) is not int for value in (*shifts, *dims)): + return None + + rank = len(input_shape) + parameters: list[tuple[int, int, int]] = [] + for shift, dim in zip(shifts, dims): + if not -rank <= dim < rank: + return None + normalized_dim = dim % rank + dim_size = int(input_shape[normalized_dim]) + parameters.append((shift % dim_size, normalized_dim, dim_size)) + return tuple(parameters) + + +def can_decompose_roll( + input_shape: Sequence[object], shifts: object, dims: object +) -> bool: + """Return whether a roll can become a nonempty slice/concat graph. + + Args: + input_shape (Sequence[object]): Shape of the roll input. + shifts (object): Roll shifts from the operator arguments. + dims (object): Roll dimensions from the operator arguments. + + Returns: + bool: True when the roll has a supported static decomposition. + + """ + parameters = get_static_roll_parameters(input_shape, shifts, dims) + return parameters is not None and any(shift != 0 for shift, _, _ in parameters) + + +class DecomposeRollPass(ArmOpTargetedPass): + """Decompose a static ``aten.roll`` into slices and concatenation. + + For each nonzero ``(shift, dim)`` pair, normalize the shift and apply: + + shift = shift % dim_size + result = cat( + ( + slice_copy(result, dim, dim_size - shift, dim_size), + slice_copy(result, dim, 0, dim_size - shift), + ), + dim, + ) + + Rewrites are applied sequentially to support multiple and repeated + dimensions. + + """ + + _passes_required_after: Set[Type[ExportPass]] = set() + target_ops = {exir_ops.edge.aten.roll.default} + + def call_operator(self, op, args, kwargs, meta, updated=False): + if op not in self.target_ops: + return super().call_operator(op, args, kwargs, meta, updated) + + input_tensor = args[0] + shifts = args[1] + dims = args[2] if len(args) > 2 else () + parameters = get_static_roll_parameters(input_tensor.data.shape, shifts, dims) + if parameters is None or not any(shift != 0 for shift, _, _ in parameters): + raise ValueError("Expected a nonempty static roll decomposition") + + result = input_tensor + for shift, dim, dim_size in parameters: + if shift == 0: + continue + split = dim_size - shift + suffix = super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (result, dim, split, dim_size, 1), + {}, + meta, + updated=True, + ) + prefix = super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (result, dim, 0, split, 1), + {}, + meta, + updated=True, + ) + result = super().call_operator( + exir_ops.edge.aten.cat.default, + ([suffix, prefix], dim), + {}, + meta, + updated=True, + ) + return result diff --git a/backends/arm/quantizer/arm_quantizer_utils.py b/backends/arm/quantizer/arm_quantizer_utils.py index 8689fbdadec..daef262ccf6 100644 --- a/backends/arm/quantizer/arm_quantizer_utils.py +++ b/backends/arm/quantizer/arm_quantizer_utils.py @@ -451,6 +451,7 @@ class SharedQspecQuantizer(Quantizer, QuantizerReporterUser): torch.ops.aten.split_copy.Tensor, torch.ops.aten.tile.default, torch.ops.aten.flip.default, + torch.ops.aten.roll.default, torch.ops.aten.index_select.default, torch.ops.aten.index_put.default, torch.ops.aten.index_put_.default, diff --git a/backends/arm/quantizer/quantization_annotator.py b/backends/arm/quantizer/quantization_annotator.py index be0c7b7b453..5cc4d46fff4 100644 --- a/backends/arm/quantizer/quantization_annotator.py +++ b/backends/arm/quantizer/quantization_annotator.py @@ -611,6 +611,7 @@ def _get_fixed_qparams_qspec( torch.ops.aten.t_copy.default, torch.ops.aten.tile.default, torch.ops.aten.flip.default, + torch.ops.aten.roll.default, torch.ops.aten.chunk.default, torch.ops.aten.contiguous.default, torch.ops.aten.upsample_bilinear2d.vec, diff --git a/backends/arm/test/ops/test_roll.py b/backends/arm/test/ops/test_roll.py new file mode 100644 index 00000000000..9d3dfd43ed4 --- /dev/null +++ b/backends/arm/test/ops/test_roll.py @@ -0,0 +1,119 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Tuple + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + TosaPipelineFP, + TosaPipelineINT, + VgfPipeline, +) + + +input_t1 = Tuple[torch.Tensor] +aten_op = "torch.ops.aten.roll.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_roll_default" + +test_data_fp = { + "bev_cyclic_shift_fp32": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.float32), + (-2, -2), + (1, 2), + ), + "bev_cyclic_shift_fp16": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.float16), + (-2, -2), + (1, 2), + ), +} + +test_data_bf16 = { + "bev_cyclic_shift_bf16": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.bfloat16), + (-2, -2), + (1, 2), + ), +} + +test_data_quant = { + "bev_cyclic_shift": lambda: ( + torch.randn(1, 8, 8, 4, dtype=torch.float32), + (-2, -2), + (1, 2), + ), +} + + +class Roll(torch.nn.Module): + def __init__(self, shifts: tuple[int, ...], dims: tuple[int, ...]) -> None: + super().__init__() + self.shifts = shifts + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.roll(x, self.shifts, self.dims) + + +@common.parametrize("test_data", test_data_fp) +def test_roll_tosa_FP(test_data) -> None: + data, shifts, dims = test_data() + pipeline = TosaPipelineFP[input_t1](Roll(shifts, dims), (data,), aten_op, exir_op) + pipeline.count_tosa_ops({"SLICE": 4, "CONCAT": 2}) + pipeline.run() + + +@common.parametrize("test_data", test_data_bf16) +def test_roll_tosa_FP_bf16(test_data) -> None: + data, shifts, dims = test_data() + pipeline = TosaPipelineFP[input_t1]( + Roll(shifts, dims), + (data,), + aten_op, + exir_op, + tosa_extensions=["bf16"], + ) + pipeline.count_tosa_ops({"SLICE": 4, "CONCAT": 2}) + pipeline.run() + + +@common.parametrize("test_data", test_data_quant) +def test_roll_tosa_INT(test_data) -> None: + data, shifts, dims = test_data() + pipeline = TosaPipelineINT[input_t1](Roll(shifts, dims), (data,), aten_op, exir_op) + pipeline.count_tosa_ops({"SLICE": 4, "CONCAT": 2}) + pipeline.run() + + +@common.parametrize("test_data", test_data_fp | test_data_bf16) +@common.SkipIfNoModelConverter +def test_roll_vgf_no_quant(test_data) -> None: + data, shifts, dims = test_data() + pipeline = VgfPipeline[input_t1]( + Roll(shifts, dims), + (data,), + aten_op, + exir_op, + quantize=False, + run_on_vulkan_runtime=False, + ) + pipeline.run() + + +@common.parametrize("test_data", test_data_quant) +@common.SkipIfNoModelConverter +def test_roll_vgf_quant(test_data) -> None: + data, shifts, dims = test_data() + pipeline = VgfPipeline[input_t1]( + Roll(shifts, dims), + (data,), + aten_op, + exir_op, + quantize=True, + run_on_vulkan_runtime=False, + ) + pipeline.run() diff --git a/backends/arm/test/passes/test_decompose_roll_pass.py b/backends/arm/test/passes/test_decompose_roll_pass.py new file mode 100644 index 00000000000..510a6ed8c09 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_roll_pass.py @@ -0,0 +1,181 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch + +from executorch.backends.arm._passes import DecomposeRollPass +from executorch.backends.arm._passes.decompose_roll_pass import can_decompose_roll +from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.util._factory import create_partitioner +from executorch.backends.arm.vgf.compile_spec import VgfCompileSpec +from executorch.exir import EdgeCompileConfig, to_edge +from executorch.exir.dialects._ops import ops as exir_ops + + +class Roll(torch.nn.Module): + def __init__(self, shifts: tuple[int, ...], dims: tuple[int, ...]) -> None: + super().__init__() + self.shifts = shifts + self.dims = dims + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.roll(x, self.shifts, self.dims) + + +roll_cases = ( + ((2, 5), (1,), (1,), 1), + ((2, 3, 5), (-2,), (-1,), 1), + ((2, 5), (12,), (1,), 1), + ((1, 8, 8, 4), (-2, -2), (1, 2), 2), + ((1, 5, 3), (1, 2), (1, 1), 2), + ((5, 4), (5, 1), (0, 1), 1), +) + + +@pytest.mark.parametrize( + "shape,shifts,dims,expected_cats", + roll_cases, + ids=( + "positive", + "negative_shift_and_dim", + "oversized_shift", + "multiple_dims", + "repeated_dim", + "mixed_zero_shift", + ), +) +def test_decompose_roll( + shape: tuple[int, ...], + shifts: tuple[int, ...], + dims: tuple[int, ...], + expected_cats: int, +) -> None: + model = Roll(shifts, dims) + inputs = (torch.randn(shape),) + eager_output = model(*inputs) + edge = to_edge( + torch.export.export(model, inputs, strict=True), + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + preserve_ops=[torch.ops.aten.roll.default], + ), + ) + + edge = edge.transform([DecomposeRollPass()]) + graph = edge.exported_program().graph + targets = [node.target for node in graph.nodes if node.op == "call_function"] + + assert exir_ops.edge.aten.roll.default not in targets + assert targets.count(exir_ops.edge.aten.slice_copy.Tensor) == 2 * expected_cats + assert targets.count(exir_ops.edge.aten.cat.default) == expected_cats + assert torch.equal(edge.exported_program().module()(*inputs), eager_output) + + +@pytest.mark.parametrize( + "shape,shifts,dims", + ( + ((2, 4), (1,), ()), + ((2, 4), (4,), (1,)), + ((2, 0), (1,), (1,)), + ((2, 4), (1, 2), (1,)), + ((2, 4), (1,), (2,)), + ), + ids=("flat", "no_op", "zero_size", "mismatched_args", "invalid_dim"), +) +def test_roll_decomposition_guards( + shape: tuple[int, ...], shifts: tuple[int, ...], dims: tuple[int, ...] +) -> None: + assert not can_decompose_roll(shape, shifts, dims) + + +@pytest.mark.parametrize( + "dtype,compile_spec", + ( + (torch.float32, TosaCompileSpec("TOSA-1.0+FP")), + (torch.float16, TosaCompileSpec("TOSA-1.0+FP")), + (torch.bfloat16, TosaCompileSpec("TOSA-1.0+FP+bf16")), + (torch.float32, VgfCompileSpec()), + ), + ids=("tosa_fp32", "tosa_fp16", "tosa_bf16", "vgf_fp32"), +) +def test_partitioner_preserves_supported_roll( + dtype: torch.dtype, compile_spec: TosaCompileSpec | VgfCompileSpec +) -> None: + model = Roll((-2, -2), (1, 2)) + exported_program = torch.export.export( + model, (torch.randn(1, 8, 8, 4, dtype=dtype),), strict=True + ) + partitioner = create_partitioner(compile_spec) + preserved_ops, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert torch.ops.aten.roll.default in preserved_ops + assert filter_fn is not None and filter_fn(roll_node) + + +def test_partitioner_does_not_preserve_flat_roll() -> None: + class FlatRoll(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.roll(x, 2) + + exported_program = torch.export.export( + FlatRoll(), (torch.randn(2, 4),), strict=True + ) + partitioner = create_partitioner(VgfCompileSpec()) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert filter_fn is not None and not filter_fn(roll_node) + + +def test_partitioner_does_not_preserve_unquantized_tosa_int_roll() -> None: + model = Roll((1,), (1,)) + exported_program = torch.export.export(model, (torch.randn(2, 4),), strict=True) + partitioner = create_partitioner(TosaCompileSpec("TOSA-1.0+INT")) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert filter_fn is not None and not filter_fn(roll_node) + + +@pytest.mark.parametrize( + "dtype,compile_spec", + ( + (torch.float64, VgfCompileSpec()), + (torch.int32, VgfCompileSpec()), + (torch.bool, VgfCompileSpec()), + (torch.bfloat16, TosaCompileSpec("TOSA-1.0+FP")), + ), + ids=("float64", "int32", "bool", "bf16_without_extension"), +) +def test_partitioner_does_not_preserve_unsupported_dtype_roll( + dtype: torch.dtype, compile_spec: TosaCompileSpec | VgfCompileSpec +) -> None: + model = Roll((1,), (1,)) + exported_program = torch.export.export( + model, (torch.zeros(2, 4, dtype=dtype),), strict=True + ) + partitioner = create_partitioner(compile_spec) + _, filter_fn = partitioner.ops_to_not_decompose(exported_program) + roll_node = next( + node + for node in exported_program.graph.nodes + if node.target == torch.ops.aten.roll.default + ) + + assert filter_fn is not None and not filter_fn(roll_node) diff --git a/backends/arm/test/quantizer/test_generic_annotater.py b/backends/arm/test/quantizer/test_generic_annotater.py index e3d32a1372a..48cfdfa0110 100644 --- a/backends/arm/test/quantizer/test_generic_annotater.py +++ b/backends/arm/test/quantizer/test_generic_annotater.py @@ -145,6 +145,17 @@ def test_flip_tosa_INT(): ) +def test_roll_tosa_INT(): + check_annotation( + SingleOpModel( + torch.roll, + (torch.randn(2, 4),), + shifts=(1, 2), + dims=(0, 1), + ), + ) + + def test_concat_tosa_INT(): check_annotation( SingleOpModel( diff --git a/backends/arm/tosa/partitioner.py b/backends/arm/tosa/partitioner.py index 96c8286f664..87e207d5450 100644 --- a/backends/arm/tosa/partitioner.py +++ b/backends/arm/tosa/partitioner.py @@ -29,14 +29,16 @@ from executorch.backends.arm._passes.decompose_large_stride_maxpool2d_pass import ( can_decompose_large_stride_maxpool2d, ) +from executorch.backends.arm._passes.decompose_roll_pass import can_decompose_roll from executorch.backends.arm._passes.decompose_unsupported_bilinear_resize_pass import ( is_exact_tosa_boundary_bilinear_downscale, ) from executorch.backends.arm.common.arm_compile_spec import ArmCompileSpec from executorch.backends.arm.common.type import ensure_type -from executorch.backends.arm.constants import DQ_OPS, Q_OPS +from executorch.backends.arm.constants import DQ_OPS, MAX_RANK, Q_OPS from executorch.backends.arm.operator_support.tosa_supported_operators import ( + is_quantized, tosa_support_factory, ) from executorch.backends.arm.tosa.backend import TOSABackend @@ -123,6 +125,51 @@ def is_node_supported( return is_exact_tosa_boundary_bilinear_downscale(node, self.tosa_spec) +def _is_decomposable_roll_node( + node: torch.fx.Node, tosa_spec: TosaSpecification +) -> bool: + """Return whether backend preprocessing can decompose a roll node.""" + if node.target not in { + torch.ops.aten.roll.default, + exir_ops.edge.aten.roll.default, + }: + return False + if ( + tosa_spec.support_integer() + and not tosa_spec.support_float() + and not is_quantized(node) + ): + return False + input_node = ensure_type(torch.fx.Node, node.args[0]) + input_tensor = get_first_fake_tensor(input_node) + if not 0 < len(input_tensor.shape) <= MAX_RANK: + return False + if input_tensor.dtype not in {torch.float16, torch.float32} and not ( + input_tensor.dtype == torch.bfloat16 and tosa_spec.support_extension("bf16") + ): + return False + + dims = node.args[2] if len(node.args) > 2 else () + return can_decompose_roll(input_tensor.shape, node.args[1], dims) + + +class DecomposableRollSupported(OperatorSupportBase): + """Accept static rolls that backend preprocessing can decompose.""" + + def __init__(self, tosa_spec: TosaSpecification) -> None: + """Initialize the check with the active TOSA specification.""" + self.tosa_spec = tosa_spec + + def is_node_supported( + self, + submodules: Mapping[str, torch.nn.Module], + node: torch.fx.Node, + ) -> bool: + """Return True when backend preprocessing can decompose the roll.""" + del submodules + return _is_decomposable_roll_node(node, self.tosa_spec) + + def _is_custom_partition_op( custom_ops: set[torch._ops.OpOverload], target: object ) -> bool: @@ -693,6 +740,7 @@ def _create_operator_support( additional_positive_checks=[ self._decomposable_resize_support, DecomposableLargeStrideMaxPool2dForU55Supported(self.tosa_spec), + DecomposableRollSupported(self.tosa_spec), ], ) @@ -782,6 +830,9 @@ def ops_to_not_decompose( # noqa: C901 ops_to_not_decompose_always = { torch.ops.aten.logit.default, } + ops_to_not_decompose_conditionally = { + torch.ops.aten.roll.default, + } ops_to_not_decompose_if_integer = { torch.ops.aten.eye.default, torch.ops.aten.linspace.default, @@ -792,6 +843,7 @@ def ops_to_not_decompose( # noqa: C901 | ops_to_not_decompose_if_quant_op | ops_to_not_decompose_if_fp | ops_to_not_decompose_if_integer + | ops_to_not_decompose_conditionally ) if not self.tosa_spec.is_U55_subset: @@ -824,6 +876,8 @@ def filter_fn(node: torch.fx.Node) -> bool: and get_first_fake_tensor(node).dtype == torch.float64 ): return False + if node.target in ops_to_not_decompose_conditionally: + return _is_decomposable_roll_node(node, self.tosa_spec) if ( self.tosa_spec.support_float() and node.target in ops_to_not_decompose_if_fp diff --git a/backends/cortex_m/test/misc/test_portable_int8.py b/backends/cortex_m/test/misc/test_portable_int8.py index 41f7f254863..ce0aeeb1de6 100644 --- a/backends/cortex_m/test/misc/test_portable_int8.py +++ b/backends/cortex_m/test/misc/test_portable_int8.py @@ -283,6 +283,12 @@ def _quantize_and_export( (torch.randn(2, 3, 4, 5), torch.randn(2, 3, 4, 5)), None, ), + "roll": OpCase( + torch.ops.aten.roll.default, + _build_module(lambda x, y: torch.ops.aten.roll.default(x, [1], [2])), + (torch.randn(2, 3, 4, 5), torch.randn(2, 3, 4, 5)), + None, + ), "index_select": OpCase( torch.ops.aten.index_select.default, _build_module( @@ -708,6 +714,7 @@ def _quantize_and_export( "where_default": "MLETORCH-1865: Properly support flaky scalar comparison ops.", "while_loop": "MLETORCH-1866: Support higher-order operators", "cond": "MLETORCH-1866: Support higher-order operators", + "roll": "MLETORCH-2563: Support roll operator", } diff --git a/docs/source/backends/arm-vgf/VGF_op_support.md b/docs/source/backends/arm-vgf/VGF_op_support.md index 7ae696a7642..092d9668419 100644 --- a/docs/source/backends/arm-vgf/VGF_op_support.md +++ b/docs/source/backends/arm-vgf/VGF_op_support.md @@ -6,7 +6,7 @@ This page lists VGF-supported PyTorch APIs and the dtype and quantization modes `8x8` means 8-bit activations and 8-bit weights. `16x8` means 16-bit activations and 8-bit weights. `8x4` means 8-bit activations and 4-bit weights. -Total supported PyTorch APIs: **154**. +Total supported PyTorch APIs: **155**. | PyTorch API | Support profile | DType | Quantization mode | | --- | --- | --- | --- | @@ -125,6 +125,7 @@ Total supported PyTorch APIs: **154**. | `torch.relu` / `torch.nn.ReLU` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.remainder` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.repeat_interleave` | FP, INT | `FP16`, `BF16`, `INT8`, `INT16`, `BOOL` | 8x8, 16x8 | +| `torch.roll` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.round` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.rsqrt` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.rsub` | FP | `FP32` | - | From f4f5e8927bbdf1ca12b556c2b22219417ee943a6 Mon Sep 17 00:00:00 2001 From: Jeremy Schoemaker Date: Thu, 10 Sep 2026 08:40:10 -0500 Subject: [PATCH 140/190] Preserve node output dtype when lifting scalars to attrs in XNNPACK quantizer (#22065) Fixes #22062 `_convert_scalars_to_attrs` in `backends/xnnpack/quantizer/xnnpack_quantizer_utils.py` lifts every scalar argument of `aten.add.Tensor` and `aten.mul.Tensor` to a buffer created with `torch.tensor(float(arg))`, which is always float32. `XNNPACKQuantizer.transform_for_annotation` runs this unconditionally, so an integer chain such as `x[:, torch.arange(4) + 0]` (the position-ids pattern emitted by Hugging Face models) has its int64 add promoted to float32, and running the `prepare_pt2e` output fails with `IndexError: tensors used as indices must be long, int, byte or bool tensors`. The fix creates the lifted constant with the node's own output dtype, `torch.tensor(args[i], dtype=n.meta["val"].dtype)`, which the function already relies on for `fake_mode`. Float behavior is unchanged. Added `test_int64_scalar_add_used_as_index` to `backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py`: it exports the index pattern above, runs `transform_for_annotation`, asserts the lifted constant keeps dtype int64, and executes the prepared module. The test fails with the float32 dtype assertion before this change and passes after. Neighboring scalar tests (`test_add_mul_scalar`, `test_add_mul_long`, `test_mul_float32_max`) still pass. flake8 and ufmt clean on both files. cc @GregoryComer @digantdesai @cbilgin @JakeStevens Co-authored-by: Jacob Stevens --- .../quantizer/xnnpack_quantizer_utils.py | 6 ++--- .../test/quantizer/test_xnnpack_quantizer.py | 24 +++++++++++++++++++ 2 files changed, 27 insertions(+), 3 deletions(-) diff --git a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py index 751388d9221..78a989bd75c 100644 --- a/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py +++ b/backends/xnnpack/quantizer/xnnpack_quantizer_utils.py @@ -1160,15 +1160,15 @@ def _convert_scalars_to_attrs(model: torch.fx.GraphModule) -> torch.fx.GraphModu prefix = "_tensor_constant_" get_new_attr_name = get_new_attr_name_with_prefix(prefix) tensor_constant_name = get_new_attr_name(model) - float_tensor = torch.tensor(float(args[i])) - model.register_buffer(tensor_constant_name, float_tensor) + scalar_tensor = torch.tensor(args[i], dtype=n.meta["val"].dtype) + model.register_buffer(tensor_constant_name, scalar_tensor) fake_mode = n.meta["val"].fake_mode with model.graph.inserting_before(n): get_attr_node = model.graph.create_node( "get_attr", tensor_constant_name, (), {} ) get_attr_node.meta["val"] = fake_mode.from_tensor( - float_tensor, static_shapes=True + scalar_tensor, static_shapes=True ) new_args.append(get_attr_node) n.args = tuple(new_args) diff --git a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py index 1e1a473dd59..27d6a8f65fb 100644 --- a/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py +++ b/backends/xnnpack/test/quantizer/test_xnnpack_quantizer.py @@ -1122,6 +1122,30 @@ def forward(self, x): node_list, ) + def test_int64_scalar_add_used_as_index(self): + """Scalars lifted to attrs must keep the op's output dtype; an int64 + add chain used as an index must not be promoted to float32.""" + + class M(torch.nn.Module): + def forward(self, x): + return x[:, torch.arange(4) + 0] + + quantizer = XNNPACKQuantizer() + quantization_config = get_symmetric_quantization_config(is_per_channel=True) + quantizer.set_global(quantization_config) + example_inputs = (torch.randn(1, 4, 5),) + m = export(M(), example_inputs, strict=True).module() + m = quantizer.transform_for_annotation(m) + lifted_constants = [ + m.get_buffer(n.target) + for n in m.graph.nodes + if n.op == "get_attr" and n.target.startswith("_tensor_constant_") + ] + self.assertEqual(len(lifted_constants), 1) + self.assertEqual(lifted_constants[0].dtype, torch.int64) + m = prepare_pt2e(m, quantizer) + m(*example_inputs) + def test_cat_same_node(self): """Ensure that concatenating the same node does not cause any unexpected behavior""" From b257c7140f82dd4316aa74ddcc0c229988af9d1e Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Thu, 10 Sep 2026 07:05:52 -0700 Subject: [PATCH 141/190] Publish CUDA 13.4 nightly wheels (#22653) ## Problem CUDA 13.4 consumers need a matching ExecuTorch wheel. The publication filter, build detection, dependencies, and compatibility environment must all support it. ## Fix Publish `cu134` when the shared generator offers it. Distinguish a missing CUDA train from an offered train with incomplete Python coverage. Add CUDA 13.4 build detection and architecture mappings. Select matching CUDA 13.4 nightly dependencies and retain them through example installation. Keep package metadata consistent with source-pinned torch and explicit torchao source builds. Preserve existing dependency selection for other CUDA versions and the CPU torchao variant on ARM. Run CUDA 13.4 compatibility with a newer driver and C++ runtime. The image's older Conda runtime aborts during thread-local cleanup after matrix multiplication. Update that runtime normally and verify the GPU result against CPU output, including normal process termination. Preserve CUDA 12.6 and 13.0 coverage. Pin the reusable workflow and its checked-out actions to the same prerequisite revision from https://github.com/pytorch/test-infra/pull/8771. That prerequisite still needs approval. Update the stale missing-weights regression to decode current metadata and exercise the legacy payload with and without constants. ## Testing The CI-script suite passed 327 tests with 15 skipped. Regressions exercise dependency selection, incorrect GPU results, and runtime-update failure. Scoped Python, shell, and workflow checks passed. The final revision passed CUDA 12.6, 13.0, and 13.4 compatibility. CUDA 13.4 completed installation, dependency checks, a GPU result comparison against CPU, and normal process termination. Both CUDA unit suites passed, including the missing-weights regressions. All ten final CUDA 13.4 wheel builds and uploads passed for Python 3.10 through 3.14 on x86_64 and ARM64. The x86_64 Python 3.12 smoke test ran delegated models with and without external weights and matched eager output. ARM64 smoke tests check packaging without GPU execution. All remaining CI checks finished without failures. The prerequisite's CI passed. Its approval is the only remaining merge-readiness gate. Assisted by Devmate. --------- Co-authored-by: PyTorch Bot --- .ci/scripts/test-cuda-build.sh | 11 +- .ci/scripts/tests/test_cu134_dependencies.py | 264 ++++++++++++++++++ .ci/scripts/tests/test_cuda_workflow.py | 114 ++++++++ .ci/scripts/tests/test_filter_cuda_matrix.py | 113 ++++++-- .ci/scripts/wheel/cuda_arch_list.sh | 4 + .github/scripts/filter_cuda_matrix.py | 47 +--- .github/workflows/cuda.yml | 17 +- .../cuda/tests/test_missing_weights_blob.py | 52 ++-- install_requirements.py | 64 ++++- install_utils.py | 1 + setup.py | 5 + 11 files changed, 592 insertions(+), 100 deletions(-) create mode 100644 .ci/scripts/tests/test_cu134_dependencies.py diff --git a/.ci/scripts/test-cuda-build.sh b/.ci/scripts/test-cuda-build.sh index e717718be66..07e2a50b1a0 100755 --- a/.ci/scripts/test-cuda-build.sh +++ b/.ci/scripts/test-cuda-build.sh @@ -49,10 +49,16 @@ import executorch print('SUCCESS: ExecuTorch imported successfully') " + python -m pip check + # Test CUDA availability and show details - python -c " + EXPECTED_CUDA_VERSION="$cuda_version" python -c " try: + import os import torch + assert torch.version.cuda == os.environ['EXPECTED_CUDA_VERSION'], ( + torch.version.cuda, os.environ['EXPECTED_CUDA_VERSION'] + ) print('INFO: PyTorch version:', torch.__version__) print('INFO: CUDA available:', torch.cuda.is_available()) @@ -68,7 +74,8 @@ try: x = torch.randn(10, 10).to(device) y = torch.randn(10, 10).to(device) z = torch.mm(x, y) - print('SUCCESS: CUDA tensor operation completed on device:', z.device) + torch.testing.assert_close(z.cpu(), x.cpu() @ y.cpu()) + print('SUCCESS: CUDA tensor operation matched CPU on device:', z.device) print('INFO: Result tensor shape:', z.shape) print('SUCCESS: ExecuTorch CUDA integration verified') diff --git a/.ci/scripts/tests/test_cu134_dependencies.py b/.ci/scripts/tests/test_cu134_dependencies.py new file mode 100644 index 00000000000..663cdb3b3bf --- /dev/null +++ b/.ci/scripts/tests/test_cu134_dependencies.py @@ -0,0 +1,264 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import ast +import importlib.util +import os +import subprocess +import sys +import unittest +from pathlib import Path +from unittest.mock import patch + +from packaging.requirements import Requirement + +ROOT = Path(__file__).resolve().parents[3] + + +def load_module(name): + spec = importlib.util.spec_from_file_location(name, ROOT / f"{name}.py") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +class TestCu134Dependencies(unittest.TestCase): + def setUp(self): + self.utils = load_module("install_utils") + self.modules = patch.dict(sys.modules, {"install_utils": self.utils}) + self.modules.start() + self.addCleanup(self.modules.stop) + self.installer = load_module("install_requirements") + + def install_commands(self, cuda, machine="x86_64", nightly=True, system="Linux"): + self.utils.determine_torch_url.cache_clear() + with ( + patch.dict(os.environ, {}, clear=True), + patch.object( + self.utils, + "_get_cuda_version", + return_value=cuda, + side_effect=RuntimeError("no nvcc") if cuda is None else None, + ), + patch.object(self.installer.platform, "machine", return_value=machine), + patch.object(self.installer.platform, "system", return_value=system), + patch.object(self.installer.sys, "platform", "linux"), + patch.object(self.installer.subprocess, "run") as run, + ): + self.installer.install_requirements(nightly) + self.installer.install_optional_example_requirements(nightly) + return [call.args[0] for call in run.call_args_list] + + def test_all_install_steps_preserve_exact_cu134_selection(self): + for machine, ao_variant in (("x86_64", "cu134"), ("aarch64", "cpu")): + with self.subTest(machine=machine): + commands = self.install_commands((13, 4), machine) + self.assertEqual(len(commands), 4) + expected = { + "torch==2.14.0.dev20260810+cu134", + "torchvision==0.29.0.dev20260811+cu134", + "torchaudio==2.11.0.dev20260811+cu134", + f"torchao==0.19.0.dev20260811+{ao_variant}", + } + for index, command in enumerate(commands): + required = ( + expected + if index >= 2 + else { + requirement + for requirement in expected + if requirement.startswith(("torch==", "torchao==")) + } + ) + self.assertTrue(required.issubset(command), command) + if index < 2: + self.assertFalse( + any( + arg.startswith(("torchvision", "torchaudio")) + for arg in command + ) + ) + self.assertIn( + "https://download.pytorch.org/whl/nightly/cu134", command + ) + self.assertNotIn( + "https://download.pytorch.org/whl/test/cu134", command + ) + self.assertNotIn("--no-deps", command) + if machine == "aarch64": + self.assertIn( + "https://download.pytorch.org/whl/nightly/cpu", command + ) + + def test_other_cuda_trains_keep_existing_pins(self): + for cuda in ((12, 6), (13, 0), (13, 2)): + for machine in ("x86_64", "aarch64"): + with self.subTest(cuda=cuda, machine=machine): + core, local, domains, examples = self.install_commands( + cuda, machine + ) + self.assertIn("torch==2.14.0", core) + self.assertIn("torchao==0.18.0.dev20260729", core) + self.assertIn("torchvision==0.29.0", domains) + self.assertIn("torchaudio==2.11.0", domains) + self.assertFalse(any("==" in arg for arg in local)) + self.assertFalse(any("==" in arg for arg in examples)) + + def test_source_pinned_torch_is_not_replaced(self): + for cuda in ((13, 2), (13, 4)): + with self.subTest(cuda=cuda): + core, _, domains, _ = self.install_commands(cuda, nightly=False) + self.assertIn("torch", core) + self.assertNotIn("torch==2.14.0.dev20260810+cu134", core) + self.assertIn("torchvision", domains) + self.assertIn("torchaudio", domains) + + def test_no_cuda_keeps_default_pins(self): + core, _, domains, _ = self.install_commands(None) + self.assertIn("torch==2.14.0", core) + self.assertIn("torchao==0.18.0.dev20260729", core) + self.assertIn("torchvision==0.29.0", domains) + self.assertIn("https://download.pytorch.org/whl/test/cpu", core) + + def test_windows_does_not_select_cu134(self): + core, _, domains, _ = self.install_commands((13, 4), system="Windows") + self.assertIn("torch==2.14.0", core) + self.assertIn("torchvision==0.29.0", domains) + self.assertIn("https://download.pytorch.org/whl/test/cpu", core) + + def test_failure_is_not_retried_with_another_cuda_train(self): + with ( + patch.object(self.utils, "_get_cuda_version", return_value=(13, 4)), + patch.object(self.installer.platform, "system", return_value="Linux"), + patch.object( + self.installer.subprocess, + "run", + side_effect=subprocess.CalledProcessError(1, "pip"), + ) as run, + ): + with self.assertRaises(subprocess.CalledProcessError): + self.installer.install_requirements(True) + self.assertEqual(run.call_count, 1) + + def torchao_requirement(self): + path = ROOT / "setup.py" + tree = ast.parse(path.read_text()) + function = next( + node + for node in tree.body + if isinstance(node, ast.FunctionDef) and node.name == "_torchao_requirement" + ) + namespace = { + "__file__": str(path), + "Path": Path, + "importlib": importlib, + "sys": sys, + "install_utils": self.utils, + } + exec( + compile(ast.Module(body=[function], type_ignores=[]), str(path), "exec"), + namespace, + ) + return namespace["_torchao_requirement"]() + + def test_package_install_preserves_source_pinned_torchao(self): + with patch.dict(sys.modules, {"install_requirements": self.installer}): + package_installer = load_module("install_executorch") + for machine in ("x86_64", "aarch64"): + with self.subTest(machine=machine): + self.utils.determine_torch_url.cache_clear() + with ( + patch.dict(os.environ, {}, clear=True), + patch.object(self.utils, "_get_cuda_version", return_value=(13, 4)), + patch.object( + self.installer.platform, "machine", return_value=machine + ), + patch.object( + self.installer.platform, "system", return_value="Linux" + ), + patch.object(self.installer.sys, "platform", "linux"), + patch.object( + sys, + "argv", + ["install_executorch", "--use-pt-pinned-commit", "--minimal"], + ), + patch.object( + package_installer, "python_is_compatible", return_value=True + ), + patch.object(package_installer, "check_and_update_submodules"), + patch.object(self.installer.subprocess, "run") as run, + ): + package_installer.main([]) + commands = [call.args[0] for call in run.call_args_list] + metadata = Requirement(self.torchao_requirement()) + self.assertEqual(len(commands), 3) + self.assertIn(".", commands[-1]) + core = commands[0] + torchao = Requirement( + next(arg for arg in core if arg.startswith("torchao==")) + ) + version = next(iter(torchao.specifier)).version + self.assertIn(version, metadata.specifier) + self.assertIn("torch", core) + self.assertFalse(any(arg.startswith("torch==") for arg in core)) + + def test_cu134_keeps_explicit_torchao_source_build(self): + with patch.dict(sys.modules, {"install_requirements": self.installer}): + package_installer = load_module("install_executorch") + for source_flag in ( + "EXECUTORCH_BUILD_KERNELS_TORCHAO", + "TORCHAO_BUILD_EXPERIMENTAL_MPS", + ): + with self.subTest(source_flag=source_flag): + self.utils.determine_torch_url.cache_clear() + with ( + patch.dict(os.environ, {source_flag: "1"}, clear=True), + patch.object(self.utils, "_get_cuda_version", return_value=(13, 4)), + patch.object( + self.installer.platform, "system", return_value="Linux" + ), + patch.object(self.installer.sys, "platform", "linux"), + patch.object(sys, "argv", ["install_executorch"]), + patch.object( + package_installer, "python_is_compatible", return_value=True + ), + patch.object(package_installer, "check_and_update_submodules"), + patch.object(self.installer.subprocess, "run") as run, + ): + package_installer.main([]) + metadata = Requirement(self.torchao_requirement()) + commands = [call.args[0] for call in run.call_args_list] + self.assertEqual(len(commands), 5) + self.assertIn("third-party/ao", commands[1]) + self.assertIn(".", commands[2]) + for command in commands: + self.assertFalse( + any(arg.startswith("torchao==") for arg in command) + ) + self.assertIn("torch==2.14.0.dev20260810+cu134", commands[-1]) + self.assertIn("0.18.0+git03ca489", metadata.specifier) + + def test_wheel_torchao_bound_matches_selected_train(self): + for cuda, expected in ( + ((13, 4), "torchao>=0.19.0.dev20260811,<0.20"), + ((13, 2), "torchao>=0.18.0.dev20260729,<0.19"), + (None, "torchao>=0.18.0.dev20260729,<0.19"), + ): + self.utils.determine_torch_url.cache_clear() + with ( + patch.object( + self.utils, + "_get_cuda_version", + return_value=cuda, + side_effect=RuntimeError("no nvcc") if cuda is None else None, + ), + patch.object(self.installer.platform, "system", return_value="Linux"), + ): + self.assertEqual(self.torchao_requirement(), expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/.ci/scripts/tests/test_cuda_workflow.py b/.ci/scripts/tests/test_cuda_workflow.py index 7eeb5e1e393..0ce37edc917 100644 --- a/.ci/scripts/tests/test_cuda_workflow.py +++ b/.ci/scripts/tests/test_cuda_workflow.py @@ -4,6 +4,9 @@ # 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 sys import unittest from pathlib import Path @@ -29,6 +32,117 @@ def _model_quant(entry): class CudaWorkflowTest(unittest.TestCase): + def test_build_matrix_preserves_existing_cuda_versions(self): + job = WORKFLOW["jobs"]["test-cuda-builds"] + self.assertEqual( + job["strategy"]["matrix"]["cuda-version"], ["12.6", "13.0", "13.4"] + ) + self.assertEqual(job["with"]["gpu-arch-version"], "${{ matrix.cuda-version }}") + + def test_cuda134_driver_uses_matching_workflow_and_action_revision(self): + job = WORKFLOW["jobs"]["test-cuda-builds"] + workflow, revision = job["uses"].split("@") + self.assertEqual( + workflow, "pytorch/test-infra/.github/workflows/linux_job_v2.yml" + ) + self.assertRegex(revision, r"^[0-9a-f]{40}$") + self.assertEqual(job["with"]["test-infra-ref"], revision) + self.assertEqual( + job["with"]["driver-version"], + "${{ matrix.cuda-version == '13.4' && '615.71.09' || '580.65.06' }}", + ) + self.assertEqual( + job["with"]["driver-download-url"], + "${{ matrix.cuda-version == '13.4' && " + "'https://download.nvidia.com/XFree86/Linux-x86_64/615.71.09/" + "NVIDIA-Linux-x86_64-615.71.09.run' || '' }}", + ) + + def test_cuda134_runtime_update_precedes_build_and_propagates_failure(self): + script = WORKFLOW["jobs"]["test-cuda-builds"]["with"]["script"] + stubs = """ +conda() { printf 'CONDA %s\n' "$*"; return "$CONDA_STATUS"; } +source() { printf 'BUILD %s\n' "$*"; } +""" + for version, conda_status in ( + ("12.6", 0), + ("13.0", 0), + ("13.4", 0), + ("13.4", 1), + ): + with self.subTest(version=version, conda_status=conda_status): + result = subprocess.run( + [ + "bash", + "-c", + stubs + script.replace("${{ matrix.cuda-version }}", version), + ], + env={**os.environ, "CONDA_STATUS": str(conda_status)}, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, conda_status, result.stderr) + expected = [] + if version == "13.4": + expected.append( + "CONDA install -y -n base -c conda-forge " + "libstdcxx-ng=16.2.0 libgcc-ng=16.2.0" + ) + if conda_status == 0: + expected.append(f"BUILD .ci/scripts/test-cuda-build.sh {version}") + self.assertEqual(result.stdout.splitlines(), expected) + + def test_cuda_probe_checks_the_result_and_torch_train(self): + script = (ROOT / ".ci/scripts/test-cuda-build.sh").read_text() + probes = [block.split('\n"', 1)[0] for block in script.split('python -c "')[1:]] + probe = next(block for block in probes if "import torch" in block) + fake_torch = """ +import os +import sys +from types import SimpleNamespace +class Tensor: + device = 'cuda' + shape = (10, 10) + def to(self, device): + return self + def cpu(self): + return self + def __matmul__(self, other): + return self +def assert_close(actual, expected): + print('RESULT CHECKED') + assert os.environ['INVALID_CUDA_RESULT'] == '0', 'CUDA result mismatch' +sys.modules['torch'] = SimpleNamespace( + __version__='test', version=SimpleNamespace(cuda='13.4'), + cuda=SimpleNamespace( + is_available=lambda: True, device_count=lambda: 1, + current_device=lambda: 0, get_device_name=lambda: 'test', + ), + device=lambda name: name, randn=lambda *args: Tensor(), + mm=lambda x, y: Tensor(), + testing=SimpleNamespace(assert_close=assert_close), +) +""" + for expected, invalid_result, succeeds in ( + ("13.4", "0", True), + ("13.0", "0", False), + ("13.4", "1", False), + ): + with self.subTest(expected=expected, invalid_result=invalid_result): + result = subprocess.run( + [sys.executable, "-c", fake_torch + probe], + env={ + **os.environ, + "EXPECTED_CUDA_VERSION": expected, + "INVALID_CUDA_RESULT": invalid_result, + }, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode == 0, succeeds, result.stdout) + if succeeds: + self.assertIn("RESULT CHECKED", result.stdout) + def test_pybind_runs_inline_for_the_expected_matrix_cells(self): job = WORKFLOW["jobs"]["test-model-cuda-e2e"] matrix = job["strategy"]["matrix"] diff --git a/.ci/scripts/tests/test_filter_cuda_matrix.py b/.ci/scripts/tests/test_filter_cuda_matrix.py index 03a849df81f..3f21762c44f 100644 --- a/.ci/scripts/tests/test_filter_cuda_matrix.py +++ b/.ci/scripts/tests/test_filter_cuda_matrix.py @@ -10,8 +10,12 @@ # with what the project can publish. Two of its comments record past bugs it now guards against, and # a regression in any of them would surface only as a broken release, so each gate is pinned here. +import contextlib import importlib.util +import io import json +import os +import subprocess import unittest from pathlib import Path from unittest import mock @@ -21,29 +25,21 @@ ROOT = Path(__file__).resolve().parents[3] -def _load_filter(): - """Load the script by path, since .github/scripts is not an importable package.""" - path = ROOT / ".github" / "scripts" / "filter_cuda_matrix.py" - spec = importlib.util.spec_from_file_location("filter_cuda_matrix", path) +def _load_module(name, path): + spec = importlib.util.spec_from_file_location(name, path) module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module -FILTER = _load_filter() +FILTER = _load_module( + "filter_cuda_matrix", ROOT / ".github" / "scripts" / "filter_cuda_matrix.py" +) +INSTALL_UTILS = _load_module("install_utils", ROOT / "install_utils.py") def _full_matrix(): - """Every supported python and CUDA pair. - - The filter refuses anything less: one gate rejects a matrix that would leave a CUDA train - unpublished, another rejects a missing python and CUDA combination. Built from the module's own - lists so it cannot go stale when either grows. - - That also means it shrinks when either list shrinks, and every gate keeps passing. Measured: - deleting cu132 and 3.13 from the filter left all sixteen cases here green. TestPublishedSets - below is what notices that, so this fixture does not have to. - """ + """Every supported pair; TestPublishedSets separately guards against shrinking the lists.""" return { "include": [ {"python_version": python, "desired_cuda": cuda} @@ -165,9 +161,6 @@ def test_unsupported_cuda_is_dropped(self): class TestGates(unittest.TestCase): def _exit_message(self, matrix, limit="false", extra=None): """The stderr text of the gate that fired, so a case can name which one it hit.""" - import contextlib - import io - argv = ["--matrix", json.dumps(matrix), "--limit-pr-builds", limit] + ( extra or [] ) @@ -206,15 +199,19 @@ def test_absent_train_is_skipped_not_fatal(self): for cuda in offered ] } - emitted = _emitted(_run(matrix)) + stdout = io.StringIO() + stderr = io.StringIO() + with contextlib.redirect_stdout(stdout), contextlib.redirect_stderr(stderr): + FILTER.main(["--matrix", json.dumps(matrix)]) + emitted = json.loads(stdout.getvalue()) + self.assertIn("the generator offered no row", stderr.getvalue()) + self.assertIn(dropped, stderr.getvalue()) published = sorted({row["desired_cuda"] for row in emitted["include"]}) self.assertEqual(published, sorted(offered)) self.assertNotIn(dropped, published) def test_dropped_train_still_publishes_the_others(self): - # The exact upstream drop this resilience is for: PyTorch stops shipping cu126, the generator - # offers only cu130 and cu132, and the release must still publish those two rather than fail - # because cu126 is gone. Skips the case cleanly if the policy no longer lists cu126. + # Losing cu126 from the generator must not block the remaining supported trains. if "cu126" not in FILTER.SUPPORTED_CUDA_VERSIONS: self.skipTest("cu126 is not a published train") survivors = [c for c in FILTER.SUPPORTED_CUDA_VERSIONS if c != "cu126"] @@ -244,6 +241,17 @@ def test_missing_combination_exits_nonzero(self): message = self._exit_message(matrix) self.assertIn("incomplete train", message) + def test_offered_train_with_only_unsupported_pythons_exits_nonzero(self): + for cuda in FILTER.SUPPORTED_CUDA_VERSIONS: + with self.subTest(cuda=cuda): + matrix = _full_matrix() + for row in matrix["include"]: + if row["desired_cuda"] == cuda: + row["python_version"] = "3.15" + message = self._exit_message(matrix) + self.assertIn("incomplete train", message) + self.assertIn(f"3.10/{cuda}", message) + def test_jetpack_not_published_exits_nonzero(self): # Refused explicitly rather than allowed to fall through to an empty result, so the reason a # reader sees is the real one. Nothing passes this flag today, which is why it had no cover. @@ -272,7 +280,66 @@ class TestPublishedSets(unittest.TestCase): """ def test_published_cuda_versions(self): - self.assertEqual(FILTER.SUPPORTED_CUDA_VERSIONS, ["cu126", "cu130", "cu132"]) + self.assertEqual( + FILTER.SUPPORTED_CUDA_VERSIONS, ["cu126", "cu130", "cu132", "cu134"] + ) + + def test_published_cuda_versions_are_supported_by_the_installer(self): + supported = { + f"cu{major}{minor}" + for major, minor in INSTALL_UTILS.SUPPORTED_CUDA_VERSIONS + } + self.assertLessEqual(set(FILTER.SUPPORTED_CUDA_VERSIONS), supported) + + def test_supported_toolkits_select_the_matching_torch_index(self): + base_url = "https://download.pytorch.org/whl/nightly" + self.addCleanup(INSTALL_UTILS._get_cuda_version.cache_clear) + self.addCleanup(INSTALL_UTILS.determine_torch_url.cache_clear) + for major, minor in INSTALL_UTILS.SUPPORTED_CUDA_VERSIONS: + with self.subTest(cuda=(major, minor)): + INSTALL_UTILS._get_cuda_version.cache_clear() + INSTALL_UTILS.determine_torch_url.cache_clear() + detected = subprocess.CompletedProcess( + args=[], + returncode=0, + stdout=f"Cuda compilation tools, release {major}.{minor}, V{major}.{minor}.0", + ) + with mock.patch.object( + INSTALL_UTILS.platform, "system", return_value="Linux" + ), mock.patch.object( + INSTALL_UTILS.subprocess, "run", return_value=detected + ): + self.assertEqual( + INSTALL_UTILS.determine_torch_url(base_url), + f"{base_url}/cu{major}{minor}", + ) + self.assertTrue(INSTALL_UTILS.is_cuda_available()) + + def test_published_cuda_versions_have_gpu_architectures(self): + script = ROOT / ".ci" / "scripts" / "wheel" / "cuda_arch_list.sh" + for machine in ("x86_64", "aarch64"): + for cuda in FILTER.SUPPORTED_CUDA_VERSIONS: + with self.subTest(machine=machine, cuda=cuda): + result = subprocess.run( + [ + "bash", + "-c", + 'uname() { printf "%s\\n" "$MACHINE"; }; ' + 'source "$1"; executorch_cuda_arch_list', + "bash", + str(script), + ], + env={ + **os.environ, + "MACHINE": machine, + "CU_VERSION": cuda, + "EXECUTORCH_BUILD_CUDA": "1", + }, + capture_output=True, + text=True, + ) + self.assertEqual(result.returncode, 0, result.stderr) + self.assertIn("8.0", result.stdout.split()) def test_published_python_versions(self): self.assertEqual( diff --git a/.ci/scripts/wheel/cuda_arch_list.sh b/.ci/scripts/wheel/cuda_arch_list.sh index 29484056633..9e74af6c644 100644 --- a/.ci/scripts/wheel/cuda_arch_list.sh +++ b/.ci/scripts/wheel/cuda_arch_list.sh @@ -33,11 +33,13 @@ # that never claimed the device. So these lists are narrower than torch at the bottom on purpose. _cuda_arch_x86_64_cu130="8.0 8.6 8.9 9.0 10.0 12.0" _cuda_arch_x86_64_cu132="${_cuda_arch_x86_64_cu130}" +_cuda_arch_x86_64_cu134="${_cuda_arch_x86_64_cu130}" # The architectures the published aarch64 PyTorch CUDA build covers, read from its own library on an ARM # machine, for the same reason as the x86_64 rows above. Includes the ARM module whose train matches. _cuda_arch_aarch64_cu130="8.0 9.0 10.0 11.0 12.0" _cuda_arch_aarch64_cu132="${_cuda_arch_aarch64_cu130}" +_cuda_arch_aarch64_cu134="${_cuda_arch_aarch64_cu130}" # The older CUDA train. # @@ -104,6 +106,7 @@ executorch_cuda_arch_list() { 126) printf '%s' "${_cuda_arch_aarch64_cu126}" ;; 130) printf '%s' "${_cuda_arch_aarch64_cu130}" ;; 132) printf '%s' "${_cuda_arch_aarch64_cu132}" ;; + 134) printf '%s' "${_cuda_arch_aarch64_cu134}" ;; *) _executorch_unknown_train "${train}" ;; esac ;; @@ -112,6 +115,7 @@ executorch_cuda_arch_list() { 126) printf '%s' "${_cuda_arch_x86_64_cu126}" ;; 130) printf '%s' "${_cuda_arch_x86_64_cu130}" ;; 132) printf '%s' "${_cuda_arch_x86_64_cu132}" ;; + 134) printf '%s' "${_cuda_arch_x86_64_cu134}" ;; *) _executorch_unknown_train "${train}" ;; esac ;; diff --git a/.github/scripts/filter_cuda_matrix.py b/.github/scripts/filter_cuda_matrix.py index d0ec2ab076a..95d7a02120b 100644 --- a/.github/scripts/filter_cuda_matrix.py +++ b/.github/scripts/filter_cuda_matrix.py @@ -48,19 +48,17 @@ # # cu126 the floor, and what Jetson devices are limited to # cu130 the generator's stable choice, and the default for accelerator consumers -# cu132 the newest, which consumers building against a current TensorRT need +# cu132 a current TensorRT build target +# cu134 the newest, which consumers building against the latest CUDA need # -# A version listed here is published only when the shared generator still offers it. When -# PyTorch stops shipping a CUDA train, its rows simply do not appear and the release skips it, -# rather than failing the whole build. So a train PyTorch drops (as it did with cu126) costs -# only that train, and a train PyTorch restores returns here with no edit. The release still -# fails if a train that IS offered comes through incomplete, which is a real build break. +# Skip wholly absent trains so an upstream removal cannot block the remaining releases. +# Offered trains must still cover every supported Python version. # # cu132 is included because omitting it would leave a published consumer row with no # ExecuTorch wheel to pair with. It is executable on a device one minor behind, since CUDA # minor versions are compatible, so a cu132 wheel has been run end to end on a CUDA 13.0 # device. The packaging properties are checked on every row regardless. -SUPPORTED_CUDA_VERSIONS: List[str] = ["cu126", "cu130", "cu132"] +SUPPORTED_CUDA_VERSIONS: List[str] = ["cu126", "cu130", "cu132", "cu134"] # Python versions to publish, stated rather than derived for the same reason the CUDA # versions are. Deriving them from the rows that survived the filter made the release @@ -190,41 +188,24 @@ def main(argv: List[str]) -> None: if args.limit_pr_builds.lower() == "true" and items: items = only_pull_request_row(items) elif items and not is_jetpack: - # A release has to publish every combination this policy advertises. Comparing the result against - # what the generator offered cannot catch anything, because both sides apply the same conditions, so - # the difference is empty by construction and the check never fires. The policy's own list is the - # thing to compare against: a CUDA version the generator stopped offering otherwise disappears from - # the release silently, and a missing job is a green check for a wheel that was never built. - # - # The generic rows only. A JetPack release advertises the single pair its own lists name rather than - # every supported CUDA version, so checking it against this list would fail a correct release. - # - # Both axes come from this policy's own lists, not from the matrix. Reading the generator's python - # axis pulled in rows this policy never builds, and deriving it from the rows that survived went - # blind to a python that disappeared from every supported train. The generator lives in another - # repository and its axes move independently of what this policy promises to publish. built = {(item["python_version"], item["desired_cuda"]) for item in items} - built_trains = {cuda for _, cuda in built} - # A train the generator offered nothing for is one PyTorch stopped shipping, not a build - # break here. Skip it and publish the rest, so one dropped train cannot take the others - # down with it. When PyTorch dropped CUDA 12.6, failing here also blocked cu130 and cu132 - # from publishing, which is the opposite of what a consumer needs. The train returns on its - # own if PyTorch ships it again, with no edit here. - absent_trains = sorted(set(SUPPORTED_CUDA_VERSIONS) - built_trains) + # Filtering out every Python row must not disguise an offered train as absent. + offered_trains = { + item["desired_cuda"] + for item in matrix.get("include", []) + if item["desired_cuda"] in SUPPORTED_CUDA_VERSIONS + } + absent_trains = sorted(set(SUPPORTED_CUDA_VERSIONS) - offered_trains) if absent_trains: print( f"the generator offered no row for {absent_trains}, so they are skipped this run; " - f"publishing {sorted(built_trains)}", + f"publishing {sorted(offered_trains)}", file=sys.stderr, ) - # A train that IS offered but missing some python versions is a real break, not an upstream - # drop: the release would ship an incomplete train, fewer wheels than promised for a version - # that is otherwise present. Checked only against the trains actually offered, so a fully - # absent train handled above does not also trip this and read as a python problem. missing = sorted( f"{python}/{cuda}" for python in SUPPORTED_PYTHON_VERSIONS - for cuda in built_trains + for cuda in offered_trains if (python, cuda) not in built ) if missing: diff --git a/.github/workflows/cuda.yml b/.github/workflows/cuda.yml index f3705104ded..0c134b30e05 100644 --- a/.github/workflows/cuda.yml +++ b/.github/workflows/cuda.yml @@ -1,6 +1,6 @@ # Test ExecuTorch CUDA Build Compatibility # This workflow tests whether ExecuTorch can be successfully built with CUDA support -# across different CUDA versions (12.6, 13.0) using the command: +# across different CUDA versions (12.6, 13.0, 13.4) using the command: # ./install_executorch.sh # # Intentionally skipped CUDA version 13.2 check due to ci image unsupported. @@ -64,10 +64,10 @@ jobs: strategy: fail-fast: false matrix: - cuda-version: ["12.6", "13.0"] + cuda-version: ["12.6", "13.0", "13.4"] name: test-executorch-cuda-build-${{ matrix.cuda-version }} - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@c3c4c4d48e97dbaaabd9b65496afd8c8d8dad713 permissions: id-token: write contents: read @@ -76,14 +76,19 @@ jobs: runner: linux.g5.4xlarge.nvidia.gpu gpu-arch-type: cuda gpu-arch-version: ${{ matrix.cuda-version }} + test-infra-ref: c3c4c4d48e97dbaaabd9b65496afd8c8d8dad713 + driver-version: ${{ matrix.cuda-version == '13.4' && '615.71.09' || '580.65.06' }} + driver-download-url: ${{ matrix.cuda-version == '13.4' && 'https://download.nvidia.com/XFree86/Linux-x86_64/615.71.09/NVIDIA-Linux-x86_64-615.71.09.run' || '' }} use-custom-docker-registry: false submodules: recursive ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | set -eux - # Test ExecuTorch CUDA build - ExecuTorch will automatically detect CUDA version - # and install the appropriate PyTorch wheel + if [ "${{ matrix.cuda-version }}" = "13.4" ]; then + # The image's older Conda runtime aborts during CUDA thread-local cleanup. + conda install -y -n base -c conda-forge 'libstdcxx-ng=16.2.0' 'libgcc-ng=16.2.0' + fi source .ci/scripts/test-cuda-build.sh "${{ matrix.cuda-version }}" # This job will fail if any of the CUDA versions fail @@ -112,7 +117,7 @@ jobs: echo "CUDA build results: ${{ needs.test-cuda-builds.result }}" exit 1 else - echo "SUCCESS: All ExecuTorch CUDA builds (12.6, 13.0) completed successfully!" + echo "SUCCESS: All ExecuTorch CUDA builds completed successfully!" fi test-models-cuda: diff --git a/backends/cuda/tests/test_missing_weights_blob.py b/backends/cuda/tests/test_missing_weights_blob.py index bbd85c29da4..410184ee238 100644 --- a/backends/cuda/tests/test_missing_weights_blob.py +++ b/backends/cuda/tests/test_missing_weights_blob.py @@ -14,10 +14,9 @@ Two payload shapes reach that code. A current export carries per-name weight metadata and goes through the weight cache, which reports a missing blob itself. -A library built before external weights carries only the two blob keys, newline -separated, and goes through the legacy path. That legacy path is the one the -check was added to, so it is the one this file exercises, by rewriting the -payload of a real export into the older shape in place. +The legacy external-weight format carries two blob keys, newline separated. +That path is the one the check was added to, so this file exercises it by +rewriting the payload of a real export into the older shape in place. """ import os @@ -27,7 +26,7 @@ import torch from executorch.backends.cuda.cuda_backend import CudaBackend from executorch.backends.cuda.cuda_partitioner import CudaPartitioner -from executorch.backends.cuda.cuda_weight_collector import CUDA_WEIGHT_CACHE_MAGIC +from executorch.backends.cuda.cuda_weight_collector import decode_cuda_aoti_metadata from executorch.exir import to_edge_transform_and_lower from executorch.exir._serialize._program import deserialize_pte_binary from torch.export import export @@ -91,10 +90,9 @@ def _rewrite_payload_as_legacy(self, path: str) -> None: ) self.assertEqual(len(payloads), 1, "expected one CUDA delegate") payload = payloads[0] - self.assertTrue( - payload.startswith(CUDA_WEIGHT_CACHE_MAGIC), - "expected the weight metadata payload this rewrite consumes", - ) + metadata = decode_cuda_aoti_metadata(payload) + self.assertEqual(len(metadata.variants), 1, "expected one compiled variant") + so_key = metadata.variants[0].so_blob_key # The payload carries a content hash, so it occurs once. offset = raw.find(payload) @@ -103,14 +101,9 @@ def _rewrite_payload_as_legacy(self, path: str) -> None: raw.find(payload, offset + 1), -1, "payload is not unique in the file" ) - # Derived from the shared library key, so it carries the library's hash - # rather than the blob's and would not resolve even if a sidecar were - # supplied. That is fine here: the point is that an unresolvable key now - # fails the load rather than binding nothing. - marker = b"_so_blob" - end = payload.index(marker) + len(marker) - so_key = payload[:end].rsplit(b"\x00", 1)[-1].decode("utf-8") - weights_key = so_key.replace("_so_blob", "_weights_blob") + # Keep the absent key short enough for a payload with no weight entries. + weights_key = "missing_weights" + self.assertNotIn(weights_key.encode(), raw) # The two keys, then zeros to keep the payload its original length. The # runtime reads the blob key as a C string, so the filler is not part of the @@ -123,6 +116,15 @@ def _rewrite_payload_as_legacy(self, path: str) -> None: blob[offset : offset + len(payload)] = legacy with open(path, "wb") as f: f.write(bytes(blob)) + with open(path, "rb") as f: + rewritten = deserialize_pte_binary(f.read()).program + rewritten_payloads = [ + bytes(rewritten.backend_delegate_data[delegate.processed.index].data) + for plan in rewritten.execution_plan + for delegate in plan.delegates + if delegate.id == "CudaBackend" + ] + self.assertEqual(rewritten_payloads, [legacy]) def test_load_reports_not_found_when_blob_is_absent(self) -> None: from executorch.runtime import Runtime @@ -149,16 +151,6 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: self._rewrite_payload_as_legacy(path) - # Without this the test would still pass if the rewrite stopped - # working, by exercising the weight cache path instead, which reports - # the same error number for the same program. - with open(path, "rb") as f: - self.assertNotIn( - CUDA_WEIGHT_CACHE_MAGIC, - f.read(), - "the rewrite left the metadata payload in place", - ) - # The blob is never supplied, so the load must fail. The runtime's # exception carries only the method name and the error number, so the # cause is asserted through the rewrite check above rather than here. @@ -171,9 +163,7 @@ def test_load_reports_not_found_when_blob_is_absent(self) -> None: def test_load_succeeds_when_the_model_has_no_constants(self) -> None: """A model with nothing to bind still loads without its weights blob. - The refusal must not fire on a model that has no constants, and that branch - has no other coverage. A model with no parameters or buffers already emits - the two-key payload this loader handles, so no rewrite is needed here. + Rewrite this export too, so it exercises the legacy constant-count check. """ with tempfile.TemporaryDirectory() as outdir: @@ -194,6 +184,8 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: with open(path, "wb") as f: lowered.to_executorch().write_to_file(f) + self._rewrite_payload_as_legacy(path) + from executorch.runtime import Runtime method = Runtime.get().load_program(path).load_method("forward") diff --git a/install_requirements.py b/install_requirements.py index c4a934a0180..3429f06dd8c 100644 --- a/install_requirements.py +++ b/install_requirements.py @@ -18,6 +18,35 @@ TORCH_URL_BASE = "https://download.pytorch.org/whl/test" TORCHAO_URL_BASE = "https://download.pytorch.org/whl/nightly" TORCHAO_NIGHTLY_VERSION = "0.18.0.dev20260729" +CU134_TORCHAO_NIGHTLY_VERSION = "0.19.0.dev20260811" +# These wheels' metadata pairs August 11 domain libraries with August 10 torch. +CU134_TORCH_PACKAGES = [ + "torch==2.14.0.dev20260810+cu134", + "torchvision==0.29.0.dev20260811+cu134", + "torchaudio==2.11.0.dev20260811+cu134", +] + + +def torchao_from_source(): + return ( + os.environ.get("EXECUTORCH_BUILD_KERNELS_TORCHAO") == "1" + or os.environ.get("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1" + ) + + +def cu134_requirements(torch_url, include_domains=False): + if not torch_url.endswith("/cu134"): + return [] + packages = list( + CU134_TORCH_PACKAGES if include_domains else CU134_TORCH_PACKAGES[:1] + ) + if not torchao_from_source(): + torchao_variant = ( + "cpu" if platform.machine().lower() in ("aarch64", "arm64") else "cu134" + ) + packages.append(f"torchao=={CU134_TORCHAO_NIGHTLY_VERSION}+{torchao_variant}") + return packages + # Since ExecuTorch often uses main-branch features of pytorch, only the nightly # pip versions will have the required features. @@ -46,6 +75,11 @@ def install_requirements(use_pytorch_nightly): # Determine the appropriate PyTorch URL based on CUDA delegate status torch_url = determine_torch_url(TORCH_URL_BASE) + cu134_packages = cu134_requirements(torch_url) + if cu134_packages: + torch_url = determine_torch_url(TORCHAO_URL_BASE) + if not use_pytorch_nightly: + cu134_packages[0] = "torch" # torchao's CUDA channel publishes x86_64 only, so asking for a CUDA build makes the pin # unsatisfiable on aarch64. Only that case is special-cased: falling back everywhere would # change which torchao a CPU x86_64 install resolves, and the CUDA build is genuinely wanted @@ -62,7 +96,7 @@ def install_requirements(use_pytorch_nightly): torchao_url = determine_torch_url(TORCHAO_URL_BASE) # pip packages needed by exir. - TORCH_PACKAGE = [ + TORCH_PACKAGE = cu134_packages or [ # Setting use_pytorch_nightly to false to test the pinned PyTorch commit. Note # that we don't need to set any version number there because they have already # been installed on CI before this step, so pip won't reinstall them @@ -91,10 +125,7 @@ def install_requirements(use_pytorch_nightly): ) LOCAL_REQUIREMENTS = [] - if ( - os.environ.get("EXECUTORCH_BUILD_KERNELS_TORCHAO") == "1" - or os.environ.get("TORCHAO_BUILD_EXPERIMENTAL_MPS") == "1" - ): + if torchao_from_source(): LOCAL_REQUIREMENTS.append("third-party/ao") if sys.platform != "win32": # TODO(larryliu0820): Setup a pypi package for this. @@ -122,6 +153,12 @@ def install_requirements(use_pytorch_nightly): # Without --no-build-isolation, setup.py can't find the torch module. "--no-build-isolation", *LOCAL_REQUIREMENTS, + *cu134_packages, + *( + ["--extra-index-url", torch_url, "--extra-index-url", torchao_url] + if cu134_packages + else [] + ), ], env=new_env, check=True, @@ -131,9 +168,21 @@ def install_requirements(use_pytorch_nightly): def install_optional_example_requirements(use_pytorch_nightly): # Determine the appropriate PyTorch URL based on CUDA delegate status torch_url = determine_torch_url(TORCH_URL_BASE) + cu134_packages = ( + cu134_requirements(torch_url, include_domains=True) + if use_pytorch_nightly + else [] + ) + if cu134_packages: + torch_url = determine_torch_url(TORCHAO_URL_BASE) + torchao_index = ( + ["--extra-index-url", f"{TORCHAO_URL_BASE}/cpu"] + if cu134_packages and platform.machine().lower() in ("aarch64", "arm64") + else [] + ) print("Installing torch domain libraries") - DOMAIN_LIBRARIES = [ + DOMAIN_LIBRARIES = cu134_packages or [ ("torchvision==0.29.0" if use_pytorch_nightly else "torchvision"), ("torchaudio==2.11.0" if use_pytorch_nightly else "torchaudio"), ] @@ -147,6 +196,7 @@ def install_optional_example_requirements(use_pytorch_nightly): *DOMAIN_LIBRARIES, "--extra-index-url", torch_url, + *torchao_index, ], check=True, ) @@ -160,8 +210,10 @@ def install_optional_example_requirements(use_pytorch_nightly): "install", "-r", "requirements-examples.txt", + *cu134_packages, "--extra-index-url", torch_url, + *torchao_index, "--upgrade-strategy", "only-if-needed", ], diff --git a/install_utils.py b/install_utils.py index 276535cf38f..4e1397413b3 100644 --- a/install_utils.py +++ b/install_utils.py @@ -21,6 +21,7 @@ (12, 6), (13, 0), (13, 2), + (13, 4), ) diff --git a/setup.py b/setup.py index 132e9bcd1e3..a00dfd69909 100644 --- a/setup.py +++ b/setup.py @@ -1106,6 +1106,11 @@ def _torchao_requirement() -> str: spec.loader.exec_module(module) version = module.TORCHAO_NIGHTLY_VERSION + if ( + install_utils.determine_torch_url(module.TORCH_URL_BASE).endswith("/cu134") + and not module.torchao_from_source() + ): + version = module.CU134_TORCHAO_NIGHTLY_VERSION major, minor = (int(part) for part in version.split(".")[:2]) return f"torchao>={version},<{major}.{minor + 1}" From 588260d3b5b0933f4e6a938aec08419097d655d5 Mon Sep 17 00:00:00 2001 From: Suryansh Sijwali <159204949+SuryanshSS1011@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:08:22 -0400 Subject: [PATCH 142/190] Check dim order in the optimized mm and bmm as the portable ones do (#22355) ### Summary Follow-up to #21866, which did the same for the optimized `layer_norm`. `opt_mm_out` and `opt_bmm_out` carry no dim order check, while portable `mm_out` and `bmm_out` each carry two: ```cpp ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(in, mat2, out), InvalidArgument, out); ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); ``` The shared `check_mm_args` and `check_bmm_args` do not supply them either, so choosing the optimized kernels quietly drops a check the portable path enforces. `opt_mm_out` passes raw data pointers to `cpublas::gemm` with the leading dimensions taken from sizes rather than strides, which describes a row-major matrix only in the default dim order. `opt_bmm_out` does the same per batch. Thus, the arithmetic also needs it. This PR adds portable guards to both the kernels, placed in the same position relative to the resize. Not reachable from export today, in the same way the `select_scatter` half of #21915 was not, since `_get_channels_last_dim_order` raises for any rank other than 4 and so a 2D `mm` or 3D `bmm` input cannot be given a non-default dim order. It is still worth adding because the optimized kernels are currently less strict than the portable ones with nothing recording that as intentional, and because #16429 asks for these layouts to be supported, at which point the optimized path would return wrong numbers where the portable path errors. ### Test plan A `NonDefaultDimOrderDies` per op. `op_mm_test.cpp` and `op_bmm_test.cpp` are both in `_optimized_kernels_test_sources`, so the same test body runs against both kernel sets, which makes the gap visible directly: | target | optimized kernels | mm and bmm | |---|---|---| | `optimized_kernels_test` | guards reverted, tests kept | both fail | | `optimized_kernels_test` | with the guards | both pass | | `portable_kernels_test` | untouched either way | both pass | All three tensors in each test share the same non-default dim order, so the same dim order check passes and only the default dim order check can reject. Both skip under `is_aten`. cc @larryliu0820 @manuelcandales @JakeStevens --- kernels/optimized/cpu/op_bmm.cpp | 5 +++++ kernels/optimized/cpu/op_mm.cpp | 5 +++++ kernels/test/op_bmm_test.cpp | 19 +++++++++++++++++++ kernels/test/op_mm_test.cpp | 16 ++++++++++++++++ 4 files changed, 45 insertions(+) diff --git a/kernels/optimized/cpu/op_bmm.cpp b/kernels/optimized/cpu/op_bmm.cpp index 171f14de399..0dd7bc40b3e 100644 --- a/kernels/optimized/cpu/op_bmm.cpp +++ b/kernels/optimized/cpu/op_bmm.cpp @@ -150,6 +150,11 @@ Tensor& opt_bmm_out( ET_KERNEL_CHECK( ctx, check_bmm_out_args(self, mat2, out), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(self, mat2, out), InvalidArgument, out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out); + static constexpr auto name = "bmm.out"; auto self_type = self.scalar_type(); diff --git a/kernels/optimized/cpu/op_mm.cpp b/kernels/optimized/cpu/op_mm.cpp index 53385a40dff..eefaf3131b7 100644 --- a/kernels/optimized/cpu/op_mm.cpp +++ b/kernels/optimized/cpu/op_mm.cpp @@ -34,6 +34,11 @@ Tensor& opt_mm_out( InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, mat2, out), InvalidArgument, out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + if (out.numel() == 0) { return out; } diff --git a/kernels/test/op_bmm_test.cpp b/kernels/test/op_bmm_test.cpp index afc4be856cf..944d8a2d48e 100644 --- a/kernels/test/op_bmm_test.cpp +++ b/kernels/test/op_bmm_test.cpp @@ -471,3 +471,22 @@ TEST_F(OpBmmOutTest, DISABLED_DynamicShapeUnbound) { Tensor ret = op_bmm_out(x, y, out); EXPECT_TENSOR_CLOSE(out, expected_result); } + +TEST_F(OpBmmOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + // All three tensors share the same non-default dim order, so the kernel's + // same dim order check passes and only the default dim order check rejects. + Tensor x = + tf.make_with_dimorder({2, 3, 4}, std::vector(24, 2), {0, 2, 1}); + Tensor y = + tf.make_with_dimorder({2, 4, 5}, std::vector(40, 3), {0, 2, 1}); + Tensor out = + tf.make_with_dimorder({2, 3, 5}, std::vector(30), {0, 2, 1}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_bmm_out(x, y, out)); +} diff --git a/kernels/test/op_mm_test.cpp b/kernels/test/op_mm_test.cpp index 05d6a7b8d7e..a2044f383f7 100644 --- a/kernels/test/op_mm_test.cpp +++ b/kernels/test/op_mm_test.cpp @@ -294,3 +294,19 @@ TEST_F(OpMmOutTest, DISABLED_DynamicShapeUnbound) { Tensor ret = op_mm_out(x, y, out); EXPECT_TENSOR_CLOSE(out, expected_result); } + +TEST_F(OpMmOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + // All three tensors share the same non-default dim order, so the kernel's + // same dim order check passes and only the default dim order check rejects. + Tensor x = tf.make_with_dimorder({3, 4}, std::vector(12, 2), {1, 0}); + Tensor y = tf.make_with_dimorder({4, 5}, std::vector(20, 3), {1, 0}); + Tensor out = tf.make_with_dimorder({3, 5}, std::vector(15), {1, 0}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_mm_out(x, y, out)); +} From 101ec9ec1a703500bdd706781219666292819ffc Mon Sep 17 00:00:00 2001 From: Suryansh Sijwali <159204949+SuryanshSS1011@users.noreply.github.com> Date: Thu, 10 Sep 2026 10:08:47 -0400 Subject: [PATCH 143/190] Check dim order in the optimized layer_norm as the portable one does (#21866) ### Summary Related to #21828 and #21865 (same underlying assumption in the optimized kernel library). `aten.native_layer_norm` returns wrong data on a channels-last input, off by 3.192 against eager PyTorch where it matches exactly on contiguous input. Nothing errors. The portable kernel is fine. It carries this, with the reason written down: ```cpp // Only support default dim order for now. // TODO: Support other dim orders. ET_KERNEL_CHECK( ctx, tensor_is_default_dim_order(input), InvalidArgument, ret_val); ``` The optimized kernel has neither that check nor the matching `tensors_have_same_dim_order`, and `native_layer_norm.out` is mapped to `torch::executor::opt_native_layer_norm_out` in `optimized.yaml`. So any build with `EXECUTORCH_BUILD_KERNELS_OPTIMIZED=ON` gets the ungated one. It needs the guard rather than stride-aware indexing, because it splits the buffer into `M` rows of `N` contiguous elements: ```cpp const size_t M = getLeadingDims(input, dim); const size_t N = getTrailingDims(input, dim) * dim_size; ``` That layout only exists in the default dim order. This copies the two checks across from the portable kernel so the two agree. ### Test plan `OpNativeLayerNormTest.NonDefaultDimOrderDies` passes a channels-last input with `out`, `mean` and `rstd` all channels-last, so the same dim order check passes and only the default dim order check can reject. `mean` and `rstd` share the input's rank with the normalized dims set to 1, which the kernel requires before it reaches any dim order check. The test file is shared by both kernel libraries, so it runs against the portable kernel too, where it already passes. --- .../optimized/cpu/op_native_layer_norm.cpp | 26 ++++++ kernels/test/op_native_layer_norm_test.cpp | 92 +++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/kernels/optimized/cpu/op_native_layer_norm.cpp b/kernels/optimized/cpu/op_native_layer_norm.cpp index 5fac9faf25e..66e2971087d 100644 --- a/kernels/optimized/cpu/op_native_layer_norm.cpp +++ b/kernels/optimized/cpu/op_native_layer_norm.cpp @@ -149,6 +149,32 @@ std::tuple opt_native_layer_norm_out( InvalidArgument, ret_val); + // Only support default dim order for now. + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(input), InvalidArgument, ret_val); + + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(input, out, mean_out, rstd_out), + InvalidArgument, + ret_val); + + if (weight.has_value()) { + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(input, weight.value()), + InvalidArgument, + ret_val); + } + + if (bias.has_value()) { + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(input, bias.value()), + InvalidArgument, + ret_val); + } + Tensor::SizesType mean_rstd_sizes[kTensorDimensionLimit]; size_t mean_rstd_ndim = 0; get_layer_norm_out_target_size( diff --git a/kernels/test/op_native_layer_norm_test.cpp b/kernels/test/op_native_layer_norm_test.cpp index e1345a10354..0504716a83a 100644 --- a/kernels/test/op_native_layer_norm_test.cpp +++ b/kernels/test/op_native_layer_norm_test.cpp @@ -452,3 +452,95 @@ TEST_F(OpNativeLayerNormTest, DynamicShapeUnbound) { test_dynamic_shape( {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); } + +TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + // mean and rstd share the input's rank with the normalized dims set to 1. + // All four are channels-last so the same-dim-order check passes and only the + // default dim order check can reject. + Tensor input = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out0 = tf.zeros_channels_last({1, 3, 2, 2}); + Tensor out1 = tf.zeros_channels_last({1, 3, 2, 1}); + Tensor out2 = tf.zeros_channels_last({1, 3, 2, 1}); + const std::vector normalized_shape = {2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_native_layer_norm_out( + input, + IntArrayRef(normalized_shape.data(), normalized_shape.size()), + exec_aten::optional(), + exec_aten::optional(), + 1e-5, + out0, + out1, + out2)); +} + +TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderWeightDies) { + TensorFactory tf; + + // weight takes the shape of normalized_shape, so it has to be rank 4 to carry + // a channels-last dim order, which makes the input rank 5. Everything else is + // default, so the same dim order check on weight shall be what rejects. + Tensor input = tf.make({2, 2, 2, 2, 2}, std::vector(32, 1)); + Tensor weight = tf.make_with_dimorder( + {2, 2, 2, 2}, std::vector(16, 1), {0, 2, 3, 1}); + Tensor out0 = tf.zeros({2, 2, 2, 2, 2}); + Tensor out1 = tf.zeros({2, 1, 1, 1, 1}); + Tensor out2 = tf.zeros({2, 1, 1, 1, 1}); + const std::vector normalized_shape = {2, 2, 2, 2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_native_layer_norm_out( + input, + IntArrayRef(normalized_shape.data(), normalized_shape.size()), + exec_aten::optional(weight), + exec_aten::optional(), + 1e-5, + out0, + out1, + out2)); +} + +TEST_F(OpNativeLayerNormTest, NonDefaultDimOrderBiasDies) { + TensorFactory tf; + + // bias takes the shape of normalized_shape, so it has to be rank 4 to carry + // a channels-last dim order, which makes the input rank 5. Everything else is + // default, so the same dim order check on bias shall be what rejects. + Tensor input = tf.make({2, 2, 2, 2, 2}, std::vector(32, 1)); + Tensor bias = tf.make_with_dimorder( + {2, 2, 2, 2}, std::vector(16, 1), {0, 2, 3, 1}); + Tensor out0 = tf.zeros({2, 2, 2, 2, 2}); + Tensor out1 = tf.zeros({2, 1, 1, 1, 1}); + Tensor out2 = tf.zeros({2, 1, 1, 1, 1}); + const std::vector normalized_shape = {2, 2, 2, 2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_native_layer_norm_out( + input, + IntArrayRef(normalized_shape.data(), normalized_shape.size()), + exec_aten::optional(), + exec_aten::optional(bias), + 1e-5, + out0, + out1, + out2)); +} From ff86d8c2f143596881492e1119b8d708a68da634 Mon Sep 17 00:00:00 2001 From: Per Held Date: Wed, 2 Sep 2026 14:06:25 +0200 Subject: [PATCH 144/190] Arm backend: Document Model Explorer performance overlays Document PTE and TOSA visualization, FVP trace artifacts, cycle comparison guidance, and current limitations. Reject performance overlays that do not select the required TOSA view. Authored with Codex. Change-Id: Ib3412dbf7d49e2c2da2f295ede1027b5201a4423 Signed-off-by: Per Held --- examples/arm/README.md | 2 + examples/arm/model-explorer.md | 137 +++++++++++++++++++++++++++++++++ examples/arm/run.sh | 4 + 3 files changed, 143 insertions(+) create mode 100644 examples/arm/model-explorer.md diff --git a/examples/arm/README.md b/examples/arm/README.md index d830356717c..8efffd5c377 100644 --- a/examples/arm/README.md +++ b/examples/arm/README.md @@ -112,6 +112,8 @@ For Cortex-M testing, use a Cortex-M target and bundled I/O: BundleIO, ETDump, profiling, semihosted files, and backend regression tests. - [ethos-u-porting-guide.md](ethos-u-porting-guide.md) - Notes for adapting the example Ethos-U runtime integration to another target. +- [model-explorer.md](model-explorer.md) - Visualize PTE and TOSA graphs and + overlay per-operator Ethos-U cycle data collected from an FVP PMU trace. - [export_standalone_tosa_graph.py](export_standalone_tosa_graph.py) - Example of exporting a standalone TOSA graph with multiple outputs. - [visualize.py](visualize.py) - Helper used by `run.sh --model_explorer` to diff --git a/examples/arm/model-explorer.md b/examples/arm/model-explorer.md new file mode 100644 index 00000000000..d42a0d04f24 --- /dev/null +++ b/examples/arm/model-explorer.md @@ -0,0 +1,137 @@ + + +# Visualize Arm models and Ethos-U performance + +The Arm example scripts can open ExecuTorch PTE or TOSA graphs in +[Model Explorer](https://github.com/google-ai-edge/model-explorer). For models +run on an Ethos-U FVP, they can also overlay the measured number of cycles on +the corresponding TOSA operators. + +The performance overlay helps identify expensive operators and compare the +NPU work performed by different versions of a model. It is generated from a +trace of an actual FVP execution, rather than from Vela's compile-time cycle +estimates. + +## Visualize a graph + +Use `--visualize_pte` to inspect the ExecuTorch program, including its delegate +calls: + +```bash +./examples/arm/run.sh \ + --model_name=mv2 \ + --target=ethos-u85-256 \ + --model_explorer \ + --visualize_pte +``` + +Use `--visualize_tosa` to inspect the TOSA graph passed to Vela: + +```bash +./examples/arm/run.sh \ + --model_name=mv2 \ + --target=ethos-u85-256 \ + --model_explorer \ + --visualize_tosa +``` + +These views serve different purposes. The PTE view shows the ExecuTorch +program around each delegate call. The TOSA view exposes the operators inside +the Ethos-U delegate and is the graph to which performance data can be mapped. + +## Overlay measured Ethos-U cycles + +Add `--perf_overlay` to a TOSA visualization: + +```bash +./examples/arm/run.sh \ + --model_name=mv2 \ + --target=ethos-u85-256 \ + --model_explorer \ + --visualize_tosa \ + --perf_overlay +``` + +`run.sh` performs the following additional steps: + +1. Enables compiler debug output so that Vela emits tables mapping command + stream offsets back to TOSA operators. +2. Enables PMU tracing when it runs the model on the FVP. +3. Combines the trace timestamps with the Vela mapping tables. +4. Adds the resulting per-operator duration data to Model Explorer as + `Duration (Cycles)`. + +With the default build root, the relevant MobileNetV2 artifacts are: + +```text +arm_test/mv2/pmu_trace.gz +arm_test/mv2/output/out_debug.xml +``` + +`pmu_trace.gz` contains the FVP trace events. `out_debug.xml` contains the Vela +debug tables required to attribute those events to TOSA operators. The two +files must come from the same compilation and execution. + +Model Explorer color-codes operators by duration. Its node-data panel also +provides aggregate values for a selected layer, which can be used to find the +parts of the delegated graph that consume the most cycles. + +## Open existing artifacts + +To reopen an existing overlay without rebuilding and rerunning the model, pass +the generated files directly to `visualize.py`: + +```bash +python3 examples/arm/visualize.py \ + --model_dir arm_test/mv2 \ + --tosa \ + --trace arm_test/mv2/pmu_trace.gz \ + --tables arm_test/mv2/output/out_debug.xml +``` + +Both `--trace` and `--tables` are required when either one is specified. + +## Compare model versions + +The overlay can expose performance improvements or regressions caused by +changes to model lowering, quantization, fusion, or Vela scheduling. Compare +the aggregate cycle count first. If the graph structure is unchanged, the +per-operator values can then show where the difference originated. + +For a meaningful comparison, keep the following fixed between runs: + +- Ethos-U target and MAC configuration +- Vela and FVP versions +- Vela system configuration and memory mode +- Input shapes and compiler options +- Quantization configuration + +When a change adds, removes, or fuses operators, node identifiers may no longer +correspond between graphs. In that case, compare totals and groups of +semantically equivalent operators instead of matching nodes only by ID. Record +the PTE or source revision and verify model outputs or accuracy alongside the +cycle results. + +The overlay measures work in the Ethos-U command stream. It is not an +end-to-end latency measurement and does not account for portable CPU operators +or all ExecuTorch and delegate overhead. Use +[ETDump](https://docs.pytorch.org/executorch/stable/etdump.html) when the total +runtime behavior is the metric of interest. + +## Current limitations + +- Performance overlays are supported for the TOSA view, not the PTE view. +- The trace parser expects the gzip-compressed JSON trace generated by the + Ethos-U FVP. +- The model directory should contain the TOSA and Vela artifacts from the same + build as the trace. +- Cycle counts are comparable only when the target and memory timing + configuration are equivalent. + +For general information about displaying node data, see the +[Model Explorer custom node data documentation](https://github.com/google-ai-edge/model-explorer/wiki/2.-User-Guide#custom-node-data). diff --git a/examples/arm/run.sh b/examples/arm/run.sh index b69f8f1c4a7..007c7c63cd2 100755 --- a/examples/arm/run.sh +++ b/examples/arm/run.sh @@ -144,6 +144,10 @@ if [ "$perf_overlay" = true ] && [ "$model_explorer" != true ]; then echo "Error: --perf_overlay requires --model_explorer" >&2 exit 1 fi +if [ "$perf_overlay" = true ] && [ "$visualize_tosa" != true ]; then + echo "Error: --perf_overlay requires --visualize_tosa" >&2 + exit 1 +fi # Cortex-M backend is an operator-library, not a delegate; force-disable # --delegate when targeting cortex-m so users don't need --no_delegate. From 06bb77f7f3b54299926ccb17f968f5186b5f78ee Mon Sep 17 00:00:00 2001 From: Per Held Date: Mon, 17 Aug 2026 11:27:49 +0200 Subject: [PATCH 145/190] Arm backend: Delegate constant tensor index Lower U55 index.Tensor operations with one constant rank-1 integer index to slices and optional concatenation, including indices on non-leading dimensions. Keep runtime, empty, mask, rank>1, multiple-index, and symbolic-shape cases on the CPU. Reject out-of-range constants during preprocessing. Authored with Codex. Signed-off-by: Per Held Change-Id: Ib16e1c800b3517d054c77cfde39639de6f8b9369 --- backends/arm/_passes/arm_pass_manager.py | 2 +- .../decompose_index_tensor_to_gather_pass.py | 88 +++++++++- .../arm/operator_support/ethos_u55_support.py | 48 +++++- .../operator_support/index_tensor_support.py | 30 +++- .../tosa_supported_operators.py | 2 + .../test/misc/test_tosa_operator_support.py | 20 +++ backends/arm/test/ops/test_index_tensor.py | 156 ++++++++++++++++++ ...t_decompose_index_tensor_to_gather_pass.py | 53 ++++++ 8 files changed, 387 insertions(+), 12 deletions(-) create mode 100644 backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 8d10cbbcf33..0587170c23c 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -608,7 +608,7 @@ def _tosa_pipeline( DecomposeStridedSliceCopyPass(), DecomposeSliceScatterPass(), AccumulateIndexPutPass(), - DecomposeIndexTensorToGatherPass(), + DecomposeIndexTensorToGatherPass(exported_program), DecomposeAdaptiveAvgPool2dPass(), DecomposeDynamicAdaptiveAvgPool2dPass(), DecomposeAvgPool2dPass(), diff --git a/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py b/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py index 93db9f9d434..c3e98c363dc 100644 --- a/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py +++ b/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py @@ -10,7 +10,11 @@ import torch from executorch.backends.arm._passes import ArmOpTargetedPass -from executorch.backends.arm._passes.arm_pass_utils import meta_without_qparams +from executorch.backends.arm._passes.arm_pass_utils import ( + get_param_tensor, + is_param_node, + meta_without_qparams, +) from executorch.backends.arm._passes.convert_expand_copy_to_repeat import ( ConvertExpandCopyToRepeatPass, ) @@ -20,6 +24,7 @@ from executorch.backends.arm._passes.replace_scalar_with_tensor_pass import ( ReplaceScalarWithTensorByProfilePass, ) +from executorch.exir import ExportedProgram from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -169,6 +174,12 @@ class DecomposeIndexTensorToGatherPass(ArmOpTargetedPass): exir_ops.edge.aten.index.Tensor, } + def __init__( + self, exported_program: ExportedProgram | None = None, *args, **kwargs + ) -> None: + super().__init__(*args, **kwargs) + self.exported_program = exported_program + @staticmethod def _shape_to_stride( values_shape: Sequence[int], @@ -245,6 +256,70 @@ def _compute_index_tensor_params(self, x, m, index_shapes): return x_data, S, W, K, C, trailing, lin_scales + def _decompose_constant_index(self, x, indices, meta): + tensor_indices = [ + (dim, index) for dim, index in enumerate(indices) if index is not None + ] + if ( + self.exported_program is None + or len(tensor_indices) != 1 + or any(not isinstance(size, int) for size in x.data.shape) + ): + return None + + indexed_dim, index_tensor = tensor_indices[0] + if not is_param_node(self.exported_program, index_tensor.node): + return None + + constant_index = get_param_tensor(self.exported_program, index_tensor.node) + if ( + constant_index is None + or constant_index.dim() != 1 + or constant_index.numel() == 0 + ): + return None + + indexed_dim_size = x.data.shape[indexed_dim] + index_values = [] + for value in constant_index.tolist(): + normalized_value = value if value >= 0 else value + indexed_dim_size + if normalized_value < 0 or normalized_value >= indexed_dim_size: + raise IndexError( + f"index {value} is out of bounds for dimension {indexed_dim} " + f"with size {indexed_dim_size}" + ) + index_values.append(normalized_value) + + if index_values == list( + range(index_values[0], index_values[0] + len(index_values)) + ): + return super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, indexed_dim, index_values[0], index_values[-1] + 1), + {}, + meta, + updated=True, + ) + + slices = [] + for index_value in index_values: + slices.append( + super().call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + (x, indexed_dim, index_value, index_value + 1), + {}, + meta, + updated=True, + ) + ) + return super().call_operator( + exir_ops.edge.aten.cat.default, + (slices, indexed_dim), + {}, + meta, + updated=True, + ) + def call_operator(self, op, args, kwargs, meta): if op not in self.target_ops: return super().call_operator(op, args, kwargs, meta) @@ -255,6 +330,17 @@ def call_operator(self, op, args, kwargs, meta): x, indices = args + tensor_indices = [index for index in indices if index is not None] + if len(tensor_indices) == 1 and tensor_indices[0].data.dtype in ( + torch.bool, + torch.uint8, + ): + return super().call_operator(op, args, kwargs, meta) + + constant_result = self._decompose_constant_index(x, indices, meta) + if constant_result is not None: + return constant_result + self._validate_tensor_indices(indices) index_shapes = [idx.data.shape for idx in indices] m = len(indices) diff --git a/backends/arm/operator_support/ethos_u55_support.py b/backends/arm/operator_support/ethos_u55_support.py index 62631893636..fc8182062c3 100644 --- a/backends/arm/operator_support/ethos_u55_support.py +++ b/backends/arm/operator_support/ethos_u55_support.py @@ -203,7 +203,6 @@ class EthosU55NotSupported(OperatorSupportBase): exir_ops.edge.aten.ne.Scalar, exir_ops.edge.aten.gather.default, # GATHER exir_ops.edge.aten.grid_sampler_2d, # GATHER - exir_ops.edge.aten.index.Tensor, # GATHER exir_ops.edge.aten.index_put.default, # SCATTER exir_ops.edge.aten.scatter.src, exir_ops.edge.aten.scatter.value, @@ -431,6 +430,53 @@ def is_node_supported( return True +class EthosU55IndexTensorCheck(OperatorSupportBase): + """Accept single constant index.Tensor cases that lower to slices.""" + + def __init__( + self, exported_program: ExportedProgram, reporter: WhyNoPartitionReporter + ): + self.exported_program = exported_program + self.reporter = reporter + + def is_node_supported( + self, submodules: typing.Mapping[str, torch.nn.Module], node: fx.Node + ) -> bool: + del submodules + if node.target != exir_ops.edge.aten.index.Tensor: + return True + + input_arg, indices_arg = node.args + input_node = typing.cast(fx.Node, input_arg) + indices = typing.cast(typing.Sequence[fx.Node | None], indices_arg) + input_shape = get_first_fake_tensor(input_node).shape + tensor_indices = [index for index in indices if index is not None] + if len(tensor_indices) != 1: + self.reporter.report_reject( + node, + "U55 index.Tensor only supports indexing along one dimension but got " + f"{len(tensor_indices)}.", + ) + return False + + index_node = tensor_indices[0] + index_shape = get_first_fake_tensor(index_node).shape + if ( + not is_param_node(self.exported_program, index_node) + or len(index_shape) != 1 + or index_shape[0] == 0 + or any(not isinstance(size, int) for size in input_shape) + ): + self.reporter.report_reject( + node, + "U55 index.Tensor requires static input shape and a nonempty " + "constant rank-1 index.", + ) + return False + + return True + + class EthosU55IndexSelectCheck(OperatorSupportBase): """Accept constant contiguous index_select cases that lower to a slice.""" diff --git a/backends/arm/operator_support/index_tensor_support.py b/backends/arm/operator_support/index_tensor_support.py index 29134fe964d..937b102ce8f 100644 --- a/backends/arm/operator_support/index_tensor_support.py +++ b/backends/arm/operator_support/index_tensor_support.py @@ -108,20 +108,32 @@ def is_node_tosa_supported( """Return True if ``aten.index.Tensor`` usage fits supported patterns. Enforces the following constraints: - - No ``None`` (unsqueeze), slice, or ellipsis before an indexing tensor. + - No ``None`` (unsqueeze), slice, or ellipsis before an indexing tensor, + except for the U55 constant-index lowering. - The value tensor element count fits in ``int32``. """ indices = node.args[1] - for index in indices: # type: ignore[union-attr] - # Usage 1 guard - if index is None: + if not tosa_spec.is_U55_subset and any( + index is None for index in indices # type: ignore[union-attr] + ): + self.reporter.report_reject( + node, + ( + "None (from slice/unsqueeze/ellipsis) before an indexing tensor" + " is not supported." + ), + ) + return False + + # The U55-specific check limits this to one constant tensor index. + for index in ( + index for index in indices if index is not None # type: ignore[union-attr] + ): + index_node = ensure_type(torch.fx.Node, index) + if get_first_fake_tensor(index_node).dtype in (torch.bool, torch.uint8): self.reporter.report_reject( - node, - ( - "None (from slice/unsqueeze/ellipsis) before an indexing tensor" - " is not supported." - ), + node, "Boolean and byte mask indices are not supported." ) return False diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 04d1e416d48..2c1ae365a8b 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -40,6 +40,7 @@ EthosU55CastCheck, EthosU55DtypeSupport, EthosU55IndexSelectCheck, + EthosU55IndexTensorCheck, EthosU55NotSupported, EthosU55ResizeCheck, EthosU55ReverseCheck, @@ -414,6 +415,7 @@ def _negative_checks( checks.append(EthosU55ResizeCheck(reporter)) checks.append(EthosU55ReverseCheck(reporter)) checks.append(EthosU55UnfoldCopyCheck(reporter)) + checks.append(EthosU55IndexTensorCheck(exported_program, reporter)) checks.append(EthosU55IndexSelectCheck(exported_program, reporter)) checks.append(EthosU55DtypeSupport(reporter)) checks.append(EthosU55CastCheck(reporter)) diff --git a/backends/arm/test/misc/test_tosa_operator_support.py b/backends/arm/test/misc/test_tosa_operator_support.py index 662b428a21f..87d2d8a4c0d 100644 --- a/backends/arm/test/misc/test_tosa_operator_support.py +++ b/backends/arm/test/misc/test_tosa_operator_support.py @@ -5,10 +5,14 @@ import pytest import torch +from executorch.backends.arm.operator_support.index_tensor_support import ( + IndexTensorSupported, +) from executorch.backends.arm.operator_support.tosa_supported_operators import ( CheckFPComparisonInputs, CheckKnownUnsupportedTOSASemantics, ) +from executorch.backends.arm.tosa import TosaSpecification from executorch.exir.backend.utils import WhyNoPartitionReporter from executorch.exir.dialects._ops import ops as exir_ops from torch._subclasses.fake_tensor import FakeTensorMode @@ -132,3 +136,19 @@ def test_rejects_argmax_with_mixed_int32_cast_and_raw_user() -> None: raw_user.meta["val"] = _fake_tensor((3,), torch.int64) assert not _checker().is_node_supported({}, node) + + +@pytest.mark.parametrize("dtype", (torch.bool, torch.uint8)) +def test_rejects_index_tensor_mask(dtype: torch.dtype) -> None: + graph = torch.fx.Graph() + x = _placeholder(graph, "x", (5, 2, 3)) + index = _placeholder(graph, "index", (5,), dtype) + node = graph.call_function(exir_ops.edge.aten.index.Tensor, (x, [index])) + node.meta["val"] = _fake_tensor((2, 2, 3)) + + checker = IndexTensorSupported( + TosaSpecification.create_from_string("TOSA-1.0+INT+u55"), + WhyNoPartitionReporter(), + ) + + assert not checker.is_node_supported({}, node) diff --git a/backends/arm/test/ops/test_index_tensor.py b/backends/arm/test/ops/test_index_tensor.py index de6a1ac5f6b..20bce650513 100644 --- a/backends/arm/test/ops/test_index_tensor.py +++ b/backends/arm/test/ops/test_index_tensor.py @@ -8,12 +8,15 @@ import torch from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.arm_tester import ArmTester from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, OpNotSupportedPipeline, TosaPipelineFP, TosaPipelineINT, VgfPipeline, ) +from executorch.exir.dialects._ops import ops as exir_ops class IndexTensorTestCommon: @@ -41,6 +44,46 @@ def forward(self, x: torch.Tensor): return x[self.index] +class ConstantIndexTensor(torch.nn.Module): + def __init__(self, indices: list[int]): + super().__init__() + self.register_buffer("index", torch.tensor(indices, dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + return x[self.index] + + +class ConstantIndexTensorDim(torch.nn.Module): + def __init__(self, dim: int, indices: list[int]): + super().__init__() + self.dim = dim + self.register_buffer("index", torch.tensor(indices, dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + indices = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + +class ConstantTensorIndex(torch.nn.Module): + def __init__(self, index: torch.Tensor): + super().__init__() + self.register_buffer("index", index) + + def forward(self, x: torch.Tensor): + return x[self.index] + + +class ConstantMultiIndexTensor(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("index_0", torch.tensor([0, 1], dtype=torch.int32)) + self.register_buffer("index_1", torch.tensor([1, 0], dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + return x[self.index_0, self.index_1] + + def test_index_tensor_tosa_FP_int64_buffer_index(): # This mirrors torchvision Swin relative_position_bias_table[index]. The # int64 get_attr must be cast before index.Tensor is decomposed to GATHER: @@ -622,3 +665,116 @@ def test_index_tensor_vgf_quant(test_data: input_params): quantize=True, ) pipeline.run() + + +@common.parametrize( + "indices", + { + "contiguous": [1, 2, 3], + "noncontiguous": [1, 3], + "descending": [3, 1], + "duplicate": [1, 1], + "negative": [-1, -3], + }, +) +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_constant(indices): + pipeline = EthosU55PipelineINT[Tuple[torch.Tensor]]( + ConstantIndexTensor(indices), + (torch.rand(5, 2, 3),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.parametrize( + "test_data", + { + "dim1_contiguous": (1, [1, 2, 3]), + "dim1_noncontiguous": (1, [1, 3]), + "dim2_descending": (2, [3, 1]), + "dim2_duplicate": (2, [1, 1]), + "dim2_negative": (2, [-1, -3]), + }, +) +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_constant_later_dim(test_data): + dim, indices = test_data + pipeline = EthosU55PipelineINT[Tuple[torch.Tensor]]( + ConstantIndexTensorDim(dim, indices), + (torch.rand(5, 5, 5),), + aten_ops=[], + exir_ops=[], + ) + pipeline.run() + + +@common.XfailIfNoCorstone300 +def test_index_tensor_u55_INT_constant_a16w8(): + pipeline = EthosU55PipelineINT[Tuple[torch.Tensor]]( + ConstantIndexTensor([1, 2, 3]), + (torch.rand(5, 2, 3),), + aten_ops=[], + exir_ops=[], + a16w8_quantization=True, + ) + pipeline.run() + + +def test_index_tensor_u55_INT_constant_empty_not_delegated(): + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor]]( + ConstantIndexTensor([]), + (torch.rand(5, 2, 3),), + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + +def test_index_tensor_u55_INT_constant_multi_not_delegated(): + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor]]( + ConstantMultiIndexTensor(), + (torch.rand(5, 2, 3),), + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + +@common.parametrize( + "index", + { + "multidimensional": torch.tensor([[0, 1]], dtype=torch.int32), + "boolean": torch.tensor([False, True, False, True, False]), + }, +) +def test_index_tensor_u55_INT_constant_shape_or_dtype_not_delegated(index): + pipeline = OpNotSupportedPipeline[Tuple[torch.Tensor]]( + ConstantTensorIndex(index), + (torch.rand(5, 2, 3),), + {IndexTensorTestCommon.exir_op: 1}, + quantize=True, + u55_subset=True, + ) + pipeline.run() + + +def test_index_tensor_u55_INT_constant_symbolic_dim_not_delegated(): + indexed_dim = torch.export.Dim("indexed_dim", min=4, max=8) + tester = ArmTester( + ConstantIndexTensor([1, 2, 3]), + (torch.rand(5, 2, 3),), + common.get_u55_compile_spec(), + dynamic_shapes={"x": {0: indexed_dim}}, + ) + tester.quantize().export().to_edge().partition() + + targets = { + node.target + for node in tester.stages[tester.cur].artifact.exported_program().graph.nodes + } + assert exir_ops.edge.aten.index.Tensor in targets + assert torch.ops.higher_order.executorch_call_delegate not in targets diff --git a/backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py b/backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py new file mode 100644 index 00000000000..2d99bca9d70 --- /dev/null +++ b/backends/arm/test/passes/test_decompose_index_tensor_to_gather_pass.py @@ -0,0 +1,53 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest +import torch +from executorch.backends.arm._passes.decompose_index_tensor_to_gather_pass import ( + DecomposeIndexTensorToGatherPass, +) +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from executorch.exir import to_edge +from torch.export import export + + +class ConstantIndexTensor(torch.nn.Module): + def __init__(self, dim: int, index: int): + super().__init__() + self.dim = dim + self.index: torch.Tensor + self.register_buffer("index", torch.tensor([index], dtype=torch.int32)) + + def forward(self, x: torch.Tensor): + indices: list[slice | torch.Tensor] = [slice(None)] * x.dim() + indices[self.dim] = self.index + return x[tuple(indices)] + + +@pytest.mark.parametrize( + "dim,index", + ( + (0, 5), + (1, -6), + ), +) +def test_constant_out_of_bounds_index_raises(dim: int, index: int): + exported_program = export( + ConstantIndexTensor(dim, index), + (torch.rand(5, 5, 5),), + ) + edge_program = to_edge(exported_program) + edge_exported_program = edge_program.exported_program() + decompose_pass = DecomposeIndexTensorToGatherPass(edge_exported_program) + + with TosaLoweringContext(TosaSpecification.create_from_string("TOSA-1.0+INT")): + with pytest.raises( + IndexError, + match=rf"index {index} is out of bounds for dimension {dim} with size 5", + ): + decompose_pass(edge_exported_program.graph_module) From 80230c90aa2548137332267e25813095224e97c5 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Thu, 10 Sep 2026 15:40:09 +0100 Subject: [PATCH 146/190] Arm backend: Fix Vulkan data graph shader module feature negotiation (#22679) Found with vulkan validation layer cc @SS-JIA @manuelcandales @digantdesai @cbilgin @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Elena Zhelezina --- backends/arm/CMakeLists.txt | 21 ++++++++ backends/arm/runtime/VGFBackend.cpp | 25 +++++++--- backends/arm/runtime/VGFVulkanFeatures.h | 34 +++++++++++++ backends/arm/runtime/targets.bzl | 1 + backends/arm/test/targets.bzl | 15 ++++++ .../arm/test/vgf_vulkan_features_test.cpp | 48 +++++++++++++++++++ 6 files changed, 138 insertions(+), 6 deletions(-) create mode 100644 backends/arm/runtime/VGFVulkanFeatures.h create mode 100644 backends/arm/test/vgf_vulkan_features_test.cpp diff --git a/backends/arm/CMakeLists.txt b/backends/arm/CMakeLists.txt index 45640e8f3dd..b7e37748c73 100644 --- a/backends/arm/CMakeLists.txt +++ b/backends/arm/CMakeLists.txt @@ -315,5 +315,26 @@ if(EXECUTORCH_BUILD_VGF) executorch_target_link_options_shared_lib(vgf_backend) + if(EXECUTORCH_BUILD_TESTS) + add_executable( + vgf_vulkan_features_test + ${EXECUTORCH_ROOT}/backends/arm/test/vgf_vulkan_features_test.cpp + ) + target_include_directories( + vgf_vulkan_features_test + PRIVATE ${_common_include_directories} ${VULKAN_HEADERS_PATH} + ${VOLK_HEADERS_PATH} + ) + target_compile_options( + vgf_vulkan_features_test PRIVATE -DUSE_VULKAN_WRAPPER -DUSE_VULKAN_VOLK + ) + if(TARGET GTest::gtest_main) + target_link_libraries(vgf_vulkan_features_test PRIVATE GTest::gtest_main) + else() + target_link_libraries(vgf_vulkan_features_test PRIVATE gtest gtest_main) + endif() + add_test(NAME vgf_vulkan_features_test COMMAND vgf_vulkan_features_test) + endif() + # end config for VGF builds endif() diff --git a/backends/arm/runtime/VGFBackend.cpp b/backends/arm/runtime/VGFBackend.cpp index c7b735376a4..e29bbe7de9f 100644 --- a/backends/arm/runtime/VGFBackend.cpp +++ b/backends/arm/runtime/VGFBackend.cpp @@ -58,6 +58,8 @@ using executorch::runtime::EventTracerEntry; // We use the platform and runtime environment provided by the Vulkan delegate #include +#include + // Dependencies for processing VGF files into Vulkan calls #include #include @@ -961,17 +963,19 @@ VkResult vkml_allocate_basics( .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, .pNext = &available_12, }; + VkPhysicalDeviceDataGraphFeaturesARM available_graph{ + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM, &available_11}; #if defined(VK_ARM_data_graph_neural_accelerator_statistics) VkPhysicalDeviceDataGraphNeuralAcceleratorStatisticsFeaturesARM available_neural_statistics{ .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_NEURAL_ACCELERATOR_STATISTICS_FEATURES_ARM, - .pNext = &available_11, + .pNext = &available_graph, .dataGraphNeuralAcceleratorStatistics = VK_FALSE, }; void* available_features_pnext = &available_neural_statistics; #else - void* available_features_pnext = &available_11; + void* available_features_pnext = &available_graph; #endif VkPhysicalDeviceFeatures2 available_2 = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_FEATURES_2, @@ -979,6 +983,17 @@ VkResult vkml_allocate_basics( }; vkGetPhysicalDeviceFeatures2(*physical_device, &available_2); + if (!vgf_data_graph_features_supported(available_graph)) { + ET_LOG( + Error, + "VGF requires VK_ARM_data_graph features dataGraph and " + "dataGraphShaderModule (reported dataGraph=%u, " + "dataGraphShaderModule=%u)", + available_graph.dataGraph, + available_graph.dataGraphShaderModule); + return VK_ERROR_FEATURE_NOT_PRESENT; + } + // Select features VkPhysicalDeviceShaderReplicatedCompositesFeaturesEXT features_c{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_REPLICATED_COMPOSITES_FEATURES_EXT, @@ -1013,10 +1028,8 @@ VkResult vkml_allocate_basics( features_tensor.shaderTensorAccess = true; features_tensor.tensors = true; features_tensor.pNext = &features_11; - VkPhysicalDeviceDataGraphFeaturesARM features_graph{ - VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM, nullptr}; - features_graph.dataGraph = true; - features_graph.pNext = &features_tensor; + VkPhysicalDeviceDataGraphFeaturesARM features_graph = + make_vgf_data_graph_features(&features_tensor); #if defined(VK_ARM_data_graph_neural_accelerator_statistics) VkPhysicalDeviceDataGraphNeuralAcceleratorStatisticsFeaturesARM features_neural_statistics{ diff --git a/backends/arm/runtime/VGFVulkanFeatures.h b/backends/arm/runtime/VGFVulkanFeatures.h new file mode 100644 index 00000000000..8e88afc8f32 --- /dev/null +++ b/backends/arm/runtime/VGFVulkanFeatures.h @@ -0,0 +1,34 @@ +/* + * Copyright 2026 Arm Limited and/or its affiliates. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace executorch { +namespace backends { +namespace vgf { + +inline VkPhysicalDeviceDataGraphFeaturesARM make_vgf_data_graph_features( + void* p_next) { + VkPhysicalDeviceDataGraphFeaturesARM features{}; + features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM; + features.pNext = p_next; + features.dataGraph = VK_TRUE; + features.dataGraphShaderModule = VK_TRUE; + return features; +} + +inline bool vgf_data_graph_features_supported( + const VkPhysicalDeviceDataGraphFeaturesARM& features) { + return features.dataGraph == VK_TRUE && + features.dataGraphShaderModule == VK_TRUE; +} + +} // namespace vgf +} // namespace backends +} // namespace executorch diff --git a/backends/arm/runtime/targets.bzl b/backends/arm/runtime/targets.bzl index 0ba2dec3994..230e95e8fba 100644 --- a/backends/arm/runtime/targets.bzl +++ b/backends/arm/runtime/targets.bzl @@ -51,6 +51,7 @@ def define_common_targets(): exported_headers = [ "VGFNeuralStatistics.h", "VGFSetup.h", + "VGFVulkanFeatures.h", ], # @lint-ignore BUCKLINT: Avoid `link_whole=True` (https://fburl.com/avoid-link-whole) link_whole = True, diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 53d6ddadd58..2597f74d9ac 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -185,3 +185,18 @@ def define_arm_tests(): "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:vulkan-headers", ], ) + + if not runtime.is_oss and _ENABLE_VGF: + runtime.cxx_test( + name = "vgf_vulkan_features_test", + srcs = ["vgf_vulkan_features_test.cpp"], + compiler_flags = [ + "-DUSE_VULKAN_WRAPPER", + "-DUSE_VULKAN_VOLK", + ], + deps = [ + "//executorch/backends/arm/runtime:vgf_backend", + "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:volk_arm", + "fbsource//third-party/vulkan-headers-1.4.343/v1.4.343/src:vulkan-headers", + ], + ) diff --git a/backends/arm/test/vgf_vulkan_features_test.cpp b/backends/arm/test/vgf_vulkan_features_test.cpp new file mode 100644 index 00000000000..149e4dcd9a9 --- /dev/null +++ b/backends/arm/test/vgf_vulkan_features_test.cpp @@ -0,0 +1,48 @@ +/* + * Copyright 2026 Arm Limited and/or its affiliates. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +namespace executorch { +namespace backends { +namespace vgf { +namespace { + +TEST(VgfVulkanFeaturesTest, EnablesDataGraphShaderModule) { + VkPhysicalDeviceTensorFeaturesARM next{}; + next.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_TENSOR_FEATURES_ARM; + + const auto features = make_vgf_data_graph_features(&next); + + EXPECT_EQ( + features.sType, + VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM); + EXPECT_EQ(features.pNext, &next); + EXPECT_EQ(features.dataGraph, VK_TRUE); + EXPECT_EQ(features.dataGraphShaderModule, VK_TRUE); +} + +// cppcheck-suppress syntaxError +TEST(VgfVulkanFeaturesTest, RequiresDataGraphShaderModuleSupport) { + VkPhysicalDeviceDataGraphFeaturesARM available{}; + available.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_DATA_GRAPH_FEATURES_ARM; + + EXPECT_FALSE(vgf_data_graph_features_supported(available)); + + available.dataGraph = VK_TRUE; + EXPECT_FALSE(vgf_data_graph_features_supported(available)); + + available.dataGraphShaderModule = VK_TRUE; + EXPECT_TRUE(vgf_data_graph_features_supported(available)); +} + +} // namespace +} // namespace vgf +} // namespace backends +} // namespace executorch From 84fd9cef7669d8fc34931dd49a4bcf15c323a2a8 Mon Sep 17 00:00:00 2001 From: Jacob Stevens Date: Thu, 10 Sep 2026 11:00:13 -0400 Subject: [PATCH 147/190] Add optimized Arm64 BF16 `_to_copy` (#22493) Summary: Add an optimized `_to_copy.out` kernel. On Arm64, contiguous FP32-to-BF16 and BF16-to-FP32 conversions process eight elements per NEON iteration and use the ExecuTorch threadpool above its grain size. Other dtype pairs, layouts, and platforms fall back to the portable implementation, which remains unchanged. Reviewed By: digantdesai Differential Revision: D118502073 Pull Request resolved: https://github.com/pytorch/executorch/pull/22493 --- kernels/optimized/cpu/op_to_copy.cpp | 184 ++++++++++++++ kernels/optimized/optimized.yaml | 5 + kernels/test/op_to_copy_test.cpp | 228 ++++++++++++++++++ kernels/test/targets.bzl | 2 +- .../executorch/build/build_variables.bzl | 2 + .../optimized/op_registration_util.bzl | 8 + 6 files changed, 428 insertions(+), 1 deletion(-) create mode 100644 kernels/optimized/cpu/op_to_copy.cpp diff --git a/kernels/optimized/cpu/op_to_copy.cpp b/kernels/optimized/cpu/op_to_copy.cpp new file mode 100644 index 00000000000..2720f46d042 --- /dev/null +++ b/kernels/optimized/cpu/op_to_copy.cpp @@ -0,0 +1,184 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +#if defined(__aarch64__) +#include + +#include +#endif + +#include +#include + +namespace torch { +namespace executor { +namespace native { + +using BFloat16 = executorch::aten::BFloat16; +using MemoryFormat = executorch::aten::MemoryFormat; +using ScalarType = executorch::aten::ScalarType; +using Tensor = executorch::aten::Tensor; + +Tensor& to_copy_out( + KernelRuntimeContext& ctx, + const Tensor& self, + bool non_blocking, + std::optional memory_format, + Tensor& out); + +namespace { + +#if defined(__aarch64__) +static_assert(sizeof(BFloat16) == sizeof(uint16_t)); +static_assert(std::is_trivially_copyable_v); + +void float_to_bfloat16_range( + const float* const input, + BFloat16* const output, + const int64_t begin, + const int64_t end) { + constexpr int64_t kVectorWidth = 8; + // Integer rounding preserves subnormals regardless of FPCR, like BFloat16. + const uint32x4_t magnitude_mask = vdupq_n_u32(0x7FFFFFFF); + const uint32x4_t infinity = vdupq_n_u32(0x7F800000); + const uint32x4_t mantissa_lsb_mask = vdupq_n_u32(1); + const uint32x4_t rounding_bias = vdupq_n_u32(0x7FFF); + const uint32x4_t canonical_nan = vdupq_n_u32(0x7FC00000); + + int64_t i = begin; +#if defined(__clang__) +#pragma unroll 4 +#elif defined(__GNUC__) +#pragma GCC unroll 4 +#endif + for (; i + kVectorWidth <= end; i += kVectorWidth) { + const uint32x4_t low_bits = vreinterpretq_u32_f32(vld1q_f32(input + i)); + const uint32x4_t high_bits = + vreinterpretq_u32_f32(vld1q_f32(input + i + 4)); + + const auto round_and_canonicalize = [&](const uint32x4_t bits) { + const uint32x4_t mantissa_lsb = + vandq_u32(vshrq_n_u32(bits, 16), mantissa_lsb_mask); + const uint32x4_t rounded = + vaddq_u32(bits, vaddq_u32(rounding_bias, mantissa_lsb)); + const uint32x4_t is_nan = + vcgtq_u32(vandq_u32(bits, magnitude_mask), infinity); + return vbslq_u32(is_nan, canonical_nan, rounded); + }; + + const uint16x4_t low = vshrn_n_u32(round_and_canonicalize(low_bits), 16); + const uint16x8_t result = + vshrn_high_n_u32(low, round_and_canonicalize(high_bits), 16); + std::memcpy(output + i, &result, sizeof(result)); + } + + for (; i < end; ++i) { + output[i] = static_cast(input[i]); + } +} + +void bfloat16_to_float_range( + const BFloat16* const input, + float* const output, + const int64_t begin, + const int64_t end) { + constexpr int64_t kVectorWidth = 8; + + int64_t i = begin; +#if defined(__clang__) +#pragma unroll 4 +#elif defined(__GNUC__) +#pragma GCC unroll 4 +#endif + for (; i + kVectorWidth <= end; i += kVectorWidth) { + uint16x8_t input_bits; + // Avoid aliasing BFloat16 storage as a NEON vector. + std::memcpy(&input_bits, input + i, sizeof(input_bits)); + const uint32x4_t low_bits = vshll_n_u16(vget_low_u16(input_bits), 16); + const uint32x4_t high_bits = vshll_high_n_u16(input_bits, 16); + vst1q_f32(output + i, vreinterpretq_f32_u32(low_bits)); + vst1q_f32(output + i + 4, vreinterpretq_f32_u32(high_bits)); + } + + for (; i < end; ++i) { + output[i] = static_cast(input[i]); + } +} + +template +bool convert_contiguous(const Tensor& self, Tensor& out) { + const auto numel = self.numel(); + if (numel == 0) { + return true; + } + + const auto* const input = self.const_data_ptr(); + auto* const output = out.mutable_data_ptr(); + const auto convert_range = [&](const auto begin, const auto end) { + if constexpr (std::is_same_v) { + float_to_bfloat16_range(input, output, begin, end); + } else { + bfloat16_to_float_range(input, output, begin, end); + } + }; + + if (numel > ::executorch::extension::internal::GRAIN_SIZE) { + return ::executorch::extension::parallel_for( + 0, numel, ::executorch::extension::internal::GRAIN_SIZE, convert_range); + } + convert_range(0, numel); + return true; +} +#endif + +} // namespace + +Tensor& opt_to_copy_out( + KernelRuntimeContext& ctx, + const Tensor& self, + bool non_blocking, + std::optional memory_format, + Tensor& out) { +#if defined(__aarch64__) + const bool float_to_bfloat16 = self.scalar_type() == ScalarType::Float && + out.scalar_type() == ScalarType::BFloat16; + const bool bfloat16_to_float = self.scalar_type() == ScalarType::BFloat16 && + out.scalar_type() == ScalarType::Float; + const bool supported_memory_format = !memory_format.has_value() || + memory_format.value() == MemoryFormat::Contiguous; + const bool can_use_optimized_kernel = + (float_to_bfloat16 || bfloat16_to_float) && !non_blocking && + supported_memory_format && tensor_is_default_dim_order(self) && + tensor_is_default_dim_order(out); + if (can_use_optimized_kernel) { + ET_KERNEL_CHECK( + ctx, + resize_tensor(out, self.sizes()) == Error::Ok, + InvalidArgument, + out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out); + + const bool success = float_to_bfloat16 + ? convert_contiguous(self, out) + : convert_contiguous(self, out); + ET_KERNEL_CHECK_MSG(ctx, success, Internal, out, "parallel_for failed"); + return out; + } +#endif + + return to_copy_out(ctx, self, non_blocking, memory_format, out); +} + +} // namespace native +} // namespace executor +} // namespace torch diff --git a/kernels/optimized/optimized.yaml b/kernels/optimized/optimized.yaml index 5a001afc7a0..1827411ff57 100644 --- a/kernels/optimized/optimized.yaml +++ b/kernels/optimized/optimized.yaml @@ -17,6 +17,11 @@ - arg_meta: null kernel_name: torch::executor::opt_log_softmax_out +- op: _to_copy.out + kernels: + - arg_meta: null + kernel_name: torch::executor::opt_to_copy_out + - op: add.out kernels: - arg_meta: null diff --git a/kernels/test/op_to_copy_test.cpp b/kernels/test/op_to_copy_test.cpp index 45b2b2f6020..a6fb6239390 100644 --- a/kernels/test/op_to_copy_test.cpp +++ b/kernels/test/op_to_copy_test.cpp @@ -7,7 +7,9 @@ */ #include +#include #include +#include #include #include @@ -22,6 +24,7 @@ #include using namespace ::testing; +using executorch::aten::BFloat16; using executorch::aten::MemoryFormat; using executorch::aten::ScalarType; using executorch::aten::Tensor; @@ -81,6 +84,22 @@ class OpToTest : public OperatorTest { const std::vector data_out; }; + static float float_from_bits(uint32_t bits) { + float value; + std::memcpy(&value, &bits, sizeof(value)); + return value; + } + + static uint32_t float_bits(float value) { + uint32_t bits; + std::memcpy(&bits, &value, sizeof(bits)); + return bits; + } + + static bool is_bfloat16_nan(BFloat16 value) { + return (value.x & 0x7FFF) > 0x7F80; + } + // Each test has different combination of input and output types. Therefore it // is a little bit mess if create template test case and custom data types for // both input data and output data. @@ -121,6 +140,81 @@ class OpToTest : public OperatorTest { } } + template < + typename INPUT_CTYPE, + ScalarType INPUT_DTYPE, + typename OUTPUT_CTYPE, + ScalarType OUTPUT_DTYPE> + void test_conversion_at_sizes( + const std::vector& sizes, + const std::vector& input_pattern, + const std::vector& expected_pattern) { + static_assert( + (std::is_same_v && + std::is_same_v) || + (std::is_same_v && + std::is_same_v), + "Only float/BFloat16 conversion pairs are supported"); + ASSERT_EQ(input_pattern.size(), expected_pattern.size()); + + TensorFactory tf_in; + TensorFactory tf_out; + for (const int32_t numel : sizes) { + SCOPED_TRACE(::testing::Message() << "numel=" << numel); + std::vector input_data; + std::vector expected_data; + input_data.reserve(numel); + expected_data.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + input_data.push_back(input_pattern[i % input_pattern.size()]); + expected_data.push_back(expected_pattern[i % expected_pattern.size()]); + } + + Tensor input = tf_in.make({numel}, input_data); + Tensor output = tf_out.zeros({numel}); + + Tensor& ret = op_to_copy_out( + input, + /*non_blocking=*/false, + executorch::aten::MemoryFormat::Contiguous, + output); + + EXPECT_EQ(&ret, &output); + const auto* const actual_data = ret.const_data_ptr(); + if (actual_data == nullptr) { + ADD_FAILURE() << "conversion returned a null data pointer"; + continue; + } + if constexpr (std::is_same_v) { + const bool is_aten = + torch::executor::testing::SupportedFeatures::get()->is_aten; + std::vector actual_bits; + std::vector expected_bits; + actual_bits.reserve(numel); + expected_bits.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + if (is_aten && is_bfloat16_nan(expected_data[i])) { + EXPECT_TRUE(is_bfloat16_nan(actual_data[i])) << "index=" << i; + continue; + } + actual_bits.push_back(actual_data[i].x); + expected_bits.push_back(expected_data[i].x); + } + EXPECT_EQ(actual_bits, expected_bits); + } else { + std::vector actual_bits; + std::vector expected_bits; + actual_bits.reserve(numel); + expected_bits.reserve(numel); + for (int32_t i = 0; i < numel; ++i) { + actual_bits.push_back(float_bits(actual_data[i])); + expected_bits.push_back(float_bits(expected_data[i])); + } + EXPECT_EQ(actual_bits, expected_bits); + } + } + } + template void test_runner_to_bool( std::vector test_case, @@ -360,6 +454,140 @@ TEST_F(OpToTest, NanInfSupported) { #undef TEST_KERNEL } +TEST_F(OpToTest, FloatToBFloat16RawBitsAtVectorAndGrainBoundaries) { + std::vector sizes; + for (int32_t size = 1; size <= 17; ++size) { + sizes.push_back(size); + } + sizes.insert(sizes.end(), {32767, 32768, 32769}); + + const std::vector input_pattern = { + float_from_bits(0x00000000), float_from_bits(0x80000000), + float_from_bits(0x00007FFF), float_from_bits(0x00008000), + float_from_bits(0x00008001), float_from_bits(0x00017FFF), + float_from_bits(0x00018000), float_from_bits(0x00018001), + float_from_bits(0x3F807FFF), float_from_bits(0x3F808000), + float_from_bits(0x3F808001), float_from_bits(0x3F817FFF), + float_from_bits(0x3F818000), float_from_bits(0x3F818001), + float_from_bits(0xBF807FFF), float_from_bits(0xBF808000), + float_from_bits(0xBF808001), float_from_bits(0xBF817FFF), + float_from_bits(0xBF818000), float_from_bits(0xBF818001), + float_from_bits(0x7F800000), float_from_bits(0xFF800000), + float_from_bits(0x7FC12345), float_from_bits(0xFFC12345), + float_from_bits(0x7FA12345), float_from_bits(0xFFA12345), + }; + const std::vector expected_pattern = { + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x8000, BFloat16::from_bits()), + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x0001, BFloat16::from_bits()), + BFloat16(0x0001, BFloat16::from_bits()), + BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0xBF80, BFloat16::from_bits()), + BFloat16(0xBF80, BFloat16::from_bits()), + BFloat16(0xBF81, BFloat16::from_bits()), + BFloat16(0xBF81, BFloat16::from_bits()), + BFloat16(0xBF82, BFloat16::from_bits()), + BFloat16(0xBF82, BFloat16::from_bits()), + BFloat16(0x7F80, BFloat16::from_bits()), + BFloat16(0xFF80, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + BFloat16(0x7FC0, BFloat16::from_bits()), + }; + + test_conversion_at_sizes< + float, + ScalarType::Float, + BFloat16, + ScalarType::BFloat16>(sizes, input_pattern, expected_pattern); +} + +#if defined(__aarch64__) +TEST_F(OpToTest, FloatToBFloat16SubnormalsIgnoreFlushToZero) { + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen conversion may flush subnormals to zero"); + + struct RestoreFpcr { + uint64_t value{}; + RestoreFpcr() { + asm volatile("mrs %0, fpcr" : "=r"(value)); + } + ~RestoreFpcr() { + asm volatile("msr fpcr, %0" : : "r"(value) : "memory"); + } + } original_fpcr; + + constexpr uint64_t kFlushToZero = uint64_t{1} << 24; + for (const bool flush_to_zero : {false, true}) { + SCOPED_TRACE(::testing::Message() << "flush_to_zero=" << flush_to_zero); + const uint64_t fpcr = flush_to_zero ? original_fpcr.value | kFlushToZero + : original_fpcr.value & ~kFlushToZero; + asm volatile("msr fpcr, %0" : : "r"(fpcr) : "memory"); + test_conversion_at_sizes< + float, + ScalarType::Float, + BFloat16, + ScalarType::BFloat16>( + {1, 7, 8, 9, 15, 16, 17}, + {float_from_bits(0x00018000), float_from_bits(0x80018000)}, + {BFloat16(0x0002, BFloat16::from_bits()), + BFloat16(0x8002, BFloat16::from_bits())}); + } +} +#endif + +TEST_F(OpToTest, BFloat16ToFloatRawBitsAtVectorAndGrainBoundaries) { + std::vector sizes; + for (int32_t size = 1; size <= 17; ++size) { + sizes.push_back(size); + } + sizes.insert(sizes.end(), {32767, 32768, 32769}); + + const std::vector input_pattern = { + BFloat16(0x0000, BFloat16::from_bits()), + BFloat16(0x8000, BFloat16::from_bits()), + BFloat16(0x3F80, BFloat16::from_bits()), + BFloat16(0x3F81, BFloat16::from_bits()), + BFloat16(0x3F82, BFloat16::from_bits()), + BFloat16(0x7F80, BFloat16::from_bits()), + BFloat16(0xFF80, BFloat16::from_bits()), + BFloat16(0x7FC1, BFloat16::from_bits()), + BFloat16(0xFFC1, BFloat16::from_bits()), + BFloat16(0x7FA1, BFloat16::from_bits()), + BFloat16(0xFFA1, BFloat16::from_bits()), + }; + const std::vector expected_pattern = { + float_from_bits(0x00000000), + float_from_bits(0x80000000), + float_from_bits(0x3F800000), + float_from_bits(0x3F810000), + float_from_bits(0x3F820000), + float_from_bits(0x7F800000), + float_from_bits(0xFF800000), + float_from_bits(0x7FC10000), + float_from_bits(0xFFC10000), + float_from_bits(0x7FA10000), + float_from_bits(0xFFA10000), + }; + + test_conversion_at_sizes< + BFloat16, + ScalarType::BFloat16, + float, + ScalarType::Float>(sizes, input_pattern, expected_pattern); +} + TEST_F(OpToTest, HardcodeFloatConvertInt) { // Hardcode input and output generated from core PyTorch // clang-format off diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 837c7327c4f..9084dd2b16d 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -352,7 +352,7 @@ def define_common_targets(): _common_op_test("op_t_copy_test", ["aten", "portable"]) _common_op_test("op_tan_test", ["aten", "portable"]) _common_op_test("op_tanh_test", ["aten", "portable"]) - _common_op_test("op_to_copy_test", ["aten", "portable"]) + _common_op_test("op_to_copy_test", ["aten", "portable", "optimized"]) _common_op_test("op_topk_test", ["aten", "portable"]) _common_op_test("op_transpose_copy_test", ["aten", "portable"]) _common_op_test("op_tril_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index c75436d3dc7..8dc8944b052 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -278,6 +278,7 @@ OPTIMIZED_KERNELS_SRCS = [ "kernels/optimized/cpu/op_native_layer_norm.cpp", "kernels/optimized/cpu/op_sub.cpp", "kernels/optimized/cpu/op_sum.cpp", + "kernels/optimized/cpu/op_to_copy.cpp", "kernels/optimized/cpu/op_where.cpp", ] @@ -320,6 +321,7 @@ OPTIMIZED_NATIVE_CPU_OPS_SRCS = [ "kernels/optimized/cpu/op_mul.cpp", "kernels/optimized/cpu/op_native_layer_norm.cpp", "kernels/optimized/cpu/op_sub.cpp", + "kernels/optimized/cpu/op_to_copy.cpp", "kernels/optimized/cpu/op_where.cpp", ] diff --git a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl index fba89adde64..25041897059 100644 --- a/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/optimized/op_registration_util.bzl @@ -327,6 +327,14 @@ OPTIMIZED_ATEN_OPS = ( "//executorch/kernels/portable/cpu/util:reduce_util", ], ), + op_target( + name = "op_to_copy", + deps = [ + "//executorch/extension/threadpool:threadpool", + "//executorch/kernels/portable/cpu:op_to_copy", + "//executorch/kernels/portable/cpu/util:copy_ops_util", + ], + ), op_target( name = "op_where", deps = [ From 9ad8f1facc0a58dfade7aee640384a7cb11ef729 Mon Sep 17 00:00:00 2001 From: Fatih Uzulmez Date: Thu, 10 Sep 2026 08:29:24 -0700 Subject: [PATCH 148/190] Fix eager eval token ID fallbacks (#22633) ## Summary - treat None-valued BOS and EOT token IDs as unavailable - fall back from BOS to EOT to EOS while preserving valid token ID 0 - add focused regression coverage for present, zero-valued, None-valued, and missing optional IDs ## Test plan - python3 -m py_compile examples/models/llama/evaluate/eager_eval.py examples/models/llama/tests/test_eager_eval.py - focused six-case fallback matrix: passed - git diff --check Fixes #22518 cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani --- examples/models/llama/evaluate/eager_eval.py | 10 +++-- examples/models/llama/tests/BUCK | 11 +++++ .../models/llama/tests/test_eager_eval.py | 42 +++++++++++++++++++ 3 files changed, 59 insertions(+), 4 deletions(-) create mode 100644 examples/models/llama/tests/test_eager_eval.py diff --git a/examples/models/llama/evaluate/eager_eval.py b/examples/models/llama/evaluate/eager_eval.py index cefb951281c..b094b2245cf 100644 --- a/examples/models/llama/evaluate/eager_eval.py +++ b/examples/models/llama/evaluate/eager_eval.py @@ -44,14 +44,16 @@ def eot_token_id(self): """ The stories model does not have an EOT token, so we use the EOS token instead. """ - if hasattr(self._tokenizer, "eot_id"): - return self._tokenizer.eot_id + eot_id = getattr(self._tokenizer, "eot_id", None) + if eot_id is not None: + return eot_id return self._tokenizer.eos_id @property def prefix_token_id(self): - if hasattr(self._tokenizer, "bos_id"): - return self._tokenizer.bos_id + bos_id = getattr(self._tokenizer, "bos_id", None) + if bos_id is not None: + return bos_id return self.eot_token_id @property diff --git a/examples/models/llama/tests/BUCK b/examples/models/llama/tests/BUCK index 74430d9e306..dc3401ced83 100644 --- a/examples/models/llama/tests/BUCK +++ b/examples/models/llama/tests/BUCK @@ -3,6 +3,17 @@ load("@fbcode_macros//build_defs:python_unittest.bzl", "python_unittest") oncall("executorch") +fbcode_target(_kind = python_unittest, + name = "test_eager_eval", + srcs = [ + "test_eager_eval.py", + ], + deps = [ + "//executorch/examples/models/llama:eval_library", + "fbsource//third-party/pypi/pytest:pytest", + ], +) + fbcode_target(_kind = python_unittest, name = "test_simple_sdpa", srcs = [ diff --git a/examples/models/llama/tests/test_eager_eval.py b/examples/models/llama/tests/test_eager_eval.py new file mode 100644 index 00000000000..bb53292b206 --- /dev/null +++ b/examples/models/llama/tests/test_eager_eval.py @@ -0,0 +1,42 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest +from types import SimpleNamespace +from typing import Optional + +import pytest + +pytest.importorskip("lm_eval", reason="requires lm-evaluation-harness") + +from executorch.examples.models.llama.evaluate.eager_eval import ( # noqa: E402 + EagerEvalWrapper, +) + + +class TestEagerEvalWrapperTokenIds(unittest.TestCase): + @staticmethod + def _wrapper(**token_ids: Optional[int]) -> EagerEvalWrapper: + # HFLM initialization loads a model and is unrelated to these properties. + wrapper = object.__new__(EagerEvalWrapper) + wrapper._tokenizer = SimpleNamespace(**token_ids) # pyre-ignore[8] + return wrapper + + def test_token_id_fallbacks(self): + cases = ( + ({"bos_id": 1, "eot_id": 2, "eos_id": 3}, 2, 1), + ({"bos_id": 0, "eot_id": 2, "eos_id": 3}, 2, 0), + ({"bos_id": None, "eot_id": 2, "eos_id": 3}, 2, 2), + ({"bos_id": None, "eot_id": 0, "eos_id": 3}, 0, 0), + ({"bos_id": None, "eot_id": None, "eos_id": 0}, 0, 0), + ({"eos_id": 0}, 0, 0), + ) + + for token_ids, expected_eot, expected_prefix in cases: + with self.subTest(token_ids=token_ids): + wrapper = self._wrapper(**token_ids) + self.assertEqual(wrapper.eot_token_id, expected_eot) + self.assertEqual(wrapper.prefix_token_id, expected_prefix) From 31f3916b17e69bcb92c5a83ab687d823f7cd7db2 Mon Sep 17 00:00:00 2001 From: Peyman Gardideh Date: Thu, 10 Sep 2026 11:33:18 -0400 Subject: [PATCH 149/190] Remove dynamic allocation from Ethos-U backend (#22531) Summary: When running ET in embedded Cortex-M we should not dynamically allocate. This `new` call added to support Cortex-A is breaking that contract. Instead, allocate ExecutionHandle on the runtime allocator and delegate the platform_state allocation/destruction to the platform. EthosUBackend_Cortex_M does not create a PlatformState so we can safely remove the delete call from platform_destory Reviewed By: rascani Differential Revision: D117747705 Pull Request resolved: https://github.com/pytorch/executorch/pull/22531 --- backends/arm/runtime/EthosUBackend.cpp | 7 ++++--- backends/arm/runtime/EthosUBackend_Cortex_M.cpp | 5 +---- 2 files changed, 5 insertions(+), 7 deletions(-) diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index e7325f96a94..079da95391d 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -125,10 +125,11 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { } MemoryAllocator* allocator = context.get_runtime_allocator(); - ExecutionHandle* handle = new (std::nothrow) ExecutionHandle(); + ExecutionHandle* handle = allocator->allocateInstance(); if (handle == nullptr) { return Error::MemoryAllocationFailed; } + new (handle) ExecutionHandle(); EXECUTORCH_PROF_START( event_tracer, @@ -138,7 +139,7 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { data, size, context.get_named_data_map(), &handle->handles); EXECUTORCH_PROF_END(event_tracer, event_tracer_local_scope); if (read_status != Error::Ok) { - delete handle; + handle->~ExecutionHandle(); return read_status; } @@ -324,7 +325,7 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { platform_destroy(exec_handle->platform_state); } - delete exec_handle; + exec_handle->~ExecutionHandle(); } private: diff --git a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp index f33a523986d..eab91382247 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp @@ -13,7 +13,6 @@ #include #include #include -#include #include @@ -60,9 +59,7 @@ PlatformState* platform_init( return nullptr; } -void platform_destroy(PlatformState* state) { - delete state; -} +void platform_destroy(PlatformState* /*state*/) {} bool needs_scratch_allocation() { return true; From d9dae2ddcf1cdcca0914422058914631ee1863a5 Mon Sep 17 00:00:00 2001 From: jathu Date: Thu, 10 Sep 2026 08:58:06 -0700 Subject: [PATCH 150/190] Don't upload golden artifacts to GitHub Actions storage (#22656) ### Summary Addresses step 4: https://docs.google.com/document/d/1CnGfiP5SU0PG0IdJWn9Ca43LS8bGlEMi3l2gjgC7k48/edit * We upload the golden artifacts to GitHub Actions storage _and_ S3. Nothing uses the GHA copy; the Android test uses the S3 one https://github.com/pytorch/executorch/blob/5b7718754e2ab4d634dc8e7ad837209345307983/extension/android/executorch_android/android_test_setup.sh#L34 So keep only the S3 upload. * `test_backend.sh` always set `GOLDEN_ARTIFACTS_DIR`, so every PR and push to main wrote goldens into the test reports of every backend, about 2 TB/day. Now only the xnnpack nightly writes them (that's where the Android pin comes from), plus a manual `save-goldens` dispatch when we need a fresh one. ### Test plan CI. On this PR `package-golden-artifacts` is skipped and the models test reports are a few MB instead of ~1.7 GB. _Authored with Claude Code._ --- .ci/scripts/test_backend.sh | 3 --- .github/workflows/_test_backend.yml | 17 +++++++++-------- .github/workflows/test-backend-xnnpack.yml | 9 +++++++++ 3 files changed, 18 insertions(+), 11 deletions(-) diff --git a/.ci/scripts/test_backend.sh b/.ci/scripts/test_backend.sh index 95c5c8e9db9..068d5adb260 100755 --- a/.ci/scripts/test_backend.sh +++ b/.ci/scripts/test_backend.sh @@ -154,9 +154,6 @@ if [[ "$FLOW" == *nxp* ]]; then export NXP_RUNNER_PATH="$(pwd)/examples/nxp/executor_runner/build/nxp_executor_runner" fi -GOLDEN_DIR="${ARTIFACT_DIR}/golden-artifacts" -export GOLDEN_ARTIFACTS_DIR="${GOLDEN_DIR}" - EXIT_CODE=0 # An Ethos-U failure captures a few hundred thousand lines of Vela operator # listings, and the runner agent throws System.OutOfMemoryException processing diff --git a/.github/workflows/_test_backend.yml b/.github/workflows/_test_backend.yml index dda19fa033f..063c0ebb161 100644 --- a/.github/workflows/_test_backend.yml +++ b/.github/workflows/_test_backend.yml @@ -46,6 +46,11 @@ on: required: false type: string default: ci-image:executorch-ubuntu-22.04-clang12 + save-goldens: + description: 'Write golden .pte/.bin files in the models suite and package them; only for regenerating the Android test fixture' + required: false + type: boolean + default: false jobs: docker-image: @@ -77,10 +82,13 @@ jobs: script: | set -eux + if [[ "${{ inputs.save-goldens }}" == "true" && "${{ matrix.suite }}" == "models" ]]; then + export GOLDEN_ARTIFACTS_DIR="${RUNNER_ARTIFACT_DIR}/golden-artifacts" + fi source .ci/scripts/test_backend.sh "${{ matrix.suite }}" "${{ matrix.flow }}" "${RUNNER_ARTIFACT_DIR}" package-golden-artifacts: - if: ${{ inputs.run-linux }} + if: ${{ inputs.run-linux && inputs.save-goldens }} needs: test-backend-linux runs-on: linux.2xlarge steps: @@ -116,13 +124,6 @@ jobs: echo "No golden artifacts found." fi - - name: Upload combined golden artifacts - uses: actions/upload-artifact@v4 - with: - name: golden-artifacts-${{ inputs.backend }} - path: golden_artifacts_*.zip - if-no-files-found: ignore - - name: Upload golden artifacts to S3 uses: seemethere/upload-artifact-s3@v5 if: ${{ hashFiles('golden_artifacts_*.zip') != '' }} diff --git a/.github/workflows/test-backend-xnnpack.yml b/.github/workflows/test-backend-xnnpack.yml index c5abb15e837..a0906af5bb2 100644 --- a/.github/workflows/test-backend-xnnpack.yml +++ b/.github/workflows/test-backend-xnnpack.yml @@ -11,6 +11,11 @@ on: - ciflow/nightly/* pull_request: workflow_dispatch: + inputs: + save-goldens: + description: 'Write and package the models-suite goldens (Android test fixture)' + type: boolean + default: false concurrency: group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref_name }}-${{ github.ref_type == 'branch' && github.sha }}-${{ github.event_name == 'workflow_dispatch' }}-${{ github.event_name == 'schedule' }} @@ -31,3 +36,7 @@ jobs: ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 run-linux: true + # The nightly keeps writing goldens: they are what the Android instrumentation + # test pins (android_test_setup.sh), and a fresh key must exist before the pin + # ages out of S3. Every other run skips them. + save-goldens: ${{ inputs.save-goldens == true || github.event_name == 'schedule' }} From 5c36c2e1e5031aabd5fc1e4a9785c214c5e2eb87 Mon Sep 17 00:00:00 2001 From: Andrew Date: Thu, 10 Sep 2026 09:24:31 -0700 Subject: [PATCH 151/190] Lower branched A16W8 convs on U55 (#22566) Summary: A16W8 convolution rewriting can expose one INT48 accumulator to multiple output `RESCALE` nodes. Regor cannot materialize INT48 on U55 and can only fuse one output quantization into each convolution, so these graphs fail compilation before scheduling. This fixes a U55 lowering regression introduced by https://github.com/pytorch/executorch/pull/21910: its A16W8 output branching can leave one INT48 convolution accumulator feeding multiple `RESCALE` paths, but Regor cannot materialize INT48 and can fold only one `RESCALE` into a convolution. Deduplicate identical complete `RESCALE` plus layout-permute heads, then clone the TOSA convolution for each remaining distinct output quantization while sharing its input, weight, and bias nodes. Mark those exact-semantics clones so later duplicate-user cleanup cannot merge the unsupported fanout back together. Add structural, duplicate-fusion barrier, and true U55 compilation regressions. This change was authored with Codex. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Differential Revision: D118736607 Pull Request resolved: https://github.com/pytorch/executorch/pull/22566 --- backends/arm/_passes/rewrite_conv_pass.py | 104 +++++++++- .../passes/test_fuse_duplicate_users_pass.py | 19 ++ .../arm/test/passes/test_rewrite_conv_pass.py | 183 +++++++++++++++++- .../transforms/fuse_duplicate_users_pass.py | 133 +++++++------ 4 files changed, 375 insertions(+), 64 deletions(-) diff --git a/backends/arm/_passes/rewrite_conv_pass.py b/backends/arm/_passes/rewrite_conv_pass.py index dc00baebaac..1ed60228582 100644 --- a/backends/arm/_passes/rewrite_conv_pass.py +++ b/backends/arm/_passes/rewrite_conv_pass.py @@ -37,11 +37,20 @@ TOSA_CONTROL_FLOW_SOURCE_NODE_META, TosaSpecialDtype, ) -from executorch.backends.arm.tosa.specification import get_context_shape_env +from executorch.backends.arm.tosa.specification import ( + get_context_shape_env, + get_context_spec, +) +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + build_node_signature, + DO_NOT_FUSE_DUPLICATE_META_KEY, +) from executorch.backends.transforms.utils import create_constant_placeholder from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.dialects.edge._ops import EdgeOpOverload from executorch.exir.pass_base import ExportPass, PassResult +from torch._ops import OpOverload from torch._subclasses.fake_tensor import FakeTensor from torch.export.graph_signature import InputKind @@ -571,6 +580,88 @@ def _insert_layout_permute( output.meta["val"] = output_fake_tensor return output, output_fake_tensor + @classmethod + def _deduplicate_a16w8_output_rescales( + cls, + graph_module: torch.fx.GraphModule, + tosa_op: torch.fx.Node, + node_order: dict[torch.fx.Node, int], + ) -> list[torch.fx.Node] | None: + """Merge only complete, canonical RESCALE-to-PERMUTE heads.""" + if any(user not in node_order for user in tosa_op.users): + return None + rescale_users = sorted(tosa_op.users, key=node_order.__getitem__) + if any( + user.target != exir_ops.backend.tosa.RESCALE.default + for user in rescale_users + ): + # RewriteConvPass creates only RESCALE users for this accumulator; + # preserve an unfamiliar future shape instead of partially rewriting it. + return None + + unique_rescales: dict[tuple[Any, ...], torch.fx.Node] = {} + deduplicated_rescales: list[torch.fx.Node] = [] + for rescale in rescale_users: + rescale_outputs = list(rescale.users) + if ( + len(rescale_outputs) != 1 + or rescale_outputs[0].target != exir_ops.edge.aten.permute_copy.default + ): + deduplicated_rescales.append(rescale) + continue + layout_permute = rescale_outputs[0] + rescale_signature = build_node_signature(rescale, positional_arg_start=1) + permute_signature = build_node_signature( + layout_permute, positional_arg_start=1 + ) + if rescale_signature is None or permute_signature is None: + deduplicated_rescales.append(rescale) + continue + signature = ( + rescale_signature, + permute_signature, + ) + canonical_permute = unique_rescales.get(signature) + if canonical_permute is not None: + # Layout permutes are inserted directly after their RESCALE, + # so the earliest RESCALE also provides a dominating permute. + layout_permute.replace_all_uses_with(canonical_permute) + graph_module.graph.erase_node(layout_permute) + graph_module.graph.erase_node(rescale) + else: + unique_rescales[signature] = layout_permute + deduplicated_rescales.append(rescale) + + return deduplicated_rescales + + def _separate_u55_a16w8_output_rescales( + self, + graph_module: torch.fx.GraphModule, + tosa_op: torch.fx.Node, + node_order: dict[torch.fx.Node, int], + ) -> None: + if len(tosa_op.users) < 2: + return + + rescale_users = self._deduplicate_a16w8_output_rescales( + graph_module, tosa_op, node_order + ) + if rescale_users is None or len(rescale_users) < 2: + return + tosa_op.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + for rescale in rescale_users[1:]: + with graph_module.graph.inserting_before(rescale): + cloned_tosa_op = create_node( + graph=graph_module.graph, + op_target=cast(OpOverload | EdgeOpOverload, tosa_op.target), + args=tosa_op.args, + kwargs=tosa_op.kwargs, + from_node=tosa_op, + inherit_qparams=True, + ) + cloned_tosa_op.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + rescale.replace_input_with(tosa_op, cloned_tosa_op) + def _insert_a16w8_output_branches( self, graph_module: torch.fx.GraphModule, @@ -787,6 +878,7 @@ def _insert_output_conversion( def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 modified = False + a16w8_tosa_ops: list[torch.fx.Node] = [] for node in graph_module.graph.nodes: if ( node.op != "call_function" @@ -1172,11 +1264,21 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult: # noqa: C901 if squeeze_view is not None: graph_module.graph.erase_node(squeeze_view) graph_module.graph.erase_node(output_conversion_node) + a16w8_tosa_ops.append(tosa_op) else: node.replace_all_uses_with(node_replacement) graph_module.graph.erase_node(node) + if a16w8_tosa_ops and get_context_spec().is_U55_subset: + node_order = { + node: index for index, node in enumerate(graph_module.graph.nodes) + } + for tosa_op in a16w8_tosa_ops: + self._separate_u55_a16w8_output_rescales( + graph_module, tosa_op, node_order + ) + if modified: graph_module.recompile() graph_module = super().call(graph_module).graph_module diff --git a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py index 893d9eefea5..831d9d91267 100644 --- a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py +++ b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py @@ -21,6 +21,9 @@ TosaLoweringContext, TosaSpecification, ) +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + DO_NOT_FUSE_DUPLICATE_META_KEY, +) from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export @@ -163,6 +166,22 @@ def test_fuse_duplicate_users_preserves_graph_order_for_representative(): assert len(_add_node_names(result.graph_module)) == 1 +def test_fuse_duplicate_users_honors_do_not_fuse_marker(): + graph_module = _graph_with_users_not_in_node_order() + marked_node = next( + node + for node in graph_module.graph.nodes + if node.target == torch.ops.aten.add.Tensor + ) + marked_node.meta[DO_NOT_FUSE_DUPLICATE_META_KEY] = True + + result = FuseDuplicateUsersPass()(graph_module) + + result.graph_module.graph.lint() + assert not result.modified + assert len(_add_node_names(result.graph_module)) == 2 + + def test_fuse_duplicate_users_keeps_identical_rescale_users(): graph_module = _graph_with_duplicate_rescale_users() diff --git a/backends/arm/test/passes/test_rewrite_conv_pass.py b/backends/arm/test/passes/test_rewrite_conv_pass.py index e1605faefa6..3928aa81e75 100644 --- a/backends/arm/test/passes/test_rewrite_conv_pass.py +++ b/backends/arm/test/passes/test_rewrite_conv_pass.py @@ -27,7 +27,10 @@ from executorch.backends.arm.test.misc.test_dw_convs_with_shared_weights import ( DWConvsModule, ) -from executorch.backends.arm.test.tester.test_pipeline import PassPipeline +from executorch.backends.arm.test.tester.test_pipeline import ( + EthosU55PipelineINT, + PassPipeline, +) from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec from executorch.backends.arm.tosa.mapping import TosaSpecialDtype from executorch.backends.arm.tosa.partitioner import TOSAPartitioner @@ -36,6 +39,9 @@ TosaSpecification, ) from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner +from executorch.backends.transforms.fuse_duplicate_users_pass import ( + build_node_signature, +) from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower from executorch.exir.dialects._ops import ops as exir_ops from torch.export import Dim, export @@ -267,7 +273,9 @@ def _get_expected_int32_scales( def _rewrite_a16w8_convs( - model: nn.Module, inputs: tuple[torch.Tensor, ...] + model: nn.Module, + inputs: tuple[torch.Tensor, ...], + tosa_spec: TosaSpecification | None = None, ) -> tuple[torch.fx.GraphModule, list[list[float]]]: """Run the passes needed to inspect rewritten A16W8 convolutions.""" exported_program = _export_quantized_a16w8(model, inputs) @@ -276,7 +284,7 @@ def _rewrite_a16w8_convs( ).exported_program() gm = _run_pre_rewrite_passes(edge_program) rewrite_pass = RewriteConvPass(edge_program) - with TosaLoweringContext(_compile_spec_int16().tosa_spec): + with TosaLoweringContext(tosa_spec or _compile_spec_int16().tosa_spec): rescale_result = InsertRescaleInt32Pass()(gm) assert rescale_result is not None expected_int32_scales = _get_expected_int32_scales( @@ -294,6 +302,29 @@ def _get_call_function_node(gm: torch.fx.GraphModule, target): raise AssertionError(f"Node with target {target} not found") +def _add_a16w8_rescale_head( + graph: torch.fx.Graph, + accumulator: torch.fx.Node, + positional_unsigned: tuple[bool, ...] = (), +) -> tuple[torch.fx.Node, torch.fx.Node]: + rescale = graph.call_function( + exir_ops.backend.tosa.RESCALE.default, + args=( + accumulator, + torch.int16, + [1.0], + 0, + 0, + *positional_unsigned, + ), + ) + layout_permute = graph.call_function( + exir_ops.edge.aten.permute_copy.default, + args=(rescale, [0, 3, 1, 2]), + ) + return rescale, layout_permute + + class ConvModule(torch.nn.Module): def __init__(self): super().__init__() @@ -510,6 +541,152 @@ def test_rewrite_conv_a16w8_mixed_consumers_restore_int16( ) +def test_rewrite_conv_rescale_signature_includes_positional_unsigned_flags() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + signed_rescale, _ = _add_a16w8_rescale_head(graph, accumulator, (False, False)) + unsigned_rescale, _ = _add_a16w8_rescale_head(graph, accumulator, (False, True)) + + assert build_node_signature( + signed_rescale, positional_arg_start=1 + ) != build_node_signature(unsigned_rescale, positional_arg_start=1) + + +def test_rewrite_conv_without_convolution_does_not_require_context() -> None: + inputs = (torch.randn(1, 4),) + edge_program = to_edge(export(nn.Identity(), inputs)).exported_program() + + result = RewriteConvPass(edge_program)(edge_program.graph_module) + + assert result is not None + assert not result.modified + + +def test_rewrite_conv_a16w8_unknown_accumulator_user_is_unchanged() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + _, first_permute = _add_a16w8_rescale_head(graph, accumulator) + _, second_permute = _add_a16w8_rescale_head(graph, accumulator) + unexpected_user = graph.call_function(torch.neg, args=(accumulator,)) + graph.output((first_permute, second_permute, unexpected_user)) + graph_module = torch.fx.GraphModule({}, graph) + nodes_before = list(graph.nodes) + users_before = {node: tuple(node.users) for node in graph.nodes} + node_order = {node: index for index, node in enumerate(graph.nodes)} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result is None + assert list(graph.nodes) == nodes_before + assert {node: tuple(node.users) for node in graph.nodes} == users_before + graph.lint() + + +def test_rewrite_conv_a16w8_multi_consumer_rescale_is_not_deduplicated() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + first_rescale, first_permute = _add_a16w8_rescale_head(graph, accumulator) + second_rescale, second_permute = _add_a16w8_rescale_head(graph, accumulator) + output = graph.output((first_permute, second_permute, second_rescale)) + graph_module = torch.fx.GraphModule({}, graph) + node_order = {node: index for index, node in enumerate(graph.nodes)} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [first_rescale, second_rescale] + assert set(second_rescale.users) == {second_permute, output} + graph.lint() + + +def test_rewrite_conv_a16w8_deduplication_uses_graph_order() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + temporary_input = graph.placeholder("temporary_input") + early_rescale, early_permute = _add_a16w8_rescale_head(graph, temporary_input) + late_rescale, late_permute = _add_a16w8_rescale_head(graph, accumulator) + early_rescale.replace_input_with(temporary_input, accumulator) + graph.output((early_permute, late_permute)) + graph_module = torch.fx.GraphModule({}, graph) + node_order = {node: index for index, node in enumerate(graph.nodes)} + + assert list(accumulator.users) == [late_rescale, early_rescale] + assert node_order[early_rescale] < node_order[late_rescale] + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result == [early_rescale] + assert late_rescale not in graph.nodes + assert late_permute not in graph.nodes + graph.lint() + + +def test_rewrite_conv_a16w8_unknown_order_user_is_unchanged() -> None: + graph = torch.fx.Graph() + accumulator = graph.placeholder("accumulator") + _, first_permute = _add_a16w8_rescale_head(graph, accumulator) + node_order = {node: index for index, node in enumerate(graph.nodes)} + _, second_permute = _add_a16w8_rescale_head(graph, accumulator) + graph.output((first_permute, second_permute)) + graph_module = torch.fx.GraphModule({}, graph) + nodes_before = list(graph.nodes) + users_before = {node: tuple(node.users) for node in graph.nodes} + + result = RewriteConvPass._deduplicate_a16w8_output_rescales( + graph_module, accumulator, node_order + ) + + assert result is None + assert list(graph.nodes) == nodes_before + assert {node: tuple(node.users) for node in graph.nodes} == users_before + graph.lint() + + +def test_rewrite_conv_a16w8_u55_separates_distinct_output_rescales() -> None: + model = A16W8MixedConsumerChain(nn.Conv2d(4, 4, 1)) + inputs = (torch.randn(1, 4, 8, 8),) + generic_graph, _ = _rewrite_a16w8_convs(model, inputs) + u55_graph, _ = _rewrite_a16w8_convs( + model, + inputs, + TosaSpecification.create_from_string("TOSA-1.0+INT+int16+int4+u55"), + ) + + conv_targets = { + exir_ops.backend.tosa.CONV2D.default, + exir_ops.backend.tosa.DEPTHWISE_CONV2D.default, + } + generic_convs = [ + node for node in generic_graph.graph.nodes if node.target in conv_targets + ] + u55_convs = [node for node in u55_graph.graph.nodes if node.target in conv_targets] + + assert len(u55_convs) == len(generic_convs) + 1 + assert all(len(conv.users) == 1 for conv in u55_convs) + assert all( + next(iter(conv.users)).target == exir_ops.backend.tosa.RESCALE.default + for conv in u55_convs + ) + + +def test_rewrite_conv_a16w8_mixed_consumers_lowers_on_u55() -> None: + inputs = (torch.randn(1, 4, 8, 8),) + pipeline = EthosU55PipelineINT[tuple[torch.Tensor]]( + A16W8MixedConsumerChain(nn.Conv2d(4, 4, 1)), + inputs, + aten_ops=[], + exir_ops=[], + run_on_fvp=False, + a16w8_quantization=True, + ) + pipeline.run() + + def test_rewrite_conv_a16w8_preserves_int32_for_int32_consumers() -> None: r"""Test that an exclusively INT32 consumer keeps the widened path. diff --git a/backends/transforms/fuse_duplicate_users_pass.py b/backends/transforms/fuse_duplicate_users_pass.py index b3989e76c94..47b6064b618 100644 --- a/backends/transforms/fuse_duplicate_users_pass.py +++ b/backends/transforms/fuse_duplicate_users_pass.py @@ -11,7 +11,76 @@ from executorch.exir.pass_base import ExportPass, PassResult from torch._ops import OpOverload from torch.fx import GraphModule, Node -from torch.fx.node import Argument, map_arg +from torch.fx.node import map_arg + + +DO_NOT_FUSE_DUPLICATE_META_KEY = "do_not_fuse_duplicate" + + +def _map_leaf_to_key(node: Node) -> str: + return node.name + + +def _to_hashable(value: Any) -> Hashable: + """Convert arbitrarily nested structures into hashable tuples.""" + + if isinstance(value, (list, tuple)): + return tuple(_to_hashable(v) for v in value) + if isinstance(value, dict): + normalized_items = [(k, _to_hashable(v)) for k, v in value.items()] + return tuple(sorted(normalized_items, key=lambda item: repr(item[0]))) + if isinstance(value, set): + hashable_values: List[Hashable] = [_to_hashable(v) for v in value] + return tuple(sorted(hashable_values, key=repr)) + if isinstance(value, slice): + return ( + "slice", + _to_hashable(value.start), + _to_hashable(value.stop), + _to_hashable(value.step), + ) + if isinstance(value, range): + return ("range", value.start, value.stop, value.step) + if isinstance(value, torch.Size): + return ("size", tuple(value)) + if isinstance(value, torch.dtype): + return ("dtype", str(value)) + if isinstance(value, torch.device): + return ("device", str(value)) + if isinstance(value, torch.memory_format): + return ("memory_format", str(value)) + if isinstance(value, torch.Tensor): + return ( + "tensor", + str(value.dtype), + tuple(value.size()), + value.device.type, + value.requires_grad, + ) + return value + + +def _get_target_key(target: Any) -> Hashable: + if isinstance(target, (EdgeOpOverload, OpOverload)): + return str(target) + return target + + +def build_node_signature( + node: Node, *, positional_arg_start: int = 0 +) -> Tuple[Hashable, ...] | None: + """Build a stable signature while ignoring leading positional operands.""" + try: + normalized_args = _to_hashable( + map_arg(node.args[positional_arg_start:], _map_leaf_to_key) + ) + normalized_kwargs = _to_hashable( + {k: map_arg(v, _map_leaf_to_key) for k, v in node.kwargs.items()} + ) + except TypeError: + return None + + return (node.op, _get_target_key(node.target), normalized_args, normalized_kwargs) class FuseDuplicateUsersPass(ExportPass): @@ -106,7 +175,7 @@ def _get_candidate_groups(self, node_order, user_nodes): if user.target in self._excluded_targets: continue - target_key = self._get_target_key(user.target) + target_key = _get_target_key(user.target) target_signature = (user.op, target_key) users_by_target.setdefault(target_signature, []).append(user) @@ -120,62 +189,6 @@ def _get_candidate_groups(self, node_order, user_nodes): return candidate_groups def _build_user_signature(self, node: Node) -> Tuple[Hashable, ...] | None: - try: - normalized_args = self._to_hashable( - map_arg(node.args, self._map_leaf_to_key) - ) - normalized_kwargs = self._to_hashable( - {k: map_arg(v, self._map_leaf_to_key) for k, v in node.kwargs.items()} - ) - except TypeError: + if node.meta.get(DO_NOT_FUSE_DUPLICATE_META_KEY, False): return None - - target_key = self._get_target_key(node.target) - - return (node.op, target_key, normalized_args, normalized_kwargs) - - def _map_leaf_to_key(self, node: Node) -> Argument: - return node.name - - def _to_hashable(self, value: Any) -> Hashable: - """Convert arbitrarily nested structures into hashable tuples.""" - - if isinstance(value, (list, tuple)): - return tuple(self._to_hashable(v) for v in value) - if isinstance(value, dict): - normalized_items = [(k, self._to_hashable(v)) for k, v in value.items()] - return tuple(sorted(normalized_items, key=lambda item: repr(item[0]))) - if isinstance(value, set): - hashable_values: List[Hashable] = [self._to_hashable(v) for v in value] - return tuple(sorted(hashable_values, key=repr)) - if isinstance(value, slice): - return ( - "slice", - self._to_hashable(value.start), - self._to_hashable(value.stop), - self._to_hashable(value.step), - ) - if isinstance(value, range): - return ("range", value.start, value.stop, value.step) - if isinstance(value, torch.Size): - return ("size", tuple(value)) - if isinstance(value, torch.dtype): - return ("dtype", str(value)) - if isinstance(value, torch.device): - return ("device", str(value)) - if isinstance(value, torch.memory_format): - return ("memory_format", str(value)) - if isinstance(value, torch.Tensor): - return ( - "tensor", - str(value.dtype), - tuple(value.size()), - value.device.type, - value.requires_grad, - ) - return value - - def _get_target_key(self, target: Any) -> Hashable: - if isinstance(target, (EdgeOpOverload, OpOverload)): - return str(target) - return target + return build_node_signature(node) From 3470baafb9ad743b95f2d7be628669c272c9838b Mon Sep 17 00:00:00 2001 From: pssrawat <34485295+pssrawat@users.noreply.github.com> Date: Thu, 10 Sep 2026 22:42:32 +0530 Subject: [PATCH 152/190] SDPA: skip the non-attendable key columns (#22112) Summary: Under causal attention a query block can only attend to `num_keys` keys, but the flash-attention custom SDPA kernel sized every key block to the full `kvSize`. The columns past `num_keys` ran through the qk gemm, the -inf fill, the softmax and the v gemm to contribute exactly zero. Clamp the block width to `num_keys`, and narrow the causal-mask guard and two `fill_stub` extents to match; the block can only shrink, and with `is_causal` false nothing changes at all. Worth -9.5% to -33% of SDPA time on multi-token prefill on an S25, and provably inert at `qSize == 1`. Reviewed By: JakeStevens Differential Revision: D117259224 Pull Request resolved: https://github.com/pytorch/executorch/pull/22112 --- extension/llm/custom_ops/op_sdpa_impl.h | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/extension/llm/custom_ops/op_sdpa_impl.h b/extension/llm/custom_ops/op_sdpa_impl.h index f6ed378ec03..7e7275c0427 100644 --- a/extension/llm/custom_ops/op_sdpa_impl.h +++ b/extension/llm/custom_ops/op_sdpa_impl.h @@ -1045,11 +1045,17 @@ void cpu_flash_attention( is_causal ? std::min(m + start_pos + qBlockSize, kvSize) : kvSize; int64_t m_start_pos = m + start_pos; auto j_kv = j / num_reps; - fill_stub(dst_data, static_cast(0), qSplitSize * headSize); + fill_stub(dst_data, static_cast(0), qBlockSize * headSize); for (int64_t n = 0; n < num_keys; n += kvSplitSize) { - int64_t kvBlockSize = std::min(kvSplitSize, kvSize - n); + // Only the first num_keys columns are causally attendable; the rest + // would be masked to -inf and contribute exactly zero, so clamping + // here skips their gemm, softmax and v-multiply. This matters for the + // leading query blocks of a prefill, where num_keys is much smaller + // than the key-cache extent. Not bit-exact: shortening the reduction + // moves the vector-lane partition, so the accumulation order changes. + int64_t kvBlockSize = std::min(kvSplitSize, num_keys - n); // Calculate scale * q @ k.T - fill_stub(qk_data, static_cast(0), qSplitSize * kvSplitSize); + fill_stub(qk_data, static_cast(0), qBlockSize * kvBlockSize); const void* q_sub_matrix_data_ptr; const void* k_sub_matrix_data_ptr; @@ -1147,10 +1153,10 @@ void cpu_flash_attention( take care of this case because the loop for (int64_t n = 0; n < num_keys; n += kvSplitSize) will exit before that. */ - if (is_causal && m_start_pos <= n + kvSplitSize) { + if (is_causal && m_start_pos <= n + kvBlockSize) { // For this fn to work k_split_size > q_split_size for (int32_t row = 0; - row < qBlockSize && (m_start_pos + row < n + (kvSplitSize - 1)); + row < qBlockSize && (m_start_pos + row < n + (kvBlockSize - 1)); ++row) { // When last_col is 0, it means that the entire row is not attended // to because m_pos is smaller than n_pos. So everything in n is for From 6960f714238167c8cf3d508256f484084859d0b6 Mon Sep 17 00:00:00 2001 From: Huy Do Date: Thu, 10 Sep 2026 10:20:02 -0700 Subject: [PATCH 153/190] Move trunk.yml, periodic.yml and nightly.yml to linux_job_v3 (#22248) Fourth of six splitting up #22107. Stacked on #22247. Also picks up the last `linux_job.yml` (v1) call site, in trunk's `test-arm-backend-zephyr`. The two arm backend jobs that raise `fs.inotify.max_user_watches` now tolerate failure, since that is a node-level setting an unprivileged pod cannot change. Authored with Claude Code. --- .ci/scripts/gather_test_models.py | 26 +++---- .github/workflows/nightly.yml | 11 ++- .github/workflows/periodic.yml | 10 ++- .github/workflows/trunk.yml | 122 +++++++++++++++++------------- 4 files changed, 99 insertions(+), 70 deletions(-) diff --git a/.ci/scripts/gather_test_models.py b/.ci/scripts/gather_test_models.py index 65fefc4073d..db0f2671594 100755 --- a/.ci/scripts/gather_test_models.py +++ b/.ci/scripts/gather_test_models.py @@ -17,23 +17,23 @@ from examples.xnnpack import MODEL_NAME_TO_OPTIONS, QuantType DEFAULT_RUNNERS = { - "linux": "linux.2xlarge", + "linux": "mt-l-x86iavx512-8-64", "macos": "macos-m1-stable", } CUSTOM_RUNNERS = { "linux": { # This one runs OOM on smaller runner, the root cause is unclear (T163016365) - "w2l": "linux.4xlarge.memory", - "ic4": "linux.4xlarge.memory", - "resnet50": "linux.4xlarge.memory", - "llava": "linux.4xlarge.memory", - "llama3_2_vision_encoder": "linux.4xlarge.memory", - "llama3_2_text_decoder": "linux.4xlarge.memory", + "w2l": "mt-l-x86iavx512-16-128", + "ic4": "mt-l-x86iavx512-16-128", + "resnet50": "mt-l-x86iavx512-16-128", + "llava": "mt-l-x86iavx512-16-128", + "llama3_2_vision_encoder": "mt-l-x86iavx512-16-128", + "llama3_2_text_decoder": "mt-l-x86iavx512-16-128", # This one causes timeout on smaller runner, the root cause is unclear (T161064121) - "dl3": "linux.4xlarge.memory", - "emformer_join": "linux.4xlarge.memory", - "emformer_predict": "linux.4xlarge.memory", - "phi_4_mini": "linux.4xlarge.memory", + "dl3": "mt-l-x86iavx512-16-128", + "emformer_join": "mt-l-x86iavx512-16-128", + "emformer_predict": "mt-l-x86iavx512-16-128", + "phi_4_mini": "mt-l-x86iavx512-16-128", } } @@ -146,7 +146,7 @@ def export_models_for_ci() -> dict[str, dict]: "build-tool": "buck2", "model": "mv3", "backend": backend, - "runner": "linux.2xlarge", + "runner": "mt-l-x86iavx512-8-64", "timeout": DEFAULT_TIMEOUT, } models["include"].append(record) @@ -175,7 +175,7 @@ def export_models_for_ci() -> dict[str, dict]: "build-tool": "cmake", "model": name, "backend": backend, - "runner": DEFAULT_RUNNERS.get(target_os, "linux.2xlarge"), + "runner": DEFAULT_RUNNERS.get(target_os, "mt-l-x86iavx512-8-64"), "timeout": DEFAULT_TIMEOUT, } diff --git a/.github/workflows/nightly.yml b/.github/workflows/nightly.yml index 922ed95ab8d..c301b61399e 100644 --- a/.github/workflows/nightly.yml +++ b/.github/workflows/nightly.yml @@ -16,6 +16,10 @@ concurrency: cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + update-pytorch-commit-hash: runs-on: ubuntu-latest environment: ${{ (github.event_name == 'schedule') && 'update-commit-hash' || '' }} @@ -50,8 +54,9 @@ jobs: timeout: 180 test-static-hf-llm-qnn-linux: + needs: docker-image name: test-static-hf-llm-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -60,8 +65,8 @@ jobs: task: [smollm2_135m] fail-fast: false with: - runner: linux.24xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-94-192 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 diff --git a/.github/workflows/periodic.yml b/.github/workflows/periodic.yml index 01bff087124..a859a9fd779 100644 --- a/.github/workflows/periodic.yml +++ b/.github/workflows/periodic.yml @@ -21,6 +21,10 @@ concurrency: permissions: read-all jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + gather-models: runs-on: ubuntu-22.04 outputs: @@ -42,17 +46,17 @@ jobs: test-models-linux: name: test-models-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read - needs: gather-models + needs: [docker-image, gather-models] strategy: matrix: ${{ fromJSON(needs.gather-models.outputs.models) }} fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: ${{ matrix.timeout }} diff --git a/.github/workflows/trunk.yml b/.github/workflows/trunk.yml index 25b2df3aa4f..f75e7e76f01 100644 --- a/.github/workflows/trunk.yml +++ b/.github/workflows/trunk.yml @@ -19,6 +19,10 @@ concurrency: cancel-in-progress: true jobs: + docker-image: + name: Resolve CI docker image + uses: ./.github/workflows/_docker-image.yml + # Emits the list of changed files for the current PR or push commit. # On PR: PR diff. On push: diff against `github.event.before`. # On events without a diff base (workflow_dispatch, tag creation, @@ -81,8 +85,12 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-arm-backend-zephyr: + needs: docker-image name: test-arm-backend-zephyr - uses: pytorch/test-infra/.github/workflows/linux_job.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main + permissions: + id-token: write + contents: read strategy: matrix: include: @@ -93,8 +101,8 @@ jobs: - { readme: zephyr/samples/mv2-ethosu/README.md, target: ethos-u85 } fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-zephyr-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-zephyr-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -119,8 +127,9 @@ jobs: --zephyr-samples-readme-path "${{ matrix.readme }}" test-models-linux-aarch64: + needs: docker-image name: test-models-linux-aarch64 - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -128,30 +137,30 @@ jobs: matrix: model: [linear, add, add_mul, ic3, ic4, mv2, mv3, resnet18, resnet50, vit, w2l, mobilebert, emformer_join, emformer_transcribe] backend: [portable, xnnpack-quantization-delegation] - runner: [linux.arm64.2xlarge] + runner: [mt-l-arm64g4-16-62] include: - model: lstm backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: mul backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: softmax backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: phi_4_mini backend: portable - runner: linux.arm64.m7g.4xlarge + runner: mt-l-arm64g4-16-62 - model: qwen2_5_1_5b backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 - model: llama3_2_vision_encoder backend: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:executorch-ubuntu-22.04-gcc11-aarch64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-gcc11-aarch64-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -223,8 +232,9 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash examples/selective_build/test_selective_build.sh "${BUILD_TOOL}" test-demo-backend-delegation: + needs: docker-image name: test-demo-backend-delegation - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -235,8 +245,8 @@ jobs: - build-tool: cmake fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} script: | @@ -250,8 +260,9 @@ jobs: PYTHON_EXECUTABLE=python bash examples/portable/scripts/test_demo_backend_delegation.sh "${BUILD_TOOL}" test-arm-backend-ethos-u: + needs: docker-image name: test-arm-backend-ethos-u - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -273,8 +284,8 @@ jobs: - test_arm_backend: test_deit_e2e_ethos_u fail-fast: false with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -290,7 +301,8 @@ jobs: # Increase number of files user can monitor to bypass buck failures. # Hopefully this is high enough for this setup. - sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024 + # Node-level, so an unprivileged pod cannot change it. + sudo sysctl fs.inotify.max_user_watches=1048576 2>/dev/null || true ARM_TEST=${{ matrix.test_arm_backend }} @@ -303,8 +315,9 @@ jobs: backends/arm/test/test_arm_backend.sh "${ARM_TEST}" test-arm-backend-vkml: + needs: docker-image name: test-arm-backend-vkml - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -317,8 +330,8 @@ jobs: - test_arm_backend: test_smaller_stories_llama_vkml fail-fast: false with: - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-24.04-arm-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-24.04-arm-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 120 @@ -333,7 +346,8 @@ jobs: # Increase number of files user can monitor to bypass buck failures. # Hopefully this is high enough for this setup. - sudo sysctl fs.inotify.max_user_watches=1048576 # 1024 * 1024 + # Node-level, so an unprivileged pod cannot change it. + sudo sysctl fs.inotify.max_user_watches=1048576 2>/dev/null || true ARM_TEST=${{ matrix.test_arm_backend }} @@ -432,9 +446,10 @@ jobs: ${CONDA_RUN} sh .ci/scripts/test_llama_torchao_lowbit.sh test-llama-runner-linux: + needs: docker-image # Test Both linux x86 and linux aarch64 name: test-llama-runner-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -442,33 +457,33 @@ jobs: matrix: dtype: [fp32] mode: [portable, xnnpack+custom] - runner: [linux.2xlarge, linux.arm64.2xlarge] + runner: [mt-l-x86iavx512-8-64, mt-l-arm64g4-16-62] docker-image: [executorch-ubuntu-22.04-clang12, executorch-ubuntu-22.04-gcc11-aarch64] include: - dtype: bf16 mode: portable - runner: linux.2xlarge + runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-clang12 - dtype: bf16 mode: portable - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - dtype: bf16 mode: custom - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 # Excluding specific runner + docker image combinations that don't make sense: - # - Excluding the ARM64 gcc image on the x86 runner (linux.2xlarge) - # - Excluding the x86 clang image on the ARM64 runner (linux.arm64.2xlarge) + # - Excluding the ARM64 gcc image on the x86 runner + # - Excluding the x86 clang image on the ARM64 runner exclude: - - runner: linux.2xlarge + - runner: mt-l-x86iavx512-8-64 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 - - runner: linux.arm64.2xlarge + - runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-clang12 fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -543,34 +558,35 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_llama.sh -model stories110M -build_tool cmake -dtype "${DTYPE}" -mode "${MODE}" test-torchao-huggingface-checkpoints: + needs: docker-image name: test-torchao-huggingface-checkpoints - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read strategy: matrix: model: [qwen3_4b, phi_4_mini, lfm2_5_1_2b] - runner: [linux.2xlarge] + runner: [mt-l-x86iavx512-8-64] docker-image: [executorch-ubuntu-22.04-clang12] backend: [xnnpack] include: - model: qwen3_4b - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 backend: torchao - model: phi_4_mini - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 backend: torchao - model: lfm2_5_1_2b - runner: linux.arm64.2xlarge + runner: mt-l-arm64g4-16-62 docker-image: executorch-ubuntu-22.04-gcc11-aarch64 backend: torchao fail-fast: false with: runner: ${{ matrix.runner }} - docker-image: ci-image:${{ matrix.docker-image }} + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:${{ matrix.docker-image }}-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -637,8 +653,9 @@ jobs: echo "::endgroup::" test-qnn-model: + needs: docker-image name: test-qnn-model - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -648,8 +665,8 @@ jobs: model: [dl3, mv3, mv2, ic4, ic3, vit, mb, w2l, conv_former] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -663,8 +680,9 @@ jobs: PYTHON_EXECUTABLE=python bash .ci/scripts/test_model.sh ${{ matrix.model }} "cmake" "qnn" test-qnn-optimum-model: + needs: docker-image name: test-qnn-optimum-model - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -674,8 +692,8 @@ jobs: model: [cvt, dit, efficientnet, focalnet, mobilevit_v1, mobilevit_v2, pvt, swin, albert, bert, distilbert, roberta] # eurobert requires transfomer >= 4.48.0, skip for now fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 @@ -740,10 +758,11 @@ jobs: PYTHON_EXECUTABLE=python ${CONDA_RUN} bash .ci/scripts/test_model.sh "${MODEL_NAME}" "${BUILD_TOOL}" "${BACKEND}" test-huggingface-transformers-xnnpack: + needs: docker-image # NB: Don't run this on fork PRs because they won't have access to the secret and would fail anyway if: ${{ !github.event.pull_request.head.repo.fork }} name: test-huggingface-transformers-xnnpack - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -763,8 +782,8 @@ jobs: fail-fast: false with: secrets-env: EXECUTORCH_HF_TOKEN - runner: linux.2xlarge.memory - docker-image: ci-image:executorch-ubuntu-22.04-clang12 + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-clang12-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 90 @@ -905,8 +924,9 @@ jobs: ${CONDA_RUN} python .ci/scripts/test_huggingface_optimum_model.py --model ${MODEL} --recipe ${RECIPE} ${QUANTIZE} test-llama-runner-qnn-linux: + needs: docker-image name: test-llama-runner-qnn-linux - uses: pytorch/test-infra/.github/workflows/linux_job_v2.yml@main + uses: pytorch/test-infra/.github/workflows/linux_job_v3.yml@main permissions: id-token: write contents: read @@ -917,8 +937,8 @@ jobs: mode: [qnn] fail-fast: false with: - runner: linux.2xlarge - docker-image: ci-image:executorch-ubuntu-22.04-qnn-sdk + runner: mt-l-x86iavx512-8-64 + docker-image: ${{ needs.docker-image.outputs.docker-registry }}/ci-image:executorch-ubuntu-22.04-qnn-sdk-${{ needs.docker-image.outputs.ci-docker-hash }} submodules: 'recursive' ref: ${{ github.event_name == 'pull_request' && github.event.pull_request.head.sha || github.sha }} timeout: 900 From 652500fa44ae0fc27b78a57a2b0ed8c3c0a55c0f Mon Sep 17 00:00:00 2001 From: Suryansh Sijwali <159204949+SuryanshSS1011@users.noreply.github.com> Date: Thu, 10 Sep 2026 13:27:58 -0400 Subject: [PATCH 154/190] Reject non-default dim order in the portable kernels that copy by block (#21865) ### Summary Follow-up to #21828, which fixes the coordinate-indexed half of the same defect. These two PRs touch no files in common and can land in either order. Five portable kernels accept a channels-last input and return wrong data without erroring. Three more share the defect but are currently masked by an unrelated check. Checked against eager PyTorch with only the memory format differing: | kernel | contiguous input | channels-last input | |---|---|---| | `aten.cat` | matches | wrong, max abs diff 4.141 | | `aten.pixel_unshuffle` | matches | wrong, max abs diff 3.941 | | `aten.slice_scatter` | matches | wrong, max abs diff 3.376 | | `aten.topk` | matches | wrong, max abs diff 2.581 | | `aten.constant_pad_nd` | matches | wrong, max abs diff 2.526 | | `aten.cumsum` | matches | rejected, status 0x12 | | `aten.split_copy` | matches | rejected, status 0x12 | | `aten.split_with_sizes_copy` | matches | rejected, status 0x12 | The check that masks the last three is `tensors_have_same_dim_order`, and it fires only because the exported graph gives them differing dim orders. Nothing guards the arithmetic itself, so the unit tests below reach those kernels directly. These kernels copy in blocks, taking `getLeadingDims` and `getTrailingDims` around the operating dim as the length of a contiguous run: ```cpp const size_t outer = getLeadingDims(out, dim); const size_t dim_stride = getTrailingDims(out, dim); ``` That run only exists in the default dim order. Under channels-last the elements after `dim` are not adjacent, so the arithmetic walks the wrong bytes. This is the same assumption as #21828, but it needs a different fix. There the kernels index element by element, so reading the tensor's strides is enough. Here the algorithm depends on contiguity itself, and supporting the layout would mean restructuring each loop. So this adds the guard that `op_addmm`, `op_bmm` and `op_avg_pool2d` already use: ```cpp ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); ``` `cat` needed it on each input as well as the output, since it had no dim order check at all. `slice_scatter` needed `src` in its same dim order check as well, since the kernel reads `src` as one contiguous run and only `input` and `out` were covered. `pixel_shuffle` already carries this exact pair of checks, so `pixel_unshuffle` missing them looks like an oversight rather than a decision. Note: `cat` and `split` appear in nearly every model, so a channels-last program that runs today would start failing instead of returning wrong numbers. That still seems better than silent corruption, but happy to restructure the loops to support the layout instead if that's preferred. `native_batch_norm` uses the same helpers but is not affected, since it checks `is_contiguous_dim_order` on the input directly. `select_scatter` and `unfold_copy` are affected and are left to a follow-up. `select_scatter` carries only the same dim order check, which passes when a rank-5 input and a rank-4 `src` are both channels-last, because `is_channels_last_dim_order` accepts 4 and 5 dims. `unfold_copy` has no dim order check at all. The same block arithmetic appears in the optimized `layer_norm`, fixed in #21866, and in the quantized `dequantize`, which #21517 covers. ### Test plan A `NonDefaultDimOrderDies` test per kernel. Each passes a uniformly channels-last set of tensors, so the existing same dim order check passes and only the new guard can reject. Reverting the eight kernel sources while keeping the tests fails all eight. This also adds `op_pixel_unshuffle_test.cpp` to `kernels/test/CMakeLists.txt`. It is registered in `targets.bzl` but missing from the CMake list, so those tests do not run in a CMake build and the new one would not either. Ten other op tests are missing the same way, and I have left those alone to keep this to one change. cc @larryliu0820 @manuelcandales @JakeStevens --- kernels/portable/cpu/op_cat.cpp | 4 ++ kernels/portable/cpu/op_constant_pad_nd.cpp | 2 + kernels/portable/cpu/op_cumsum.cpp | 2 + kernels/portable/cpu/op_pixel_unshuffle.cpp | 5 ++ kernels/portable/cpu/op_slice_scatter.cpp | 5 +- kernels/portable/cpu/op_split_copy.cpp | 2 + .../portable/cpu/op_split_with_sizes_copy.cpp | 2 + kernels/portable/cpu/op_topk.cpp | 8 ++++ kernels/test/CMakeLists.txt | 1 + kernels/test/op_cat_test.cpp | 18 +++++++ kernels/test/op_constant_pad_nd_test.cpp | 18 +++++++ kernels/test/op_cumsum_test.cpp | 15 ++++++ kernels/test/op_pixel_unshuffle_test.cpp | 31 ++++++++++++ kernels/test/op_slice_scatter_test.cpp | 38 +++++++++++++++ kernels/test/op_split_copy_test.cpp | 19 ++++++++ .../test/op_split_with_sizes_copy_test.cpp | 26 ++++++++++ kernels/test/op_topk_test.cpp | 48 +++++++++++++++++++ 17 files changed, 243 insertions(+), 1 deletion(-) diff --git a/kernels/portable/cpu/op_cat.cpp b/kernels/portable/cpu/op_cat.cpp index ab15d5249df..5d5d930e29c 100644 --- a/kernels/portable/cpu/op_cat.cpp +++ b/kernels/portable/cpu/op_cat.cpp @@ -28,6 +28,10 @@ Tensor& cat_out( ET_KERNEL_CHECK(ctx, check_cat_args(tensors, dim, out), InvalidArgument, out); + // check_cat_args already requires every input to share out's dim order, so + // checking out is enough to force all of them to the default. + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(out), InvalidArgument, out); + Tensor::SizesType expected_out_size[kTensorDimensionLimit]; size_t expected_out_dim = 0; get_cat_out_target_size(tensors, dim, expected_out_size, &expected_out_dim); diff --git a/kernels/portable/cpu/op_constant_pad_nd.cpp b/kernels/portable/cpu/op_constant_pad_nd.cpp index 0f287a5ac53..0947a14dc2d 100644 --- a/kernels/portable/cpu/op_constant_pad_nd.cpp +++ b/kernels/portable/cpu/op_constant_pad_nd.cpp @@ -252,6 +252,8 @@ Tensor& constant_pad_nd_out( ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + // resize out tensor for dynamic shapes ET_KERNEL_CHECK_MSG( ctx, diff --git a/kernels/portable/cpu/op_cumsum.cpp b/kernels/portable/cpu/op_cumsum.cpp index 5023be7b694..3b7abcbed63 100644 --- a/kernels/portable/cpu/op_cumsum.cpp +++ b/kernels/portable/cpu/op_cumsum.cpp @@ -103,6 +103,8 @@ Tensor& cumsum_out( ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(self, out), InvalidArgument, out); + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(self), InvalidArgument, out); + ET_KERNEL_CHECK( ctx, resize_tensor(out, self.sizes()) == Error::Ok, InvalidArgument, out); diff --git a/kernels/portable/cpu/op_pixel_unshuffle.cpp b/kernels/portable/cpu/op_pixel_unshuffle.cpp index 68d7bbbc27a..c6b69a56a04 100644 --- a/kernels/portable/cpu/op_pixel_unshuffle.cpp +++ b/kernels/portable/cpu/op_pixel_unshuffle.cpp @@ -81,6 +81,11 @@ Tensor& pixel_unshuffle_out( InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + // @lint-ignore CLANGTIDY facebook-hte-CArray Tensor::SizesType expected_out_size[kTensorDimensionLimit]; size_t expected_out_dim = 0; diff --git a/kernels/portable/cpu/op_slice_scatter.cpp b/kernels/portable/cpu/op_slice_scatter.cpp index 29c4ff7ab90..5a59df8d25b 100644 --- a/kernels/portable/cpu/op_slice_scatter.cpp +++ b/kernels/portable/cpu/op_slice_scatter.cpp @@ -42,7 +42,10 @@ Tensor& slice_scatter_out( out); ET_KERNEL_CHECK( - ctx, tensors_have_same_dim_order(input, out), InvalidArgument, out); + ctx, tensors_have_same_dim_order(input, src, out), InvalidArgument, out); + + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(input), InvalidArgument, out); if (input.numel() == 0) { return out; diff --git a/kernels/portable/cpu/op_split_copy.cpp b/kernels/portable/cpu/op_split_copy.cpp index fdc89727897..0f97dc76345 100644 --- a/kernels/portable/cpu/op_split_copy.cpp +++ b/kernels/portable/cpu/op_split_copy.cpp @@ -49,6 +49,8 @@ void split_copy_Tensor_out( for (size_t i = 0; i < out.size(); ++i) { ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(input, out[i]), InvalidArgument, ); + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(out[i]), InvalidArgument, ); } const size_t leading_dims = getLeadingDims(input, dim); diff --git a/kernels/portable/cpu/op_split_with_sizes_copy.cpp b/kernels/portable/cpu/op_split_with_sizes_copy.cpp index c99a7fb6815..0353e048b9e 100644 --- a/kernels/portable/cpu/op_split_with_sizes_copy.cpp +++ b/kernels/portable/cpu/op_split_with_sizes_copy.cpp @@ -43,6 +43,8 @@ void split_with_sizes_copy_out( for (const auto i : c10::irange(out.size())) { ET_KERNEL_CHECK( ctx, tensors_have_same_dim_order(in, out[i]), InvalidArgument, ); + ET_KERNEL_CHECK( + ctx, tensor_is_default_dim_order(out[i]), InvalidArgument, ); } // If out is empty, then nothing needs to be done after checking the args. diff --git a/kernels/portable/cpu/op_topk.cpp b/kernels/portable/cpu/op_topk.cpp index 3082bc94662..7bda44fccd6 100644 --- a/kernels/portable/cpu/op_topk.cpp +++ b/kernels/portable/cpu/op_topk.cpp @@ -170,6 +170,14 @@ std::tuple topk_values( ET_KERNEL_CHECK( ctx, check_topk_args(in, k, dim, values, indices), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, + tensors_have_same_dim_order(in, values, indices), + InvalidArgument, + out); + + ET_KERNEL_CHECK(ctx, tensor_is_default_dim_order(in), InvalidArgument, out); + if (dim < 0) { dim += nonzero_dim(in); } diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index da9fbc2f55b..a8e703cb5aa 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -261,6 +261,7 @@ set(all_test_sources "op_pdist_forward_test.cpp" "op_permute_copy_test.cpp" "op_pixel_shuffle_test.cpp" + "op_pixel_unshuffle_test.cpp" "op_prod_test.cpp" "op_rand_test.cpp" "op_randn_test.cpp" diff --git a/kernels/test/op_cat_test.cpp b/kernels/test/op_cat_test.cpp index d3bda1e8abd..cfdd6e7a426 100644 --- a/kernels/test/op_cat_test.cpp +++ b/kernels/test/op_cat_test.cpp @@ -464,3 +464,21 @@ TEST_F(OpCatOutTest, DynamicShapeUnbound) { op_cat_out(x, 0, out); EXPECT_TENSOR_EQ(out, expected); } + +TEST_F(OpCatOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor x = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor y = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros_channels_last({1, 6, 2, 2}); + std::vector inputs = {x, y}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_cat_out(TensorList(inputs.data(), inputs.size()), 1, out)); +} diff --git a/kernels/test/op_constant_pad_nd_test.cpp b/kernels/test/op_constant_pad_nd_test.cpp index 7bd908e0ecb..00a3b5bbef9 100644 --- a/kernels/test/op_constant_pad_nd_test.cpp +++ b/kernels/test/op_constant_pad_nd_test.cpp @@ -484,3 +484,21 @@ TEST_F(OpConstantPadNDOutTest, IncorrectOutputShapeFail) { } GENERATE_SCALAR_OVERFLOW_TESTS(OpConstantPadNDOutTest) + +TEST_F(OpConstantPadNDOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor self = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros_channels_last({1, 3, 2, 4}); + const std::vector padding = {1, 1}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_constant_pad_nd_out( + self, IntArrayRef(padding.data(), padding.size()), 0.0, out)); +} diff --git a/kernels/test/op_cumsum_test.cpp b/kernels/test/op_cumsum_test.cpp index 8ddc197217b..bbf333db4c6 100644 --- a/kernels/test/op_cumsum_test.cpp +++ b/kernels/test/op_cumsum_test.cpp @@ -286,3 +286,18 @@ TEST_F(OpCumSumOutTest, DISABLED_DynamicShapeUnbound) { Tensor ret = op_cumsum_out(x, 1, ScalarType::Float, out); EXPECT_TENSOR_CLOSE(out, expected_result); } + +TEST_F(OpCumSumOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor in = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros_channels_last({1, 3, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_cumsum_out(in, 1, ScalarType::Float, out)); +} diff --git a/kernels/test/op_pixel_unshuffle_test.cpp b/kernels/test/op_pixel_unshuffle_test.cpp index 21bed318b9c..9cb4498aabf 100644 --- a/kernels/test/op_pixel_unshuffle_test.cpp +++ b/kernels/test/op_pixel_unshuffle_test.cpp @@ -9,6 +9,7 @@ #include // Declares the operator #include #include +#include #include #include #include @@ -126,3 +127,33 @@ TEST_F(OpPixelUnshuffleOutTest, NegativeUpscaleFactorDies) { // Using a negative upscale factor should exit with an error code. ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, -3, out)); } + +TEST_F(OpPixelUnshuffleOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor a = tf.channels_last_like(tf.make( + {1, 1, 4, 4}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + Tensor out = tf.zeros_channels_last({1, 4, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, 2, out)); +} + +TEST_F(OpPixelUnshuffleOutTest, MixedDimOrderDies) { + TensorFactory tf; + + // Only out has a non-default dim order, so the same dim order check shall be + // what rejects it. + Tensor a = tf.make( + {1, 1, 4, 4}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + Tensor out = tf.zeros_channels_last({1, 4, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE(context_, op_pixel_unshuffle_out(a, 2, out)); +} diff --git a/kernels/test/op_slice_scatter_test.cpp b/kernels/test/op_slice_scatter_test.cpp index 309f2b8b5f7..256e6e261f5 100644 --- a/kernels/test/op_slice_scatter_test.cpp +++ b/kernels/test/op_slice_scatter_test.cpp @@ -884,3 +884,41 @@ TEST_F(OpSliceScatterTensorOutTest, LargeEndValue) { EXPECT_TENSOR_EQ(ret, out); EXPECT_TENSOR_EQ(ret, expected); } + +TEST_F(OpSliceScatterTensorOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor input = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor src = tf.zeros_channels_last({1, 1, 2, 2}); + Tensor out = tf.zeros_channels_last({1, 3, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, op_slice_scatter_out(input, src, 1, 0, 1, 1, out)); +} + +TEST_F(OpSliceScatterTensorOutTest, MixedDimOrderDies) { + TensorFactory tf; + + // Only src has a non-default dim order, so the same dim order check shall be + // what rejects it. + Tensor input = + tf.make({1, 3, 2, 4}, {0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, + 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23}); + Tensor src = tf.channels_last_like( + tf.make({1, 3, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12})); + Tensor out = tf.zeros({1, 3, 2, 4}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_slice_scatter_out( + input, src, /*dim=*/3, /*start=*/0, /*end=*/2, /*step=*/1, out)); +} diff --git a/kernels/test/op_split_copy_test.cpp b/kernels/test/op_split_copy_test.cpp index 34df2c749ff..7714770d383 100644 --- a/kernels/test/op_split_copy_test.cpp +++ b/kernels/test/op_split_copy_test.cpp @@ -576,3 +576,22 @@ TEST_F(OpSplitCopyTensorOutTest, DISABLED_DynamicShapeUnbound) { test_dynamic_shape( {1, 1}, torch::executor::TensorShapeDynamism::DYNAMIC_UNBOUND); } + +TEST_F(OpSplitCopyTensorOutTest, NonDefaultDimOrderDies) { + TensorFactory tf; + + Tensor input = tf.channels_last_like(tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + std::vector outs = { + tf.zeros_channels_last({1, 2, 2, 2}), + tf.zeros_channels_last({1, 2, 2, 2})}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_split_copy_tensor_out( + input, 2, 1, TensorList(outs.data(), outs.size()))); +} diff --git a/kernels/test/op_split_with_sizes_copy_test.cpp b/kernels/test/op_split_with_sizes_copy_test.cpp index cc81ffff19d..f33b622490f 100644 --- a/kernels/test/op_split_with_sizes_copy_test.cpp +++ b/kernels/test/op_split_with_sizes_copy_test.cpp @@ -9,6 +9,7 @@ #include // Declares the operator #include #include +#include #include #include #include @@ -115,3 +116,28 @@ TEST_F(OpSplitWithSizesCopyOutTest, DynamicShape) { test_tensor_shape_dynamism( executorch::aten::TensorShapeDynamism::DYNAMIC_BOUND); } + +TEST_F(OpSplitWithSizesCopyOutTest, NonDefaultDimOrderDies) { + torch::executor::testing::TensorFactory + tf; + + executorch::aten::Tensor self = tf.channels_last_like(tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + std::vector outs = { + tf.zeros_channels_last({1, 2, 2, 2}), + tf.zeros_channels_last({1, 2, 2, 2})}; + const std::vector split_sizes = {2, 2}; + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + ET_EXPECT_KERNEL_FAILURE( + context_, + op_split_with_sizes_copy_out( + self, + executorch::aten::ArrayRef( + split_sizes.data(), split_sizes.size()), + 1, + executorch::aten::TensorList(outs.data(), outs.size()))); +} diff --git a/kernels/test/op_topk_test.cpp b/kernels/test/op_topk_test.cpp index 17c7141d12d..331d3c276eb 100644 --- a/kernels/test/op_topk_test.cpp +++ b/kernels/test/op_topk_test.cpp @@ -8,6 +8,8 @@ #include // Declares the operator #include +#include +#include #include #include #include @@ -173,3 +175,49 @@ TEST_F(OpTopkValuesTest, NonPartialSort) { EXPECT_TENSOR_EQ(indices, indices_expected); } } + +TEST_F(OpTopkValuesTest, NonDefaultDimOrderDies) { + TensorFactory tf; + TensorFactory tf_long; + + Tensor in = tf.channels_last_like(tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16})); + Tensor values = tf.zeros_channels_last({1, 2, 2, 2}); + Tensor indices = tf_long.zeros_channels_last({1, 2, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + TempMemoryAllocator allocator = TempMemoryAllocator(); + executorch::ET_RUNTIME_NAMESPACE::KernelRuntimeContext context( + nullptr, &allocator); + torch::executor::aten::topk_outf( + context, in, 2, 1, true, true, values, indices); + + EXPECT_NE(context.failure_state(), torch::executor::Error::Ok); +} + +TEST_F(OpTopkValuesTest, MixedDimOrderDies) { + TensorFactory tf; + TensorFactory tf_long; + + // Only values has a non-default dim order, so the same dim order check shall + // be what rejects it. + Tensor in = tf.make( + {1, 4, 2, 2}, {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16}); + Tensor values = tf.zeros_channels_last({1, 2, 2, 2}); + Tensor indices = tf_long.zeros({1, 2, 2, 2}); + + ET_SKIP_IF( + torch::executor::testing::SupportedFeatures::get()->is_aten, + "ATen kernel can handle non-default dim order"); + + TempMemoryAllocator allocator = TempMemoryAllocator(); + executorch::ET_RUNTIME_NAMESPACE::KernelRuntimeContext context( + nullptr, &allocator); + torch::executor::aten::topk_outf( + context, in, 2, 1, true, true, values, indices); + + EXPECT_NE(context.failure_state(), torch::executor::Error::Ok); +} From 719e8514afbb9fee9003a41b53724ff2b5509392 Mon Sep 17 00:00:00 2001 From: Hyungkeun Park Date: Fri, 11 Sep 2026 02:40:16 +0900 Subject: [PATCH 155/190] Bound the dynamic-qdq traceback in XNNPACK ChannelsLastTaggedReshapePass (#21637) ### Summary The dynamic-quant branch of `ChannelsLastTaggedReshapePass.input_to_nhwc` traces back over the q/dq wrapper so the NHWC copy is inserted ahead of the quantize, keeping the `x -> q -> dq -> conv` chain XNNPACK matches intact. The loop stops on "`args[0]` is not a Node" rather than on "this is not a q/dq node", so it does not stop at the quantized tensor and continues into ordinary compute: ```python while getattr(input_node, "args", None) and isinstance( input_node.args[0], torch.fx.Node ): input_node = input_node.args[0] ``` The `input_node.replace_all_uses_with(input_node_nhwc)` that follows is then applied from wherever the walk landed, rewriting consumers the pass never reasoned about. That shows up two ways: If the walk stops on an intermediate op, lowering succeeds but the serialized graph is inconsistent, and only the first `execute()` reports it: ``` [XNNExecutor.cpp:133] Internal Error: Propagating input shapes failed with code: xnn_status_invalid_parameter [method.cpp:1421] CALL_DELEGATE execute failed at instruction 3: 0x1 ``` The failing partition contains an elementwise op whose input and output dims disagree, next to a correct one in the sibling branch: ``` [0] XNNStaticTranspose (0)[1,32,160,160] -> (1)[1,160,160,32] [1] XNNSigmoid (1)[1,160,160,32] -> (2)[1,32,160,160] [2] XNNStaticTranspose (2)[1,32,160,160] -> (3)[1,160,160,32] [3] XNNSigmoid (5)[1,32,160,160] -> (4)[1,32,160,160] ``` If the walk reaches a non-4D constant it fails earlier, during the pass: ``` RuntimeError: required rank 4 tensor to use channels_last format While executing _to_copy(%b_mul_10_const_input, memory_format=channels_last) ``` `b_mul_10_const_input` is a `(256, 1, 1)` per-channel constant. A placeholder has no `args`, so the walk stops there and tries to convert it. `can_be_converted_to_nhwc` does check rank 4, but this path never calls it. Both were found on w8a8 dynamic-quantized vision models, a YOLOX detector for the first and a SAM image encoder for the second. ### Fix Restrict the walk to q/dq nodes. `dq -> q -> source` is two hops and the source is not a q/dq node, so the walk stops at the source, which is the node it was aiming for, and the `x -> _to_copy -> q -> dq -> conv` ordering is unchanged. `is_quant` is already exported from `backends/xnnpack/utils/quant_utils.py` alongside the `is_dynamic_qdq` this file imports. Instrumenting the loop on the YOLOX graph: 83 invocations, every one with exactly 2 q/dq hops, and 69 of them (83%) walking past that, up to 26 hops. Bounding the walk leaves the delegate count at 16 either way and removes 16 `XNNStaticTranspose` nodes (414 to 398 total), so the overshoot was not buying larger fused partitions. ### Test `test_dq_conv2d_eltwise_source_channels_last_tagged_reshape_pass` builds the smallest graph that triggers it: a dynamically quantized conv whose input is a `sigmoid` reading the placeholder, so the conv sees `sigmoid -> q -> dq`. It asserts the sigmoid keeps its placeholder input and that the channels-last copy sits on the sigmoid's output. Without the fix the pass produces `x -> _to_copy(channels_last) -> sigmoid -> q -> dq -> conv` and the test fails with `AssertionError: 'call_function' != 'placeholder'`; with it the order is `x -> sigmoid -> _to_copy(channels_last) -> q -> dq -> conv`. The assertion is structural because `channels_last` does not change eager results, so `run_method_and_compare_outputs` alone cannot catch this. ``` pytest backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py -k eltwise_source ``` The full file passes (22 tests). Separately, 10 w8a8 dynamic models that already lowered and executed correctly before this change (googlenet, inception_v3, efficientnet_b4, wideresnet50, sesr_m5, mobile_vit_s, swin_t, vit_b_16, quicksrnet_small, squeezenet1_0) were re-lowered with the grouped partitioner and executed: 10/10 pass. The two models above now lower and execute, with output shapes matching a per-op-partitioned reference build. cc @GregoryComer @digantdesai @cbilgin @JakeStevens --- .../channels_last_tagged_reshape_pass.py | 38 ++- .../test_channels_last_tagged_reshape.py | 217 ++++++++++++++++++ 2 files changed, 247 insertions(+), 8 deletions(-) diff --git a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py index c74c35532b0..5238c9689ee 100644 --- a/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py +++ b/backends/xnnpack/_passes/channels_last_tagged_reshape_pass.py @@ -7,11 +7,13 @@ # pyre-unsafe from enum import Enum -from typing import Optional, Tuple +from typing import List, Optional, Tuple import torch from executorch.backends.xnnpack._passes.xnnpack_pass import XNNPACKPass from executorch.backends.xnnpack.utils.quant_utils import ( + extract_qdq_affine_op_args_for_decomposed_ops, + is_affine_qdq, is_dequant, is_dynamic_qdq, is_tagged_as_implicit_q_dq, @@ -365,6 +367,23 @@ def input_dim_order( else ChannelsLastTaggedReshapePass.is_nhwc_node(input_node) ) + @staticmethod + def redirect_dynamic_chain_to_nhwc( + input_node: torch.fx.Node, + input_node_nhwc: torch.fx.Node, + chain: List[torch.fx.Node], + ) -> None: + """Point the traversed chain's quantize and the choose_qparams feeding it + at the NHWC copy; consumers outside the chain keep the source. + """ + quantize = chain[-1] + quantize_args = quantize.args + if is_affine_qdq(quantize): + quantize_args = extract_qdq_affine_op_args_for_decomposed_ops(quantize) + qparam = quantize_args[1].args[0] + quantize.replace_input_with(input_node, input_node_nhwc) + qparam.replace_input_with(input_node, input_node_nhwc) + def input_to_nhwc( self, graph_module: torch.fx.GraphModule, @@ -409,12 +428,15 @@ def input_to_nhwc( # Check if input uses dynamic quantization is_dynamic_input = is_dynamic_qdq(input_node) + dynamic_chain = [] if is_dynamic_input: - # Trace back to original source node. Stop if args[0] is not - # a Node (e.g., immutable_list from cat). - while getattr(input_node, "args", None) and isinstance( + # Trace back over this consumer's own q/dq chain to the source + # node, so the copy lands ahead of the quantize, and remember the + # chain: only it is redirected below. + while is_dynamic_qdq(input_node) and isinstance( input_node.args[0], torch.fx.Node ): + dynamic_chain.append(input_node) input_node = input_node.args[0] with graph_module.graph.inserting_after(input_node): @@ -427,10 +449,10 @@ def input_to_nhwc( # Use static method for consistency ChannelsLastTaggedReshapePass.mark_as_nhwc_node(input_node_nhwc) - if is_dynamic_input: - # Replace downstream input_nodes with NHWC node - input_node.replace_all_uses_with(input_node_nhwc) - input_node_nhwc.args = (input_node,) + if dynamic_chain: + self.redirect_dynamic_chain_to_nhwc( + input_node, input_node_nhwc, dynamic_chain + ) self.insert_copy_and_assign_partner_nodes_quantization_sensitive( graph_module=graph_module, diff --git a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py index adf1c694b22..bfe2ce862c9 100644 --- a/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py +++ b/backends/xnnpack/test/passes/test_channels_last_tagged_reshape.py @@ -22,6 +22,7 @@ from executorch.backends.xnnpack.test.tester import Quantize, RunPasses, Tester from executorch.backends.xnnpack.utils.quant_utils import ( is_dequant, + is_dynamic_qdq, is_quant, is_tagged_as_implicit_q_dq, ) @@ -364,6 +365,222 @@ def test_dq_conv2d_channels_last_tagged_reshape_pass(self) -> None: .run_method_and_compare_outputs() ) + class EltwiseConv2dDynamicQuant(torch.nn.Module): + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 10, 3) + + def forward(self, x): + return self.conv(torch.sigmoid(x)) + + def test_dq_conv2d_eltwise_source_channels_last_tagged_reshape_pass(self) -> None: + # The conv's input is sigmoid -> q -> dq. Stepping past the q/dq pair leaves + # the sigmoid reading NHWC while its own output stays NCHW, which XNNPACK + # only rejects at runtime. + tester = ( + Tester(self.EltwiseConv2dDynamicQuant().eval(), (torch.randn(1, 3, 8, 8),)) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + artifact = tester.get_artifact(StageType.RUN_PASSES) + graph_module = artifact.exported_program().graph_module + sigmoid_nodes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.sigmoid.default + ] + self.assertEqual(len(sigmoid_nodes), 1) + sigmoid = sigmoid_nodes[0] + + # The sigmoid keeps its NCHW input and the copy sits on its output instead. + self.assertEqual(sigmoid.args[0].op, "placeholder") + copies = [ + user + for user in sigmoid.users + if user.target == exir_ops.edge.aten._to_copy.default + and user.kwargs.get("memory_format") == torch.channels_last + ] + self.assertEqual(len(copies), 1) + + tester.run_method_and_compare_outputs() + + class SiLUStemSharedConv2dDynamicQuant(torch.nn.Module): + """A SiLU stem ahead of the first convolution, as detection backbones have. + + Two producers here have more than one consumer: the input activation feeds + both the sigmoid and the mul, and the SiLU output feeds both the quantized + convolution and the graph output. Both are reachable by the blanket + ``replace_all_uses_with`` in the dynamic-quant branch of ``input_to_nhwc``. + """ + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x): + act = x * torch.sigmoid(x) + return self.conv(act), act + + def test_dq_conv2d_silu_stem_shared_channels_last_tagged_reshape_pass(self) -> None: + tester = ( + Tester( + self.SiLUStemSharedConv2dDynamicQuant().eval(), + (torch.randn(1, 3, 16, 16),), + ) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ) + muls = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.mul.Tensor + ] + self.assertEqual(len(muls), 1) + silu = muls[0] + + # The walk must stop at the SiLU output rather than run on to the input + # activation, which would leave the mul and the sigmoid reading NHWC while + # their own outputs stay NCHW. + for arg in silu.all_input_nodes: + self.assertNotEqual(arg.target, exir_ops.edge.aten._to_copy.default) + + # The SiLU output is converted once, for the convolution only. The graph + # output keeps the unconverted node. + copies = [ + user + for user in silu.users + if user.target == exir_ops.edge.aten._to_copy.default + and user.kwargs.get("memory_format") == torch.channels_last + ] + self.assertEqual(len(copies), 1) + output_node = next( + node for node in graph_module.graph.nodes if node.op == "output" + ) + self.assertIn(silu, output_node.args[0]) + + tester.run_method_and_compare_outputs() + + class SiblingBranchConv2dDynamicQuant(torch.nn.Module): + """A producer feeding both a quantized conv and an ordinary op.""" + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 8, 3, padding=1) + + def forward(self, x): + act = torch.sigmoid(x) + return self.conv(act), torch.tanh(act) + + def test_dq_conv2d_sibling_branch_channels_last_tagged_reshape_pass(self) -> None: + tester = ( + Tester( + self.SiblingBranchConv2dDynamicQuant().eval(), + (torch.randn(1, 3, 16, 16),), + ) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ) + tanhs = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.aten.tanh.default + ] + self.assertEqual(len(tanhs), 1) + + # Only the quantize wrapper moved to the NHWC copy. + self.assertEqual(tanhs[0].args[0].target, exir_ops.edge.aten.sigmoid.default) + + tester.run_method_and_compare_outputs() + + class SharedLinearConvDynamicQuant(torch.nn.Module): + """Two dynamically quantized siblings sharing one source; only the conv + wants NHWC. + """ + + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(8, 8) + self.conv = torch.nn.Conv2d(3, 4, 1) + + def forward(self, x): + act = torch.sigmoid(x) + # Keep linear before conv so it is processed first. + return self.linear(act), self.conv(act) + + def test_dq_shared_linear_conv_channels_last_tagged_reshape_pass(self) -> None: + tester = ( + Tester( + self.SharedLinearConvDynamicQuant().eval(), + (torch.randn(1, 3, 8, 8),), + ) + .quantize( + Quantize( + quantization_config=get_symmetric_quantization_config( + is_dynamic=True + ) + ) + ) + .export() + .to_edge() + .run_passes(self.PassStage) + ) + + graph_module = ( + tester.get_artifact(StageType.RUN_PASSES).exported_program().graph_module + ) + quantizes = [ + node + for node in graph_module.graph.nodes + if is_dynamic_qdq(node) and is_quant(node) + ] + self.assertEqual(len(quantizes), 2) + for quantize in quantizes: + consumer = next(iter(next(iter(quantize.users)).users)) + source = quantize.args[0].target + qparam_source = quantize.args[1].args[0].args[0].target + self.assertEqual(source, qparam_source) + if consumer.target == exir_ops.edge.aten.convolution.default: + # The conv's chain reads the NHWC copy. + self.assertEqual(source, exir_ops.edge.aten._to_copy.default) + else: + # The linear's chain keeps the source. + self.assertEqual(source, exir_ops.edge.aten.sigmoid.default) + + tester.run_method_and_compare_outputs() + class ConvAddConvOutput(torch.nn.Module): def __init__(self): super().__init__() From 4cfe494a1c24a6e5f5b40fba886a349fd8a44982 Mon Sep 17 00:00:00 2001 From: roman-janik-nxp Date: Thu, 10 Sep 2026 20:45:33 +0200 Subject: [PATCH 156/190] NXP backend: Add support for Cortex-M backend benchmarking (#22228) ### Summary Adds support for benchmarking Cortex-M backend in NXP backend test pipeline. ### Test plan Tested on `backends/nxp/tests/generic_tests/test_cifarnet.py`. This feature is for benchmarking only, so no automatic tests were added. cc @robert-kalmar @JakeStevens @digantdesai @rascani --- backends/nxp/tests/cortex_m_benchmarking.py | 79 ++++++ backends/nxp/tests/executorch_pipeline.py | 6 +- backends/nxp/tests/nsys_testing.py | 280 ++++++++++---------- backends/nxp/tests/utils.py | 166 ++++++++++++ 4 files changed, 381 insertions(+), 150 deletions(-) create mode 100644 backends/nxp/tests/cortex_m_benchmarking.py diff --git a/backends/nxp/tests/cortex_m_benchmarking.py b/backends/nxp/tests/cortex_m_benchmarking.py new file mode 100644 index 00000000000..7859e9e2c3b --- /dev/null +++ b/backends/nxp/tests/cortex_m_benchmarking.py @@ -0,0 +1,79 @@ +# Copyright 2025-2026 Arm Limited and/or its affiliates. +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from executorch.backends.cortex_m.target_config import CortexM, CortexMTargetConfig +from executorch.backends.cortex_m.test.tester import CortexMQuantize, CortexMTester +from executorch.backends.nxp.tests.utils import ( + process_input_sample, + process_output_sample, + read_prepared_samples, + store_results, +) +from executorch.backends.test.harness.stages import StageType + + +class CortexMNXPBenchmarkTester(CortexMTester): + def __init__( + self, + module, + example_inputs, + target_config: CortexMTargetConfig | None = None, + timeout: int = 120, + ): + target_config = target_config or CortexMTargetConfig( + cpu=CortexM.M33 + ) # set default to M33 for NXP boards + super().__init__(module, example_inputs, target_config, timeout) + + def run_benchmark( + self, + calibration_samples, + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ): + quantization_stage = CortexMQuantize(calibration_samples=calibration_samples) + + self.quantize(quantization_stage) + self.export() + self.to_edge() + self.run_passes() + self.to_executorch() + self.serialize() + self.run_program( + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ) + + return self.stages[StageType.SERIALIZE].executorch_program_manager + + def run_program( + self, + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ): + all_outputs = [] + + for input_samples in read_prepared_samples(testing_dataset_dir, input_spec): + current_input_samples = process_input_sample(input_spec, input_samples) + + # Run the model. + output = self.stages[StageType.SERIALIZE].run_artifact( + *current_input_samples + ) + current_outputs = process_output_sample(output, output_spec) + all_outputs.append(current_outputs) + + # Store all the results. + store_results(all_outputs, cpu_results_dir, npu_results_dir) diff --git a/backends/nxp/tests/executorch_pipeline.py b/backends/nxp/tests/executorch_pipeline.py index 896bef9a3d5..0567f62b101 100644 --- a/backends/nxp/tests/executorch_pipeline.py +++ b/backends/nxp/tests/executorch_pipeline.py @@ -155,7 +155,7 @@ def _nested( return _nested -def _get_example_input( +def get_example_input( input_spec: tuple[ModelInputSpec, ...], ) -> tuple[torch.Tensor, ...]: example_input = [] @@ -199,7 +199,7 @@ def to_quantized_edge_program( get_quantizer_fn = partial(get_default_quantizer, _neutron_target_spec, use_qat) input_spec = to_model_input_spec(input_spec) calibration_inputs = get_calibration_inputs_fn(input_spec) - example_input = _get_example_input(input_spec) + example_input = get_example_input(input_spec) # Make sure the model is in the evaluation mode. model.eval() @@ -311,7 +311,7 @@ def to_edge_program( model: nn.Module, input_spec: Iterable[ModelInputSpec] | tuple[int, ...] | list[tuple[int, ...]], ) -> EdgeProgramManager: - example_input = _get_example_input(to_model_input_spec(input_spec)) + example_input = get_example_input(to_model_input_spec(input_spec)) # Make sure the model is in the evaluation mode. model.eval() diff --git a/backends/nxp/tests/nsys_testing.py b/backends/nxp/tests/nsys_testing.py index 6ec0d821bae..c7364967b6b 100644 --- a/backends/nxp/tests/nsys_testing.py +++ b/backends/nxp/tests/nsys_testing.py @@ -16,14 +16,8 @@ from os import environ, mkdir from typing import Callable, Iterable -import numpy as np import torch import yaml -from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order -from executorch.backends.nxp.backend.ir.converter.conversion import translator -from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( - torch_type_to_numpy_type, -) from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.tests.config_importer import test_config @@ -34,6 +28,7 @@ ) from executorch.backends.nxp.tests.executorch_pipeline import ( get_calibration_inputs_fn_from_dataset_dir, + get_example_input, ModelInputSpec, to_edge_program, to_model_input_spec, @@ -46,8 +41,14 @@ AllCloseOutputComparator, ) from executorch.backends.nxp.tests.outputs_dir_importer import outputs_dir -from executorch.backends.nxp.tests.utils import save_pte_program, store_txt_input_tensor - +from executorch.backends.nxp.tests.utils import ( + process_input_sample, + process_output_sample, + read_prepared_samples, + save_pte_program, + store_results, + store_txt_input_tensor, +) from executorch.devtools.visualization.visualization_utils import ( visualize_with_clusters, ) @@ -72,6 +73,7 @@ class ReferenceModel(Enum): # QUANTIZED_ATEN_PYTHON = 2 # Not implemented. # FLOAT_ATEN_PYTHON = 3 # Not implemented. FLOAT_PYTORCH_PYTHON = 4 + QUANTIZED_CORTEX_M = 5 def _get_dataset_cli_args(input_spec: list[ModelInputSpec], testing_dataset_dir): @@ -251,102 +253,6 @@ def _save_non_quantized_fp32_executorch_program( return non_quantized_program.exported_program() -def read_prepared_samples( - dataset_dir: str, input_spec: list[ModelInputSpec] -) -> list[tuple[np.ndarray, ...]]: - """Read numpy arrays generated by a `DatasetCreator`. - - :param dataset_dir: Directory containing the generated samples - :param input_spec: List of ModelInputSpec defining the shape and type of each input - - :return: List of tuples, where each tuple contains numpy arrays for one sample - """ - all_samples = [] - - # Multi-input: samples are in numbered subdirectories - if len(input_spec) > 1: - sample_dirs = sorted( - [ - d - for d in os.listdir(dataset_dir) - if os.path.isdir(os.path.join(dataset_dir, d)) - ] - ) - - for sample_name in sample_dirs: - sample_dir = os.path.join(dataset_dir, sample_name) - current_samples = [] - - for spec_idx, spec in enumerate(input_spec): - bin_file_path = os.path.join( - sample_dir, f"{str(spec_idx).zfill(2)}.bin" - ) - sample_vector = np.fromfile( - bin_file_path, dtype=torch_type_to_numpy_type(spec.dtype) - ).reshape(spec.shape) - current_samples.append(sample_vector) - - all_samples.append(tuple(current_samples)) - - # Single-input: binary files are directly in dataset_dir - else: - bin_files = sorted([f for f in os.listdir(dataset_dir) if f.endswith(".bin")]) - - for bin_file in bin_files: - bin_file_path = os.path.join(dataset_dir, bin_file) - sample_vector = np.fromfile( - bin_file_path, dtype=torch_type_to_numpy_type(input_spec[0].dtype) - ).reshape(input_spec[0].shape) - all_samples.append((sample_vector,)) - - return all_samples - - -def store_results( - results: list[tuple[np.ndarray, ...]], output_dir: str, reference_dir: str -): - """Store a list of output arrays in the directory structure matching the reference directory. - - :param results: List of tuples, where each tuple contains numpy arrays (outputs for one sample) - :param output_dir: Directory where results will be stored - - Directory structure created matches reference_dir: - output_dir/ - ├── sample_0/ - │ ├── 0000.bin - │ └── 0001.bin - ├── some_other_sample/ - │ ├── 0000.bin - │ └── 0001.bin - """ - os.makedirs(output_dir, exist_ok=True) - - # Get subdirectories from reference directory - sample_dirs = sorted( - [ - d - for d in os.listdir(reference_dir) - if os.path.isdir(os.path.join(reference_dir, d)) - ] - ) - - assert len(sample_dirs) == len( - results - ), f"Number of samples ({len(results)}) must match number of subdirectories in reference_dir ({len(sample_dirs)})" - - for _sample_idx, (sample_name, sample_outputs) in enumerate( - zip(sample_dirs, results) - ): - sample_dir = os.path.join(output_dir, sample_name) - os.makedirs(sample_dir, exist_ok=True) - - # Store each output tensor - for output_idx, output_array in enumerate(sample_outputs): - bin_file_name = f"{str(output_idx).zfill(4)}.bin" - bin_file_path = os.path.join(sample_dir, bin_file_name) - output_array.tofile(bin_file_path) - - def _run_python_program( model: torch.nn.Module | GraphModule, testing_dataset_dir, @@ -371,53 +277,102 @@ def _run_python_program( all_outputs = [] for input_samples in read_prepared_samples(testing_dataset_dir, input_spec): - current_input_samples = [] - for spec, sample in zip(input_spec, input_samples, strict=True): - match spec.dim_order: - case torch.contiguous_format: - # Use the data as is, just turn it into a PyTorch tensor. - sample = torch.tensor(sample) - - case torch.channels_last: - # The tensor data was stored by the DatasetCreator as channels last (NHWC), but it was now - # incorrectly parsed as contiguous/channels first (NCHW). Transpose it to channels last to preserve - # the semantics. - channels_last_shape = translator.dims_to_channels_last( - list(spec.shape) - ) - sample = np.moveaxis(sample.reshape(channels_last_shape), -1, 1) - sample = torch.tensor(sample).to(memory_format=torch.channels_last) - - case _: - raise ValueError(f"Unsupported dim_order: {spec.dim_order}") - - current_input_samples.append(sample) + current_input_samples = process_input_sample(input_spec, input_samples) # Run the model. output = model(*current_input_samples) - if isinstance(output, torch.Tensor): - output = (output,) + current_outputs = process_output_sample(output, output_spec) + all_outputs.append(current_outputs) - current_outputs = [] + # Store all the results. + store_results(all_outputs, cpu_results_dir, npu_results_dir) - for o, o_spec in zip(output, output_spec, strict=True): - dim_order = list(o_spec.dim_order()) # ExecuTorch dim order. - rank = len(o_spec.shape) - if dim_order == list(range(rank)): # Contiguous dim order. - current_outputs.append(o.detach().numpy()) - elif is_channels_last_dim_order(dim_order): # Channels last dim order. - # The NPU variant outputs channels last (NHWC). We need to convert the CPU output to match. - o = o.detach().numpy().reshape(o_spec.shape) - current_outputs.append(np.moveaxis(o, 1, -1)) +def _run_cortex_m_program( + model: torch.nn.Module | GraphModule, + test_dir, + test_name, + calibration_dataset_dir, + testing_dataset_dir, + input_spec: list[ModelInputSpec], + output_spec: list[torch.Tensor], + cpu_results_dir, + npu_results_dir, +): + """Run a model with Cortex-M backend with channels last inputs. + + :param model: Any PyTorch/ExecuTorch model runnable with Cortex-M with channels last inputs. + :param test_dir: Directory for saving test artifacts. + :param test_name: Name of the test. + :param calibration_dataset_dir: Directory containing calibration data. + :param testing_dataset_dir: Directory containing testing data. The samples have to be channels last (NHWC) for 4D tensors. + The format must match the input_spec.dim_order. + :param input_spec: List of ModelInputSpec defining the shape, type, and dimension order of each input. + :param output_spec: List of output tensor specifications. + :param cpu_results_dir: Directory where CPU results will be stored. The structure will match the existing structure + of `npu_results_dir`. + :param npu_results_dir: Directory where NPU results are already stored, to serve as reference directory structure + for `cpu_results_dir`. + """ + # Assert Cortex-M dependencies are available + assert_cortex_m() - else: - raise ValueError(f"Unsupported dim_order: {o_spec.dim_order}") + from executorch.backends.nxp.tests.cortex_m_benchmarking import ( + CortexMNXPBenchmarkTester, + ) - all_outputs.append(current_outputs) + numpy_samples = read_prepared_samples(calibration_dataset_dir, input_spec) + calibration_samples = [ + process_input_sample(input_spec, sample) for sample in numpy_samples + ] - # Store all the results. - store_results(all_outputs, cpu_results_dir, npu_results_dir) + example_inputs = get_example_input(input_spec) + + tester = CortexMNXPBenchmarkTester( + model, + example_inputs, + ) + cortex_m_delegated_program = tester.run_benchmark( + calibration_samples, + input_spec, + output_spec, + testing_dataset_dir, + cpu_results_dir, + npu_results_dir, + ) + save_pte_program( + cortex_m_delegated_program, test_name + "_cortex_m_delegated", test_dir + ) + + +def assert_cortex_m(): + # Follow backends/cortex_m/README.md to install the required dependencies. + # Build Arm executor runner with target="cortex-m33": + # ./backends/cortex_m/test/build_test_runner.sh --target="cortex-m33" + # FVP Corstone-300 simulator needs to be added to PATH. + + import sysconfig + + suffix = sysconfig.get_config_var( + "EXT_SUFFIX" + ) # e.g. ".cpython-312-x86_64-linux-gnu.so" + cmsis_nn_lib_path = os.path.join( + PROJECT_DIR, "backends/cortex_m/library/_cmsis_nn", f"cmsis_nn{suffix}" + ) + + assert os.path.exists( + cmsis_nn_lib_path + ), "CMSIS-NN lib is not available, check if ET is built correctly." + fvp_simulator_path = os.path.join( + PROJECT_DIR, + "examples/arm/arm-scratch/FVP-corstone300/models/Linux64_GCC-9.3/FVP_Corstone_SSE-300_Ethos-U55", + ) + assert os.path.exists(fvp_simulator_path), "Arm FVP Corstone-300 is not installed." + arm_executor_runner_path = os.path.join( + PROJECT_DIR, + "arm_test/arm_semihosting_executor_runner_corstone-300_cortex-m33/arm_executor_runner", + ) + assert os.path.exists(arm_executor_runner_path), "Arm executor is not installed." def assert_NSYS(): @@ -590,11 +545,44 @@ def lower_run_compare( npu_results_dir, ) + case ReferenceModel.QUANTIZED_CORTEX_M: + if use_qat: + raise ValueError( + "Flag use_qat is not applicable to QUANTIZED_CORTEX_M reference model " + "as it doesn't support QAT. Run with use_qat=False." + ) + if remove_quant_io_ops: + raise ValueError( + "Flag remove_quant_io_ops is not applicable to QUANTIZED_CORTEX_M reference model " + "as it works with float data only. Run with remove_quant_io_ops=False." + ) + if any( + spec.dim_order != torch.channels_last + for spec in input_spec + if len(spec.shape) == 4 + ): + raise ValueError( + "Cortex-M backend supports only channel last dim order for 4D inputs." + ) + + model_to_delegate_cortex_m = deepcopy(model) + + # Lower to quantized Cortex-M program and run on Arm simulator. + _run_cortex_m_program( + model_to_delegate_cortex_m, + test_dir, + test_name, + calibration_dataset_dir, + testing_dataset_dir, + input_spec, + output_spec, + cpu_results_dir, + npu_results_dir, + ) + case _: raise ValueError(f"Unsupported reference model: `{reference_model}`.") - output_tensor_spec = _get_program_output_spec(delegated_program) - if logging.root.isEnabledFor(logging.DEBUG): _generate_txt_test_data( calibration_dataset_dir, testing_dataset_dir, list(input_spec) @@ -602,9 +590,7 @@ def lower_run_compare( dump_debug_test_summary(test_name, test_dir) npu_results_dir = os.path.join(test_dir, "results_npu") cpu_results_dir = os.path.join(test_dir, "results_cpu") - output_comparator.compare_results( - cpu_results_dir, npu_results_dir, output_tensor_spec - ) + output_comparator.compare_results(cpu_results_dir, npu_results_dir, output_spec) def lower_run_compare_ptq_qat( diff --git a/backends/nxp/tests/utils.py b/backends/nxp/tests/utils.py index 00b7c364a31..ef0fe956b4c 100644 --- a/backends/nxp/tests/utils.py +++ b/backends/nxp/tests/utils.py @@ -11,7 +11,11 @@ import numpy as np +import torch +from executorch.backends.nxp.backend.edge_helper import is_channels_last_dim_order + from executorch.backends.nxp.backend.ir.converter.conversion.translator import ( + dims_to_channels_last, torch_type_to_numpy_type, ) from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec @@ -64,3 +68,165 @@ def store_txt_input_tensor( def archive_test_dir(test_dir: str): shutil.make_archive(test_dir, "zip", test_dir) + + +def read_prepared_samples( + dataset_dir: str, input_spec: list[ModelInputSpec] +) -> list[tuple[np.ndarray, ...]]: + """Read numpy arrays generated by a `DatasetCreator`. + + :param dataset_dir: Directory containing the generated samples + :param input_spec: List of ModelInputSpec defining the shape and type of each input + + :return: List of tuples, where each tuple contains numpy arrays for one sample + """ + all_samples = [] + + # Multi-input: samples are in numbered subdirectories + if len(input_spec) > 1: + sample_dirs = sorted( + [ + d + for d in os.listdir(dataset_dir) + if os.path.isdir(os.path.join(dataset_dir, d)) + ] + ) + + for sample_name in sample_dirs: + sample_dir = os.path.join(dataset_dir, sample_name) + current_samples = [] + + for spec_idx, spec in enumerate(input_spec): + bin_file_path = os.path.join( + sample_dir, f"{str(spec_idx).zfill(2)}.bin" + ) + sample_vector = np.fromfile( + bin_file_path, dtype=torch_type_to_numpy_type(spec.dtype) + ).reshape(spec.shape) + current_samples.append(sample_vector) + + all_samples.append(tuple(current_samples)) + + # Single-input: binary files are directly in dataset_dir + else: + bin_files = sorted([f for f in os.listdir(dataset_dir) if f.endswith(".bin")]) + + for bin_file in bin_files: + bin_file_path = os.path.join(dataset_dir, bin_file) + sample_vector = np.fromfile( + bin_file_path, dtype=torch_type_to_numpy_type(input_spec[0].dtype) + ).reshape(input_spec[0].shape) + all_samples.append((sample_vector,)) + + return all_samples + + +def store_results( + results: list[tuple[np.ndarray, ...]], output_dir: str, reference_dir: str +): + """Store a list of output arrays in the directory structure matching the reference directory. + + :param results: List of tuples, where each tuple contains numpy arrays (outputs for one sample) + :param output_dir: Directory where results will be stored + + Directory structure created matches reference_dir: + output_dir/ + ├── sample_0/ + │ ├── 0000.bin + │ └── 0001.bin + ├── some_other_sample/ + │ ├── 0000.bin + │ └── 0001.bin + """ + os.makedirs(output_dir, exist_ok=True) + + # Get subdirectories from reference directory + sample_dirs = sorted( + [ + d + for d in os.listdir(reference_dir) + if os.path.isdir(os.path.join(reference_dir, d)) + ] + ) + + assert len(sample_dirs) == len( + results + ), f"Number of samples ({len(results)}) must match number of subdirectories in reference_dir ({len(sample_dirs)})" + + for _sample_idx, (sample_name, sample_outputs) in enumerate( + zip(sample_dirs, results) + ): + sample_dir = os.path.join(output_dir, sample_name) + os.makedirs(sample_dir, exist_ok=True) + + # Store each output tensor + for output_idx, output_array in enumerate(sample_outputs): + bin_file_name = f"{str(output_idx).zfill(4)}.bin" + bin_file_path = os.path.join(sample_dir, bin_file_name) + output_array.tofile(bin_file_path) + + +def process_input_sample( + input_spec: list[ModelInputSpec], input_samples: tuple[np.ndarray, ...] +) -> list[torch.Tensor]: + """Process input samples by converting them to PyTorch tensors with correct dimension order. + + :param input_spec: List of ModelInputSpec defining the shape, type, and dimension order of each input + :param input_samples: Tuple of numpy arrays representing one sample + + :return: List of PyTorch tensors with correct dimension order + """ + current_input_samples = [] + for spec, sample in zip(input_spec, input_samples, strict=True): + match spec.dim_order: + case torch.contiguous_format: + # Use the data as is, just turn it into a PyTorch tensor. + sample = torch.tensor(sample) + + case torch.channels_last: + # The tensor data was stored by the DatasetCreator as channels last (NHWC), but it was now + # incorrectly parsed as contiguous/channels first (NCHW). Transpose it to channels last to preserve + # the semantics. + channels_last_shape = dims_to_channels_last(list(spec.shape)) + sample = np.moveaxis(sample.reshape(channels_last_shape), -1, 1) + sample = torch.tensor(sample).to(memory_format=torch.channels_last) + + case _: + raise ValueError(f"Unsupported dim_order: {spec.dim_order}") + + current_input_samples.append(sample) + + return current_input_samples + + +def process_output_sample( + output: tuple[torch.Tensor, ...] | torch.Tensor, output_spec: list[torch.Tensor] +) -> list[np.ndarray]: + """Process output tensors by converting them to numpy arrays with correct dimension order. + + :param output: Model output - either a single tensor or tuple of tensors + :param output_spec: List of output tensor specifications + + :return: List of numpy arrays with correct dimension order matching NPU output format + """ + + if isinstance(output, torch.Tensor): + output = (output,) + + current_outputs = [] + + for o, o_spec in zip(output, output_spec, strict=True): + dim_order = list(o_spec.dim_order()) # ExecuTorch dim order. + rank = len(o_spec.shape) + if dim_order == list(range(rank)): # Contiguous dim order. + current_outputs.append(o.detach().numpy()) + + elif is_channels_last_dim_order(dim_order): # Channels last dim order. + # The NPU variant outputs channels last (NHWC). We need to convert the CPU output to match. + o = o.detach().numpy().reshape(o_spec.shape) + current_outputs.append(np.moveaxis(o, 1, -1)) + + else: + raise ValueError(f"Unsupported dim_order: {o_spec.dim_order()}") + + return current_outputs From 49905c7aa30f28f0e51297acacff1bfb60008f9b Mon Sep 17 00:00:00 2001 From: roman-janik-nxp Date: Thu, 10 Sep 2026 21:14:43 +0200 Subject: [PATCH 157/190] NXP backend: Add support for AdaptiveAvgPool1d, fix AdaptiveAvgPool1d tests, add check for 4D pool ops inputs, refactor tests (#22361) ### Summary Adds support for AdaptiveAvgPool1d, fixes tests for AdaptiveAvgPool1d. Adds check for 4D pool ops inputs. Refactors `ops_aliases.py` usage, all tests now use this file instead os `exir_ops`. Add # noinspection PyUnusedImports for pytest import where it was missing. Moves tests with name `test_*` from `backends/nxp/tests` into subdirs `generic_tests` and `models`. Renames `backends/nxp/tests/models.py` file to `simple_models.py` to not confuse with new package `backends/nxp/tests/models`. Changes affects almost all test files. ### Test plan Unit tests provided. cc @robert-kalmar @JakeStevens @digantdesai @rascani --- .../adaptive_avg_pool_2d_converter.py | 5 ++ .../ops_converters/avg_pool_2d_converter.py | 6 ++ .../max_pool2d_with_indices_converter.py | 6 +- backends/nxp/backend/ops_aliases.py | 14 +++ .../convert_reshaping_nodes_to_view.py | 25 +++--- ...operator_into_separate_qdq_cluster_pass.py | 85 ++++++++++--------- ...ditional_quantize_dequantize_nodes_pass.py | 27 ++++-- .../remove_as_strided_copy_nodes.py | 6 +- .../edge_passes/remove_io_quant_ops_pass.py | 15 ++-- backends/nxp/quantizer/neutron_quantizer.py | 6 +- backends/nxp/quantizer/patterns.py | 11 ++- ...add_batch_size_for_3d_input_pool_2d_ops.py | 4 +- .../tests/generic_tests/test_aot_example.py | 1 + .../generic_tests/test_batch_norm_fusion.py | 4 +- .../test_context_sensitive_delegation.py | 23 +++-- .../test_convert_1d_conv_to_2d.py | 19 +++-- .../generic_tests/test_convert_div_to_mul.py | 4 +- .../test_convert_scalar_to_attr.py | 10 +-- .../tests/generic_tests/test_debug_results.py | 4 +- .../test_decompose_split_to_slices.py | 6 +- .../test_fold_redundant_qdq.py | 3 +- .../test_fuse_batch_norm_single_user.py | 0 .../tests/generic_tests/test_gru_splitting.py | 2 + .../tests/generic_tests/test_integration.py | 2 +- .../generic_tests/test_kernel_selection.py | 4 +- ...st_move_activation_before_concatenation.py | 26 ++++-- .../generic_tests/test_neutron_backend.py | 5 +- .../test_neutron_backend_executor.py | 5 +- .../test_neutron_compiler_manager.py | 2 +- .../test_node_format_inference.py | 2 +- .../generic_tests/test_operator_selector.py | 2 +- .../test_per_channel_conversion.py | 19 ++--- .../nxp/tests/generic_tests/test_profiling.py | 5 +- .../generic_tests/test_qdq_clustering_conv.py | 2 +- .../test_quantized_input_data.py | 2 +- .../nxp/tests/generic_tests/test_quantizer.py | 36 +++++--- .../tests/generic_tests/test_recipe_export.py | 2 +- .../generic_tests/test_removing_dead_code.py | 2 + .../test_split_group_convolution.py | 10 +-- .../test_adaptive_avg_pool2d_converter.py | 77 ++++++++++------- .../test_add_tensor_converter.py | 5 +- .../node_converter/test_addmm_converter.py | 2 +- .../test_avg_pool2d_converter.py | 8 +- .../node_converter/test_bmm_converter.py | 4 +- .../node_converter/test_clone_converter.py | 23 ++--- .../test_constant_pad_nd_converter.py | 2 +- .../node_converter/test_conv_converter.py | 7 +- .../test_hardswish_converter.py | 4 +- .../node_converter/test_hardtanh_converter.py | 5 +- .../test_max_pool_2d_converter.py | 8 +- .../node_converter/test_maximum_converter.py | 5 +- .../node_converter/test_minimum_converter.py | 5 +- .../node_converter/test_mm_converter.py | 2 +- .../test_mul_tensor_converter.py | 5 +- .../node_converter/test_pad_converter.py | 4 +- .../node_converter/test_prelu_converter.py | 6 +- .../node_converter/test_relu_converter.py | 8 +- .../test_slice_copy_tensor_converter.py | 4 +- .../node_converter/test_softmax_converter.py | 4 +- .../test_sub_tensor_converter.py | 5 +- .../node_converter/test_tanh_converter.py | 2 +- .../test_view_copy_converter.py | 2 + .../test_convert_reshaping_nodes_to_view.py | 5 +- .../tests/ir/edge_passes/test_edge_passes.py | 33 +++---- .../ir/edge_passes/test_linear_bn_fusing.py | 12 +-- .../test_remove_io_quant_ops_pass.py | 2 +- backends/nxp/tests/models/__init__.py | 0 .../test_cifarnet.py | 1 + .../test_mlperf_tiny_image_classification.py | 5 +- .../nxp/tests/{models.py => simple_models.py} | 10 +++ backends/nxp/tests/use_qat.py | 1 + docs/source/backends/nxp/nxp-quantization.md | 1 + docs/source/backends/nxp/op-support.csv | 3 +- 73 files changed, 424 insertions(+), 258 deletions(-) rename backends/nxp/tests/{ => generic_tests}/test_convert_1d_conv_to_2d.py (98%) rename backends/nxp/tests/{ => generic_tests}/test_fold_redundant_qdq.py (97%) rename backends/nxp/tests/{ => generic_tests}/test_fuse_batch_norm_single_user.py (100%) create mode 100644 backends/nxp/tests/models/__init__.py rename backends/nxp/tests/{generic_tests => models}/test_cifarnet.py (98%) rename backends/nxp/tests/{models.py => simple_models.py} (99%) diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py index bd812cccb76..d4029ebebce 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/adaptive_avg_pool_2d_converter.py @@ -6,6 +6,7 @@ import executorch.backends.nxp.backend.ir.lib.tflite.Padding as tflPadding import torch +from executorch.backends.nxp.backend.edge_helper import input_rank from executorch.backends.nxp.backend.ir.converter.conversion import common from executorch.backends.nxp.backend.ir.converter.node_converter import ( CustomDelegationOptions, @@ -45,6 +46,10 @@ def _is_supported_in_IR( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + # The input must be 4D. + if input_rank(node, 0) != 4: + return False + input_size = node.args[0].meta["val"].shape output_size = node.args[1] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py index b3157ab4c4b..5791edcfac9 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/avg_pool_2d_converter.py @@ -5,6 +5,8 @@ import numpy as np import torch + +from executorch.backends.nxp.backend.edge_helper import input_rank from executorch.backends.nxp.backend.ir.converter.conversion import ( aten_translator, common, @@ -36,6 +38,10 @@ def _is_supported_in_IR( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + # The input must be 4D. + if input_rank(node, 0) != 4: + return False + n_args = len(node.args) padding = node.args[3] if n_args >= 4 else [0, 0] diff --git a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py index c8d24ea34a6..507de4283c1 100644 --- a/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py +++ b/backends/nxp/backend/ir/converter/node_converters/ops_converters/max_pool2d_with_indices_converter.py @@ -7,7 +7,7 @@ import numpy as np import torch -from executorch.backends.nxp.backend.edge_helper import try_get_arg +from executorch.backends.nxp.backend.edge_helper import input_rank, try_get_arg from executorch.backends.nxp.backend.ir.converter.conversion import ( aten_translator, common, @@ -42,6 +42,10 @@ def _is_supported_in_IR( parameters_mapping: dict[str, Parameter], custom_delegation_options: CustomDelegationOptions, ) -> bool: + # The input must be 4D. + if input_rank(node, 0) != 4: + return False + kernel_size, stride, padding, dilation, ceil_mode = ( MaxPool2DWithIndicesConverter._get_node_args(node) ) diff --git a/backends/nxp/backend/ops_aliases.py b/backends/nxp/backend/ops_aliases.py index a9eb04e1960..fea772eb8ee 100644 --- a/backends/nxp/backend/ops_aliases.py +++ b/backends/nxp/backend/ops_aliases.py @@ -17,6 +17,7 @@ AddTensor = exir_ops.edge.aten.add.Tensor Amax = exir_ops.edge.aten.amax.default Amin = exir_ops.edge.aten.amin.default +AsStridedCopy = exir_ops.edge.aten.as_strided_copy.default AvgPool2D = exir_ops.edge.aten.avg_pool2d.default BMM = exir_ops.edge.aten.bmm.default Cat = exir_ops.edge.aten.cat.default @@ -27,6 +28,9 @@ Convolution = exir_ops.edge.aten.convolution.default DequantizePerChannel = exir_ops.edge.quantized_decomposed.dequantize_per_channel.default DequantizePerTensor = exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default +DequantizePerTensorTensor = ( + exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor +) ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate Exp = exir_ops.edge.aten.exp.default GetItem = operator.getitem @@ -35,6 +39,7 @@ HardTanh = exir_ops.edge.aten.hardtanh.default HardTanh_ = exir_ops.edge.aten.hardtanh_.default LeakyRelu = exir_ops.edge.aten.leaky_relu.default +Linear = exir_ops.edge.aten.linear.default Log = exir_ops.edge.aten.log.default MM = exir_ops.edge.aten.mm.default Maximum = exir_ops.edge.aten.maximum.default @@ -43,12 +48,17 @@ MeanDim = exir_ops.edge.aten.mean.dim Minimum = exir_ops.edge.aten.minimum.default MulTensor = exir_ops.edge.aten.mul.Tensor +NativebatchNormLegitNoStats = exir_ops.edge.aten._native_batch_norm_legit.no_stats +NativebatchNormLegitNoTraining = ( + exir_ops.edge.aten._native_batch_norm_legit_no_training.default +) Neg = exir_ops.edge.aten.neg.default Pad = exir_ops.edge.aten.pad.default PermuteCopy = exir_ops.edge.aten.permute_copy.default Prelu = exir_ops.edge.aten.prelu.default QuantizePerChannel = exir_ops.edge.quantized_decomposed.quantize_per_channel.default QuantizePerTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.default +QuantizePerTensorTensor = exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor Relu = exir_ops.edge.aten.relu.default Rsqrt = exir_ops.edge.aten.rsqrt.default Sigmoid = exir_ops.edge.aten.sigmoid.default @@ -56,6 +66,9 @@ SliceCopy = exir_ops.edge.aten.slice_copy.Tensor Softmax = exir_ops.edge.aten._softmax.default Squeeze = exir_ops.edge.aten.squeeze.default +SqueezeCopy = exir_ops.edge.aten.squeeze_copy.default +SqueezeCopyDim = exir_ops.edge.aten.squeeze_copy.dim +SqueezeCopyDims = exir_ops.edge.aten.squeeze_copy.dims SqueezeDim = exir_ops.edge.aten.squeeze.dim SqueezeDims = exir_ops.edge.aten.squeeze.dims SubTensor = exir_ops.edge.aten.sub.Tensor @@ -63,6 +76,7 @@ Tanh = exir_ops.edge.aten.tanh.default Tanh_ = exir_ops.edge.aten.tanh_.default Unsqueeze = exir_ops.edge.aten.unsqueeze.default +UnsqueezeCopy = exir_ops.edge.aten.unsqueeze_copy.default UpsampleBilinear2D = exir_ops.edge.aten.upsample_bilinear2d.vec UpsampleNearest2D = exir_ops.edge.aten.upsample_nearest2d.vec ViewCopy = exir_ops.edge.aten.view_copy.default diff --git a/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py b/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py index a82c823f593..ad296b165b2 100644 --- a/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py +++ b/backends/nxp/edge_passes/convert_reshaping_nodes_to_view.py @@ -6,8 +6,14 @@ import torch +from executorch.backends.nxp.backend.ops_aliases import ( + SqueezeCopy, + SqueezeCopyDim, + SqueezeCopyDims, + UnsqueezeCopy, + ViewCopy, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass -from executorch.exir.dialects._ops import ops as exir_ops from torch._subclasses import FakeTensor, FakeTensorMode from torch.fx import GraphModule, Node from torch.fx.passes.infra.pass_base import PassResult @@ -43,25 +49,20 @@ class ConvertReshapingNodesToViewPass(NeutronEdgePass): @staticmethod def _is_squeeze(node_: Node) -> bool: return node_.op == "call_function" and ( - node_.target == exir_ops.edge.aten.squeeze_copy.dim - or node_.target == exir_ops.edge.aten.squeeze_copy.dims - or node_.target == exir_ops.edge.aten.squeeze_copy.default + node_.target == SqueezeCopyDim + or node_.target == SqueezeCopyDims + or node_.target == SqueezeCopy ) @staticmethod def _is_unsqueeze(node_: Node) -> bool: - return ( - node_.op == "call_function" - and node_.target == exir_ops.edge.aten.unsqueeze_copy.default - ) + return node_.op == "call_function" and node_.target == UnsqueezeCopy def _create_view_copy_node(self, *view_args) -> Node: - view_target = exir_ops.edge.aten.view_copy.default + view_target = ViewCopy view_node = self.graph_module.graph.call_function(view_target, view_args) - view_node.meta["source_fn_stack"] = [ - (view_node.name, exir_ops.edge.aten.view_copy.default) - ] + view_node.meta["source_fn_stack"] = [(view_node.name, ViewCopy)] x_val = view_args[0].meta["val"] with FakeTensorMode() as mode: diff --git a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py index 9ac888addb9..dc9cc28a25e 100644 --- a/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py +++ b/backends/nxp/edge_passes/move_auxiliary_operator_into_separate_qdq_cluster_pass.py @@ -3,36 +3,35 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import operator - import torch -from executorch.backends.nxp.backend.ops_aliases import PermuteCopy + +from executorch.backends.nxp.backend.ops_aliases import ( + AdaptiveAvgPool2D, + AddMM, + AvgPool2D, + Clone, + CloneDimOrder, + Convolution, + DequantizePerTensor, + GetItem, + HardTanh, + MaxPool2DWithIndices, + MM, + PermuteCopy, + QuantizePerTensor, + Relu, + Sigmoid, + SqueezeCopyDims, + Tanh, + UnsqueezeCopy, + ViewCopy, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass from executorch.backends.nxp.neutron_partitioner import QDQClusterRecognizer - -# noinspection PyProtectedMember -from executorch.exir.dialects._ops import ops as exir_ops from torch.fx import Node from torch.fx.passes.infra.pass_base import PassResult -# Operator aliases for better readability. -AddMM = exir_ops.edge.aten.addmm.default -AvgPool2D = exir_ops.edge.aten.avg_pool2d.default -MaxPool2D = exir_ops.edge.aten.max_pool2d_with_indices.default -Conv = exir_ops.edge.aten.convolution.default -Clone = exir_ops.edge.aten.clone.default -CloneDimOrder = exir_ops.edge.dim_order_ops._clone_dim_order.default -Getitem = operator.getitem -HardTanh = exir_ops.edge.aten.hardtanh.default -MM = exir_ops.edge.aten.mm.default -Relu = exir_ops.edge.aten.relu.default -Sigmoid = exir_ops.edge.aten.sigmoid.default -SqueezeCopy = exir_ops.edge.aten.squeeze_copy.dims -Tanh = exir_ops.edge.aten.tanh.default -UnsqueezeCopy = exir_ops.edge.aten.unsqueeze_copy.default -ViewCopy = exir_ops.edge.aten.view_copy.default - def insert_qdq_pair_after_node( graph: torch.fx.Graph, anchor: torch.fx.Node, q_params: tuple @@ -41,7 +40,7 @@ def insert_qdq_pair_after_node( with graph.inserting_after(anchor): quantize_op = graph.create_node( op="call_function", - target=exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, + target=QuantizePerTensor, args=(), # Will be added later. ) quantize_op.meta = anchor.meta @@ -50,7 +49,7 @@ def insert_qdq_pair_after_node( with graph.inserting_after(quantize_op): dequantize_op = graph.create_node( op="call_function", - target=exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, + target=DequantizePerTensor, args=(quantize_op,) + q_params, ) dequantize_op.meta = quantize_op.meta @@ -65,8 +64,7 @@ def _is_dequantize(node_: Node) -> bool: return ( hasattr(node_, "op") and node_.op == "call_function" - and node_.target - == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default + and node_.target == DequantizePerTensor ) @@ -74,8 +72,7 @@ def _is_quantize(node_: Node) -> bool: return ( hasattr(node_, "op") and node_.op == "call_function" - and node_.target - == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default + and node_.target == QuantizePerTensor ) @@ -117,7 +114,7 @@ class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): PermuteCopy, ], ViewCopy: [Clone, CloneDimOrder], - Conv: [ + Convolution: [ ViewCopy, # For 1D conv ], # AvgPool1D is represented in edge as Unsqueeze -> AvgPool2D -> Squeeze. The reshaping nodes must be moved out @@ -126,9 +123,15 @@ class MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): ViewCopy, UnsqueezeCopy, ], - # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> Getitem -> Squeeze. The reshaping nodes must be moved out + # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> GetItem -> Squeeze. The reshaping nodes must be moved out # of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. - MaxPool2D: [ + MaxPool2DWithIndices: [ + ViewCopy, + UnsqueezeCopy, + ], + # AdaptiveAvgPool1D is represented in edge as Unsqueeze -> AdaptiveAvgPool2D -> Squeeze. The reshaping nodes + # must be moved out of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. + AdaptiveAvgPool2D: [ ViewCopy, UnsqueezeCopy, ], @@ -222,7 +225,7 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): Sigmoid, Tanh, ], - Conv: [ + Convolution: [ HardTanh, Relu, Sigmoid, @@ -234,13 +237,19 @@ class MoveTrailingAuxiliaryOperatorIntoSeparateQDQClusterPass(NeutronEdgePass): # of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. AvgPool2D: [ ViewCopy, - SqueezeCopy, + SqueezeCopyDims, ], - # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> Getitem -> Squeeze. The reshaping nodes must be moved out + # MaxPool1D is represented in edge as Unsqueeze -> MaxPool2D -> GetItem -> Squeeze. The reshaping nodes must be moved out # of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. - Getitem: [ + GetItem: [ + ViewCopy, + SqueezeCopyDims, + ], + # AdaptiveAvgPool1D is represented in edge as Unsqueeze -> AdaptiveAvgPool2D -> Squeeze. The reshaping nodes + # must be moved out of the cluster. Instead of [Un]squeeze, ViewCopy can be used as well. + AdaptiveAvgPool2D: [ ViewCopy, - SqueezeCopy, + SqueezeCopyDims, ], } @@ -278,7 +287,7 @@ def run(self, graph_module: torch.fx.GraphModule) -> PassResult: # satisfy the requirements of the `QDQClusterRecognizer`. actual_main_cluster_node = ( main_cluster_node - if main_cluster_node.target != Getitem + if main_cluster_node.target != GetItem else main_cluster_node.args[0] ) cluster = QDQClusterRecognizer().get_qdq_cluster(actual_main_cluster_node) diff --git a/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py b/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py index 4edcc0b0e97..549b77f6eca 100644 --- a/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py +++ b/backends/nxp/edge_passes/remove_additional_quantize_dequantize_nodes_pass.py @@ -7,9 +7,18 @@ import torch from executorch.backends.nxp.backend.edge_helper import get_quantization_parameters_for +from executorch.backends.nxp.backend.ops_aliases import ( + Cat, + DequantizePerChannel, + DequantizePerTensor, + DequantizePerTensorTensor, + PermuteCopy, + QuantizePerChannel, + QuantizePerTensor, + QuantizePerTensorTensor, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass from executorch.backends.nxp.neutron_partitioner import QDQClusterRecognizer -from executorch.exir.dialects._ops import ops as exir_ops from torch.fx.passes.infra.pass_base import PassResult @@ -36,15 +45,15 @@ class RemoveAdditionalQDQClustersPass(NeutronEdgePass): """ qdq_per_channel_nodes = ( - exir_ops.edge.quantized_decomposed.dequantize_per_channel.default, - exir_ops.edge.quantized_decomposed.quantize_per_channel.default, + DequantizePerChannel, + QuantizePerChannel, ) qdq_per_tensor_nodes = ( - exir_ops.edge.quantized_decomposed.quantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.quantize_per_tensor.tensor, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, - exir_ops.edge.quantized_decomposed.dequantize_per_tensor.tensor, + QuantizePerTensor, + QuantizePerTensorTensor, + DequantizePerTensor, + DequantizePerTensorTensor, ) def run(self, graph_module: torch.fx.GraphModule) -> PassResult: @@ -55,8 +64,8 @@ def run(self, graph_module: torch.fx.GraphModule) -> PassResult: for cluster in qdq_clusterer.cluster_map.values(): # For now, enable only permute_copy and cat. if cluster.compute_node.target not in [ - exir_ops.edge.aten.permute_copy.default, - exir_ops.edge.aten.cat.default, + PermuteCopy, + Cat, ]: continue diff --git a/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py b/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py index 1ed9a9f607c..d2257a453ac 100644 --- a/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py +++ b/backends/nxp/edge_passes/remove_as_strided_copy_nodes.py @@ -4,8 +4,8 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +from executorch.backends.nxp.backend.ops_aliases import AsStridedCopy, MeanDim from executorch.backends.nxp.edge_passes.neutron_edge_pass import NeutronEdgePass -from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.passes import dead_code_elimination_pass from torch.fx import GraphModule from torch.fx.passes.infra.pass_base import PassResult @@ -18,12 +18,12 @@ def __init__(self): def gen_pattern_as_strided_copy(self, graph_module: GraphModule): # Unedited method taken from `backends/samsung/_passes/remove_useless_ops.py`. for node in list(graph_module.graph.nodes): # noqa: C416 - if node.target != exir_ops.edge.aten.mean.dim: + if node.target != MeanDim: continue if len(node.users) != 1: continue successor = list(node.users.keys())[0] - if successor.target != exir_ops.edge.aten.as_strided_copy.default: + if successor.target != AsStridedCopy: continue is_pattern = True count = 0 diff --git a/backends/nxp/edge_passes/remove_io_quant_ops_pass.py b/backends/nxp/edge_passes/remove_io_quant_ops_pass.py index a87eac7360c..bb2ae14c397 100644 --- a/backends/nxp/edge_passes/remove_io_quant_ops_pass.py +++ b/backends/nxp/edge_passes/remove_io_quant_ops_pass.py @@ -5,8 +5,11 @@ import torch +from executorch.backends.nxp.backend.ops_aliases import ( + DequantizePerTensor, + QuantizePerTensor, +) from executorch.exir import EdgeProgramManager -from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass from executorch.exir.passes.quantize_io_pass import QuantizeInputs, QuantizeOutputs from torch.fx.passes.infra.pass_base import PassResult @@ -37,10 +40,7 @@ def _get_quantizable_input_indices(self): raise ValueError(f"Input {input_index} has more than one users") quantize = next(iter(target_placeholder.users)) - if ( - quantize.target - != exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - ): + if quantize.target != QuantizePerTensor: continue inputs_to_quantization.append(input_index) @@ -59,10 +59,7 @@ def _get_quantizable_output_indices(self): user_outputs = list(outputs[0].args[0]) for output_index, user_output in enumerate(user_outputs): - if ( - user_output.target - != exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ): + if user_output.target != DequantizePerTensor: continue outputs_to_quantization.append(output_index) diff --git a/backends/nxp/quantizer/neutron_quantizer.py b/backends/nxp/quantizer/neutron_quantizer.py index 1f1d26ce229..7f186b5f769 100644 --- a/backends/nxp/quantizer/neutron_quantizer.py +++ b/backends/nxp/quantizer/neutron_quantizer.py @@ -13,7 +13,8 @@ from executorch.backends.nxp.quantizer.patterns import ( AbsPattern, ActivationsConcatClusterPattern, - AdaptiveAvgPoolPattern, + AdaptiveAvgPool1DPattern, + AdaptiveAvgPool2DPattern, AddmmPattern, AddTensorPattern, AmaxPattern, @@ -267,7 +268,8 @@ def __init__(self, neutron_target_spec: NeutronTargetSpec, is_qat: bool = False) super().__init__( [ OpQuantizer(AbsPattern(is_qat=is_qat), static_qconfig), - OpQuantizer(AdaptiveAvgPoolPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(AdaptiveAvgPool1DPattern(is_qat=is_qat), static_qconfig), + OpQuantizer(AdaptiveAvgPool2DPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddTensorPattern(is_qat=is_qat), static_qconfig), OpQuantizer(AddmmPattern(self, is_qat=is_qat), static_fc_qconfig), OpQuantizer(AmaxPattern(is_qat=is_qat), static_qconfig), diff --git a/backends/nxp/quantizer/patterns.py b/backends/nxp/quantizer/patterns.py index 3bbd2bed54e..4ac4b777cba 100644 --- a/backends/nxp/quantizer/patterns.py +++ b/backends/nxp/quantizer/patterns.py @@ -279,7 +279,16 @@ def partition_types(self): return [torch.ops.aten.abs.default] -class AdaptiveAvgPoolPattern(SharedSpecPattern): +class AdaptiveAvgPool1DPattern(SharedSpecPattern): + """ + Quantizer for AdaptiveAvgPool1D operator. + """ + + def partition_types(self): + return [torch.ops.aten.adaptive_avg_pool1d.default] + + +class AdaptiveAvgPool2DPattern(SharedSpecPattern): """ Quantizer for AdaptiveAvgPool2D operator. """ diff --git a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py index 679e7365577..a3eae135144 100644 --- a/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py +++ b/backends/nxp/tests/generic_tests/test_add_batch_size_for_3d_input_pool_2d_ops.py @@ -29,12 +29,12 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( AdaptiveAvgPool2dModule, AvgPool2dModule, MaxPool2dModule, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_aot_example.py b/backends/nxp/tests/generic_tests/test_aot_example.py index 1f8dc410917..63e56ddd4bf 100644 --- a/backends/nxp/tests/generic_tests/test_aot_example.py +++ b/backends/nxp/tests/generic_tests/test_aot_example.py @@ -2,6 +2,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import contextlib import os import subprocess diff --git a/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py b/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py index 64cd7ce83f0..85c02e0d53d 100644 --- a/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py +++ b/backends/nxp/tests/generic_tests/test_batch_norm_fusion.py @@ -6,6 +6,8 @@ from copy import deepcopy import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( @@ -27,7 +29,7 @@ graph_contains_any_of_ops, OverrideTargetSupportCheck, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( ConvBatchNormModule, LinearBatchNormModule, ) diff --git a/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py index 677312f9483..0f931075927 100644 --- a/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py +++ b/backends/nxp/tests/generic_tests/test_context_sensitive_delegation.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -13,12 +15,15 @@ from executorch.backends.nxp.backend.ir.converter.node_converters.ops_converters import ( ViewCopyConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddTensor, + Cat, + ExecutorchDelegateCall, + SubTensor, + ViewCopy, +) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.exir.dialects._ops import ops as exir_ops - -# noinspection PyProtectedMember -ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate class SingleViewCopyModule(torch.nn.Module): @@ -70,7 +75,7 @@ def test_single_view_copy_partition(): ep = to_quantized_edge_program(module, input_shape).exported_program() # Make sure the `view_copy` was not delegated. - assert graph_contains_any_of_ops(ep.graph, [exir_ops.edge.aten.view_copy.default]) + assert graph_contains_any_of_ops(ep.graph, [ViewCopy]) assert not graph_contains_any_of_ops(ep.graph, [ExecutorchDelegateCall]) @@ -114,8 +119,8 @@ def test_noop_partitions__concatenate_one_tensor_and_add_zeros(): assert graph_contains_any_of_ops( ep.graph, [ - exir_ops.edge.aten.cat.default, - exir_ops.edge.aten.add.Tensor, + Cat, + AddTensor, ], ) @@ -158,8 +163,8 @@ def test_noop_partitions__add_sub(): assert graph_contains_any_of_ops( ep.graph, [ - exir_ops.edge.aten.add.Tensor, - exir_ops.edge.aten.sub.Tensor, + AddTensor, + SubTensor, ], ) diff --git a/backends/nxp/tests/test_convert_1d_conv_to_2d.py b/backends/nxp/tests/generic_tests/test_convert_1d_conv_to_2d.py similarity index 98% rename from backends/nxp/tests/test_convert_1d_conv_to_2d.py rename to backends/nxp/tests/generic_tests/test_convert_1d_conv_to_2d.py index b7db6fbc46e..d4a07d86fab 100644 --- a/backends/nxp/tests/test_convert_1d_conv_to_2d.py +++ b/backends/nxp/tests/generic_tests/test_convert_1d_conv_to_2d.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.aten_passes.neutron_aten_pass_manager import ( @@ -14,6 +16,10 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + ExecutorchDelegateCall, +) from executorch.backends.nxp.tests.executorch_pipeline import ( neutron_target_spec, to_quantized_edge_program, @@ -22,8 +28,10 @@ convert_run_compare, graph_contains_any_of_ops, ) -from executorch.backends.nxp.tests.models import Conv1dModule, ConvTranspose1dModule -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.simple_models import ( + Conv1dModule, + ConvTranspose1dModule, +) from torch import nn from torch.export import ExportedProgram @@ -46,9 +54,6 @@ def reseed_model_per_test_run(): AtenHardtanh = torch.ops.aten.hardtanh.default AtenBatchNorm = torch.ops.aten.batch_norm.default -EdgeConvolution = exir_ops.edge.aten.convolution.default -ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate - @pytest.mark.parametrize( "input_shape, kernel_size, stride, padding, dilation, groups, bias", @@ -301,7 +306,7 @@ def test_convert_conv_1d_to_conv2d_full_pipeline( # Make sure `edge.aten.convolution.default` is in the model. assert graph_contains_any_of_ops( exported_program.graph, - [EdgeConvolution], + [Convolution], ) example_input = (np.random.random(input_shape).astype(np.float32) * 50).astype( @@ -378,7 +383,7 @@ def test_convert_conv_1d_to_conv2d_transp_full_pipeline( # Make sure `edge.aten.convolution.default` is in the model. assert graph_contains_any_of_ops( exported_program.graph, - [EdgeConvolution], + [Convolution], ) example_input = (np.random.random(input_shape).astype(np.float32) * 50).astype( diff --git a/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py index 044467ca223..89879226c3a 100644 --- a/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py +++ b/backends/nxp/tests/generic_tests/test_convert_div_to_mul.py @@ -18,11 +18,11 @@ from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( NonstaticDivLinearModel, StaticDivLinearModel, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py index c68cde0c23b..be214b77568 100644 --- a/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py +++ b/backends/nxp/tests/generic_tests/test_convert_scalar_to_attr.py @@ -23,15 +23,15 @@ from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( - AddScalarModule, - MulScalarModule, - SubScalarModule, -) from executorch.backends.nxp.tests.nsys_testing import ( AllCloseOutputComparator, lower_run_compare, ) +from executorch.backends.nxp.tests.simple_models import ( + AddScalarModule, + MulScalarModule, + SubScalarModule, +) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_debug_results.py b/backends/nxp/tests/generic_tests/test_debug_results.py index ccf5c13b501..387acd69b2c 100644 --- a/backends/nxp/tests/generic_tests/test_debug_results.py +++ b/backends/nxp/tests/generic_tests/test_debug_results.py @@ -7,17 +7,19 @@ import os import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier -from executorch.backends.nxp.tests.models import AddTensorModule, AvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import ( get_test_name, lower_run_compare, OUTPUTS_DIR, ) +from executorch.backends.nxp.tests.simple_models import AddTensorModule, AvgPool2dModule @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py index 3c3c29f95b5..cddf2feff34 100644 --- a/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py +++ b/backends/nxp/tests/generic_tests/test_decompose_split_to_slices.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -16,12 +18,12 @@ from executorch.backends.nxp.tests.executorch_pipeline import neutron_target_spec from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( GRUModel, SplitWithSections, SplitWithSize, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/test_fold_redundant_qdq.py b/backends/nxp/tests/generic_tests/test_fold_redundant_qdq.py similarity index 97% rename from backends/nxp/tests/test_fold_redundant_qdq.py rename to backends/nxp/tests/generic_tests/test_fold_redundant_qdq.py index e79d29b1166..7da9986ad7b 100644 --- a/backends/nxp/tests/test_fold_redundant_qdq.py +++ b/backends/nxp/tests/generic_tests/test_fold_redundant_qdq.py @@ -6,13 +6,12 @@ import torch +from executorch.backends.nxp.backend.ops_aliases import ExecutorchDelegateCall from executorch.backends.nxp.edge_passes.fold_redundant_qdq_pass import ( FoldRedundantDequantizeQuantizePass, ) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -ExecutorchDelegateCall = torch.ops.higher_order.executorch_call_delegate - class ConvDropoutConvModule(torch.nn.Module): """Two conv clusters separated by an eval-mode dropout (an identity). diff --git a/backends/nxp/tests/test_fuse_batch_norm_single_user.py b/backends/nxp/tests/generic_tests/test_fuse_batch_norm_single_user.py similarity index 100% rename from backends/nxp/tests/test_fuse_batch_norm_single_user.py rename to backends/nxp/tests/generic_tests/test_fuse_batch_norm_single_user.py diff --git a/backends/nxp/tests/generic_tests/test_gru_splitting.py b/backends/nxp/tests/generic_tests/test_gru_splitting.py index 297f9677fb2..03f3f4a2947 100644 --- a/backends/nxp/tests/generic_tests/test_gru_splitting.py +++ b/backends/nxp/tests/generic_tests/test_gru_splitting.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch diff --git a/backends/nxp/tests/generic_tests/test_integration.py b/backends/nxp/tests/generic_tests/test_integration.py index f0cb25548ec..869aed432f8 100644 --- a/backends/nxp/tests/generic_tests/test_integration.py +++ b/backends/nxp/tests/generic_tests/test_integration.py @@ -12,7 +12,7 @@ from executorch.backends.nxp.tests.executorch_pipeline import ( to_quantized_executorch_program, ) -from executorch.backends.nxp.tests.models import ConvFCSoftmaxModule +from executorch.backends.nxp.tests.simple_models import ConvFCSoftmaxModule from executorch.devtools.backend_debug import get_delegation_info from executorch.examples.nxp.experimental.cifar_net.cifar_net import CifarNet diff --git a/backends/nxp/tests/generic_tests/test_kernel_selection.py b/backends/nxp/tests/generic_tests/test_kernel_selection.py index 732913c0485..2dd498881c7 100644 --- a/backends/nxp/tests/generic_tests/test_kernel_selection.py +++ b/backends/nxp/tests/generic_tests/test_kernel_selection.py @@ -7,11 +7,13 @@ import eiq_neutron_sdk import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( AdaptiveAvgPool2dConvModule, Conv2dReLUMaxPoolModule, ) diff --git a/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py b/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py index 6aa07dbba8d..64a7a5e162d 100644 --- a/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py +++ b/backends/nxp/tests/generic_tests/test_move_activation_before_concatenation.py @@ -19,6 +19,15 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Cat, + Convolution, + HardTanh, + Relu, + Sigmoid, + Tanh, +) from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.quantizer.utils import calibrate_and_quantize from executorch.backends.nxp.tests.executorch_pipeline import ( @@ -33,21 +42,20 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.backends.nxp.tests.models import get_activation -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.simple_models import get_activation from parameterized import parameterized from torch import nn from torch.export import ExportedProgram from torch.fx import GraphModule concat_cluster_ops = [ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.convolution.default, - exir_ops.edge.aten.hardtanh.default, - exir_ops.edge.aten.relu.default, - exir_ops.edge.aten.sigmoid.default, - exir_ops.edge.aten.tanh.default, - exir_ops.edge.aten.cat.default, + AddMM, + Convolution, + HardTanh, + Relu, + Sigmoid, + Tanh, + Cat, ] diff --git a/backends/nxp/tests/generic_tests/test_neutron_backend.py b/backends/nxp/tests/generic_tests/test_neutron_backend.py index 867b585ef64..f2eaf097b14 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_backend.py +++ b/backends/nxp/tests/generic_tests/test_neutron_backend.py @@ -4,7 +4,10 @@ # LICENSE file in the root directory of this source tree. from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dModule, LinearSoftmaxModule +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + LinearSoftmaxModule, +) def test_neutron_backend__single_conv_model(): diff --git a/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py b/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py index 06a95142b1a..6dbbf177b7d 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py +++ b/backends/nxp/tests/generic_tests/test_neutron_backend_executor.py @@ -21,7 +21,10 @@ TFLiteExecutor, ToNHWCPreprocess, ) -from executorch.backends.nxp.tests.models import Conv2dModule, ConvFCSoftmaxModule +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + ConvFCSoftmaxModule, +) from torch.export import ExportedProgram diff --git a/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py b/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py index fab2c8afc06..2b3b798dd67 100644 --- a/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py +++ b/backends/nxp/tests/generic_tests/test_neutron_compiler_manager.py @@ -10,7 +10,7 @@ NeutronCompilerManager, ) from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import LinearModule +from executorch.backends.nxp.tests.simple_models import LinearModule def test_conv2d_neutron_conversion__prefetching(mocker): diff --git a/backends/nxp/tests/generic_tests/test_node_format_inference.py b/backends/nxp/tests/generic_tests/test_node_format_inference.py index 5206a029e01..8d6ac52627d 100644 --- a/backends/nxp/tests/generic_tests/test_node_format_inference.py +++ b/backends/nxp/tests/generic_tests/test_node_format_inference.py @@ -20,7 +20,7 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( Conv2dModule, MaxPool2dModule, SoftmaxModule, diff --git a/backends/nxp/tests/generic_tests/test_operator_selector.py b/backends/nxp/tests/generic_tests/test_operator_selector.py index ca301daf738..71c30c73fac 100644 --- a/backends/nxp/tests/generic_tests/test_operator_selector.py +++ b/backends/nxp/tests/generic_tests/test_operator_selector.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree.import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dModule +from executorch.backends.nxp.tests.simple_models import Conv2dModule def test_operator_selector_mechanism(): diff --git a/backends/nxp/tests/generic_tests/test_per_channel_conversion.py b/backends/nxp/tests/generic_tests/test_per_channel_conversion.py index af9ef08057b..25a4b28224b 100644 --- a/backends/nxp/tests/generic_tests/test_per_channel_conversion.py +++ b/backends/nxp/tests/generic_tests/test_per_channel_conversion.py @@ -12,6 +12,10 @@ from executorch.backends.nxp.backend.edge_program_converter import ( EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + Convolution, + DequantizePerChannel, +) from executorch.backends.nxp.quantizer.neutron_quantizer import ( act_qspec, NeutronAtenQuantizer, @@ -29,8 +33,7 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.backends.nxp.tests.models import Conv2dModule -from executorch.exir.dialects._ops import ops as exir_ops +from executorch.backends.nxp.tests.simple_models import Conv2dModule from parameterized import parameterized from torch import fx @@ -172,16 +175,10 @@ def test_per_channel_convolution(self, _, use_qat: bool): conv_nodes = [ node for node in exported_program.graph.nodes - if node.target == exir_ops.edge.aten.convolution.default + if node.target == Convolution ] assert len(conv_nodes) == 1 conv_node = conv_nodes[0] - assert ( - conv_node.args[1].target - == exir_ops.edge.quantized_decomposed.dequantize_per_channel.default - ) - assert ( - conv_node.args[2].target - == exir_ops.edge.quantized_decomposed.dequantize_per_channel.default - ) + assert conv_node.args[1].target == DequantizePerChannel + assert conv_node.args[2].target == DequantizePerChannel diff --git a/backends/nxp/tests/generic_tests/test_profiling.py b/backends/nxp/tests/generic_tests/test_profiling.py index 7e147dc989c..09ae6c079ed 100644 --- a/backends/nxp/tests/generic_tests/test_profiling.py +++ b/backends/nxp/tests/generic_tests/test_profiling.py @@ -2,6 +2,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import ast import logging import os @@ -9,13 +10,14 @@ from typing import Any, Union import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier from executorch.backends.nxp.tests.model_output_comparator import ( NumericalStatsOutputComparator, ) -from executorch.backends.nxp.tests.models import AvgPool2dModule, SoftmaxModule from executorch.backends.nxp.tests.nsys_testing import ( get_test_name, lower_run_compare, @@ -27,6 +29,7 @@ get_neutron_driver_version, get_neutron_kernel_kinds, ) +from executorch.backends.nxp.tests.simple_models import AvgPool2dModule, SoftmaxModule from executorch.devtools.inspector._inspector import Inspector from executorch.examples.models.mlperf_tiny import ( DeepAutoEncoder, diff --git a/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py b/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py index 6db55347452..b0a25c68dd3 100644 --- a/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py +++ b/backends/nxp/tests/generic_tests/test_qdq_clustering_conv.py @@ -4,7 +4,7 @@ # LICENSE file in the root directory of this source tree. from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dModule +from executorch.backends.nxp.tests.simple_models import Conv2dModule def test_conv2d_partitioner(): diff --git a/backends/nxp/tests/generic_tests/test_quantized_input_data.py b/backends/nxp/tests/generic_tests/test_quantized_input_data.py index bd2a6056a0b..23a39ac11de 100644 --- a/backends/nxp/tests/generic_tests/test_quantized_input_data.py +++ b/backends/nxp/tests/generic_tests/test_quantized_input_data.py @@ -9,12 +9,12 @@ from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import AvgPool2dModule, MulTensorModule from executorch.backends.nxp.tests.nsys_testing import ( lower_run_compare, OUTPUTS_DIR, ReferenceModel, ) +from executorch.backends.nxp.tests.simple_models import AvgPool2dModule, MulTensorModule def test__single_quantized_inputs(mocker, request): diff --git a/backends/nxp/tests/generic_tests/test_quantizer.py b/backends/nxp/tests/generic_tests/test_quantizer.py index 6180d2fd9ae..886301fb149 100644 --- a/backends/nxp/tests/generic_tests/test_quantizer.py +++ b/backends/nxp/tests/generic_tests/test_quantizer.py @@ -9,8 +9,10 @@ from copy import deepcopy import executorch.backends.nxp.tests.executorch_pipeline as executorch_pipeline -import executorch.backends.nxp.tests.models as models +import executorch.backends.nxp.tests.simple_models as models import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -18,6 +20,18 @@ EdgeProgramToIRConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + AddMM, + Convolution, + HardTanh, + MM, + NativebatchNormLegitNoStats, + NativebatchNormLegitNoTraining, + Relu, + Sigmoid, + Tanh, +) + from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.tests.executorch_pipeline import ( neutron_target_spec, @@ -31,8 +45,6 @@ ToChannelLastPreprocess, ) -from executorch.exir.dialects._ops import ops as exir_ops - requires_tflite = pytest.mark.skipif( tflite is None, reason="tensorflow/tflite not available" ) @@ -50,13 +62,13 @@ ) fuse_activation_ops = [ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.mm.default, - exir_ops.edge.aten.convolution.default, - exir_ops.edge.aten.hardtanh.default, - exir_ops.edge.aten.relu.default, - exir_ops.edge.aten.sigmoid.default, - exir_ops.edge.aten.tanh.default, + AddMM, + MM, + Convolution, + HardTanh, + Relu, + Sigmoid, + Tanh, ] @@ -74,8 +86,8 @@ ] batch_norm_ops = ( - exir_ops.edge.aten._native_batch_norm_legit.no_stats, - exir_ops.edge.aten._native_batch_norm_legit_no_training.default, + NativebatchNormLegitNoStats, + NativebatchNormLegitNoTraining, torch.ops.aten._native_batch_norm_legit_no_training.default, torch.ops.aten.batch_norm.default, torch.ops.aten.native_batch_norm.default, diff --git a/backends/nxp/tests/generic_tests/test_recipe_export.py b/backends/nxp/tests/generic_tests/test_recipe_export.py index 2403f2e62a7..b853a5aa560 100644 --- a/backends/nxp/tests/generic_tests/test_recipe_export.py +++ b/backends/nxp/tests/generic_tests/test_recipe_export.py @@ -35,7 +35,7 @@ graph_contains_any, graph_contains_any_of_ops, ) -from executorch.backends.nxp.tests.models import ConvBatchNormModule +from executorch.backends.nxp.tests.simple_models import ConvBatchNormModule from executorch.backends.transforms.quantize_fused_convbn_bias_pass import ( QuantizeFusedConvBnBiasAtenPass, ) diff --git a/backends/nxp/tests/generic_tests/test_removing_dead_code.py b/backends/nxp/tests/generic_tests/test_removing_dead_code.py index 8b3a979f412..7c01be7832d 100644 --- a/backends/nxp/tests/generic_tests/test_removing_dead_code.py +++ b/backends/nxp/tests/generic_tests/test_removing_dead_code.py @@ -6,6 +6,8 @@ import unittest import numpy as np + +# noinspection PyUnusedImports import pytest import torch diff --git a/backends/nxp/tests/generic_tests/test_split_group_convolution.py b/backends/nxp/tests/generic_tests/test_split_group_convolution.py index 12d2f193f57..a6bd3fafdc9 100644 --- a/backends/nxp/tests/generic_tests/test_split_group_convolution.py +++ b/backends/nxp/tests/generic_tests/test_split_group_convolution.py @@ -15,6 +15,7 @@ from executorch.backends.nxp.aten_passes.split_group_convolution import ( SplitGroupConvolution, ) +from executorch.backends.nxp.backend.ops_aliases import Cat, Convolution from executorch.backends.nxp.neutron_partitioner import NeutronPartitioner from executorch.backends.nxp.nxp_backend import generate_neutron_compile_spec from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer @@ -26,13 +27,12 @@ to_quantized_edge_program, ) from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( Conv1dModule, Conv2dModule, Conv3dModule, ) from executorch.exir import EdgeCompileConfig, EdgeProgramManager -from executorch.exir.dialects._ops import ops as exir_ops from executorch.extension.export_util import export_to_edge from parameterized import parameterized from torch.fx import GraphModule @@ -129,7 +129,7 @@ def test_split_group_convolution__2d( assert nodes[-5].name == "lowered_module_0" assert not graph_contains_any_of_ops( ep.graph, - [exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.cat.default], + [Convolution, Cat], ) @parameterized.expand( @@ -206,7 +206,7 @@ def test_split_group_convolution__1d( assert nodes[-5].name == "lowered_module_0" assert not graph_contains_any_of_ops( ep.graph, - [exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.cat.default], + [Convolution, Cat], ) @parameterized.expand( @@ -303,5 +303,5 @@ def test_split_group_convolution__applied_by_default(self, _, is_qat: bool): assert nodes[-5].name == "lowered_module_0" assert not graph_contains_any_of_ops( ep.graph, - [exir_ops.edge.aten.convolution.default, exir_ops.edge.aten.cat.default], + [Convolution, Cat], ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py index 65e2b82555c..75093aba530 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_adaptive_avg_pool2d_converter.py @@ -11,6 +11,7 @@ from executorch.backends.nxp.backend.ops_aliases import ( AdaptiveAvgPool2D, ExecutorchDelegateCall, + ViewCopy, ) from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator @@ -20,8 +21,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import AdaptiveAvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + AdaptiveAvgPool1dModule, + AdaptiveAvgPool2dModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -42,6 +46,11 @@ class TestAdaptiveAvgPool2D: (2, 3), id="H != W, non multiples of num_macs, batch != 1.", ), + pytest.param( + (2, 3, 10, 15), + (5, 5), + id="H != W, non multiples of num_macs, batch != 1, fixed fail.", + ), ], ) def test__basic_nsys_inference( @@ -54,9 +63,8 @@ def test__basic_nsys_inference( expected_non_delegated_ops={}, ) - output_comparator = AllCloseOutputComparator( - 7.84e-3 - ) # Accept small error due to Neutron bug (AIR-14585). + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. lower_run_compare( model, @@ -66,32 +74,7 @@ def test__basic_nsys_inference( RandomDatasetCreator(low=-1, high=1), output_comparator=output_comparator, use_qat=use_qat, - ) - - @pytest.mark.xfail( - strict=True, - reason="Known Neutron bad compute issue. Will be fixed in Neutron SW 3.1.2.", - ) - def test__know_neutron_issue(self, mocker, request): - input_shape = (2, 3, 10, 15) - output_size = (5, 5) - model = AdaptiveAvgPool2dModule(output_size) - graph_verifier = DetailedGraphVerifier( - mocker, - expected_delegated_ops={AdaptiveAvgPool2D: 1}, - expected_non_delegated_ops={}, - ) - - # Use high tolerance so we notice when the issue is fixed. - output_comparator = AllCloseOutputComparator(7.8e-3) - - lower_run_compare( - model, - input_shape, - graph_verifier, - request, - RandomDatasetCreator(low=-1, high=1), - output_comparator=output_comparator, + remove_quant_io_ops=remove_quant_io_ops, ) def test__kernel_size_and_stride_limit(self, mocker, request): @@ -110,9 +93,8 @@ def test__kernel_size_and_stride_limit(self, mocker, request): expected_non_delegated_ops={}, ) - output_comparator = AllCloseOutputComparator( - 7.9e-3 - ) # Accept small error due to Neutron bug (AIR-14585). + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. lower_run_compare( model, @@ -121,6 +103,7 @@ def test__kernel_size_and_stride_limit(self, mocker, request): request, RandomDatasetCreator(low=-1, high=1), output_comparator=output_comparator, + remove_quant_io_ops=remove_quant_io_ops, ) def test__kernel_size_and_stride_limit_exceeded(self): @@ -140,3 +123,31 @@ def test__kernel_size_and_stride_limit_exceeded(self): delegated_ep.graph, [ExecutorchDelegateCall] ) assert graph_contains_any_of_ops(delegated_ep.graph, [AdaptiveAvgPool2D]) + + +class TestAdaptiveAvgPool1DTo2D: + + # Just a basic test to verify that the operator gets extended to the 2D variant correctly. + def test__basic_nsys_inference(self, mocker, request, use_qat): + input_shape = (2, 4, 6) # The old flow limited the batch size to 1. + output_size = (3,) + model = AdaptiveAvgPool1dModule(output_size) + graph_verifier = DetailedGraphVerifier( + mocker, + expected_delegated_ops={AdaptiveAvgPool2D: 1, ViewCopy: 2}, + expected_non_delegated_ops={}, + ) + + remove_quant_io_ops = True # Use quantized dataset. + output_comparator = AllCloseOutputComparator(atol=1) # Allow single bit error. + + lower_run_compare( + model, + input_shape, + graph_verifier, + request, + RandomDatasetCreator(low=-1, high=1), + output_comparator=output_comparator, + use_qat=use_qat, + remove_quant_io_ops=remove_quant_io_ops, + ) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py index 237d15aefbd..5c3b0dd0100 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_add_tensor_converter.py @@ -25,8 +25,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import AddTensorModule, MaxPoolAddTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + AddTensorModule, + MaxPoolAddTensorModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py index 6ebb8ac128c..713a618ea53 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_addmm_converter.py @@ -20,8 +20,8 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator -from executorch.backends.nxp.tests.models import AddmmModule, LinearModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import AddmmModule, LinearModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py index c132c9e509d..bd28e8df6e6 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_avg_pool2d_converter.py @@ -17,8 +17,8 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import AvgPool2dModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import AvgPool2dModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 @@ -112,10 +112,10 @@ def test__stride_limit_exceeded(self): assert graph_contains_any_of_ops(delegated_ep.graph, [AvgPool2D]) -class TestAvgPool1D: +class TestAvgPool1DTo2D: # Just a basic test to verify that the operator gets extended to the 2D variant correctly. - def test__basic_nsys_inference(self, mocker, request): + def test__basic_nsys_inference(self, mocker, request, use_qat): input_shape = (2, 4, 6) # The old flow limited the batch size to 1. model = AvgPool1DModule() graph_verifier = DetailedGraphVerifier( @@ -124,4 +124,4 @@ def test__basic_nsys_inference(self, mocker, request): expected_non_delegated_ops={}, ) - lower_run_compare(model, input_shape, graph_verifier, request) + lower_run_compare(model, input_shape, graph_verifier, request, use_qat=use_qat) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py index ee46ad94d44..5bb7a080b23 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_bmm_converter.py @@ -21,11 +21,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( BatchMatMulMaxPoolModel, BatchMatMulModel, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py index 1238e31e246..aaf6a6e35d3 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_clone_converter.py @@ -2,6 +2,7 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. + import itertools import unittest @@ -16,6 +17,11 @@ PermuteCopyConverter, ) from executorch.backends.nxp.backend.node_format_inference import NodeFormatInference +from executorch.backends.nxp.backend.ops_aliases import ( + Clone, + CloneDimOrder, + PermuteCopy, +) from executorch.backends.nxp.edge_passes.move_auxiliary_operator_into_separate_qdq_cluster_pass import ( MoveLeadingAuxiliaryOperatorIntoSeparateQDQClusterPass, ) @@ -45,7 +51,6 @@ ToChannelLastPreprocess, ) from executorch.exir import EdgeCompileConfig -from executorch.exir.dialects._ops import ops as exir_ops from executorch.extension.export_util.utils import export_to_edge from parameterized import parameterized from torch import nn @@ -144,8 +149,8 @@ def setUpClass(cls): @staticmethod def _node_is_clone(node) -> bool: clone_ops = [ - exir_ops.edge.aten.clone.default, - exir_ops.edge.dim_order_ops._clone_dim_order.default, + Clone, + CloneDimOrder, ] def target_can_be_clone(node): @@ -215,8 +220,8 @@ def test_conv_dropout_no_quant( has_clone = graph_contains_any_of_ops( graph=edge_program.graph, ops=[ - exir_ops.edge.aten.clone.default, - exir_ops.edge.dim_order_ops._clone_dim_order.default, + Clone, + CloneDimOrder, ], ) @@ -287,7 +292,7 @@ def test_clone__to_contiguous_format(self): ) # Make sure the `aten.clone` was inserted as expected. nodes = list(edge_program_manager.exported_program().graph.nodes) - assert nodes[9].target == exir_ops.edge.dim_order_ops._clone_dim_order.default + assert nodes[9].target == CloneDimOrder assert nodes[9].kwargs["dim_order"] == [0, 1, 2, 3] # Move the `clone` out of the cluster with the `view_copy`. @@ -340,9 +345,7 @@ def _unsupported_target(*_): ep = to_quantized_edge_program(model, input_shape).exported_program() nodes = list(ep.graph.nodes) - assert not graph_contains_any_of_ops( - ep.graph, [exir_ops.edge.aten.clone.default] - ) + assert not graph_contains_any_of_ops(ep.graph, [Clone]) assert nodes[3].name == "executorch_call_delegate" - assert nodes[5].target == exir_ops.edge.aten.permute_copy.default + assert nodes[5].target == PermuteCopy assert nodes[7].name == "executorch_call_delegate_1" diff --git a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py index 730334932d7..b4de3c4984f 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_constant_pad_nd_converter.py @@ -14,8 +14,8 @@ ) from executorch.backends.nxp.backend.ops_aliases import ConstantPadND, Convolution from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import PadConvModule, PadModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import PadConvModule, PadModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py index f0fb0fdb7b4..b802309d8e0 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_conv_converter.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.backend.ops_aliases import ( @@ -15,12 +17,15 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dModule, Conv2dTransposedModule from executorch.backends.nxp.tests.nsys_testing import ( AllCloseOutputComparator, lower_run_compare, ReferenceModel, ) +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + Conv2dTransposedModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py index 76aaab7520e..e0f626a2140 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_hardswish_converter.py @@ -21,12 +21,12 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( ConvHardswishModule, HardswishModule, LinearHardswishModule, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py index 1199609bb17..dc71f5df921 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_hardtanh_converter.py @@ -27,8 +27,11 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dWithActivation, HardTanhModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + Conv2dWithActivation, + HardTanhModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py index 7d3769c896d..cfc4b46022a 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_max_pool_2d_converter.py @@ -8,13 +8,13 @@ # noinspection PyUnusedImports import pytest import torch + from executorch.backends.nxp.backend.ops_aliases import ( ExecutorchDelegateCall, GetItem, MaxPool2DWithIndices, ViewCopy, ) - from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier @@ -158,10 +158,10 @@ def test__padding_to_kernel_ratio_exceeded(self): to_quantized_edge_program(model, input_shape) -class TestMaxPool1D: +class TestMaxPool1DTo2D: # Just a basic test to verify that the operator gets extended to the 2D variant correctly. - def test__basic_nsys_inference__view_not_delegated(self, mocker, request): + def test__basic_nsys_inference__view_not_delegated(self, mocker, request, use_qat): input_shape = (2, 4, 6) # The old flow limited the batch size to 1. model = MaxPool1DModule() @@ -171,4 +171,4 @@ def test__basic_nsys_inference__view_not_delegated(self, mocker, request): expected_non_delegated_ops={}, ) - lower_run_compare(model, input_shape, graph_verifier, request) + lower_run_compare(model, input_shape, graph_verifier, request, use_qat=use_qat) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py index 349a624242a..01df6211dfc 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_maximum_converter.py @@ -25,8 +25,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaximumModule, MaxPoolMaximumModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + MaximumModule, + MaxPoolMaximumModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py index 7c74fd4a1c9..727ff914ceb 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_minimum_converter.py @@ -25,8 +25,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaxPoolMinimumModule, MinimumModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + MaxPoolMinimumModule, + MinimumModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py index 405e955a5ca..a8997c11ef5 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mm_converter.py @@ -12,8 +12,8 @@ from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier, Operator -from executorch.backends.nxp.tests.models import LinearModule, MmModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import LinearModule, MmModule from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py index b031ab7e47c..414cbecad85 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_mul_tensor_converter.py @@ -25,8 +25,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaxPoolMulTensorModule, MulTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + MaxPoolMulTensorModule, + MulTensorModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py index f5c0114c81e..d7bbe0f21c6 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_pad_converter.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch @@ -12,8 +14,8 @@ ) from executorch.backends.nxp.backend.ops_aliases import Convolution, Pad from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import PadConvModule, PadModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import PadConvModule, PadModule @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py index 358320c30c7..58a820f1a7f 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_prelu_converter.py @@ -28,14 +28,14 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( + +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( ConvPReLUModule, LinearPReLUModule, PReLUModule, TwoPartitionPReLUModel, ) - -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare from torch.export import ExportedProgram from executorch.backends.nxp.tests.use_qat import * # noqa F403 from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program diff --git a/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py index 4bf9263e93f..d44fc1480db 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_relu_converter.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.backend.edge_program_converter import exir_ops @@ -21,8 +23,12 @@ from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program from executorch.backends.nxp.tests.executors import graph_contains_any_of_ops from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dModule, LinearModule, ReLUModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + Conv2dModule, + LinearModule, + ReLUModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py index 4a66c000505..00b89c859ef 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_slice_copy_tensor_converter.py @@ -21,11 +21,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( SliceTensorConvModule, SliceTensorModule, ) -from executorch.backends.nxp.tests.nsys_testing import lower_run_compare @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py index 8cec44b5274..f893cc1e394 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_softmax_converter.py @@ -4,6 +4,8 @@ # LICENSE file in the root directory of this source tree. import numpy as np + +# noinspection PyUnusedImports import pytest import torch from executorch.backends.nxp.backend.ops_aliases import ( @@ -21,8 +23,8 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import SoftmaxModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import SoftmaxModule @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py index 6e88bc270e4..54d0b3c6cc2 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_sub_tensor_converter.py @@ -25,8 +25,11 @@ from executorch.backends.nxp.tests.model_output_comparator import ( AllCloseOutputComparator, ) -from executorch.backends.nxp.tests.models import MaxPoolSubTensorModule, SubTensorModule from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import ( + MaxPoolSubTensorModule, + SubTensorModule, +) from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py index 2795f607494..71b6ced4ed2 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_tanh_converter.py @@ -11,8 +11,8 @@ from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import Conv2dWithActivation from executorch.backends.nxp.tests.nsys_testing import lower_run_compare +from executorch.backends.nxp.tests.simple_models import Conv2dWithActivation from executorch.backends.nxp.tests.use_qat import * # noqa F403 diff --git a/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py b/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py index 200ffb6fc99..009e1abc142 100644 --- a/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py +++ b/backends/nxp/tests/ir/converter/node_converter/test_view_copy_converter.py @@ -6,6 +6,8 @@ from typing import Sequence import numpy as np + +# noinspection PyUnusedImports import pytest import torch diff --git a/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py b/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py index c1850f625dd..82cc3f63526 100644 --- a/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py +++ b/backends/nxp/tests/ir/edge_passes/test_convert_reshaping_nodes_to_view.py @@ -9,11 +9,14 @@ from executorch.backends.nxp.backend.ops_aliases import AddTensor, ViewCopy from executorch.backends.nxp.tests.dataset_creator import RandomDatasetCreator from executorch.backends.nxp.tests.graph_verifier import DetailedGraphVerifier -from executorch.backends.nxp.tests.models import SqueezeAddModel, UnsqueezeAddModel from executorch.backends.nxp.tests.nsys_testing import ( AllCloseOutputComparator, lower_run_compare, ) +from executorch.backends.nxp.tests.simple_models import ( + SqueezeAddModel, + UnsqueezeAddModel, +) @pytest.fixture(autouse=True) diff --git a/backends/nxp/tests/ir/edge_passes/test_edge_passes.py b/backends/nxp/tests/ir/edge_passes/test_edge_passes.py index 9fa56989cf0..cbd2b2ff48e 100644 --- a/backends/nxp/tests/ir/edge_passes/test_edge_passes.py +++ b/backends/nxp/tests/ir/edge_passes/test_edge_passes.py @@ -21,6 +21,12 @@ PermuteCopyConverter, ViewCopyConverter, ) +from executorch.backends.nxp.backend.ops_aliases import ( + DequantizePerTensor, + PermuteCopy, + QuantizePerTensor, + ViewCopy, +) from executorch.backends.nxp.edge_passes.neutron_edge_pass_manager import ( NeutronEdgePassManager, ) @@ -42,13 +48,12 @@ EdgeProgramExecutor, OverrideTargetSupportCheck, ) -from executorch.backends.nxp.tests.models import ( +from executorch.backends.nxp.tests.simple_models import ( Conv2dModule, ConvActivationModule, ConvFCFCSoftmaxModuleWithoutReshape, LinearActivationModule, ) -from executorch.exir.dialects._ops import ops as exir_ops from executorch.extension.export_util.utils import export_to_edge from parameterized import parameterized from torch.export import ExportedProgram @@ -56,10 +61,7 @@ def _is_view_copy(node_: Node) -> bool: - return ( - node_.op == "call_function" - and node_.target == exir_ops.edge.aten.view_copy.default - ) + return node_.op == "call_function" and node_.target == ViewCopy def _find_view_copy_node_indices(graph_nodes: list[Node]) -> list[int]: @@ -352,16 +354,10 @@ def test_remove_additional_quantize_dequantize_nodes_pass(self): ) nodes = list(edge_program_with_qdq_cluster.graph.nodes) assert len(nodes) == 10 - assert ( - nodes[5].target - == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ) - assert nodes[6].target == exir_ops.edge.aten.permute_copy.default + assert nodes[5].target == DequantizePerTensor + assert nodes[6].target == PermuteCopy assert "cluster" in nodes[6].meta - assert ( - nodes[7].target - == exir_ops.edge.quantized_decomposed.quantize_per_tensor.default - ) + assert nodes[7].target == QuantizePerTensor # Run pass for removal of additional QDQ nodes and compute in non-float types where possible edge_program_manager = edge_program_manager.transform( @@ -373,12 +369,9 @@ def test_remove_additional_quantize_dequantize_nodes_pass(self): nodes = list(edge_program_without_qdq_cluster.graph.nodes) assert len(nodes) == 8 assert nodes[4].name == "getitem" - assert nodes[5].target == exir_ops.edge.aten.permute_copy.default + assert nodes[5].target == PermuteCopy assert "cluster" not in nodes[5].meta - assert ( - nodes[6].target - == exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default - ) + assert nodes[6].target == DequantizePerTensor edge_program_executor_without_qdq_cluster = EdgeProgramExecutor( edge_program_without_qdq_cluster diff --git a/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py b/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py index aadef8c7731..8a41ee7221a 100644 --- a/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py +++ b/backends/nxp/tests/ir/edge_passes/test_linear_bn_fusing.py @@ -3,7 +3,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import executorch.backends.nxp.tests.models as models +import executorch.backends.nxp.tests.simple_models as models import numpy as np import pytest import torch @@ -21,6 +21,7 @@ batch_norm_target_ops, is_batch_norm, ) +from executorch.backends.nxp.backend.ops_aliases import AddMM, Linear from executorch.backends.nxp.quantizer.neutron_quantizer import NeutronQuantizer from executorch.backends.nxp.tests.executorch_pipeline import ( get_random_calibration_inputs, @@ -34,7 +35,6 @@ ToChannelFirstPreprocess, ToChannelLastPreprocess, ) -from executorch.exir.dialects._ops import ops as exir_ops from torch.export import export, ExportedProgram from torchao.quantization.pt2e.prepare import _is_activation_post_process_node from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_qat_pt2e @@ -243,8 +243,8 @@ def test_linear_bn_full_qat_pipeline_conversion( assert not graph_contains_any_of_ops( graph=edge_program.graph, ops=[ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.linear.default, + AddMM, + Linear, ] + batch_norm_target_ops, ) @@ -303,8 +303,8 @@ def test_incompatible_linear_bn_not_fused(mocker, input_shape, linear_bias, bn_e assert graph_contains_any_of_ops( graph=edge_program.graph, ops=[ - exir_ops.edge.aten.addmm.default, - exir_ops.edge.aten.linear.default, + AddMM, + Linear, ], ) assert graph_contains_any_of_ops( diff --git a/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py b/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py index ef669897b51..b24b3e4a555 100644 --- a/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py +++ b/backends/nxp/tests/ir/edge_passes/test_remove_io_quant_ops_pass.py @@ -10,7 +10,7 @@ import torch from executorch.backends.nxp.tests.executorch_pipeline import to_quantized_edge_program -from executorch.backends.nxp.tests.models import Conv2dReLUModule +from executorch.backends.nxp.tests.simple_models import Conv2dReLUModule from executorch.examples.nxp.experimental.cifar_net.cifar_net import CifarNet from executorch.exir import ExecutorchBackendConfig from executorch.exir.passes.quantize_io_pass import get_config_method_name diff --git a/backends/nxp/tests/models/__init__.py b/backends/nxp/tests/models/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/nxp/tests/generic_tests/test_cifarnet.py b/backends/nxp/tests/models/test_cifarnet.py similarity index 98% rename from backends/nxp/tests/generic_tests/test_cifarnet.py rename to backends/nxp/tests/models/test_cifarnet.py index 6db8ebc9a03..c1d1cd91502 100644 --- a/backends/nxp/tests/generic_tests/test_cifarnet.py +++ b/backends/nxp/tests/models/test_cifarnet.py @@ -5,6 +5,7 @@ import os.path +# noinspection PyUnusedImports import pytest import torch diff --git a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py index 782f7c13a9a..741b68aa1fd 100644 --- a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py +++ b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py @@ -6,7 +6,11 @@ from functools import partial import numpy as np + +# noinspection PyUnusedImports +import pytest import torch + from executorch.backends.nxp.tests.dataset_creator import ( FromCalibrationDataDatasetCreator, ) @@ -23,7 +27,6 @@ ReferenceModel, ) from executorch.backends.nxp.tests.use_qat import * # noqa F403 -import pytest from executorch.examples.nxp.models.mlperf_tiny.image_classification.mlperf_tiny_image_classification import ( MLPerfTinyImageClassification, ) diff --git a/backends/nxp/tests/models.py b/backends/nxp/tests/simple_models.py similarity index 99% rename from backends/nxp/tests/models.py rename to backends/nxp/tests/simple_models.py index eaa5e57cee6..85158b77625 100644 --- a/backends/nxp/tests/models.py +++ b/backends/nxp/tests/simple_models.py @@ -456,6 +456,16 @@ def forward(self, x): return self.avg_pool(x) +class AdaptiveAvgPool1dModule(torch.nn.Module): + def __init__(self, output_size): + super().__init__() + + self.adaptive_avg_pool = torch.nn.AdaptiveAvgPool1d(output_size=output_size) + + def forward(self, x): + return self.adaptive_avg_pool(x) + + class AdaptiveAvgPool2dModule(torch.nn.Module): def __init__(self, output_size): super().__init__() diff --git a/backends/nxp/tests/use_qat.py b/backends/nxp/tests/use_qat.py index 7a63270ae21..c7996f8aad1 100644 --- a/backends/nxp/tests/use_qat.py +++ b/backends/nxp/tests/use_qat.py @@ -3,6 +3,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +# noinspection PyUnusedImports import pytest diff --git a/docs/source/backends/nxp/nxp-quantization.md b/docs/source/backends/nxp/nxp-quantization.md index ef038e47f44..3ba39fda0bb 100644 --- a/docs/source/backends/nxp/nxp-quantization.md +++ b/docs/source/backends/nxp/nxp-quantization.md @@ -10,6 +10,7 @@ The Neutron delegate supports the following quantization schemes: - Static quantization with 8-bit symmetric weights and 8-bit asymmetric activations (via the PT2E quantization flow), per-tensor granularity. - Following operators are supported at this moment: - `aten.abs.default` + - `aten.adaptive_avg_pool1d.default` - `aten.adaptive_avg_pool2d.default` - `aten.add.Tensor` - `aten.addmm.default` diff --git a/docs/source/backends/nxp/op-support.csv b/docs/source/backends/nxp/op-support.csv index b6aa870e31a..368b981face 100644 --- a/docs/source/backends/nxp/op-support.csv +++ b/docs/source/backends/nxp/op-support.csv @@ -1,6 +1,7 @@ Operator,Compute DType,Quantization,Constraints aten.abs.default,int8,static int8, -aten._adaptive_avg_pool2d.default,int8,static int8,"ceil_mode=False, count_include_pad=False, divisor_override=False" +aten._adaptive_avg_pool2d.default,int8,static int8,"Must be representable by avg_pool2D" +aten.adaptive_avg_pool1d.default,int8,static int8,"Must be representable by avg_pool2D" aten.addmm.default,int8,static int8,2D tensor only aten.add.Tensor,int8,static int8,"alpha = 1" aten.amax.default,int8,static int8, From 5efa92515cabe8646589ce45f7ea5a5aa8ea486f Mon Sep 17 00:00:00 2001 From: Scott Roy <161522778+metascroy@users.noreply.github.com> Date: Thu, 10 Sep 2026 12:59:11 -0700 Subject: [PATCH 158/190] Add options for logit slicing to HF export (#22655) Publishes model metadata (max context length, max sequence length, vocabulary size, activation dtype, and logits-to-keep mode) as PTE constant methods and reads it back through a single typed reader in extension/llm/runner/model_metadata.h. The Python writer moves to a shared module, extension/llm/export/model_metadata.py, and the method-name strings live in one place per language (constants.h on the C++ side) so the writer and reader cannot drift apart. The reader requires the core fields and rejects malformed values; the MLX example runners use it to bound prefill, default the KV-cache storage dtype to the model's activation dtype, and check the published vocabulary size against the forward output. Exports can request full, last, or runtime-selected logits, honored by both the single-sequence and batched runners, and the C++ field prefill_chunk_size is renamed to max_seq_len to match the constant it serializes as. --------- Co-authored-by: kiymetakdemir --- .github/workflows/mlx.yml | 5 +- backends/mlx/examples/llm/README.md | 4 +- backends/mlx/examples/llm/dflash/export.py | 14 +- backends/mlx/examples/llm/dflash/run.py | 3 +- backends/mlx/examples/llm/export_llm_hf.py | 153 +++++++++--- backends/mlx/examples/llm/run_llm_batched.cpp | 59 ++--- backends/mlx/examples/llm/run_llm_hf.cpp | 131 +++++++++-- backends/mlx/examples/llm/run_llm_hf.py | 43 +++- backends/mlx/examples/llm/runner_utils.h | 9 + backends/mlx/examples/llm/runtime_meta.py | 15 +- backends/mlx/llm/exportable.py | 164 ++++++++++++- backends/mlx/llm/hf_attention.py | 15 +- extension/llm/batching/module_executor.cpp | 221 ++++++++++++++---- extension/llm/batching/module_executor.h | 5 +- extension/llm/export/BUCK | 18 ++ extension/llm/export/model_metadata.py | 94 ++++++++ extension/llm/runner/constants.h | 2 + extension/llm/runner/model_metadata.h | 184 +++++++++++++++ extension/llm/runner/targets.bzl | 1 + extension/llm/runner/test/CMakeLists.txt | 34 +++ .../llm/runner/test/export_model_metadata.py | 78 +++++++ .../llm/runner/test/test_model_metadata.cpp | 126 ++++++++++ 22 files changed, 1218 insertions(+), 160 deletions(-) create mode 100644 extension/llm/export/model_metadata.py create mode 100644 extension/llm/runner/model_metadata.h create mode 100644 extension/llm/runner/test/export_model_metadata.py create mode 100644 extension/llm/runner/test/test_model_metadata.cpp diff --git a/.github/workflows/mlx.yml b/.github/workflows/mlx.yml index e762110693f..734af8b329f 100644 --- a/.github/workflows/mlx.yml +++ b/.github/workflows/mlx.yml @@ -853,7 +853,7 @@ jobs: # DFlash speculative decoding: target and draft exported as two methods of one # .pte. Also the only coverage that the constants the export publishes - # (get_max_ctx_len / get_prefill_chunk_size / get_max_block_len / + # (get_max_context_len / get_max_seq_len / get_max_block_len / # get_mask_token_id) match what # the runner derives from them -- the runner takes no capacity or block flags. test-mlx-dflash: @@ -919,7 +919,7 @@ jobs: BLOCK_SIZE=$(${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; from executorch.backends.mlx.examples.llm.dflash.model import load_dflash_config; print(load_dflash_config(snapshot_download('${DRAFT_MODEL}', allow_patterns=['*.json'])).block_size)") MASK_TOKEN_ID=$(${CONDA_RUN} python -c "from huggingface_hub import snapshot_download; from executorch.backends.mlx.examples.llm.dflash.model import load_dflash_config; print(load_dflash_config(snapshot_download('${DRAFT_MODEL}', allow_patterns=['*.json'])).mask_token_id)") echo "draft checkpoint block_size: ${BLOCK_SIZE} mask_token_id: ${MASK_TOKEN_ID}" - ${CONDA_RUN} python -c "from executorch.runtime import Runtime, Verification; from executorch.backends.mlx.examples.llm.runtime_meta import read_const_int as r; p = Runtime.get().load_program('${PTE}', verification=Verification.Minimal); got = {n: r(p, n) for n in ['get_max_ctx_len', 'get_prefill_chunk_size', 'get_max_block_len', 'get_mask_token_id']}; want = {'get_max_ctx_len': ${MAX_CTX_LEN}, 'get_prefill_chunk_size': ${PREFILL_CHUNK_SIZE}, 'get_max_block_len': ${BLOCK_SIZE}, 'get_mask_token_id': ${MASK_TOKEN_ID}}; print('published:', got); assert got == want, f'mismatch: {got} != {want}'; assert {'draft', 'target'} <= set(p.method_names), p.method_names; print('Success: constants match and both methods present')" + ${CONDA_RUN} python -c "from executorch.runtime import Runtime, Verification; from executorch.backends.mlx.examples.llm.runtime_meta import read_const_int as r; p = Runtime.get().load_program('${PTE}', verification=Verification.Minimal); got = {n: r(p, n) for n in ['get_max_context_len', 'get_max_seq_len', 'get_max_block_len', 'get_mask_token_id']}; want = {'get_max_context_len': ${MAX_CTX_LEN}, 'get_max_seq_len': ${PREFILL_CHUNK_SIZE}, 'get_max_block_len': ${BLOCK_SIZE}, 'get_mask_token_id': ${MASK_TOKEN_ID}}; print('published:', got); assert got == want, f'mismatch: {got} != {want}'; assert {'draft', 'target'} <= set(p.method_names), p.method_names; print('Success: constants match and both methods present')" echo "::endgroup::" echo "::group::Run DFlash speculative decoding" @@ -1064,6 +1064,7 @@ jobs: --model-id "${MODEL_ID}" \ --output /tmp/${MODEL_NAME}_offgraph.pte \ --use-offgraph-cache \ + --logits-to-keep selected \ --max-ctx-len 1024 \ --dtype bf16 \ --qlinear 4w diff --git a/backends/mlx/examples/llm/README.md b/backends/mlx/examples/llm/README.md index 21c6f72f548..8969a3be3ab 100644 --- a/backends/mlx/examples/llm/README.md +++ b/backends/mlx/examples/llm/README.md @@ -96,7 +96,7 @@ pip install -U "transformers @ git+https://github.com/huggingface/transformers.g | `--use-custom-sdpa` | `False` | Use MLX custom SDPA (`mlx::custom_sdpa`) | | `--use-custom-kv-cache` | `False` | Use MLX custom KV cache (`mlx::kv_cache_update`) | | `--use-offgraph-cache` | `False` | Use the off-graph KV cache (`kvcache::update_and_attend`); replaces the two flags above | -| `--prefill-chunk-size` | `512` | Max tokens per forward step. Bounds the traced `seq_len` dimension and is published as `get_prefill_chunk_size` for the runner. It is also the largest single cache write, so a ring layer is sized `window + chunk - 1`; it may not exceed the sliding window or the context length. Ignored on the optimum-executorch path, which owns its own `seq_len` bound | +| `--prefill-chunk-size` | `512` | Max tokens per forward step. Bounds the traced `seq_len` dimension and is published as `get_max_seq_len` for the runner. It is also the largest single cache write, so a ring layer is sized `window + chunk - 1`; it may not exceed the sliding window or the context length. Ignored on the optimum-executorch path, which owns its own `seq_len` bound | Off-graph exports keep no cache in the `.pte`, so the pybindings `run_llm_hf` cannot run them — use [`mlx_run_llm_hf`](#mlx_run_llm_hf-c) below, which builds @@ -231,7 +231,7 @@ default. | `--temperature` | `0` | Sampling temperature; 0 is greedy argmax, which is what makes two `.pte` files comparable | | `--chat` | `llama3` | Chat template: `llama3`, `gemma`, `gemma4`, or `0` to disable | | `--kv-max-capacity` | `0` | Off-graph: history the cache may hold. Setting it selects the off-graph path | -| `--kv-storage-dtype` | `bf16` | Off-graph: KV storage dtype (`bf16`, `fp16`, `fp32`) | +| `--kv-storage-dtype` | PTE activation dtype | Off-graph: optional KV storage override (`bf16`, `fp16`, `fp32`); defaults to the PTE's `get_activation_dtype`, which is required, so a `.pte` exported before this metadata must be re-exported | | `--kv-initial-capacity` | `-1` | Off-graph: starting pool size; grows by doubling up to capacity | | `--kv-windows` | *(model's own)* | Off-graph: attention pattern override, e.g. `512` | | `--interactive` | `false` | Multi-turn chat on stdin; off-graph only | diff --git a/backends/mlx/examples/llm/dflash/export.py b/backends/mlx/examples/llm/dflash/export.py index 72baca5cc45..fea62c55e89 100644 --- a/backends/mlx/examples/llm/dflash/export.py +++ b/backends/mlx/examples/llm/dflash/export.py @@ -15,7 +15,6 @@ from pathlib import Path import torch - from executorch.backends.mlx.examples.llm.dflash.adapters import get_adapter from executorch.backends.mlx.examples.llm.dflash.model import ( DFlashDraftModel, @@ -151,7 +150,6 @@ def main(): block_dim = Dim("block_len", min=2, max=block_size) import torch.fx.experimental._config as fx_config - from executorch.backends.mlx.examples.llm.dflash.cache import DFlashDraftKVCache class DFlashCachedDraftModel(torch.nn.Module): @@ -195,6 +193,7 @@ def forward(self, tokens, new_ctx, cache_position): from executorch.backends.mlx.examples.llm.export_llm_hf import ( build_hf_exported_program, + model_constant_methods, ) print( @@ -202,7 +201,7 @@ def forward(self, tokens, new_ctx, cache_position): f"and quant {qlinear}/{qembedding} g={qlinear_group_size}/{qembedding_group_size} " f"max_ctx_len {max_ctx_len} prefill_chunk_size {prefill_chunk_size}..." ) - target_exported, prefill_chunk_size = build_hf_exported_program( + target_exported, prefill_chunk_size, vocab_size = build_hf_exported_program( model_id=args.target_model, revision=None, max_ctx_len=max_ctx_len, @@ -226,8 +225,13 @@ def forward(self, tokens, new_ctx, cache_position): from executorch.exir.passes import MemoryPlanningPass constant_methods = { - "get_max_ctx_len": max_ctx_len, - "get_prefill_chunk_size": prefill_chunk_size, + **model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep="full", + activation_dtype=args.dtype, + vocab_size=vocab_size, + max_seq_len=prefill_chunk_size, + ), "get_max_block_len": block_size, "get_mask_token_id": draft_config.mask_token_id, } diff --git a/backends/mlx/examples/llm/dflash/run.py b/backends/mlx/examples/llm/dflash/run.py index 31d900b5730..660c95f87cd 100644 --- a/backends/mlx/examples/llm/dflash/run.py +++ b/backends/mlx/examples/llm/dflash/run.py @@ -24,7 +24,6 @@ import time import torch - from executorch.backends.mlx.examples.llm.runtime_meta import ( apply_chat_template, chunked_prefill, @@ -112,7 +111,7 @@ def resolve_limits(program, pte_path, n_draft_arg): max_ctx_len, prefill_chunk_size = read_model_limits(program) if max_ctx_len is None: raise ValueError( - f"{pte_path} publishes no get_max_ctx_len; re-export it with " + f"{pte_path} publishes no get_max_context_len; re-export it with " "dflash/export.py." ) if prefill_chunk_size is None: diff --git a/backends/mlx/examples/llm/export_llm_hf.py b/backends/mlx/examples/llm/export_llm_hf.py index 825fcb36365..25f2915a329 100644 --- a/backends/mlx/examples/llm/export_llm_hf.py +++ b/backends/mlx/examples/llm/export_llm_hf.py @@ -42,12 +42,43 @@ from typing import Optional import torch +from executorch.extension.llm.export.model_metadata import ( + model_vocab_size, + write_activation_dtype, + write_logits_to_keep_mode, + write_max_context_len, + write_max_seq_len, + write_vocab_size, +) FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" logging.basicConfig(level=logging.INFO, format=FORMAT) logger = logging.getLogger(__name__) +def model_constant_methods( + *, + max_context_len: int, + logits_to_keep: str, + activation_dtype: str, + vocab_size: int, + max_seq_len: Optional[int] = None, +) -> dict[str, int]: + """Build the full metadata set an MLX HF export publishes. + + ``max_seq_len`` is the largest single forward step this export traces + (serialized as ``get_max_seq_len``); ``max_context_len`` is the KV-cache + capacity (``get_max_context_len``). Composes the shared per-constant writers. + """ + return { + **write_max_context_len(max_context_len), + **write_vocab_size(vocab_size), + **write_activation_dtype(activation_dtype), + **write_logits_to_keep_mode(logits_to_keep), + **write_max_seq_len(max_seq_len), + } + + def resolve_prefill_chunk_size( prefill_chunk_size: Optional[int], max_ctx_len: int, @@ -108,6 +139,7 @@ def _export_with_optimum( dtype=dtype_str, max_seq_len=max_ctx_len, ) + vocab_size = model_vocab_size(exportable.model) from executorch.backends.mlx.llm.quantization import quantize_model_ @@ -137,7 +169,15 @@ def _export_with_optimum( # optimum drives its own torch.export call, so it owns the seq-len bound. constant_methods = dict(exportable.metadata) - constant_methods["get_max_ctx_len"] = max_ctx_len + constant_methods.update( + model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep="full", + activation_dtype=dtype, + vocab_size=vocab_size, + max_seq_len=max_ctx_len, + ) + ) edge_program = exir.to_edge_transform_and_lower( exported_progs, @@ -172,15 +212,18 @@ def build_hf_exported_program( qembedding_group_size: Optional[int] = None, tap_layers: Optional[list[int]] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ): """Build the torch.export program for an HF model with custom MLX components. - Returns ``(exported_program, resolved_prefill_chunk_size)``. The resolved chunk + Returns ``(exported_program, resolved_prefill_chunk_size, vocab_size)``. The resolved chunk is the traced ``seq_len`` upper bound and is what callers should publish as - ``get_prefill_chunk_size``. + ``get_max_seq_len``. """ + from executorch.backends.mlx.llm.exportable import LogitsToKeepMode from transformers import AutoModelForCausalLM + logits_to_keep_mode = int(LogitsToKeepMode.from_value(logits_to_keep)) torch_dtype_map = { "fp32": torch.float32, "fp16": torch.float16, @@ -206,6 +249,7 @@ def build_hf_exported_program( if attn_implementation: load_kwargs["attn_implementation"] = attn_implementation model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) + vocab_size = model_vocab_size(model) # Check if model uses sliding window attention. Multimodal configs like # Gemma 4 keep transformer attributes under text_config. @@ -240,6 +284,7 @@ def build_hf_exported_program( model=model, max_cache_len=effective_cache_len, tap_layers=tap_layers, + logits_to_keep_mode=logits_to_keep_mode, ) if use_custom_kv_cache: @@ -283,19 +328,30 @@ def build_hf_exported_program( # prefill_chunk_size is the largest single step: it bounds the traced seq_len # and sizes the ring buffer as window + chunk - 1. seq_len_dim = torch.export.Dim("seq_length_dim", max=prefill_chunk_size) + export_kwargs = { + "input_ids": example_input_ids, + "cache_position": example_cache_position, + } dynamic_shapes = { "input_ids": {1: seq_len_dim}, "cache_position": {0: seq_len_dim}, } + if logits_to_keep == "selected": + logits_example_length = min(seq_length, prefill_chunk_size) + export_kwargs["logits_to_keep"] = torch.arange( + logits_example_length, dtype=torch.int64 + ) + dynamic_shapes["logits_to_keep"] = ( + {0: torch.export.Dim("logits_to_keep_dim", min=1, max=prefill_chunk_size)} + if prefill_chunk_size > 1 + else None + ) with torch.no_grad(): exported_program = torch.export.export( exportable, args=(), - kwargs={ - "input_ids": example_input_ids, - "cache_position": example_cache_position, - }, + kwargs=export_kwargs, dynamic_shapes=dynamic_shapes, strict=True, ) @@ -304,7 +360,7 @@ def build_hf_exported_program( for sym, constraint in exported_program.range_constraints.items(): logger.info(f" Range constraint: {sym}: {constraint}") - return exported_program, prefill_chunk_size + return exported_program, prefill_chunk_size, vocab_size def _export_with_custom_components( @@ -322,6 +378,7 @@ def _export_with_custom_components( qembedding_group_size: Optional[int] = None, tap_layers: Optional[list[int]] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ) -> None: """Export using direct HF model with custom MLX components.""" import executorch.exir as exir @@ -331,7 +388,7 @@ def _export_with_custom_components( from executorch.exir.capture._config import ExecutorchBackendConfig from executorch.exir.passes import MemoryPlanningPass - exported_program, prefill_chunk_size = build_hf_exported_program( + exported_program, prefill_chunk_size, vocab_size = build_hf_exported_program( model_id=model_id, revision=revision, max_ctx_len=max_ctx_len, @@ -345,6 +402,7 @@ def _export_with_custom_components( qembedding_group_size=qembedding_group_size, tap_layers=tap_layers, prefill_chunk_size=prefill_chunk_size, + logits_to_keep=logits_to_keep, ) logger.info("Delegating to MLX backend...") @@ -353,10 +411,13 @@ def _export_with_custom_components( _skip_dim_order=True, ) - constant_methods = { - "get_max_ctx_len": max_ctx_len, - "get_prefill_chunk_size": prefill_chunk_size, - } + constant_methods = model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep=logits_to_keep, + activation_dtype=dtype, + vocab_size=vocab_size, + max_seq_len=prefill_chunk_size, + ) edge_program = exir.to_edge_transform_and_lower( {"forward": exported_program}, @@ -389,6 +450,7 @@ def _export_with_offgraph_cache( qlinear_group_size: Optional[int] = None, qembedding_group_size: Optional[int] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ) -> None: """Export using the off-graph KV cache op (kvcache::update_and_attend).""" import executorch.exir as exir @@ -422,6 +484,7 @@ def _export_with_offgraph_cache( if revision is not None: load_kwargs["revision"] = revision model = AutoModelForCausalLM.from_pretrained(model_id, **load_kwargs) + vocab_size = model_vocab_size(model) model.eval() from executorch.backends.mlx.llm.quantization import quantize_model_ @@ -436,7 +499,7 @@ def _export_with_offgraph_cache( and not no_tie_word_embeddings, ) - exportable = OffGraphExportWrapper(model) + exportable = OffGraphExportWrapper(model, logits_to_keep) from executorch.backends.mlx.llm.cache import resolve_hf_cache_layout @@ -451,14 +514,21 @@ def _export_with_offgraph_cache( prefill_chunk_size, max_ctx_len, min(sliding) if sliding else None ) - kv_metadata = { - "get_n_caches": len(layer_types), - "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32), - "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32), - "get_windows": torch.tensor(cache_windows, dtype=torch.int32), - "get_prefill_chunk_size": prefill_chunk_size, - "get_max_ctx_len": max_ctx_len, - } + kv_metadata = model_constant_methods( + max_context_len=max_ctx_len, + logits_to_keep=logits_to_keep, + activation_dtype=dtype, + vocab_size=vocab_size, + max_seq_len=prefill_chunk_size, + ) + kv_metadata.update( + { + "get_n_caches": len(layer_types), + "get_kv_heads": torch.tensor(cache_kv_heads, dtype=torch.int32), + "get_head_dims": torch.tensor(cache_head_dims, dtype=torch.int32), + "get_windows": torch.tensor(cache_windows, dtype=torch.int32), + } + ) logger.info( f"KV cache layout: {len(layer_types)} caches, " f"{sum(1 for w in cache_windows if w)} sliding (window {sliding_window})" @@ -470,19 +540,30 @@ def _export_with_offgraph_cache( example_cache_position = torch.arange(seq_length, dtype=torch.long) seq_len_dim = torch.export.Dim("seq_length_dim", max=prefill_chunk_size) + export_kwargs = { + "input_ids": example_input_ids, + "cache_position": example_cache_position, + } dynamic_shapes = { "input_ids": {1: seq_len_dim}, "cache_position": {0: seq_len_dim}, } + if logits_to_keep == "selected": + logits_example_length = min(seq_length, prefill_chunk_size) + export_kwargs["logits_to_keep"] = torch.arange( + logits_example_length, dtype=torch.int64 + ) + dynamic_shapes["logits_to_keep"] = ( + {0: torch.export.Dim("logits_to_keep_dim", min=1, max=prefill_chunk_size)} + if prefill_chunk_size > 1 + else None + ) with torch.no_grad(): exported_program = torch.export.export( exportable, args=(), - kwargs={ - "input_ids": example_input_ids, - "cache_position": example_cache_position, - }, + kwargs=export_kwargs, dynamic_shapes=dynamic_shapes, strict=True, ) @@ -534,6 +615,7 @@ def export_llama_hf( qembedding_group_size: Optional[int] = None, tap_layers: Optional[list[int]] = None, prefill_chunk_size: Optional[int] = None, + logits_to_keep: str = "full", ) -> None: if use_offgraph_cache: if use_custom_sdpa or use_custom_kv_cache: @@ -554,12 +636,19 @@ def export_llama_hf( qlinear_group_size=qlinear_group_size, qembedding_group_size=qembedding_group_size, prefill_chunk_size=prefill_chunk_size, + logits_to_keep=logits_to_keep, ) - elif use_custom_sdpa or use_custom_kv_cache or tap_layers is not None: + elif ( + use_custom_sdpa + or use_custom_kv_cache + or tap_layers is not None + or logits_to_keep != "full" + ): logger.info( f"Using custom components: sdpa={use_custom_sdpa}, " f"kv_cache={use_custom_kv_cache}, tap_layers={tap_layers}, " - f"prefill_chunk_size={prefill_chunk_size}" + f"prefill_chunk_size={prefill_chunk_size}, " + f"logits_to_keep={logits_to_keep}" ) _export_with_custom_components( model_id=model_id, @@ -576,6 +665,7 @@ def export_llama_hf( qembedding_group_size=qembedding_group_size, tap_layers=tap_layers, prefill_chunk_size=prefill_chunk_size, + logits_to_keep=logits_to_keep, ) else: logger.info("Using optimum-executorch pipeline (no custom components)") @@ -628,6 +718,12 @@ def main(): help="Comma-separated layer indices whose hidden states are concatenated and returned alongside logits. E.g. '1,9,17,25,33'", ) parser.add_argument("--use-offgraph-cache", action="store_true", default=False) + parser.add_argument( + "--logits-to-keep", + choices=("full", "last", "selected"), + default="full", + help="Logits output: full sequence, last token, or runtime-selected positions.", + ) parser.add_argument( "--prefill-chunk-size", type=int, @@ -658,6 +754,7 @@ def main(): qembedding_group_size=args.qembedding_group_size, tap_layers=tap_layers, prefill_chunk_size=args.prefill_chunk_size, + logits_to_keep=args.logits_to_keep, ) diff --git a/backends/mlx/examples/llm/run_llm_batched.cpp b/backends/mlx/examples/llm/run_llm_batched.cpp index 92d9e6cd7e7..c7333b3a79c 100644 --- a/backends/mlx/examples/llm/run_llm_batched.cpp +++ b/backends/mlx/examples/llm/run_llm_batched.cpp @@ -33,6 +33,7 @@ #include #include #include +#include #include #include @@ -42,8 +43,9 @@ DEFINE_string(out_prefix, "gen", "Output files are _.txt"); DEFINE_int32(max_session_tokens, 2048, "Maximum tokens retained per session"); DEFINE_string( kv_storage_dtype, - "bf16", - "KV storage dtype: bf16, fp16, or fp32"); + "", + "Override KV storage dtype with bf16, fp16, or fp32. Defaults to the PTE " + "activation dtype, or bf16 when metadata is absent."); DEFINE_int32( kv_initial_capacity, -1, @@ -65,14 +67,13 @@ DEFINE_string( "Chat template: llama3, gemma, gemma4, or 0 for raw text"); namespace batching = ::executorch::extension::llm::batching; +using ::executorch::backends::mlx::examples::llm::resolve_kv_storage_dtype; using ::executorch::backends::mlx::examples::llm::resolve_stop_tokens; using ::executorch::backends::mlx::examples::llm::StopTokens; -using ::executorch::backends::mlx::examples::llm::storage_dtype; using ::executorch::backends::mlx::examples::llm::wrap_turn; using ::executorch::extension::Module; using ::executorch::extension::llm::TextStream; using ::executorch::runtime::Error; -using ::executorch::runtime::Result; namespace { @@ -150,26 +151,6 @@ const char* reason_name(const std::optional& reason) { return "unknown"; } -Result> optional_const_int( - Module& module, - const char* name) { - const auto methods = module.method_names(); - if (!methods.ok()) { - return methods.error(); - } - if (methods->count(name) == 0) { - return std::optional{}; - } - const auto result = module.execute(name); - if (!result.ok()) { - return result.error(); - } - if (result->size() != 1 || !result->at(0).isInt()) { - return Error::InvalidProgram; - } - return std::optional{result->at(0).toInt()}; -} - void submit_prompt( batching::Session& session, const tokenizers::Tokenizer& tokenizer, @@ -269,11 +250,6 @@ int main(int argc, char** argv) { std::cerr << "too many prompts" << std::endl; return 1; } - const int kv_dtype = storage_dtype(FLAGS_kv_storage_dtype); - if (kv_dtype < 0) { - std::cerr << "--kv_storage_dtype must be bf16, fp16, or fp32" << std::endl; - return 1; - } if (FLAGS_max_new_tokens > FLAGS_max_session_tokens) { std::cerr << "--max_new_tokens exceeds --max_session_tokens" << std::endl; return 1; @@ -292,16 +268,27 @@ int main(int argc, char** argv) { return 1; } - const auto model_max_context = optional_const_int(*module, "get_max_ctx_len"); - if (!model_max_context.ok()) { - std::cerr << "could not read get_max_ctx_len" << std::endl; + const auto activation_dtype = + ::executorch::extension::llm::read_activation_dtype(*module); + if (!activation_dtype.ok()) { + std::cerr << "could not read model metadata" << std::endl; + return 1; + } + const int kv_dtype = + resolve_kv_storage_dtype(FLAGS_kv_storage_dtype, *activation_dtype); + if (kv_dtype < 0) { + std::cerr << "--kv_storage_dtype must be bf16, fp16, or fp32" << std::endl; + return 1; + } + const auto max_context_length = + ::executorch::extension::llm::read_max_context_length(*module); + if (!max_context_length.ok()) { + std::cerr << "could not read model metadata" << std::endl; return 1; } - if (*model_max_context && - (**model_max_context <= 0 || - FLAGS_max_session_tokens > **model_max_context)) { + if (FLAGS_max_session_tokens > *max_context_length) { std::cerr << "--max_session_tokens " << FLAGS_max_session_tokens - << " exceeds the model context limit " << **model_max_context + << " exceeds the model context limit " << *max_context_length << std::endl; return 1; } diff --git a/backends/mlx/examples/llm/run_llm_hf.cpp b/backends/mlx/examples/llm/run_llm_hf.cpp index 925c1360cf3..4c6a55c875a 100644 --- a/backends/mlx/examples/llm/run_llm_hf.cpp +++ b/backends/mlx/examples/llm/run_llm_hf.cpp @@ -31,6 +31,7 @@ #include #include #include +#include #include #include #include @@ -80,8 +81,9 @@ DEFINE_int32( "only choose policy."); DEFINE_string( kv_storage_dtype, - "bf16", - "Off-graph: KV storage dtype, bf16|fp16|fp32."); + "", + "Off-graph: override KV storage dtype with bf16|fp16|fp32. Defaults to " + "the PTE activation dtype, or bf16 when metadata is absent."); DEFINE_int32( kv_initial_capacity, -1, @@ -101,12 +103,18 @@ DEFINE_bool( false, "Run once before measuring, to absorb JIT and pool growth."); +using ::executorch::backends::mlx::examples::llm::resolve_kv_storage_dtype; using ::executorch::backends::mlx::examples::llm::resolve_stop_tokens; using ::executorch::backends::mlx::examples::llm::StopTokens; -using ::executorch::backends::mlx::examples::llm::storage_dtype; using ::executorch::backends::mlx::examples::llm::wrap_turn; using ::executorch::extension::make_tensor_ptr; using ::executorch::extension::Module; +using ::executorch::extension::llm::check_vocab_size; +using ::executorch::extension::llm::LogitsToKeepMode; +using ::executorch::extension::llm::read_activation_dtype; +using ::executorch::extension::llm::read_logits_to_keep_mode; +using ::executorch::extension::llm::read_max_seq_len; +using ::executorch::extension::llm::read_vocab_size; using ::executorch::extension::llm::TextStream; using ::executorch::runtime::Error; @@ -157,6 +165,53 @@ std::optional const_int(Module& module, const char* name) { return r->at(0).toInt(); } +// The sampler (sample_from_logits) fatally aborts on any other dtype, so an +// unsupported logits type must be rejected at startup rather than at inference. +bool is_supported_logits_type(::executorch::aten::ScalarType type) { + using ScalarType = ::executorch::aten::ScalarType; + return type == ScalarType::Float || type == ScalarType::Half || + type == ScalarType::BFloat16 || type == ScalarType::UInt16; +} + +bool validate_forward_abi( + Module& module, + LogitsToKeepMode logits_to_keep_mode, + std::int64_t& vocab_size) { + const auto meta = module.method_meta("forward"); + if (!meta.ok()) { + std::cerr << "Forward metadata is unavailable" << std::endl; + return false; + } + // The runner feeds tokens + positions, plus a selector in Selected mode; a + // mismatch means the published logits mode disagrees with the traced graph. + const std::size_t expected_inputs = + logits_to_keep_mode == LogitsToKeepMode::Selected ? 3 : 2; + if (meta->num_inputs() != expected_inputs) { + std::cerr << "Forward must take " << expected_inputs + << " inputs for its logits-to-keep mode, got " + << meta->num_inputs() << std::endl; + return false; + } + // The logits output's last dim is the observed vocab width, cross-checked + // against the published get_vocab_size by the caller; its dtype must be one + // the sampler supports. + if (meta->num_outputs() == 0) { + std::cerr << "Forward publishes no logits output" << std::endl; + return false; + } + const auto logits = meta->output_tensor_meta(0); + if (!logits.ok() || logits->sizes().size() < 2 || + logits->sizes()[logits->sizes().size() - 1] <= 0 || + !is_supported_logits_type(logits->scalar_type())) { + std::cerr << "Forward logits must have a sampler-supported dtype and shape " + "[..., vocab]" + << std::endl; + return false; + } + vocab_size = logits->sizes()[logits->sizes().size() - 1]; + return true; +} + std::optional> const_ints(Module& module, const char* name) { const auto r = module.execute(name); if (!r.ok() || r->empty() || !r->at(0).isTensor()) { @@ -333,15 +388,35 @@ int main(int argc, char** argv) { std::cerr << "Failed to load " << pte << std::endl; return 1; } - const auto published_prefill_chunk = - const_int(module, "get_prefill_chunk_size"); - if (!published_prefill_chunk || *published_prefill_chunk <= 0 || - *published_prefill_chunk > std::numeric_limits::max()) { - std::cerr << "Invalid or missing get_prefill_chunk_size in " << pte + const auto logits_to_keep_mode_result = read_logits_to_keep_mode(module); + if (!logits_to_keep_mode_result.ok()) { + std::cerr << "Invalid model metadata in " << pte << std::endl; + return 1; + } + const LogitsToKeepMode logits_to_keep_mode = *logits_to_keep_mode_result; + std::int64_t output_vocab_size = 0; + if (!validate_forward_abi(module, logits_to_keep_mode, output_vocab_size)) { + return 1; + } + const auto published_vocab_size = read_vocab_size(module); + if (!published_vocab_size.ok()) { + std::cerr << "Invalid get_vocab_size in " << pte << std::endl; + return 1; + } + const auto vocab_size_result = + check_vocab_size(*published_vocab_size, output_vocab_size); + if (!vocab_size_result.ok()) { + std::cerr << "Invalid get_vocab_size for the forward output in " << pte << std::endl; return 1; } - const int prefill_chunk = static_cast(*published_prefill_chunk); + const std::int32_t vocab_size = *vocab_size_result; + const auto max_seq_len = read_max_seq_len(module); + if (!max_seq_len.ok()) { + std::cerr << "Invalid or missing get_max_seq_len in " << pte << std::endl; + return 1; + } + const int prefill_chunk = static_cast(*max_seq_len); StopTokens stop_tokens; if (!resolve_stop_tokens(*tokenizer, module, chat, stop_tokens)) { std::cerr << "Could not resolve stop tokens for --chat=" << chat @@ -392,14 +467,37 @@ int main(int argc, char** argv) { auto in = make_tensor_ptr({1, (int)ids.size()}, std::vector(ids)); auto cp = make_tensor_ptr({(int)pos.size()}, std::vector(pos)); - auto out = module.execute("forward", {in, cp}); + auto out = [&]() -> ::executorch::runtime::Result< + std::vector<::executorch::runtime::EValue>> { + if (logits_to_keep_mode == LogitsToKeepMode::Selected) { + auto selector = make_tensor_ptr( + {1}, + std::vector{static_cast(ids.size() - 1)}); + return module.execute("forward", {in, cp, selector}); + } + return module.execute("forward", {in, cp}); + }(); if (!out.ok()) { throw std::runtime_error("execute failed"); } + if (out->empty() || !out->at(0).isTensor()) { + throw std::runtime_error("forward returned no logits"); + } const auto& logits = out->at(0).toTensor(); + const int64_t actual_vocab_size = logits.dim() == 0 + ? 0 + : static_cast(logits.size(logits.dim() - 1)); + const int64_t expected_rows = + logits_to_keep_mode == LogitsToKeepMode::Full + ? static_cast(ids.size()) + : 1; + if (logits.dim() != 3 || logits.size(0) != 1 || + logits.size(1) != expected_rows || + actual_vocab_size != vocab_size) { + throw std::runtime_error("forward returned an invalid logits shape"); + } if (!sampler) { - sampler.emplace( - static_cast(logits.size(logits.dim() - 1)), temperature); + sampler.emplace(vocab_size, temperature); } stats.on_sampling_begin(); const int32_t tok = @@ -632,11 +730,16 @@ int main(int argc, char** argv) { /*run_prefill_chunk=*/prefill_chunk); } + const auto activation_dtype = read_activation_dtype(module); + if (!activation_dtype.ok()) { + std::cerr << "Invalid get_activation_dtype in " << pte << std::endl; + return 1; + } cache::CacheConfig cfg{}; cfg.capacity = kv_capacity; - cfg.kv_dtype = storage_dtype(kv_dtype); + cfg.kv_dtype = resolve_kv_storage_dtype(kv_dtype, *activation_dtype); if (cfg.kv_dtype < 0) { - std::cerr << "Invalid --kv-storage-dtype: " << kv_dtype + std::cerr << "Invalid --kv-storage-dtype override: " << kv_dtype << " (bf16|fp16|fp32)" << std::endl; return 1; } diff --git a/backends/mlx/examples/llm/run_llm_hf.py b/backends/mlx/examples/llm/run_llm_hf.py index e73da846623..79081884c2c 100644 --- a/backends/mlx/examples/llm/run_llm_hf.py +++ b/backends/mlx/examples/llm/run_llm_hf.py @@ -25,14 +25,18 @@ import time import torch - from executorch.backends.mlx.examples.llm.runtime_meta import ( apply_chat_template, chunked_prefill, get_eos_token_ids, load_text_processor, + read_const_int, read_model_limits, ) +from executorch.extension.llm.export.model_metadata import ( + LOGITS_TO_KEEP_MODE_METHOD, + LOGITS_TO_KEEP_MODES, +) from executorch.runtime import Runtime, Verification FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -40,11 +44,12 @@ logger = logging.getLogger(__name__) -def _get_max_input_seq_len(program) -> int: - """Inspect the .pte program metadata to determine the max input_ids seq len. +def _forward_input_seq_len(program) -> int: + """The forward's traced token-input width -- what set_inputs will accept. - Fallback for .pte files exported before get_prefill_chunk_size existed. - Returns the static seq-len dimension of the first input tensor (input_ids). + 1 for a static token-by-token export (e.g. optimum's static cache), or the + dynamic upper bound for a chunked-prefill export. This is authoritative: + feeding more tokens than this per step fails set_inputs. """ meta = program.metadata("forward") input_ids_info = meta.input_tensor_meta(0) @@ -67,9 +72,31 @@ def run_inference( et_runtime = Runtime.get() program = et_runtime.load_program(pte_path, verification=Verification.Minimal) - max_ctx_len, prefill_chunk_size = read_model_limits(program) - if prefill_chunk_size is None: - prefill_chunk_size = _get_max_input_seq_len(program) + # This pybindings runner only feeds tokens and positions. A model exported + # with --logits-to-keep selected takes a third runtime selector input, so + # its forward cannot be invoked here; use the C++ runner (mlx_run_llm_hf). + if ( + read_const_int(program, LOGITS_TO_KEEP_MODE_METHOD) + == LOGITS_TO_KEEP_MODES["selected"] + ): + raise ValueError( + "This .pte was exported with --logits-to-keep selected, which needs " + "a runtime-supplied logits selector input that run_llm_hf.py does " + "not provide. Run it with the C++ runner mlx_run_llm_hf, or " + "re-export with --logits-to-keep full or last." + ) + + max_ctx_len, declared_max_seq_len = read_model_limits(program) + # The forward only accepts up to its traced token width, so clamp the + # declared step to it: optimum's static export takes 1 token/forward while + # its get_max_seq_len is the context length, and feeding more crashes + # set_inputs. A chunked-prefill export reports the two as equal. + input_seq_len = _forward_input_seq_len(program) + prefill_chunk_size = ( + min(declared_max_seq_len, input_seq_len) + if declared_max_seq_len is not None + else input_seq_len + ) logger.info( f"Model limits: max_ctx_len={max_ctx_len}, " f"prefill_chunk_size={prefill_chunk_size}" diff --git a/backends/mlx/examples/llm/runner_utils.h b/backends/mlx/examples/llm/runner_utils.h index 6e8ac83c1d9..6455dff7830 100644 --- a/backends/mlx/examples/llm/runner_utils.h +++ b/backends/mlx/examples/llm/runner_utils.h @@ -102,6 +102,15 @@ inline int storage_dtype(const std::string& name) { return -1; } +inline int resolve_kv_storage_dtype( + const std::string& override_name, + ::executorch::aten::ScalarType activation_dtype) { + if (!override_name.empty()) { + return storage_dtype(override_name); + } + return static_cast(activation_dtype); +} + } // namespace llm } // namespace examples } // namespace mlx diff --git a/backends/mlx/examples/llm/runtime_meta.py b/backends/mlx/examples/llm/runtime_meta.py index 7e50b198861..d38c1d8eee2 100644 --- a/backends/mlx/examples/llm/runtime_meta.py +++ b/backends/mlx/examples/llm/runtime_meta.py @@ -6,8 +6,8 @@ """Shared runtime helpers for the MLX LLM example runners. -Exports publish their limits as constant methods (``get_max_ctx_len``, -``get_prefill_chunk_size``) so runners do not have to be told what a .pte +Exports publish their limits as constant methods (``get_max_context_len``, +``get_max_seq_len``) so runners do not have to be told what a .pte supports. This mirrors ``const_int`` in run_llm_hf.cpp. Prompt handling (processor loading, chat templating, EOS lookup) lives here too: @@ -21,6 +21,11 @@ import torch +from executorch.extension.llm.export.model_metadata import ( + MAX_CONTEXT_LEN_METHOD, + MAX_SEQ_LEN_METHOD, +) + logger = logging.getLogger(__name__) @@ -37,10 +42,10 @@ def read_const_int(program, name: str) -> Optional[int]: def read_model_limits(program) -> Tuple[Optional[int], Optional[int]]: - """Return (max_ctx_len, prefill_chunk_size) as published by the export.""" + """Return (max_context_len, max_seq_len) as published by the export.""" return ( - read_const_int(program, "get_max_ctx_len"), - read_const_int(program, "get_prefill_chunk_size"), + read_const_int(program, MAX_CONTEXT_LEN_METHOD), + read_const_int(program, MAX_SEQ_LEN_METHOD), ) diff --git a/backends/mlx/llm/exportable.py b/backends/mlx/llm/exportable.py index 4660247e27b..174d5ae7ea9 100644 --- a/backends/mlx/llm/exportable.py +++ b/backends/mlx/llm/exportable.py @@ -21,7 +21,8 @@ """ import logging -from typing import List, Optional, Sequence +from enum import IntEnum +from typing import List, Optional, Sequence, Union import torch from transformers.integrations.executorch import ( @@ -32,6 +33,50 @@ logger = logging.getLogger(__name__) +class LogitsToKeepMode(IntEnum): + FULL = 0 + LAST = 1 + SELECTED = 2 + + @classmethod + def from_value(cls, value: Union["LogitsToKeepMode", str, int]): + if isinstance(value, str): + try: + return cls[value.upper()] + except KeyError as error: + raise ValueError(f"Unsupported logits-to-keep mode: {value}") from error + return cls(value) + + +class _LogitsToKeepMixin: + logits_to_keep_mode: LogitsToKeepMode + + def _resolve_logits_to_keep( + self, logits_to_keep: Optional[torch.LongTensor] + ) -> Union[int, torch.LongTensor]: + if self.logits_to_keep_mode == LogitsToKeepMode.SELECTED: + if logits_to_keep is None: + raise ValueError("selected logits-to-keep requires an index tensor") + if logits_to_keep.dtype != torch.int64 or logits_to_keep.dim() != 1: + raise ValueError("logits_to_keep must be an int64[K] tensor") + return logits_to_keep + return int(self.logits_to_keep_mode) + + def _logits_to_keep_kwargs( + self, logits_to_keep: Optional[torch.LongTensor] + ) -> dict: + if self.logits_to_keep_mode == LogitsToKeepMode.FULL: + return {} + return {"logits_to_keep": self._resolve_logits_to_keep(logits_to_keep)} + + def _sync_cache_position(self, cache, cache_position) -> None: + if cache_position is None or not hasattr(cache, "layers"): + return + for layer in cache.layers: + if hasattr(layer, "cumulative_length"): + layer.cumulative_length.copy_(cache_position[0]) + + class _HiddenTapMixin: """Shared tapping logic - expects self.layer_ids and self.model to exist.""" @@ -41,8 +86,78 @@ def _tap_hidden(self, outs): return torch.cat(captured, dim=-1) +class TorchExportableModuleWithStaticCacheAndLogitsToKeep( + _LogitsToKeepMixin, TorchExportableModuleWithStaticCache +): + def __init__( + self, + model, + batch_size: Optional[int] = None, + max_cache_len: Optional[int] = None, + device: Optional[torch.device] = None, + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, + ): + super().__init__( + model, batch_size=batch_size, max_cache_len=max_cache_len, device=device + ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, + ): + self._sync_cache_position(self.static_cache, cache_position) + return self.model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + attention_mask=None, + past_key_values=self.static_cache, + use_cache=True, + **self._logits_to_keep_kwargs(logits_to_keep), + ).logits + + +class TorchExportableModuleWithHybridCacheAndLogitsToKeep( + _LogitsToKeepMixin, TorchExportableModuleWithHybridCache +): + def __init__( + self, + model, + batch_size: Optional[int] = None, + max_cache_len: Optional[int] = None, + device: Optional[torch.device] = None, + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, + ): + super().__init__( + model, batch_size=batch_size, max_cache_len=max_cache_len, device=device + ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) + + def forward( + self, + input_ids: Optional[torch.LongTensor] = None, + inputs_embeds: Optional[torch.Tensor] = None, + cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, + ): + self._sync_cache_position(self.cache, cache_position) + return self.model( + input_ids=input_ids, + inputs_embeds=inputs_embeds, + cache_position=cache_position, + attention_mask=None, + past_key_values=self.cache, + use_cache=True, + **self._logits_to_keep_kwargs(logits_to_keep), + ).logits + + class TorchExportableModuleWithStaticCacheAndHidden( - _HiddenTapMixin, TorchExportableModuleWithStaticCache + _HiddenTapMixin, _LogitsToKeepMixin, TorchExportableModuleWithStaticCache ): def __init__( self, @@ -51,10 +166,12 @@ def __init__( max_cache_len: Optional[int] = None, device: Optional[torch.device] = None, layer_ids: Sequence[int] = (), + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, ): super().__init__( model, batch_size=batch_size, max_cache_len=max_cache_len, device=device ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) if not layer_ids: raise ValueError("layer_ids must be non-empty") self.layer_ids: List[int] = list(layer_ids) @@ -64,7 +181,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.Tensor] = None, cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, ): + self._sync_cache_position(self.static_cache, cache_position) outs = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, @@ -73,6 +192,7 @@ def forward( past_key_values=self.static_cache, use_cache=True, output_hidden_states=True, + **self._logits_to_keep_kwargs(logits_to_keep), ) hidden = self._tap_hidden(outs) if hasattr(outs, "logits"): @@ -81,7 +201,7 @@ def forward( class TorchExportableModuleWithHybridCacheAndHidden( - _HiddenTapMixin, TorchExportableModuleWithHybridCache + _HiddenTapMixin, _LogitsToKeepMixin, TorchExportableModuleWithHybridCache ): def __init__( self, @@ -90,10 +210,12 @@ def __init__( max_cache_len: Optional[int] = None, device: Optional[torch.device] = None, layer_ids: Sequence[int] = (), + logits_to_keep_mode: LogitsToKeepMode = LogitsToKeepMode.FULL, ): super().__init__( model, batch_size=batch_size, max_cache_len=max_cache_len, device=device ) + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) if not layer_ids: raise ValueError("layer_ids must be non-empty") self.layer_ids: List[int] = list(layer_ids) @@ -103,7 +225,9 @@ def forward( input_ids: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.Tensor] = None, cache_position: Optional[torch.Tensor] = None, + logits_to_keep: Optional[torch.LongTensor] = None, ): + self._sync_cache_position(self.cache, cache_position) outs = self.model( input_ids=input_ids, inputs_embeds=inputs_embeds, @@ -112,6 +236,7 @@ def forward( past_key_values=self.cache, use_cache=True, output_hidden_states=True, + **self._logits_to_keep_kwargs(logits_to_keep), ) hidden = self._tap_hidden(outs) if hasattr(outs, "logits"): @@ -124,6 +249,7 @@ def create_hf_exportable( max_cache_len: int, tap_layers: Optional[Sequence[int]] = None, batch_size: int = 1, + logits_to_keep_mode: Union[LogitsToKeepMode, str, int] = LogitsToKeepMode.FULL, ): """Factory: picks static vs hybrid and hidden-tapping vs plain. @@ -132,6 +258,7 @@ def create_hf_exportable( max_cache_len: cache capacity tap_layers: optional layer indices to tap and concat as second output batch_size: batch size for cache init + logits_to_keep_mode: full, last, or selected logits selection Returns: An exportable module with .model attribute pointing to HF model @@ -139,6 +266,7 @@ def create_hf_exportable( """ text_config = model.config.get_text_config() sliding_window = getattr(text_config, "sliding_window", None) + logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) if sliding_window is not None: if tap_layers is not None: @@ -150,12 +278,23 @@ def create_hf_exportable( batch_size=batch_size, max_cache_len=max_cache_len, layer_ids=tap_layers, + logits_to_keep_mode=logits_to_keep_mode, + ) + if logits_to_keep_mode == LogitsToKeepMode.FULL: + logger.info("Creating TorchExportableModuleWithHybridCache wrapper...") + return TorchExportableModuleWithHybridCache( + model=model, + batch_size=batch_size, + max_cache_len=max_cache_len, ) - logger.info("Creating TorchExportableModuleWithHybridCache wrapper...") - return TorchExportableModuleWithHybridCache( + logger.info( + f"Creating hybrid-cache wrapper with {logits_to_keep_mode.name.lower()} logits..." + ) + return TorchExportableModuleWithHybridCacheAndLogitsToKeep( model=model, batch_size=batch_size, max_cache_len=max_cache_len, + logits_to_keep_mode=logits_to_keep_mode, ) else: if tap_layers is not None: @@ -167,12 +306,23 @@ def create_hf_exportable( batch_size=batch_size, max_cache_len=max_cache_len, layer_ids=tap_layers, + logits_to_keep_mode=logits_to_keep_mode, ) - logger.info("Creating TorchExportableModuleWithStaticCache wrapper...") - return TorchExportableModuleWithStaticCache( + if logits_to_keep_mode == LogitsToKeepMode.FULL: + logger.info("Creating TorchExportableModuleWithStaticCache wrapper...") + return TorchExportableModuleWithStaticCache( + model=model, + batch_size=batch_size, + max_cache_len=max_cache_len, + ) + logger.info( + f"Creating static-cache wrapper with {logits_to_keep_mode.name.lower()} logits..." + ) + return TorchExportableModuleWithStaticCacheAndLogitsToKeep( model=model, batch_size=batch_size, max_cache_len=max_cache_len, + logits_to_keep_mode=logits_to_keep_mode, ) diff --git a/backends/mlx/llm/hf_attention.py b/backends/mlx/llm/hf_attention.py index c8f34b11d5e..ff8fc3ff9a3 100644 --- a/backends/mlx/llm/hf_attention.py +++ b/backends/mlx/llm/hf_attention.py @@ -41,8 +41,8 @@ from typing import Callable, Optional, Tuple, Union import executorch.backends.mlx.custom_ops as _mlx_custom_ops # noqa: F401 - import torch +from executorch.backends.mlx.llm.exportable import _LogitsToKeepMixin, LogitsToKeepMode def mlx_sdpa_with_start_pos_forward( @@ -219,8 +219,8 @@ def register_mlx_attention(name: str = "mlx") -> None: ) -class OffGraphExportWrapper(torch.nn.Module): - """forward(input_ids, cache_position) -> logits, with no in-graph cache. +class OffGraphExportWrapper(_LogitsToKeepMixin, torch.nn.Module): + """Export logits with no in-graph cache. The analog of TorchExportableModuleWithStaticCache for the off-graph op: runs the model with use_cache=False so each attention layer sees only this @@ -228,12 +228,16 @@ class OffGraphExportWrapper(torch.nn.Module): signature the runner drives. """ - def __init__(self, model: torch.nn.Module): + def __init__(self, model: torch.nn.Module, logits_to_keep_mode="full"): super().__init__() self.model = model + self.logits_to_keep_mode = LogitsToKeepMode.from_value(logits_to_keep_mode) def forward( - self, input_ids: torch.Tensor, cache_position: torch.Tensor + self, + input_ids: torch.Tensor, + cache_position: torch.Tensor, + logits_to_keep: Optional[torch.LongTensor] = None, ) -> torch.Tensor: # Single sequence: the op takes [q_len, n_dims] positions and the # attention function reads position_ids[0], so a batch would be placed @@ -249,6 +253,7 @@ def forward( position_ids=cache_position.unsqueeze(0), use_cache=False, past_key_values=None, + **self._logits_to_keep_kwargs(logits_to_keep), ).logits diff --git a/extension/llm/batching/module_executor.cpp b/extension/llm/batching/module_executor.cpp index 522a928e508..a2e920445af 100644 --- a/extension/llm/batching/module_executor.cpp +++ b/extension/llm/batching/module_executor.cpp @@ -9,7 +9,6 @@ #include #include -#include #include #include #include @@ -33,6 +32,12 @@ using ::executorch::runtime::Result; namespace { +bool is_supported_logits_type(::executorch::aten::ScalarType type) { + using ScalarType = ::executorch::aten::ScalarType; + return type == ScalarType::Float || type == ScalarType::Half || + type == ScalarType::BFloat16 || type == ScalarType::UInt16; +} + // Constant methods carry no delegate, so the layout reads with the program // loaded and the method not. Sizing is the caller's and is left unset. Result config_from_program(Module& module) { @@ -211,7 +216,8 @@ ModuleExecutor::ModuleExecutor( std::string backend_id, std::string method, std::int32_t vocab_size, - int max_step_tokens) + int max_step_tokens, + LogitsToKeepMode logits_to_keep_mode) : install_guard_(cache), module_(std::move(module)), ctl_(cache->as()), @@ -220,7 +226,8 @@ ModuleExecutor::ModuleExecutor( backend_id_(std::move(backend_id)), method_(std::move(method)), vocab_size_(vocab_size), - max_step_tokens_(max_step_tokens) {} + max_step_tokens_(max_step_tokens), + logits_to_keep_mode_(logits_to_keep_mode) {} ModuleExecutor::~ModuleExecutor() = default; @@ -247,6 +254,32 @@ Result> ModuleExecutor::create( return load_error; } + const auto max_context_length = read_max_context_length(*module); + if (!max_context_length.ok()) { + ET_LOG(Error, "ModuleExecutor: the program's metadata is malformed"); + return max_context_length.error(); + } + if (max_session_tokens > *max_context_length) { + ET_LOG( + Error, + "ModuleExecutor: max session tokens %d exceeds model context length %" PRId64, + max_session_tokens, + *max_context_length); + return Error::InvalidArgument; + } + const auto logits_mode_result = read_logits_to_keep_mode(*module); + if (!logits_mode_result.ok()) { + ET_LOG(Error, "ModuleExecutor: the program's metadata is malformed"); + return logits_mode_result.error(); + } + const LogitsToKeepMode logits_mode = *logits_mode_result; + if (logits_mode == LogitsToKeepMode::Last) { + ET_LOG( + Error, + "ModuleExecutor: logits-to-keep mode last is incompatible with batched execution"); + return Error::NotSupported; + } + auto cfg = config_from_program(*module); if (!cfg.ok()) { return cfg.error(); @@ -271,6 +304,83 @@ Result> ModuleExecutor::create( return meta.error(); } + const std::size_t expected_inputs = + logits_mode == LogitsToKeepMode::Selected ? 3 : 2; + if (meta->num_inputs() != expected_inputs) { + ET_LOG( + Error, + "ModuleExecutor: %s expects %zu inputs for its logits mode, got %zu", + method.c_str(), + expected_inputs, + meta->num_inputs()); + return Error::InvalidProgram; + } + + const auto tokens_info = meta->input_tensor_meta(0); + const auto positions_info = meta->input_tensor_meta(1); + if (!tokens_info.ok() || !positions_info.ok()) { + ET_LOG(Error, "ModuleExecutor: %s inputs must be tensors", method.c_str()); + return Error::InvalidProgram; + } + const auto token_sizes = tokens_info->sizes(); + const auto position_sizes = positions_info->sizes(); + if (tokens_info->scalar_type() != ::executorch::aten::ScalarType::Long || + token_sizes.size() != 2 || token_sizes[0] != 1 || token_sizes[1] <= 0 || + positions_info->scalar_type() != ::executorch::aten::ScalarType::Long || + position_sizes.size() != 1 || position_sizes[0] != token_sizes[1]) { + ET_LOG( + Error, + "ModuleExecutor: %s must take Long[1, T] tokens and Long[T] positions", + method.c_str()); + return Error::InvalidProgram; + } + if (logits_mode == LogitsToKeepMode::Selected) { + const auto selector_info = meta->input_tensor_meta(2); + if (!selector_info.ok() || + selector_info->scalar_type() != ::executorch::aten::ScalarType::Long || + selector_info->sizes().size() != 1) { + ET_LOG( + Error, + "ModuleExecutor: %s selected logits selector must be rank-one Long", + method.c_str()); + return Error::InvalidProgram; + } + } + + if (meta->num_outputs() == 0) { + ET_LOG(Error, "ModuleExecutor: %s publishes no outputs", method.c_str()); + return Error::InvalidProgram; + } + const auto logits_info = meta->output_tensor_meta(0); + if (!logits_info.ok()) { + ET_LOG(Error, "ModuleExecutor: %s has no logits metadata", method.c_str()); + return logits_info.error(); + } + const auto logits_sizes = logits_info->sizes(); + if (logits_sizes.size() < 2 || logits_sizes[logits_sizes.size() - 1] <= 0 || + !is_supported_logits_type(logits_info->scalar_type())) { + ET_LOG( + Error, + "ModuleExecutor: %s logits must have supported dtype and shape [..., vocab]", + method.c_str()); + return Error::InvalidProgram; + } + const auto published_vocab_size = read_vocab_size(*module); + if (!published_vocab_size.ok()) { + ET_LOG( + Error, "ModuleExecutor: invalid get_vocab_size for %s", method.c_str()); + return published_vocab_size.error(); + } + const auto vocab_size = check_vocab_size( + *published_vocab_size, logits_sizes[logits_sizes.size() - 1]); + if (!vocab_size.ok()) { + ET_LOG( + Error, + "ModuleExecutor: invalid get_vocab_size for %s output width", + method.c_str()); + return vocab_size.error(); + } + std::string backend_id; for (std::size_t i = 0; i < meta->num_backends(); ++i) { const auto name = meta->get_backend_name(i); @@ -311,36 +421,6 @@ Result> ModuleExecutor::create( return Error::InvalidType; } - if (meta->num_outputs() == 0) { - ET_LOG(Error, "ModuleExecutor: %s publishes no outputs", method.c_str()); - return Error::InvalidProgram; - } - const auto logits_info = meta->output_tensor_meta(0); - if (!logits_info.ok()) { - ET_LOG(Error, "ModuleExecutor: %s has no logits metadata", method.c_str()); - return logits_info.error(); - } - if (logits_info->sizes().empty()) { - ET_LOG(Error, "ModuleExecutor: %s has no logits shape", method.c_str()); - return Error::InvalidProgram; - } - const auto logits_sizes = logits_info->sizes(); - - const auto tokens_info = meta->input_tensor_meta(0); - if (!tokens_info.ok()) { - ET_LOG( - Error, - "ModuleExecutor: %s has no token input metadata", - method.c_str()); - return tokens_info.error(); - } - if (tokens_info->sizes().empty()) { - ET_LOG( - Error, "ModuleExecutor: %s has no token input shape", method.c_str()); - return Error::InvalidProgram; - } - const auto tokens_sizes = tokens_info->sizes(); - return std::unique_ptr(new ModuleExecutor( std::move(module), std::move(cache), @@ -348,8 +428,9 @@ Result> ModuleExecutor::create( max_session_tokens, std::move(backend_id), std::move(method), - logits_sizes[logits_sizes.size() - 1], - tokens_sizes[tokens_sizes.size() - 1])); + *vocab_size, + token_sizes[1], + logits_mode)); } bool ModuleExecutor::initialize() { @@ -443,7 +524,32 @@ bool ModuleExecutor::execute(const BatchInput& batch, BatchOutput& out) { {n}, std::vector( step->positions.begin() + off, step->positions.begin() + off + n)); - auto result = module_->execute(method_, {tokens, positions}); + std::vector selector_values; + std::vector selected_inputs; + if (logits_to_keep_mode_ == LogitsToKeepMode::Selected) { + for (std::size_t i = 0; i < step->logit_indices.size(); ++i) { + const int row = step->logit_indices[i]; + if (row >= off && row < off + n) { + selector_values.push_back(row - off); + selected_inputs.push_back(i); + } + } + if (selector_values.empty()) { + selector_values.push_back(n - 1); + } + } + + const int expected_rows = logits_to_keep_mode_ == LogitsToKeepMode::Selected + ? static_cast(selector_values.size()) + : n; + auto result = [&]() -> Result> { + if (logits_to_keep_mode_ == LogitsToKeepMode::Selected) { + auto selector = + make_tensor_ptr({expected_rows}, std::move(selector_values)); + return module_->execute(method_, {tokens, positions, selector}); + } + return module_->execute(method_, {tokens, positions}); + }(); if (!result.ok()) { ET_LOG( Error, @@ -458,18 +564,43 @@ bool ModuleExecutor::execute(const BatchInput& batch, BatchOutput& out) { } // Non-const: the sampler reduces each row in place. Each is read once. auto logits = result->at(0).toTensor(); + if (logits.dim() < 2 || logits.size(logits.dim() - 1) != vocab_size_ || + logits.numel() != + static_cast(expected_rows) * vocab_size_) { + ET_LOG( + Error, + "ModuleExecutor: %s returned invalid logits shape for %d rows and vocab %d", + method_.c_str(), + expected_rows, + vocab_size_); + return false; + } - for (std::size_t i = 0; i < batch.inputs.size(); ++i) { - const int row = step->logit_indices[i]; - if (row < off || row >= off + n) { - continue; // another slice's row, or a chunk whose prediction is dropped + if (logits_to_keep_mode_ == LogitsToKeepMode::Selected) { + for (std::size_t row = 0; row < selected_inputs.size(); ++row) { + const std::size_t input_index = selected_inputs[row]; + const SessionId session = batch.inputs[input_index].sid; + const std::optional token = + sample_row(logits, static_cast(row), session); + if (!token) { + return false; + } + out.outputs[input_index] = Output{session, {*token}}; } - const SessionId session = batch.inputs[i].sid; - const std::optional token = sample_row(logits, row - off, session); - if (!token) { - return false; + } else { + for (std::size_t i = 0; i < batch.inputs.size(); ++i) { + const int row = step->logit_indices[i]; + if (row < off || row >= off + n) { + continue; // another slice's row, or a dropped chunk prediction + } + const SessionId session = batch.inputs[i].sid; + const std::optional token = + sample_row(logits, row - off, session); + if (!token) { + return false; + } + out.outputs[i] = Output{session, {*token}}; } - out.outputs[i] = Output{session, {*token}}; } } return true; diff --git a/extension/llm/batching/module_executor.h b/extension/llm/batching/module_executor.h index ec249b3ecfa..e0832f5fbc6 100644 --- a/extension/llm/batching/module_executor.h +++ b/extension/llm/batching/module_executor.h @@ -23,6 +23,7 @@ #include #include #include +#include #include #include #include // ET_EXPERIMENTAL @@ -108,7 +109,8 @@ class ET_EXPERIMENTAL ModuleExecutor : public Executor { std::string backend_id, std::string method, std::int32_t vocab_size, - int max_step_tokens); + int max_step_tokens, + LogitsToKeepMode logits_to_keep_mode); // Draw the token an input produced from its row of `logits`, which the // session's sampler consumes in place. @@ -127,6 +129,7 @@ class ET_EXPERIMENTAL ModuleExecutor : public Executor { // The method's logits width, so a sampler can be built by its policy. std::int32_t vocab_size_; int max_step_tokens_; + LogitsToKeepMode logits_to_keep_mode_; SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids std::unordered_map sessions_; diff --git a/extension/llm/export/BUCK b/extension/llm/export/BUCK index 7ce87cb3e73..d956bca2258 100644 --- a/extension/llm/export/BUCK +++ b/extension/llm/export/BUCK @@ -58,6 +58,24 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "model_metadata", + srcs = [ + "model_metadata.py", + ], + _is_external_target = True, + base_module = "executorch.extension.llm.export", + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + "//executorch/extension/llm/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:scalar_type", + ], +) + fbcode_target(_kind = runtime.python_library, name = "int4", srcs = [ diff --git a/extension/llm/export/model_metadata.py b/extension/llm/export/model_metadata.py new file mode 100644 index 00000000000..2057fe910df --- /dev/null +++ b/extension/llm/export/model_metadata.py @@ -0,0 +1,94 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Per-constant model-metadata writers for exported LLM programs. + +Exports publish their metadata as PTE constant methods; the shared typed C++ +readers in ``extension/llm/runner/model_metadata.h`` consume them. The method +names below are the single Python-side source of truth and must stay in sync +with the constants in ``extension/llm/runner/constants.h``. + +Each ``write_*`` returns the ``{name: value}`` for one constant. Producers +compose the set they publish (the MLX export's ``model_constant_methods`` builds +the full set); the backend-neutral runner test composes them too. This module +deliberately depends only on ``torch`` (and a lazy ``ScalarType`` import), so +those consumers pull in no backend. +""" + +from typing import Optional + +import torch + +# Constant-method names. Keep in sync with extension/llm/runner/constants.h. +MAX_CONTEXT_LEN_METHOD = "get_max_context_len" +MAX_SEQ_LEN_METHOD = "get_max_seq_len" +VOCAB_SIZE_METHOD = "get_vocab_size" +ACTIVATION_DTYPE_METHOD = "get_activation_dtype" +LOGITS_TO_KEEP_MODE_METHOD = "get_logits_to_keep_mode" + +# Serialized logits-to-keep modes. Keep in sync with LogitsToKeepMode in +# extension/llm/runner/model_metadata.h. +LOGITS_TO_KEEP_MODES = {"full": 0, "last": 1, "selected": 2} + + +def _require_positive(name: str, value: int) -> dict[str, int]: + if value <= 0: + raise ValueError(f"Invalid value for {name}: {value}") + return {name: value} + + +def write_max_context_len(max_context_len: int) -> dict[str, int]: + """max_context_len -> get_max_context_len (the KV-cache capacity).""" + return _require_positive(MAX_CONTEXT_LEN_METHOD, max_context_len) + + +def write_vocab_size(vocab_size: int) -> dict[str, int]: + """vocab_size -> get_vocab_size.""" + return _require_positive(VOCAB_SIZE_METHOD, vocab_size) + + +def write_max_seq_len(max_seq_len: Optional[int]) -> dict[str, int]: + """max_seq_len -> get_max_seq_len (largest single forward step); optional.""" + if max_seq_len is None: + return {} + return _require_positive(MAX_SEQ_LEN_METHOD, max_seq_len) + + +def write_activation_dtype(activation_dtype: str) -> dict[str, int]: + """activation_dtype name -> get_activation_dtype (ExecuTorch ScalarType).""" + from executorch.exir.scalar_type import ScalarType + + table = { + "fp16": ScalarType.HALF, + "fp32": ScalarType.FLOAT, + "bf16": ScalarType.BFLOAT16, + } + try: + return {ACTIVATION_DTYPE_METHOD: int(table[activation_dtype])} + except KeyError as error: + raise ValueError(f"Unsupported activation dtype: {activation_dtype}") from error + + +def write_logits_to_keep_mode(logits_to_keep: str) -> dict[str, int]: + """logits_to_keep name -> get_logits_to_keep_mode.""" + try: + return {LOGITS_TO_KEEP_MODE_METHOD: LOGITS_TO_KEEP_MODES[logits_to_keep]} + except KeyError as error: + raise ValueError( + f"Unsupported logits-to-keep mode: {logits_to_keep}" + ) from error + + +def model_vocab_size(model: torch.nn.Module) -> int: + """Return the model's actual output vocabulary width.""" + output_embeddings = model.get_output_embeddings() + if output_embeddings is None or not hasattr(output_embeddings, "weight"): + raise ValueError("Model has no output embedding weight") + vocab_size = int(output_embeddings.weight.shape[0]) + if vocab_size <= 0: + raise ValueError(f"Invalid vocabulary size: {vocab_size}") + return vocab_size diff --git a/extension/llm/runner/constants.h b/extension/llm/runner/constants.h index 77ea7c9fd91..c9d062174b8 100644 --- a/extension/llm/runner/constants.h +++ b/extension/llm/runner/constants.h @@ -16,6 +16,8 @@ inline constexpr auto kEosIds = "get_eos_ids"; inline constexpr auto kMaxSeqLen = "get_max_seq_len"; inline constexpr auto kMaxContextLen = "get_max_context_len"; inline constexpr auto kVocabSize = "get_vocab_size"; +inline constexpr auto kActivationDtype = "get_activation_dtype"; +inline constexpr auto kLogitsToKeepMode = "get_logits_to_keep_mode"; inline constexpr auto kUseKVCache = "use_kv_cache"; inline constexpr auto kUseSDPAWithKVCache = "use_sdpa_with_kv_cache"; diff --git a/extension/llm/runner/model_metadata.h b/extension/llm/runner/model_metadata.h new file mode 100644 index 00000000000..43f5d396ee3 --- /dev/null +++ b/extension/llm/runner/model_metadata.h @@ -0,0 +1,184 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +enum class LogitsToKeepMode : std::int64_t { + Full = 0, + Last = 1, + Selected = 2, +}; + +// Readers for the metadata a program publishes as constant methods. Each +// required field's reader returns Error::InvalidProgram (and logs which method) +// if it is absent or malformed (non-positive size, unknown enum). A genuinely +// optional field added later should read through detail::read_int_method +// (nullopt when absent) so old programs stay readable; there are none today. +// vocab_size is returned as published (int64); check_vocab_size() narrows it to +// int32 after cross-checking the forward output. + +namespace detail { + +// Read a named int constant method: nullopt if absent, the value if present, +// Error::InvalidProgram if it does not evaluate to a single int. +inline runtime::Result> read_int_method( + Module& module, + const char* name) { + const auto names = ET_UNWRAP(module.method_names()); + if (names.count(name) == 0) { + return std::optional{}; + } + const auto result = module.execute(name); + if (!result.ok()) { + return result.error(); + } + ET_CHECK_OR_RETURN_ERROR( + result->size() == 1 && result->at(0).isInt(), + InvalidProgram, + "metadata %s must evaluate to a single int", + name); + return std::optional{result->at(0).toInt()}; +} + +// A required int constant that must be present and positive. +inline runtime::Result read_required_positive_int( + Module& module, + const char* name) { + const auto value = ET_UNWRAP(read_int_method(module, name)); + ET_CHECK_OR_RETURN_ERROR( + value.has_value(), InvalidProgram, "metadata %s is required", name); + ET_CHECK_OR_RETURN_ERROR( + *value > 0, + InvalidProgram, + "metadata %s must be positive, got %" PRId64, + name, + *value); + return *value; +} + +} // namespace detail + +// One reader per constant: the name, its encoding, and its validation together. +// Each rejection logs which constant method was at fault. + +inline runtime::Result read_max_context_length(Module& module) { + return detail::read_required_positive_int(module, kMaxContextLen); +} + +inline runtime::Result read_vocab_size(Module& module) { + return detail::read_required_positive_int(module, kVocabSize); +} + +inline runtime::Result read_activation_dtype(Module& module) { + const auto value = + ET_UNWRAP(detail::read_int_method(module, kActivationDtype)); + ET_CHECK_OR_RETURN_ERROR( + value.has_value(), + InvalidProgram, + "metadata %s is required", + kActivationDtype); + switch (*value) { + case static_cast(aten::ScalarType::Half): + return aten::ScalarType::Half; + case static_cast(aten::ScalarType::Float): + return aten::ScalarType::Float; + case static_cast(aten::ScalarType::BFloat16): + return aten::ScalarType::BFloat16; + default: + ET_LOG( + Error, + "metadata %s has unsupported value %" PRId64, + kActivationDtype, + *value); + return runtime::Error::InvalidProgram; + } +} + +inline runtime::Result read_logits_to_keep_mode( + Module& module) { + const auto value = + ET_UNWRAP(detail::read_int_method(module, kLogitsToKeepMode)); + ET_CHECK_OR_RETURN_ERROR( + value.has_value(), + InvalidProgram, + "metadata %s is required", + kLogitsToKeepMode); + switch (*value) { + case static_cast(LogitsToKeepMode::Full): + return LogitsToKeepMode::Full; + case static_cast(LogitsToKeepMode::Last): + return LogitsToKeepMode::Last; + case static_cast(LogitsToKeepMode::Selected): + return LogitsToKeepMode::Selected; + default: + ET_LOG( + Error, + "metadata %s has unsupported value %" PRId64, + kLogitsToKeepMode, + *value); + return runtime::Error::InvalidProgram; + } +} + +inline runtime::Result read_max_seq_len(Module& module) { + const auto value = + ET_UNWRAP(detail::read_required_positive_int(module, kMaxSeqLen)); + // Consumers narrow this to int for chunked prefill, so reject a value that + // would truncate rather than letting the cast overflow. + ET_CHECK_OR_RETURN_ERROR( + value <= std::numeric_limits::max(), + InvalidProgram, + "metadata %s %" PRId64 " exceeds the maximum forward step %d", + kMaxSeqLen, + value, + std::numeric_limits::max()); + return value; +} + +// Check the published vocab size against the model's actual forward output +// width: reject a disagreement or an out-of-int32 width, and hand back the +// (now int32) value the sampler takes. +inline runtime::Result check_vocab_size( + std::int64_t published_vocab_size, + std::int64_t output_vocab_size) { + ET_CHECK_OR_RETURN_ERROR( + output_vocab_size > 0 && + output_vocab_size <= std::numeric_limits::max(), + InvalidProgram, + "forward output vocab width %" PRId64 " is out of range", + output_vocab_size); + ET_CHECK_OR_RETURN_ERROR( + published_vocab_size == output_vocab_size, + InvalidProgram, + "published %s %" PRId64 " disagrees with forward output width %" PRId64, + kVocabSize, + published_vocab_size, + output_vocab_size); + // Equal to output_vocab_size, already checked to fit int32. + return static_cast(published_vocab_size); +} + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/targets.bzl b/extension/llm/runner/targets.bzl index 9af2597b4f2..c66bec170fe 100644 --- a/extension/llm/runner/targets.bzl +++ b/extension/llm/runner/targets.bzl @@ -35,6 +35,7 @@ def define_common_targets(): runtime.cxx_library( name = "stats" + aten_suffix, exported_headers = [ + "model_metadata.h", "stats.h", "util.h", ], diff --git a/extension/llm/runner/test/CMakeLists.txt b/extension/llm/runner/test/CMakeLists.txt index 942cd8d75ec..d43b297fa52 100644 --- a/extension/llm/runner/test/CMakeLists.txt +++ b/extension/llm/runner/test/CMakeLists.txt @@ -19,6 +19,7 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs test_generation_config.cpp + test_model_metadata.cpp test_text_llm_runner.cpp test_text_prefiller.cpp test_text_decoder_runner.cpp @@ -33,9 +34,42 @@ if(APPLE) list(APPEND _test_srcs lsan_stub.cpp) endif() +set(_metadata_pte_files + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_full.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_last.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_selected.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_context.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_prefill.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_vocab.pte" + "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_missing.pte" +) +add_custom_command( + OUTPUT ${_metadata_pte_files} + # Run against the installed executorch (which stages exir/_serialize/*.fbs at + # install time). Do NOT prepend the source tree to PYTHONPATH: it shadows the + # installed package with a source copy that lacks the staged .fbs, breaking + # to_executorch() under a non-editable install. + COMMAND + ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py + --outdir ${CMAKE_CURRENT_BINARY_DIR} + WORKING_DIRECTORY ${EXECUTORCH_ROOT} + DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py + ${EXECUTORCH_ROOT}/extension/llm/export/model_metadata.py +) +add_custom_target( + generated_model_metadata_test_files DEPENDS ${_metadata_pte_files} +) + et_cxx_test( test_runner SOURCES ${_test_srcs} EXTRA_LIBS executorch extension_llm_runner ) +add_dependencies(test_runner generated_model_metadata_test_files) +set_property( + TEST test_runner + PROPERTY + ENVIRONMENT + "ET_MODEL_METADATA_FULL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_full.pte;ET_MODEL_METADATA_LAST_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_last.pte;ET_MODEL_METADATA_SELECTED_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_selected.pte;ET_MODEL_METADATA_INVALID_CONTEXT_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_context.pte;ET_MODEL_METADATA_INVALID_PREFILL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_prefill.pte;ET_MODEL_METADATA_INVALID_VOCAB_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_vocab.pte;ET_MODEL_METADATA_MISSING_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_missing.pte" +) # Override sanitizer to this issue: # https://github.com/abseil/abseil-cpp/issues/841 Root issue: diff --git a/extension/llm/runner/test/export_model_metadata.py b/extension/llm/runner/test/export_model_metadata.py new file mode 100644 index 00000000000..71e0255aa6e --- /dev/null +++ b/extension/llm/runner/test/export_model_metadata.py @@ -0,0 +1,78 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import argparse +from pathlib import Path + +import torch + +from executorch.exir import to_edge +from executorch.extension.llm.export.model_metadata import ( + write_activation_dtype, + write_logits_to_keep_mode, + write_max_context_len, + write_max_seq_len, + write_vocab_size, +) + + +class Identity(torch.nn.Module): + """Minimal model used to serialize metadata fixtures.""" + + def forward(self, value: torch.Tensor) -> torch.Tensor: + """Return the input unchanged.""" + return value + + +def all_methods(logits_to_keep: str, activation_dtype: str) -> dict[str, int]: + """Compose the full metadata set from the individual per-constant writers.""" + return { + **write_max_context_len(4096), + **write_vocab_size(128256), + **write_activation_dtype(activation_dtype), + **write_logits_to_keep_mode(logits_to_keep), + **write_max_seq_len(512), + } + + +def main() -> None: + """Generate model metadata PTE fixtures.""" + parser = argparse.ArgumentParser() + parser.add_argument("--outdir", required=True) + args = parser.parse_args() + + output_dir = Path(args.outdir) + output_dir.mkdir(parents=True, exist_ok=True) + exported = torch.export.export(Identity(), (torch.ones(1),), strict=True) + for mode, dtype in ( + ("full", "fp32"), + ("last", "fp16"), + ("selected", "bf16"), + ): + program = to_edge( + exported, + constant_methods=all_methods(mode, dtype), + ).to_executorch() + (output_dir / f"ModelMetadata_{mode}.pte").write_bytes(program.buffer) + + for name, invalid_field in ( + ("invalid_context", "get_max_context_len"), + ("invalid_prefill", "get_max_seq_len"), + ("invalid_vocab", "get_vocab_size"), + ): + methods = all_methods("full", "fp32") + methods[invalid_field] = 0 + program = to_edge(exported, constant_methods=methods).to_executorch() + (output_dir / f"ModelMetadata_{name}.pte").write_bytes(program.buffer) + + # No constant methods: exercises rejection of missing required fields. + missing_program = to_edge(exported).to_executorch() + (output_dir / "ModelMetadata_missing.pte").write_bytes(missing_program.buffer) + + +if __name__ == "__main__": + main() diff --git a/extension/llm/runner/test/test_model_metadata.cpp b/extension/llm/runner/test/test_model_metadata.cpp new file mode 100644 index 00000000000..256f915bd14 --- /dev/null +++ b/extension/llm/runner/test/test_model_metadata.cpp @@ -0,0 +1,126 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +#include +#include + +#include + +namespace { + +using ::executorch::extension::Module; +using ::executorch::extension::llm::check_vocab_size; +using ::executorch::extension::llm::LogitsToKeepMode; +using ::executorch::extension::llm::read_activation_dtype; +using ::executorch::extension::llm::read_logits_to_keep_mode; +using ::executorch::extension::llm::read_max_context_length; +using ::executorch::extension::llm::read_max_seq_len; +using ::executorch::extension::llm::read_vocab_size; +using ::executorch::runtime::Error; + +std::unique_ptr load_fixture(const char* environment_variable) { + const char* path = std::getenv(environment_variable); + EXPECT_NE(path, nullptr); + auto module = std::make_unique(path); + EXPECT_EQ(module->load(), Error::Ok); + return module; +} + +struct ModeCase { + const char* environment_variable; + LogitsToKeepMode expected_mode; + ::executorch::aten::ScalarType expected_dtype; +}; + +class ModelMetadataTest : public ::testing::TestWithParam {}; + +TEST_P(ModelMetadataTest, ReadsPythonExportedConstants) { + auto module = load_fixture(GetParam().environment_variable); + + const auto max_context_length = read_max_context_length(*module); + ASSERT_TRUE(max_context_length.ok()); + EXPECT_EQ(*max_context_length, 4096); + + const auto max_seq_len = read_max_seq_len(*module); + ASSERT_TRUE(max_seq_len.ok()); + EXPECT_EQ(*max_seq_len, 512); + + const auto vocab_size = read_vocab_size(*module); + ASSERT_TRUE(vocab_size.ok()); + EXPECT_EQ(*vocab_size, 128256); + + const auto activation_dtype = read_activation_dtype(*module); + ASSERT_TRUE(activation_dtype.ok()); + EXPECT_EQ(*activation_dtype, GetParam().expected_dtype); + + const auto logits_to_keep_mode = read_logits_to_keep_mode(*module); + ASSERT_TRUE(logits_to_keep_mode.ok()); + EXPECT_EQ(*logits_to_keep_mode, GetParam().expected_mode); +} + +TEST(ModelMetadataTest, RejectsNonPositiveSizes) { + { + auto module = load_fixture("ET_MODEL_METADATA_INVALID_CONTEXT_PATH"); + const auto value = read_max_context_length(*module); + ASSERT_FALSE(value.ok()); + EXPECT_EQ(value.error(), Error::InvalidProgram); + } + { + auto module = load_fixture("ET_MODEL_METADATA_INVALID_PREFILL_PATH"); + const auto value = read_max_seq_len(*module); + ASSERT_FALSE(value.ok()); + EXPECT_EQ(value.error(), Error::InvalidProgram); + } + { + auto module = load_fixture("ET_MODEL_METADATA_INVALID_VOCAB_PATH"); + const auto value = read_vocab_size(*module); + ASSERT_FALSE(value.ok()); + EXPECT_EQ(value.error(), Error::InvalidProgram); + } +} + +TEST(ModelMetadataTest, ChecksVocabAgainstForwardOutput) { + auto matched = check_vocab_size(128256, 128256); + ASSERT_TRUE(matched.ok()); + EXPECT_EQ(*matched, 128256); + + auto mismatched = check_vocab_size(128256, 128000); + ASSERT_FALSE(mismatched.ok()); + EXPECT_EQ(mismatched.error(), Error::InvalidProgram); +} + +TEST(ModelMetadataTest, RejectsMissingRequiredFields) { + auto module = load_fixture("ET_MODEL_METADATA_MISSING_PATH"); + EXPECT_FALSE(read_max_context_length(*module).ok()); + EXPECT_FALSE(read_max_seq_len(*module).ok()); + EXPECT_FALSE(read_vocab_size(*module).ok()); + EXPECT_FALSE(read_activation_dtype(*module).ok()); + EXPECT_FALSE(read_logits_to_keep_mode(*module).ok()); +} + +INSTANTIATE_TEST_SUITE_P( + RoundTrip, + ModelMetadataTest, + ::testing::Values( + ModeCase{ + "ET_MODEL_METADATA_FULL_PATH", + LogitsToKeepMode::Full, + ::executorch::aten::ScalarType::Float}, + ModeCase{ + "ET_MODEL_METADATA_LAST_PATH", + LogitsToKeepMode::Last, + ::executorch::aten::ScalarType::Half}, + ModeCase{ + "ET_MODEL_METADATA_SELECTED_PATH", + LogitsToKeepMode::Selected, + ::executorch::aten::ScalarType::BFloat16})); + +} // namespace From 03f41d2031b9a3a21e105c1c45049ba7c9ffa313 Mon Sep 17 00:00:00 2001 From: Ivan Xu Date: Thu, 10 Sep 2026 14:51:48 -0700 Subject: [PATCH 159/190] Vulkan backend: Add missing include of algorithm ## Problem `backends/vulkan/runtime/graph/containers/SharedObject.cpp` and `backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp` call `std::find` (and `Squeeze.cpp` also `std::rotate`) without including ``. This compiled previously because the declarations leaked in transitively through other headers, but GCC 15's libstdc++ header reorganization removed that transitive path, breaking the build with: ``` error: no matching function for call to 'find(...)' note: candidate 1: ... std::find(istreambuf_iterator<...>, ...) ``` The remaining candidate is a narrow `istreambuf_iterator`-only overload pulled in via an unrelated header (``/`` chain), not the generic `` overload these call sites actually need. This is the same class of issue as #15220 (missing `` include in the Vulkan backend) and relates to the GCC 15 CI build failures that led to disabling `executorch-ubuntu-26.04-gcc15` in favor of gcc14 (#20304, tracked in #19917). ## Fix Add `#include ` to both files, matching the precedent set in `backends/vulkan/runtime/utils/StorageUtils.h` (#15220). ## Test Plan Built the `vulkan_backend` CMake target with GCC 15 (`gcc (Ubuntu 15.x)`), which fails to compile on `main` today with the errors above. With this fix, `cmake --build cmake-out --target vulkan_backend` completes cleanly. Swept the rest of `backends/vulkan/runtime` for the same pattern (`std::find`/`sort`/`rotate`/etc. without ``) and found no other occurrences outside these two files. cc @SS-JIA @manuelcandales @digantdesai @cbilgin --- backends/vulkan/runtime/graph/containers/SharedObject.cpp | 2 ++ backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp | 2 ++ 2 files changed, 4 insertions(+) diff --git a/backends/vulkan/runtime/graph/containers/SharedObject.cpp b/backends/vulkan/runtime/graph/containers/SharedObject.cpp index 10ddd6f2ca3..7ea59d59e2e 100644 --- a/backends/vulkan/runtime/graph/containers/SharedObject.cpp +++ b/backends/vulkan/runtime/graph/containers/SharedObject.cpp @@ -10,6 +10,8 @@ #include +#include + namespace vkcompute { bool SharedObject::has_user(const ValueRef idx) const { diff --git a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp index a78456538bf..bec078564b4 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Squeeze.cpp @@ -14,6 +14,8 @@ #include #include +#include + namespace vkcompute { void add_squeeze_copy_dims_node( From 1d87fa9d9209afd59a8f9c89743ffe40acac4ccb Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Fri, 11 Sep 2026 00:23:56 +0200 Subject: [PATCH 160/190] [ET-VK] Serialize non-finite floats using flatc's spelling (#22307) ### Summary Fixes #22305. The Vulkan graph is serialized by dumping it to JSON with Python's `json` module and handing that to `flatc`. The two disagree on how to spell the non-finite floats: Python emits `Infinity` / `-Infinity` / `NaN`, and `flatc` accepts none of them. Any model whose graph carries a non-finite scalar therefore fails to lower, with the `-inf` fill value of a transformer attention mask being the common source. `all-MiniLM-L6-v2` and CLIP ViT-B/32 both partition cleanly and then die with: ``` schema.json:1: 6012: error: cannot parse value starting with: - ``` which names a byte offset inside a deleted temporary file and points at nothing in the model. `flatc` does accept `inf` and `-inf` (checked against flatc 24.3.25; `nan`, `NaN`, `infinity` are all rejected), so the infinities round-trip exactly. The decompile direction needs the inverse rewrite, since `flatc` writes bare `inf` tokens that `json.load` will not parse; string literals are stepped over so a shader or key name containing "inf" is left alone. FlatBuffers JSON has no spelling for NaN at all, so that case raises with a message naming the problem rather than emitting JSON `flatc` cannot read. ### Test plan Three new cases in `backends/vulkan/test/test_serialization.py`: an exact round-trip of `-inf` / `inf` / a finite value, the NaN error, and a string literal containing "inf" surviving the decode rewrite untouched. ``` pytest backends/vulkan/test/test_serialization.py 7 passed ``` Verified end to end that `sentence-transformers/all-MiniLM-L6-v2` now lowers through `VulkanPartitioner` (14 delegate blobs, 3/3 runs at `PYTHONHASHSEED=0`), where before it failed 3/3. cc @SS-JIA @manuelcandales @digantdesai @cbilgin --- .../serialization/vulkan_graph_serialize.py | 70 ++++++++++++++++- backends/vulkan/test/test_serialization.py | 76 +++++++++++++++++++ 2 files changed, 144 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/serialization/vulkan_graph_serialize.py b/backends/vulkan/serialization/vulkan_graph_serialize.py index 96f944560a8..81de183021b 100644 --- a/backends/vulkan/serialization/vulkan_graph_serialize.py +++ b/backends/vulkan/serialization/vulkan_graph_serialize.py @@ -11,6 +11,7 @@ import importlib.resources as _resources import json import os +import re import tempfile from dataclasses import dataclass from typing import ClassVar, List @@ -27,8 +28,72 @@ from executorch.exir._serialize._flatbuffer import _flatc_compile, _flatc_decompile +# Python's json module spells the non-finite floats "Infinity" / "-Infinity" / +# "NaN"; flatc spells the infinities "inf" / "-inf" and rejects Python's +# spelling, so both directions need translating. A graph carries a non-finite +# scalar whenever the model does -- the -inf fill value of a transformer +# attention mask is the common case -- and without this the failure surfaces as +# a flatc byte offset into a temporary file rather than anything pointing at +# the graph. +# +# The rewrite runs over the serialized text rather than over the encoder's +# chunks: json only emits a float as a chunk of its own inside an object, and +# inside a list the chunk carries the delimiter with it ("[-Infinity"), so +# matching whole chunks silently missed every DoubleList. +_JSON_STRING_RE = re.compile(r'"(?:[^"\\]|\\.)*"') +_PY_NONFINITE_RE = re.compile(r"(? str: + """Apply ``sub`` to everything in ``text`` that is not a JSON string. + + String literals are copied through untouched, so a shader name or a string + value that happens to read "inf" is never rewritten. + """ + out = [] + last = 0 + for m in _JSON_STRING_RE.finditer(text): + out.append(sub(text[last : m.start()])) + out.append(m.group(0)) + last = m.end() + out.append(sub(text[last:])) + return "".join(out) + + +def _python_json_to_flatc_json(text: str) -> str: + """Rewrite json's ``Infinity`` tokens into the ``inf`` flatc accepts.""" + if "Infinity" not in text and "NaN" not in text: + return text + + def replace(m: "re.Match[str]") -> str: + if m.group(1) == "NaN": + raise ValueError( + "Cannot serialize a NaN float value into a Vulkan graph: " + "flatc rejects every spelling of NaN for a value inside a " + "union, and every float in the Vulkan schema is a member of " + "the VkValue union." + ) + return "inf" + + return _rewrite_outside_strings( + text, lambda segment: _PY_NONFINITE_RE.sub(replace, segment) + ) + + +def _flatc_json_to_python_json(text: str) -> str: + """Rewrite flatc's bare ``inf`` tokens so json.loads accepts them.""" + if "inf" not in text: + return text + return _rewrite_outside_strings( + text, lambda segment: _FLATC_INF_RE.sub("Infinity", segment) + ) + + def convert_to_flatbuffer(vk_graph: VkGraph) -> bytes: - vk_graph_json = json.dumps(vk_graph, cls=_DataclassEncoder) + vk_graph_json = _python_json_to_flatc_json( + json.dumps(vk_graph, cls=_DataclassEncoder) + ) with tempfile.TemporaryDirectory() as d: schema_path = os.path.join(d, "schema.fbs") @@ -63,7 +128,8 @@ def flatbuffer_to_vk_graph(flatbuffers: bytes) -> VkGraph: json_path = os.path.join(d, "schema.json") with open(json_path, "rb") as output_file: - return _json_to_dataclass(json.load(output_file), VkGraph) + raw = output_file.read().decode("utf-8") + return _json_to_dataclass(json.loads(_flatc_json_to_python_json(raw)), VkGraph) def extract_vk_flatbuffer(data: bytes) -> bytes: diff --git a/backends/vulkan/test/test_serialization.py b/backends/vulkan/test/test_serialization.py index 71a6980635a..513150ba9fe 100644 --- a/backends/vulkan/test/test_serialization.py +++ b/backends/vulkan/test/test_serialization.py @@ -19,6 +19,8 @@ ) from executorch.backends.vulkan.serialization.vulkan_graph_schema import ( + Double, + DoubleList, IntList, OperatorCall, String, @@ -269,3 +271,77 @@ def test_serialize_deserialize_vkgraph(self) -> None: out_vk_graph = flatbuffer_to_vk_graph(bs) self.assertEqual(in_vk_graph, out_vk_graph) + + def _round_trip(self, values, chain=None) -> VkGraph: + in_vk_graph = VkGraph( + version="1", + chain=chain if chain is not None else [], + values=values, + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + out_vk_graph = flatbuffer_to_vk_graph(convert_to_flatbuffer(in_vk_graph)) + self.assertEqual(in_vk_graph, out_vk_graph) + return out_vk_graph + + def test_serialize_deserialize_non_finite_scalars(self) -> None: + # Python's json module spells the infinities "Infinity" / "-Infinity" + # while flatc spells them "inf" / "-inf" and rejects Python's spelling, + # so both directions need translating. A graph picks up a non-finite + # scalar whenever the model has one -- the -inf fill value of a + # transformer attention mask being the usual source. + self._round_trip( + [ + VkValue(value=Double(double_val=float("-inf"))), + VkValue(value=Double(double_val=float("inf"))), + VkValue(value=Double(double_val=1.5)), + ] + ) + + def test_serialize_deserialize_non_finite_floats_in_list(self) -> None: + # json only emits a float as a chunk of its own inside an object; in a + # list the chunk carries the delimiter with it, so a rewrite that works + # on the scalar above can still miss every element of a DoubleList. + self._round_trip( + [ + VkValue(value=DoubleList(items=[float("-inf")])), + VkValue( + value=DoubleList(items=[1.5, float("inf"), 2.5, float("-inf")]) + ), + VkValue(value=DoubleList(items=[])), + ] + ) + + def test_serialize_nan_float_raises(self) -> None: + # flatc rejects nan, NaN and Nan alike for a value inside a union, and + # every float in the Vulkan schema is a member of the VkValue union, so + # report it here rather than emitting JSON that flatc cannot read. + for value in ( + Double(double_val=float("nan")), + DoubleList(items=[1.0, float("nan")]), + ): + vk_graph = VkGraph( + version="1", + chain=[], + values=[VkValue(value=value)], + input_ids=[], + output_ids=[], + constants=[], + shaders=[], + ) + with self.assertRaisesRegex(ValueError, "NaN"): + convert_to_flatbuffer(vk_graph) + + def test_serialize_deserialize_leaves_strings_alone(self) -> None: + # The token rewrites run over the serialized JSON, so they must not + # reach into string literals in either direction. + self._round_trip( + [ + VkValue(value=String(string_val="value: inf, -inf")), + VkValue(value=String(string_val="Infinity NaN nan")), + VkValue(value=String(string_val='quoted "inf" and \\ inf')), + ], + chain=[OperatorCall(node_id=1, name="inf_shader", args=[])], + ) From 2698ba313e4f72f84f77151780c4287d3a6d0af9 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 9 Sep 2026 20:37:29 -0700 Subject: [PATCH 161/190] [ET-VK][q8ta-conv] Optimize narrow q8ta convolution workgroups Use narrow local workgroups for regular and depthwise q8ta convolutions whose output width underfills existing workgroups. This reduces inactive invocations in SceneX small-spatial layers. Authored with Codex. Differential Revision: [D119274189](https://our.internmc.facebook.com/intern/diff/D119274189/) ghstack-source-id: 427402958 Pull-Request: https://github.com/pytorch/executorch/pull/22667 --- .../runtime/graph/ops/impl/Q8taConv2d.cpp | 19 ++++- .../runtime/graph/ops/impl/Q8taConv2dDW.cpp | 20 +++++ .../test/custom_ops/test_q8ta_conv2d.cpp | 73 ++++++++++++++--- .../test/custom_ops/test_q8ta_conv2d_dw.cpp | 79 +++++++++++++++++-- 4 files changed, 170 insertions(+), 21 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp index 70171a04820..262858b76fb 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp @@ -111,6 +111,7 @@ GlobalWorkGrid pick_q8ta_conv2d_gwg( (void)shader; (void)resize_args; + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); const uint32_t W = graph->size_at(-1, output); @@ -133,6 +134,7 @@ GlobalWorkGrid pick_q8ta_conv2d_gwg( * tensor dimensions. Uses experimentation results: * - {4, 2, 8} for medium tensors: +57% improvement on 81x81 * - {8, 1, 8} for very large tensors: best baseline performance + * - {2, 1, 32} or {4, 1, 16} for narrow output widths * - {64, 1, 1} for narrow channel dimensions: minimize inactive invocations */ LocalWorkGroup pick_q8ta_conv2d_lwg( @@ -144,14 +146,13 @@ LocalWorkGroup pick_q8ta_conv2d_lwg( (void)shader; (void)resize_args; + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); - - // Get actual tensor dimensions for adaptive sizing - const uint32_t H = graph->size_at(-2, output); + const uint32_t output_height = graph->size_at(-2, output); // For very large tensors (H >= 100 and large x/z), use {8, 1, 8} // This configuration performed best for 128x128 tensors in experiments - if (H >= 100 && gwg[0u] >= 24 && gwg[2u] >= 24) { + if (output_height >= 100 && gwg[0u] >= 24 && gwg[2u] >= 24) { return LocalWorkGroup(8u, 1u, 8u); } @@ -161,6 +162,16 @@ LocalWorkGroup pick_q8ta_conv2d_lwg( return LocalWorkGroup(4u, 2u, 8u); } + if (gwg[0u] == 2u && gwg[2u] >= 32u) { + return LocalWorkGroup(2u, 1u, 32u); + } + + // LWG x oversubscribes the 3 global groups here; safe only because the + // shader early-returns out-of-bounds invocations. + if (gwg[0u] == 3u && gwg[2u] >= 16u) { + return LocalWorkGroup(4u, 1u, 16u); + } + // For tensors with sufficient x and z dimensions, use square configuration if (gwg[0u] >= 6 && gwg[2u] >= 6) { return LocalWorkGroup(8u, 1u, 8u); diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp index 182e8d684d2..a504d9ab3f6 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dDW.cpp @@ -29,6 +29,7 @@ GlobalWorkGrid pick_q8ta_conv2d_dw_gwg( (void)shader; (void)resize_args; + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); const uint32_t W = graph->size_at(-1, output); @@ -46,6 +47,15 @@ GlobalWorkGrid pick_q8ta_conv2d_dw_gwg( kTiledWorkGrid); } +/** + * Picks a local workgroup size for q8ta_conv2d_dw with adaptive sizing based + * on tensor dimensions. Uses experimentation results: + * - {2, 1, 32} or {4, 1, 16} for narrow output widths + * + * Unlike the regular conv picker, there is no medium-tensor branch shadowing + * gwg[0] == 4, so the second narrow branch matches 3..4 (the conv picker's + * {4, 2, 8} branch claims gwg[0] >= 4 first, leaving only == 3 reachable). + */ LocalWorkGroup pick_q8ta_conv2d_dw_lwg( ComputeGraph* graph, const vkapi::ShaderInfo& shader, @@ -57,6 +67,16 @@ LocalWorkGroup pick_q8ta_conv2d_dw_lwg( (void)args; (void)resize_args; + if (gwg[0u] == 2u && gwg[2u] >= 32u) { + return LocalWorkGroup(2u, 1u, 32u); + } + + // LWG x oversubscribes when gwg[0] is 3; safe only because the shader + // early-returns out-of-bounds invocations. + if (gwg[0u] >= 3u && gwg[0u] <= 4u && gwg[2u] >= 16u) { + return LocalWorkGroup(4u, 1u, 16u); + } + // Some inactive invocations are okay; set 6 as the threshold to use the // a square wg size. if (gwg[0u] >= 6 && gwg[2u] >= 6) { diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index 308f2eac86a..51da65da57c 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -279,6 +280,40 @@ static TestCase create_test_case_from_config( return test_case; } +static std::vector generate_narrow_workgroup_test_cases() { + std::vector test_cases; + std::vector configs = { + {OutInChannels(64, 32), + InputSize2D(9, 9), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + {OutInChannels(128, 32), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1}, + }; + + for (auto& config : configs) { + const bool is_performance = config.channels.out > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C, + /*impl_selector=*/"general")); + } + return test_cases; +} + // Generate easy test cases for quantized conv2d operation (for debugging) std::vector generate_quantized_conv2d_easy_cases() { std::vector test_cases; @@ -627,6 +662,12 @@ static std::vector generate_quantized_conv2d_test_cases() { } } + auto narrow_workgroup_cases = generate_narrow_workgroup_test_cases(); + test_cases.insert( + test_cases.end(), + narrow_workgroup_cases.begin(), + narrow_workgroup_cases.end()); + return test_cases; } @@ -1037,6 +1078,7 @@ int main(int argc, char* argv[]) { !can_use_unsigned_pw_dot(adapter, kMaxUnsignedDotAccumulatorBytes + 1)); std::string im2col_impl_selector; + bool narrow_workgroups_only = false; for (int i = 1; i < argc; ++i) { const std::string arg(argv[i]); if (arg == "--im2col-path=signed") { @@ -1045,11 +1087,19 @@ int main(int argc, char* argv[]) { im2col_impl_selector = "im2col_unsigned"; } else if (arg == "--im2col-path=auto") { im2col_impl_selector = "im2col_auto"; + } else if (arg == "--narrow-workgroups-only") { + narrow_workgroups_only = true; } else { std::cerr << "Unknown argument: " << arg << std::endl; return 2; } } + if (narrow_workgroups_only && !im2col_impl_selector.empty()) { + std::cerr + << "Narrow-workgroup and im2col path selectors are mutually exclusive" + << std::endl; + return 2; + } set_debugging(false); set_print_output(false); #ifdef DEBUG_MODE @@ -1067,18 +1117,23 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; - // Execute test cases using the new framework with custom FLOP calculator - const auto test_case_generator = [im2col_impl_selector]() { - return im2col_impl_selector.empty() - ? generate_quantized_conv2d_test_cases() - : generate_im2col_unsigned_test_cases(im2col_impl_selector); - }; - auto results = execute_test_cases( #ifdef DEBUG_MODE - generate_quantized_conv2d_easy_cases, + std::function()> test_case_generator = + generate_quantized_conv2d_easy_cases; #else - test_case_generator, + std::function()> test_case_generator = + [im2col_impl_selector]() { + return im2col_impl_selector.empty() + ? generate_quantized_conv2d_test_cases() + : generate_im2col_unsigned_test_cases(im2col_impl_selector); + }; #endif + if (narrow_workgroups_only) { + test_case_generator = generate_narrow_workgroup_test_cases; + } + + auto results = execute_test_cases( + test_case_generator, quantized_conv2d_flop_calculator, "QuantizedConv2dQ8ToQ8To", /*warmup_runs = */ 1, diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp index 2dbb4909adb..c573b8b5d39 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d_dw.cpp @@ -22,6 +22,7 @@ using namespace executorch::vulkan::prototyping; using namespace vkcompute; static constexpr int64_t kRefDimSizeLimit = 100; +static constexpr int64_t kRefOperationLimit = 2 * 1024 * 1024; // Utility function to create a test case from a Conv2dConfig for depthwise // convolution @@ -271,6 +272,43 @@ std::vector generate_quantized_conv2d_dw_easy_cases() { return test_cases; } +std::vector generate_quantized_conv2d_dw_narrow_workgroup_cases() { + std::vector test_cases; + std::vector configs = { + {OutInChannels(128, 128), + InputSize2D(7, 7), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 128}, + {OutInChannels(64, 64), + InputSize2D(9, 9), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 64}, + {OutInChannels(64, 64), + InputSize2D(13, 13), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 64}, + }; + + for (auto& config : configs) { + const bool is_performance = config.channels.out > kRefDimSizeLimit; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = make_test_case_name( + config, is_performance, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + config, vkapi::kFloat, utils::kTexture3D, utils::kPackedInt8_4C)); + } + return test_cases; +} + // Generate test cases for quantized depthwise conv2d operation std::vector generate_quantized_conv2d_dw_test_cases() { std::vector test_cases; @@ -439,6 +477,13 @@ std::vector generate_quantized_conv2d_dw_test_cases() { } } + auto narrow_workgroup_cases = + generate_quantized_conv2d_dw_narrow_workgroup_cases(); + test_cases.insert( + test_cases.end(), + narrow_workgroup_cases.begin(), + narrow_workgroup_cases.end()); + return test_cases; } @@ -498,10 +543,14 @@ void conv2d_q8ta_q8csw_q8to_dw_reference_impl(TestCase& test_case) { int64_t dilation_w = dilation_data[1]; int64_t groups = groups_spec.get_int_value(); - // Skip for large tensors since computation time will be extremely slow - if (N > kRefDimSizeLimit || C_in > kRefDimSizeLimit || - H_in > kRefDimSizeLimit || W_in > kRefDimSizeLimit || - C_out > kRefDimSizeLimit) { + // Skip large tensors only when the reference would be expensive: each + // output element costs K_h * K_w MACs (one input channel per output + // channel), so large-dim cases with few total operations still validate. + const int64_t reference_operations = N * C_out * H_out * W_out * K_h * K_w; + const bool has_large_dimension = N > kRefDimSizeLimit || + C_in > kRefDimSizeLimit || H_in > kRefDimSizeLimit || + W_in > kRefDimSizeLimit || C_out > kRefDimSizeLimit; + if (has_large_dimension && reference_operations > kRefOperationLimit) { throw std::invalid_argument( "One or more dimensions exceed the allowed limit for reference implementation."); } @@ -680,13 +729,27 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; - // Execute test cases using the new framework with custom FLOP calculator - auto results = execute_test_cases( + bool narrow_workgroups_only = false; + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--narrow-workgroups-only") { + narrow_workgroups_only = true; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } + } #ifdef DEBUG_MODE - generate_quantized_conv2d_dw_easy_cases, + auto test_case_generator = generate_quantized_conv2d_dw_easy_cases; #else - generate_quantized_conv2d_dw_test_cases, + auto test_case_generator = generate_quantized_conv2d_dw_test_cases; #endif + if (narrow_workgroups_only) { + test_case_generator = generate_quantized_conv2d_dw_narrow_workgroup_cases; + } + + auto results = execute_test_cases( + test_case_generator, quantized_conv2d_dw_flop_calculator, "QuantizedDepthwiseInt8Conv2d", /*warmup_runs = */ 1, From d96e3f885e1a78859704de87860b3c6510e7c5cb Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 9 Sep 2026 20:37:34 -0700 Subject: [PATCH 162/190] [ET-VK][q8ta-conv] Bound q8ta im2col scratch memory Large or batched q8ta convolutions materialize the complete im2col tensor, which can require hundreds of MiB. Stream flattened batch/output-height rows through one reusable 16 MiB scratch buffer. Preserve the original one-dispatch path when full scratch fits, and fall back to direct convolution when one row cannot fit. The cap was selected from 4/8/16/32 MiB device sweeps. 16 MiB is the smallest cap without extra tiling for the representative 8.6 MiB operator and cuts Edits Saliency uint8 VMA allocation by at least 129 MiB. Consolidation (same commit): the full-fit case is the single-tile instance of the streaming path, not a separate path. The scratch is always `[1, K, rows_per_tile, align_up_4(W)]`, the full/streaming dispatch fork and `Q8taIm2ColMode` are gone, and the pointwise shader selects flat-tile vs batched-activation addressing with a `use_flat_tile` specialization constant instead of a codegen variant (8 -> 4 PW variants). The standalone 1x1 pointwise path is unchanged. Review hardening (same commit): row tiles are fixed at build time, so dynamic batch/height growth past the tiled rows has no dispatches. Each im2col/PW resize now fail-fasts against a build-time `max_im2col_rows` bound (`num_tiles * rows_per_tile`) instead of leaving outputs stale; shrinkage only ever lowers the total below the bound. Also fixes `-Wshadow` errors in `ComputeGraph.cpp` exposed by the new test target on Apple builds, and documents `ScopedAdapterCapabilityOverride` as not thread-safe with deleted move operations. Authored with Codex. Differential Revision: [D119274190](https://our.internmc.facebook.com/intern/diff/D119274190/) ghstack-source-id: 427402963 Pull-Request: https://github.com/pytorch/executorch/pull/22668 --- .../vulkan/runtime/graph/ComputeGraph.cpp | 14 +- .../graph/ops/glsl/q8ta_conv2d_pw.glsl | 34 ++- .../runtime/graph/ops/glsl/q8ta_im2col.glsl | 26 +- .../runtime/graph/ops/impl/Q8taConv2d.h | 41 ++- .../graph/ops/impl/Q8taConv2dIm2Col.cpp | 279 +++++++++++++----- .../runtime/graph/ops/impl/Q8taConv2dPW.cpp | 136 +++++++-- backends/vulkan/runtime/vk_api/Adapter.cpp | 11 + backends/vulkan/runtime/vk_api/Adapter.h | 16 + .../vk_api/AdapterCapabilityOverrides.h | 61 ++++ .../test/custom_ops/impl/TestQ8taConv2d.cpp | 28 +- .../q8ta_conv2d_stream_plan_test.cpp | 178 +++++++++++ backends/vulkan/test/custom_ops/targets.bzl | 10 + .../test/custom_ops/test_q8ta_conv2d.cpp | 265 ++++++++++++++++- 13 files changed, 958 insertions(+), 141 deletions(-) create mode 100644 backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h create mode 100644 backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp diff --git a/backends/vulkan/runtime/graph/ComputeGraph.cpp b/backends/vulkan/runtime/graph/ComputeGraph.cpp index 74407cb00bc..f84eac26d49 100644 --- a/backends/vulkan/runtime/graph/ComputeGraph.cpp +++ b/backends/vulkan/runtime/graph/ComputeGraph.cpp @@ -123,16 +123,16 @@ TmpTensor::~TmpTensor() { } int64_t TmpTensor::get_sobj_idx() { - int64_t sobj_idx; + int64_t idx; // If no available temporary shared objects, request a new one to be created if (graph_p->tmp_shared_object_idxs_.empty()) { - sobj_idx = graph_p->shared_objects_.size(); + idx = graph_p->shared_objects_.size(); } else { // Get the first available shared object idx - sobj_idx = graph_p->tmp_shared_object_idxs_.top(); + idx = graph_p->tmp_shared_object_idxs_.top(); graph_p->tmp_shared_object_idxs_.pop(); } - return sobj_idx; + return idx; } // @@ -1212,9 +1212,9 @@ void ComputeGraph::prepack() { shared_object.bind_users(this); } // Make sure all remaining tensors have allocations - for (int i = 0; i < values_.size(); i++) { - if (values_.at(i).isTensor()) { - create_dedicated_allocation_for(i); + for (int value_idx = 0; value_idx < values_.size(); value_idx++) { + if (values_.at(value_idx).isTensor()) { + create_dedicated_allocation_for(value_idx); } } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl index b68497a96be..6e258cbf2ac 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_conv2d_pw.glsl @@ -63,6 +63,7 @@ layout(push_constant) uniform restrict Block { int output_zp; int K4_per_group; int OC4_per_group; + int stream_row_offset; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -73,6 +74,9 @@ ${layout_declare_spec_const(C, "int", "activation_type", "0")} // Layout specialization constants ${layout_declare_spec_const(C, "int", "outp_layout", "CONTIG_LAYOUT_INT")} ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} +// Row-tile input (im2col scratch) vs batched activation input. Uniform per +// dispatch; declared last so existing constant ids are unchanged. +${layout_declare_spec_const(C, "int", "use_flat_tile", "0")} int compute_outp_buffer_idx( const int w_block_idx, @@ -107,9 +111,29 @@ void main() { const int W4 = div_up_4(int(outp.sizes[0][0])); const int H = int(outp.sizes[0][1]); const int OC4 = div_up_4(int(outp.sizes[0][2])); - const int hn = int(gl_GlobalInvocationID.z); - const int n = hn / H; - const int oh = hn % H; + const int local_row_idx = int(gl_GlobalInvocationID.z); + int n; + int oh; + int input_n; + int input_h; + if (use_flat_tile == 1) { + if (local_row_idx >= int(inp.sizes[0][1])) { + return; + } + const int global_row_idx = stream_row_offset + local_row_idx; + if (global_row_idx >= int(outp.sizes[0][3]) * H) { + return; + } + n = global_row_idx / H; + oh = global_row_idx % H; + input_n = 0; + input_h = local_row_idx; + } else { + n = local_row_idx / H; + oh = local_row_idx % H; + input_n = n; + input_h = oh; + } // Bounds check in block space if (ow_block_idx >= W4 || @@ -152,8 +176,8 @@ void main() { // Compute initial input tile index with group offset // For grouped im2col, each group's K range starts at group_idx * K4_per_group // For non-grouped (groups=1), group_idx is always 0 so offset is 0 - int input_idx = n * inp_n_stride - + oh * inp_h_stride + int input_idx = input_n * inp_n_stride + + input_h * inp_h_stride + ow_block_idx * inp_w_stride + group_idx * K4_per_group; diff --git a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl index b58035f59a5..a278ce394d1 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/q8ta_im2col.glsl @@ -33,6 +33,7 @@ ${layout_declare_spec_const(C, "int", "inp_layout", "CONTIG_LAYOUT_INT")} layout(push_constant) uniform restrict Block { int zp; + int stream_row_offset; }; layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -79,24 +80,31 @@ void main() { const int im2col_W4 = div_up_4(im2col_sizes.x); const int im2col_H = im2col_sizes.y; const int im2col_Z4 = div_up_4(im2col_sizes.z); - const int im2col_N = im2col_sizes.w; - // im2col block index from linear output buffer index + // im2col block index from linear output buffer index. The scratch holds one + // row tile, so local rows are rebased onto the global batch*height rows. const int c4_idx = out_buf_idx % im2col_Z4; const int row = out_buf_idx / im2col_Z4; const int w4_idx = row % im2col_W4; const int hn_idx = row / im2col_W4; + const int local_row_idx = hn_idx; const int h_idx = hn_idx % im2col_H; - const int n_idx = hn_idx / im2col_H; + const int output_H = + (input_sizes.y + 2 * conv2d_params.padding.y - + conv2d_params.dilation.y * (conv2d_params.kernel_size.y - 1) - 1) / + conv2d_params.stride.y + + 1; + const int global_row_idx = stream_row_offset + local_row_idx; + const int n_idx = global_row_idx / output_H; // out of bounds check - if (w4_idx >= im2col_W4 || h_idx >= im2col_H || - c4_idx >= im2col_Z4 || n_idx >= im2col_N) { + if (w4_idx >= im2col_W4 || local_row_idx >= im2col_H || + c4_idx >= im2col_Z4 || n_idx >= input_sizes.w) { return; } const int im2col_w = mul_4(w4_idx); - const int im2col_h = h_idx; + const int im2col_h = global_row_idx % output_H; const int im2col_k = mul_4(c4_idx); const int group_idx = im2col_k / conv2d_params.K_per_group; @@ -167,9 +175,9 @@ void main() { input_Z4, zp_packed)); - // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1) - const int buffer_idx = n_idx * int(im2col_outp.strides[0][3]) - + h_idx * int(im2col_outp.strides[0][1]) + // store_packed_int8_output_tile (with TILE_M4=1, TILE_N4=1). The scratch + // has a single batch, so every tile writes from row 0 of the same buffer. + const int buffer_idx = h_idx * int(im2col_outp.strides[0][1]) + w4_idx * int(im2col_outp.strides[0][0]) + c4_idx; diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h index 690868e9f8c..e48d1e32a08 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.h @@ -14,6 +14,34 @@ namespace vkcompute { +inline constexpr int64_t kQ8taConv2dIm2ColScratchBudgetBytes = 16 * 1024 * 1024; +inline constexpr int64_t kQ8taConv2dMaxRowsPerTile = 65535; + +struct Q8taConv2dStreamPlan final { + int64_t aligned_out_width; + int64_t rows_per_tile; + int64_t num_tiles; + int64_t scratch_bytes; + bool feasible; +}; + +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan( + int64_t batch, + int64_t flattened_kernel_size, + int64_t out_height, + int64_t out_width, + int64_t scratch_budget_bytes); + +// max_buffer_bytes is Adapter::max_buffer_numel(), which returns +// maxStorageBufferRange in bytes (not elements) — directly comparable with +// the byte-denominated scratch budget. +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan_for_device( + int64_t batch, + int64_t flattened_kernel_size, + int64_t out_height, + int64_t out_width, + uint64_t max_buffer_bytes); + enum class ActivationType : uint32_t { kNone = 0, kRelu = 1, @@ -129,7 +157,10 @@ void add_q8ta_conv2d_pw_node( const ValueRef kernel_size = kDummyValueRef, const ValueRef stride = kDummyValueRef, const ValueRef padding = kDummyValueRef, - const ValueRef dilation = kDummyValueRef); + const ValueRef dilation = kDummyValueRef, + const bool is_im2col = false, + const ValueRef stream_row_offset_ref = kDummyValueRef, + const ValueRef max_im2col_rows_ref = kDummyValueRef); constexpr int64_t kMaxUnsignedDotAccumulatorBytes = 33025; @@ -161,7 +192,9 @@ void add_q8ta_im2col_node( const ValueRef groups, const ValueRef packed_int8_output, const ValueRef packed_int8_im2col, - const int32_t zp); + const int32_t zp, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref = kDummyValueRef); void q8ta_conv2d_im2col(ComputeGraph& graph, const std::vector& args); @@ -170,6 +203,10 @@ void q8ta_conv2d_im2col_impl( bool use_unsigned_dot, const std::vector& args); +void q8ta_conv2d_general( + ComputeGraph& graph, + const std::vector& args); + // Transposed convolution void q8ta_conv2d_transposed( diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp index 24538b8736a..062e5ea1d8d 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dIm2Col.cpp @@ -16,8 +16,68 @@ #include #include +#include +#include + namespace vkcompute { +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan( + const int64_t batch, + const int64_t flattened_kernel_size, + const int64_t out_height, + const int64_t out_width, + const int64_t scratch_budget_bytes) { + Q8taConv2dStreamPlan plan{}; + if (batch <= 0 || flattened_kernel_size <= 0 || out_height <= 0 || + out_width <= 0 || scratch_budget_bytes <= 0) { + return plan; + } + + constexpr int64_t kAlignment = 4; + if (out_width > std::numeric_limits::max() - (kAlignment - 1)) { + return plan; + } + plan.aligned_out_width = + (out_width + kAlignment - 1) / kAlignment * kAlignment; + if (flattened_kernel_size > + std::numeric_limits::max() / plan.aligned_out_width) { + return plan; + } + const int64_t bytes_per_row = flattened_kernel_size * plan.aligned_out_width; + if (bytes_per_row > scratch_budget_bytes || + batch > std::numeric_limits::max() / out_height) { + return plan; + } + + const int64_t total_rows = batch * out_height; + if (total_rows > std::numeric_limits::max()) { + return plan; + } + const int64_t max_rows_per_tile = std::min( + {total_rows, + scratch_budget_bytes / bytes_per_row, + kQ8taConv2dMaxRowsPerTile}); + plan.num_tiles = total_rows / max_rows_per_tile + + static_cast(total_rows % max_rows_per_tile != 0); + plan.rows_per_tile = total_rows / plan.num_tiles + + static_cast(total_rows % plan.num_tiles != 0); + plan.scratch_bytes = plan.rows_per_tile * bytes_per_row; + plan.feasible = true; + return plan; +} + +Q8taConv2dStreamPlan make_q8ta_conv2d_stream_plan_for_device( + const int64_t batch, + const int64_t flattened_kernel_size, + const int64_t out_height, + const int64_t out_width, + const uint64_t max_buffer_bytes) { + const int64_t scratch_budget = static_cast(std::min( + kQ8taConv2dIm2ColScratchBudgetBytes, max_buffer_bytes)); + return make_q8ta_conv2d_stream_plan( + batch, flattened_kernel_size, out_height, out_width, scratch_budget); +} + // // Shader dispatch utilities // @@ -28,21 +88,41 @@ GlobalWorkGrid pick_q8ta_im2col_gwg( const std::vector& args, const std::vector& resize_args) { (void)shader; - (void)resize_args; - + VK_CHECK_COND(graph != nullptr); const ValueRef im2col_output = args.at(0).refs.at(0); - - const uint32_t N = graph->size_at(-4, im2col_output); const uint32_t K = graph->size_at(-3, im2col_output); - const uint32_t H = graph->size_at(-2, im2col_output); + const uint32_t rows_per_tile = graph->size_at(-2, im2col_output); const uint32_t W = graph->size_at(-1, im2col_output); + const ValueRef input = resize_args.at(0); + const ValueRef kernel_size = resize_args.at(1); + const ValueRef stride = resize_args.at(2); + const ValueRef padding = resize_args.at(3); + const ValueRef dilation = resize_args.at(4); + const int64_t row_offset = graph->extract_scalar(resize_args.at(6)); + + const std::vector input_sizes = graph->sizes_of(input); + const int64_t batch = utils::val_at(-4, input_sizes); + const std::vector out_hw = calc_out_sizes_hw( + *graph, + input_sizes, + kernel_size, + /*kernel_size_only=*/true, + {stride, padding, dilation, dilation}, + /*transposed=*/false); + const int64_t total_rows = batch * out_hw.at(0); + if (row_offset >= total_rows) { + return graph->create_linear_gwg(0u); + } + const uint32_t live_rows = utils::safe_downcast( + std::min(rows_per_tile, total_rows - row_offset)); + const uint32_t K4 = utils::div_up_4(K); const uint32_t W4 = utils::div_up_4(W); // Each thread handles one 4x4 block in the output - return graph->create_linear_gwg( - utils::safe_downcast(static_cast(K4) * W4 * H * N)); + return graph->create_linear_gwg(utils::safe_downcast( + static_cast(K4) * W4 * live_rows)); } LocalWorkGroup pick_q8ta_im2col_lwg( @@ -102,19 +182,18 @@ std::vector calculate_q8ta_im2col_sizes( // Resize // -// resize_args = { input, kernel_size, stride, padding, dilation, groups } +// resize_args = { input, kernel_size, stride, padding, dilation, groups, +// row_offset, max_im2col_rows } // -// The im2col scratch tensor is [N, K, H_out, align_up_4(W_out)] where K (the -// flattened conv window, channel/kernel-derived) is shape-independent and -// H_out/W_out are the conv output spatial dims. The downstream PW GEMM that -// consumes this scratch is resized separately (it preserves H/W). Without this, -// the scratch freezes at the build-time upper bound and feeds garbage rows into -// the GEMM. Recompute H_out/W_out from the CURRENT input (NOT the conv output -// tensor, which may itself still be frozen at this point in the resize order). +// The scratch tensor is [1, K, rows_per_tile, align_up_4(W_out)]. K and +// rows_per_tile are fixed; only W_out tracks the current input shape. +// Batch/height growth past max im2col rows has no dispatches, +// so fail fast instead of leaving outputs stale. void resize_q8ta_im2col_node( ComputeGraph* graph, const std::vector& args, const std::vector& resize_args) { + VK_CHECK_COND(graph != nullptr); const ValueRef im2col_out = args.at(0).refs.at(0); const ValueRef in = resize_args.at(0); const ValueRef kernel_size = resize_args.at(1); @@ -122,11 +201,11 @@ void resize_q8ta_im2col_node( const ValueRef padding = resize_args.at(3); const ValueRef dilation = resize_args.at(4); const ValueRef groups = resize_args.at(5); + const int64_t max_im2col_rows = + graph->extract_scalar(resize_args.at(7)); const std::vector in_sizes = graph->sizes_of(in); - const int64_t batch = utils::val_at(-4, in_sizes); - - // Conv output H/W from the current input. + // Conv output width from the current input. const std::vector out_hw = calc_out_sizes_hw( *graph, in_sizes, @@ -134,7 +213,6 @@ void resize_q8ta_im2col_node( /*kernel_size_only=*/true, {stride, padding, dilation, dilation}, /*transposed=*/false); - const int64_t out_height = out_hw.at(0); const int64_t out_width = out_hw.at(1); // K (flattened conv window) is shape-independent — recompute from channels + @@ -149,7 +227,15 @@ void resize_q8ta_im2col_node( const int64_t K = flattened_kernel_len * groups_val; const int64_t W = utils::align_up_4(out_width); - graph->virtual_resize(im2col_out, {batch, K, out_height, W}); + const int64_t rows_per_tile = graph->size_at(-2, im2col_out); + + const int64_t batch = utils::val_at(-4, in_sizes); + const int64_t out_height = out_hw.at(0); + VK_CHECK_COND( + batch * out_height <= max_im2col_rows, + "q8ta im2col resize grew past max im2col rows"); + + graph->virtual_resize(im2col_out, {1, K, rows_per_tile, W}); } // @@ -166,7 +252,9 @@ void add_q8ta_im2col_node( const ValueRef groups, const ValueRef packed_int8_output, const ValueRef packed_int8_im2col, - const int32_t zp) { + const int32_t zp, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref) { // Validate packed dim info for input and output tensors VK_CHECK_COND(q8ta_conv2d_check_packed_dim_info( graph.packed_dim_info_of(packed_int8_input))); @@ -195,8 +283,14 @@ void add_q8ta_im2col_node( graph.buffer_meta_ubo(packed_int8_input), graph.create_params_buffer(conv_params)}; + VK_CHECK_COND(stream_row_offset_ref != kDummyValueRef); + VK_CHECK_COND(max_im2col_rows_ref != kDummyValueRef); + const int32_t stream_row_offset = utils::safe_downcast( + graph.extract_scalar(stream_row_offset_ref)); + std::vector push_constants = { PushConstantDataInfo(&zp, sizeof(zp)), + PushConstantDataInfo(&stream_row_offset, sizeof(stream_row_offset)), }; // Build spec constants: apply_bias + layout constants (for generic shader) @@ -212,6 +306,19 @@ void add_q8ta_im2col_node( // spec_constants.append(graph.hashed_layout_of(packed_int8_im2col)); // } + // resize_args = { input, kernel_size, stride, padding, dilation, groups, + // row_offset, max_im2col_rows }. The grid picker reads the + // row offset at index 6; append-only. + std::vector resize_args = { + packed_int8_input, + kernel_size, + stride, + padding, + dilation, + groups, + stream_row_offset_ref, + max_im2col_rows_ref}; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), @@ -225,10 +332,7 @@ void add_q8ta_im2col_node( push_constants, // Specialization Constants spec_constants, - // Resize args: { input, kernel_size, stride, padding, dilation, groups } - {packed_int8_input, kernel_size, stride, padding, dilation, groups}, - // Resizing Logic: recompute the im2col scratch dims from the current - // input + resize_args, resize_q8ta_im2col_node)); } @@ -258,6 +362,25 @@ void q8ta_conv2d_im2col_impl( const ValueRef activation = args.at(idx++); const ValueRef packed_int8_output = args.at(idx++); + const std::vector full_im2col_sizes = calculate_q8ta_im2col_sizes( + &graph, packed_int8_input, packed_int8_output, kernel_size, groups); + const Q8taConv2dStreamPlan stream_plan = + make_q8ta_conv2d_stream_plan_for_device( + full_im2col_sizes.at(0), + full_im2col_sizes.at(1), + full_im2col_sizes.at(2), + full_im2col_sizes.at(3), + graph.max_buffer_numel()); + if (!stream_plan.feasible) { + q8ta_conv2d_general(graph, args); + return; + } + VK_CHECK_COND( + !use_unsigned_dot || + graph.size_at(-1, weight_data) <= + kMaxUnsignedDotAccumulatorBytes, + "Unsigned q8ta im2col convolution exceeds the accumulator bound"); + QuantizationConfig weight_quant_config(8, kPerChannel, {}); // Prepack weight using linear weight packing (for im2col approach) @@ -287,11 +410,15 @@ void q8ta_conv2d_im2col_impl( uint32_t activation_type_val = static_cast( activation_type_from_string(graph.extract_string(activation))); - // Calculate im2col output sizes - std::vector im2col_sizes = calculate_q8ta_im2col_sizes( - &graph, packed_int8_input, packed_int8_output, kernel_size, groups); + // One fixed-size scratch buffer is reused across all row tiles; the full + // fit is a single tile. Interleaved write/read dispatches insert the + // barrier before the next tile overwrites it. + const std::vector im2col_sizes = { + 1, + full_im2col_sizes.at(1), + stream_plan.rows_per_tile, + stream_plan.aligned_out_width}; - // Create temporary tensor for im2col output (4W4C layout) TmpTensor packed_int8_im2col( &graph, im2col_sizes, @@ -300,51 +427,59 @@ void q8ta_conv2d_im2col_impl( utils::kPackedInt8_4W4C); int32_t zp = graph.extract_scalar(input_zp); - - // Step 1: Perform im2col transformation - add_q8ta_im2col_node( - graph, - packed_int8_input, - kernel_size, - stride, - padding, - dilation, - groups, - packed_int8_output, - packed_int8_im2col, - zp); - - // Step 2: Perform pointwise convolution on the im2col result const int32_t groups_val = graph.extract_scalar(groups); - VK_CHECK_COND( - !use_unsigned_dot || - graph.size_at(-1, weight_data) <= - kMaxUnsignedDotAccumulatorBytes, - "Unsigned q8ta im2col convolution exceeds the accumulator bound"); - add_q8ta_conv2d_pw_node( - graph, - use_unsigned_dot, - packed_int8_im2col, - input_scale, - input_zp, - packed_weight, - packed_weight_sums, - packed_weight_scales, - output_scale, - output_zp, - bias_data, - packed_bias, - activation_type_val, - packed_int8_output, - groups_val, - // Original activation + conv geometry so the PW output H/W is recomputed - // from the true conv result, not the width-padded im2col scratch. - packed_int8_input, - kernel_size, - stride, - padding, - dilation); + // Row tiles are fixed at build time: dynamic growth past max im2col rows + // has no dispatches, so each resize fails fast below instead of leaving + // outputs stale. Shrinkage only ever lowers the total below this bound. + const ValueRef max_im2col_rows_ref = graph.add_scalar( + stream_plan.num_tiles * stream_plan.rows_per_tile); + + for (int64_t tile = 0; tile < stream_plan.num_tiles; ++tile) { + const int64_t row_offset = tile * stream_plan.rows_per_tile; + // One row-offset scalar per tile feeds both nodes' push constants (via + // re-extraction) and resize args. + const ValueRef row_offset_ref = graph.add_scalar(row_offset); + + add_q8ta_im2col_node( + graph, + packed_int8_input, + kernel_size, + stride, + padding, + dilation, + groups, + packed_int8_output, + packed_int8_im2col, + zp, + row_offset_ref, + max_im2col_rows_ref); + + add_q8ta_conv2d_pw_node( + graph, + use_unsigned_dot, + packed_int8_im2col, + input_scale, + input_zp, + packed_weight, + packed_weight_sums, + packed_weight_scales, + output_scale, + output_zp, + bias_data, + packed_bias, + activation_type_val, + packed_int8_output, + groups_val, + packed_int8_input, + kernel_size, + stride, + padding, + dilation, + /*is_im2col=*/true, + row_offset_ref, + max_im2col_rows_ref); + } } void q8ta_conv2d_im2col( diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp index 217afddb78f..454a115498b 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dPW.cpp @@ -14,19 +14,20 @@ #include #include +#include + namespace vkcompute { // // Shader dispatch utilities // -GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( +GlobalWorkGrid pick_q8ta_conv2d_pw_gwg_impl( ComputeGraph* graph, - const vkapi::ShaderInfo& shader, const std::vector& args, - const std::vector& resize_args) { - (void)shader; - (void)resize_args; + const std::vector& resize_args, + const bool is_im2col_tile) { + VK_CHECK_COND(graph != nullptr); const ValueRef output = args.at(0).refs.at(0); @@ -35,23 +36,50 @@ GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( const uint32_t C = graph->size_at(-3, output); const uint32_t N = graph->size_at(-4, output); - // Each thread covers a 4-width x 4-channel output block. - // Tile constants must match TILE_M4 / TILE_N4 in q8ta_conv2d_pw.glsl. - constexpr uint32_t TILE_N4 = 1; - constexpr uint32_t TILE_M4 = 1; - + // Each thread covers a 4-width x 4-channel output block, matching TILE_M4 / + // TILE_N4 (= 1) in q8ta_conv2d_pw.glsl. const uint32_t C4 = utils::div_up_4(C); const uint32_t W4 = utils::div_up_4(W); - // Global workgroup size: - // x = output channels / (TILE_N4 * 4) = C4 / TILE_N4 = C4 - // y = width / (TILE_M4 * 4) = W4 / TILE_M4 = W4 - // z = height * batch - return GlobalWorkGrid( - {utils::div_up(C4, TILE_N4), - utils::div_up(W4, TILE_M4), - utils::safe_downcast(static_cast(H) * N)}, - kTiledWorkGrid); + uint32_t z; + if (is_im2col_tile) { + // The bound input is one [1, K, rows, W] scratch tile; the grid covers + // the live tile rows after the tile offset. + const ValueRef input = args.at(1).refs.at(0); + const uint32_t rows_per_tile = graph->size_at(-2, input); + const int64_t row_offset = + graph->extract_scalar(resize_args.at(5)); + const int64_t total_rows = static_cast(N) * H; + if (row_offset >= total_rows) { + return GlobalWorkGrid({0u, 0u, 0u}, kTiledWorkGrid); + } + z = utils::safe_downcast( + std::min(rows_per_tile, total_rows - row_offset)); + } else { + // The bound input is the batched activation; the grid covers every + // output row. + z = utils::safe_downcast(static_cast(H) * N); + } + return GlobalWorkGrid({C4, W4, z}, kTiledWorkGrid); +} + +GlobalWorkGrid pick_q8ta_conv2d_pw_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + (void)resize_args; + return pick_q8ta_conv2d_pw_gwg_impl(graph, args, resize_args, false); +} + +GlobalWorkGrid pick_q8ta_conv2d_pw_streaming_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + return pick_q8ta_conv2d_pw_gwg_impl(graph, args, resize_args, true); } LocalWorkGroup pick_q8ta_conv2d_pw_lwg( @@ -218,14 +246,16 @@ void resize_q8ta_conv2d_pw_node( graph->virtual_resize(out, new_sizes); } -// resize_args = { conv_input, kernel_size, stride, padding, dilation } +// resize_args = { conv_input, kernel_size, stride, padding, dilation, +// row_offset, max_im2col_rows }. The grid picker reads the row +// offset at index 5; append-only. // // im2col-path PW conv. Here the PW node's bound input is the im2col scratch -// tensor sized {K, H_out, align_up_4(W_out)} — its width is rounded up to a -// multiple of 4 for texel alignment, so it must NOT be used to size the output. -// Recompute the TRUE conv H_out/W_out from the ORIGINAL activation + conv -// geometry, exactly as resize_q8ta_conv2d_node does. N/C are shape-independent -// and stay as currently allocated. +// tensor sized {1, K, rows, align_up_4(W)} (one row tile) — its width is +// rounded up to a multiple of 4 for texel alignment, so it must NOT be used +// to size the output. +// Recompute the true conv N/H_out/W_out from the original activation + conv +// geometry, exactly as resize_q8ta_conv2d_node does. C is shape-independent. void resize_q8ta_conv2d_pw_im2col_node( ComputeGraph* graph, const std::vector& args, @@ -236,8 +266,13 @@ void resize_q8ta_conv2d_pw_im2col_node( const ValueRef stride = resize_args.at(2); const ValueRef padding = resize_args.at(3); const ValueRef dilation = resize_args.at(4); + // Row tiles are fixed at build time: fail fast on growth past max im2col + // rows instead of leaving outputs stale. + const int64_t max_im2col_rows = + graph->extract_scalar(resize_args.at(6)); const std::vector in_sizes = graph->sizes_of(conv_input); + const int64_t batch = utils::val_at(-4, in_sizes); const std::vector out_hw = calc_out_sizes_hw( *graph, @@ -247,8 +282,13 @@ void resize_q8ta_conv2d_pw_im2col_node( {stride, padding, dilation, dilation}, /*transposed=*/false); + VK_CHECK_COND( + batch * out_hw.at(0) <= max_im2col_rows, + "q8ta im2col resize grew past max im2col rows"); + std::vector new_sizes = graph->sizes_of(out); const size_t ndim = new_sizes.size(); + new_sizes.at(ndim - 4) = utils::val_at(-4, in_sizes); new_sizes.at(ndim - 2) = out_hw.at(0); new_sizes.at(ndim - 1) = out_hw.at(1); graph->virtual_resize(out, new_sizes); @@ -278,7 +318,10 @@ void add_q8ta_conv2d_pw_node( const ValueRef kernel_size, const ValueRef stride, const ValueRef padding, - const ValueRef dilation) { + const ValueRef dilation, + const bool is_im2col, + const ValueRef stream_row_offset_ref, + const ValueRef max_im2col_rows_ref) { VK_CHECK_COND(q8ta_conv2d_check_4w4c_packed_dim_info( graph.packed_dim_info_of(packed_int8_input))); VK_CHECK_COND(q8ta_conv2d_check_packed_dim_info( @@ -304,6 +347,16 @@ void add_q8ta_conv2d_pw_node( int32_t output_zp_val = graph.extract_scalar(output_zp); uint32_t apply_bias = graph.val_is_none(bias_data) ? 0u : 1u; + // The tile offset lives in the graph scalar; re-extract it here so the + // shader push constant and the grid picker read one value. Standalone + // dispatches have no tile and push 0. + int32_t stream_row_offset = 0; + if (is_im2col) { + VK_CHECK_COND(stream_row_offset_ref != kDummyValueRef); + VK_CHECK_COND(max_im2col_rows_ref != kDummyValueRef); + stream_row_offset = utils::safe_downcast( + graph.extract_scalar(stream_row_offset_ref)); + } std::vector push_constants = { PushConstantDataInfo(&input_scale_val, sizeof(input_scale_val)), PushConstantDataInfo(&input_zp_val, sizeof(input_zp_val)), @@ -311,8 +364,14 @@ void add_q8ta_conv2d_pw_node( PushConstantDataInfo(&output_zp_val, sizeof(output_zp_val)), PushConstantDataInfo(&K4_per_group, sizeof(K4_per_group)), PushConstantDataInfo(&OC4_per_group, sizeof(OC4_per_group)), + PushConstantDataInfo(&stream_row_offset, sizeof(stream_row_offset)), }; + // The im2col path consumes one flat scratch tile per dispatch; the + // standalone 1x1 path reads its batched activation input directly. The + // addressing is selected by spec constant, so both share one shader. + const uint32_t use_flat_tile = is_im2col ? 1u : 0u; + const bool use_hw_dot = graph.context()->adapter_ptr()->supports_int8_dot_product(); std::string kernel_name; @@ -327,6 +386,13 @@ void add_q8ta_conv2d_pw_node( } else { kernel_name = use_hw_dot ? "q8ta_conv2d_pw" : "q8ta_conv2d_pw_fallback"; } + if (!use_unsigned_dot) { + // Signed PW kernels are only codegen'd for texture weights; a buffer + // weight here would fail kernel lookup at dispatch, so fail fast. + VK_CHECK_COND( + graph.storage_type_of(packed_weight) != utils::kBuffer, + "Signed q8ta pointwise convolution requires texture weights"); + } add_dtype_suffix(kernel_name, graph.dtype_of(packed_weight_scales)); vkapi::ParamsBindList param_buffers = { @@ -338,6 +404,8 @@ void add_q8ta_conv2d_pw_node( activation_type, graph.hashed_layout_of(packed_int8_output), graph.hashed_layout_of(packed_int8_input), + // Appended last to match the use_flat_tile declaration order. + use_flat_tile, }; // The im2col path passes the original activation + conv geometry so the @@ -347,18 +415,28 @@ void add_q8ta_conv2d_pw_node( // output matches directly. std::vector resize_args; ExecuteNode::ResizeFunction resize_fn; - if (conv_input == kDummyValueRef) { + if (!is_im2col) { resize_args = {packed_int8_input}; resize_fn = resize_q8ta_conv2d_pw_node; } else { - resize_args = {conv_input, kernel_size, stride, padding, dilation}; + resize_args = { + conv_input, + kernel_size, + stride, + padding, + dilation, + stream_row_offset_ref, + max_im2col_rows_ref}; resize_fn = resize_q8ta_conv2d_pw_im2col_node; } + const auto pick_gwg = + is_im2col ? pick_q8ta_conv2d_pw_streaming_gwg : pick_q8ta_conv2d_pw_gwg; + graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - pick_q8ta_conv2d_pw_gwg, + pick_gwg, pick_q8ta_conv2d_pw_lwg, {{packed_int8_output, vkapi::kWrite}, {{packed_int8_input, diff --git a/backends/vulkan/runtime/vk_api/Adapter.cpp b/backends/vulkan/runtime/vk_api/Adapter.cpp index 3d9acae8975..eea61a8fe3f 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.cpp +++ b/backends/vulkan/runtime/vk_api/Adapter.cpp @@ -350,6 +350,17 @@ Adapter::~Adapter() { } } +ScopedAdapterCapabilityOverride::ScopedAdapterCapabilityOverride( + Adapter* adapter, + AdapterCapabilityOverrides overrides) + : adapter_(adapter), previous_(adapter->capability_overrides_) { + adapter_->capability_overrides_ = overrides; +} + +ScopedAdapterCapabilityOverride::~ScopedAdapterCapabilityOverride() { + adapter_->capability_overrides_ = previous_; +} + Adapter::Queue Adapter::request_queue() { // Lock the mutex as multiple threads can request a queue at the same time std::lock_guard lock(queue_usage_mutex_); diff --git a/backends/vulkan/runtime/vk_api/Adapter.h b/backends/vulkan/runtime/vk_api/Adapter.h index 968ca798369..0c2784ae473 100644 --- a/backends/vulkan/runtime/vk_api/Adapter.h +++ b/backends/vulkan/runtime/vk_api/Adapter.h @@ -17,6 +17,8 @@ #include +#include + #include namespace vkcompute { @@ -78,6 +80,8 @@ class Adapter final { VkQueue handle; }; + friend class ScopedAdapterCapabilityOverride; + private: // Use a mutex to manage queue usage info since // it can be accessed from multiple threads @@ -102,6 +106,9 @@ class Adapter final { // Miscellaneous bool linear_tiling_3d_enabled_; bool owns_device_; + // Test-only capability overrides; empty unless a ScopedCapabilityOverride + // is live. + AdapterCapabilityOverrides capability_overrides_; public: // Physical Device metadata @@ -232,6 +239,9 @@ class Adapter final { } inline bool supports_int8_dot_product() const { + if (capability_overrides_.int8_dot_product.has_value()) { + return *capability_overrides_.int8_dot_product; + } #ifdef ETVK_FORCE_NO_EXTENSIONS return false; #endif @@ -244,6 +254,9 @@ class Adapter final { } inline bool accelerates_signed_packed4x8_dot() const { + if (capability_overrides_.signed_packed4x8_dot.has_value()) { + return *capability_overrides_.signed_packed4x8_dot; + } #ifdef ETVK_FORCE_NO_EXTENSIONS return false; #endif @@ -258,6 +271,9 @@ class Adapter final { } inline bool accelerates_unsigned_packed4x8_dot() const { + if (capability_overrides_.unsigned_packed4x8_dot.has_value()) { + return *capability_overrides_.unsigned_packed4x8_dot; + } #ifdef ETVK_FORCE_NO_EXTENSIONS return false; #endif diff --git a/backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h b/backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h new file mode 100644 index 00000000000..f4b6c6d9907 --- /dev/null +++ b/backends/vulkan/runtime/vk_api/AdapterCapabilityOverrides.h @@ -0,0 +1,61 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { +namespace vkapi { + +class Adapter; + +// Test-only device capability overrides. Each set field replaces the +// corresponding Adapter query, simulating a weaker device so fallback paths +// can be covered on capable hardware. Forcing a capability the device +// lacks is undefined behavior and typically fails at pipeline creation. +// Overrides apply to every graph built against the Adapter while set; +// hold a ScopedAdapterCapabilityOverride to bound the scope to one test. +// Not thread-safe: while set, every query on the Adapter observes the +// override, so hold it only around single-test build plus execute with no +// concurrent adapter use. +struct AdapterCapabilityOverrides { + std::optional int8_dot_product; + std::optional signed_packed4x8_dot; + std::optional unsigned_packed4x8_dot; + + static AdapterCapabilityOverrides without_dot_product_support() { + AdapterCapabilityOverrides overrides; + overrides.int8_dot_product = false; + overrides.signed_packed4x8_dot = false; + overrides.unsigned_packed4x8_dot = false; + return overrides; + } +}; + +class ScopedAdapterCapabilityOverride final { + public: + ScopedAdapterCapabilityOverride( + Adapter* adapter, + AdapterCapabilityOverrides overrides); + ~ScopedAdapterCapabilityOverride(); + ScopedAdapterCapabilityOverride(const ScopedAdapterCapabilityOverride&) = + delete; + ScopedAdapterCapabilityOverride& operator=( + const ScopedAdapterCapabilityOverride&) = delete; + ScopedAdapterCapabilityOverride(ScopedAdapterCapabilityOverride&&) = delete; + ScopedAdapterCapabilityOverride& operator=( + ScopedAdapterCapabilityOverride&&) = delete; + + private: + Adapter* adapter_; + AdapterCapabilityOverrides previous_; +}; + +} // namespace vkapi +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp index e97ef8d61e0..250ea51746d 100644 --- a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp @@ -20,7 +20,8 @@ namespace { void assert_im2col_kernel_selection( ComputeGraph& graph, const bool expect_unsigned, - const bool expect_buffer_weights) { + const bool expect_buffer_weights, + const bool expect_fallback_kernel = false) { const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); std::string expected_execute; std::string expected_prepack; @@ -33,9 +34,13 @@ void assert_im2col_kernel_selection( : "pack_q8_linear_weight_unsigned_texture2d"; } else { VK_CHECK_COND(!expect_buffer_weights); - expected_execute = adapter->supports_int8_dot_product() - ? "q8ta_conv2d_pw_float" - : "q8ta_conv2d_pw_fallback_float"; + if (expect_fallback_kernel) { + expected_execute = "q8ta_conv2d_pw_fallback_float"; + } else { + expected_execute = adapter->supports_int8_dot_product() + ? "q8ta_conv2d_pw_float" + : "q8ta_conv2d_pw_fallback_float"; + } expected_prepack = "pack_q8_linear_weight_texture2d"; } @@ -311,7 +316,20 @@ void test_q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { groups, activation, packed_int8_output}; - if (impl_selector == "im2col" || impl_selector == "im2col_unsigned" || + if (impl_selector == "im2col_fallback") { + // Simulate a device without dot-product support so the fallback + // kernel is selected on any hardware. + vkapi::ScopedAdapterCapabilityOverride no_dot_support( + graph.context()->adapter_ptr(), + vkapi::AdapterCapabilityOverrides::without_dot_product_support()); + q8ta_conv2d_im2col_impl(graph, /*use_unsigned_dot=*/false, conv_args); + assert_im2col_kernel_selection( + graph, + /*expect_unsigned=*/false, + /*expect_buffer_weights=*/false, + /*expect_fallback_kernel=*/true); + } else if ( + impl_selector == "im2col" || impl_selector == "im2col_unsigned" || impl_selector == "im2col_auto") { const vkapi::Adapter* const adapter = graph.context()->adapter_ptr(); bool expect_unsigned = impl_selector == "im2col_unsigned"; diff --git a/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp b/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp new file mode 100644 index 00000000000..6f634a83c27 --- /dev/null +++ b/backends/vulkan/test/custom_ops/q8ta_conv2d_stream_plan_test.cpp @@ -0,0 +1,178 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include + +#include + +namespace vkcompute { +namespace { + +constexpr int64_t kEightMiB = 8 * 1024 * 1024; +constexpr int64_t kSixteenMiB = 16 * 1024 * 1024; + +TEST(Q8taConv2dStreamPlanTest, SplitsFirstSceneXConvolution) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/60, + /*flattened_kernel_size=*/576, + /*out_height=*/20, + /*out_width=*/26, + kEightMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.aligned_out_width, 28); + EXPECT_EQ(plan.rows_per_tile, 400); + EXPECT_EQ(plan.num_tiles, 3); + EXPECT_EQ(plan.scratch_bytes, 6451200); +} + +TEST(Q8taConv2dStreamPlanTest, SplitsSecondSceneXConvolution) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/60, + /*flattened_kernel_size=*/1152, + /*out_height=*/10, + /*out_width=*/13, + kEightMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.aligned_out_width, 16); + EXPECT_EQ(plan.rows_per_tile, 300); + EXPECT_EQ(plan.num_tiles, 2); + EXPECT_EQ(plan.scratch_bytes, 5529600); +} + +TEST(Q8taConv2dStreamPlanTest, UsesOneTileWhenFullScratchFits) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/2, + /*flattened_kernel_size=*/288, + /*out_height=*/7, + /*out_width=*/7, + kEightMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 14); + EXPECT_EQ(plan.num_tiles, 1); + EXPECT_EQ(plan.scratch_bytes, 32256); +} + +TEST(Q8taConv2dStreamPlanTest, SelectsFullFitBelowProductionBudget) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/10, + /*flattened_kernel_size=*/288, + /*out_height=*/30, + /*out_width=*/99, + kSixteenMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 300); + EXPECT_EQ(plan.num_tiles, 1); + EXPECT_EQ(plan.scratch_bytes, 8640000); +} + +TEST(Q8taConv2dStreamPlanTest, SelectsStreamingAboveProductionBudget) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/10, + /*flattened_kernel_size=*/576, + /*out_height=*/30, + /*out_width=*/99, + kSixteenMiB); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 150); + EXPECT_EQ(plan.num_tiles, 2); + EXPECT_EQ(plan.scratch_bytes, 8640000); +} + +TEST(Q8taConv2dStreamPlanTest, RejectsOneRowLargerThanBudget) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/1, + /*flattened_kernel_size=*/16384, + /*out_height=*/1, + /*out_width=*/513, + kEightMiB); + + EXPECT_FALSE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 0); + EXPECT_EQ(plan.num_tiles, 0); + EXPECT_EQ(plan.scratch_bytes, 0); +} + +TEST(Q8taConv2dStreamPlanTest, RejectsOutputWidthAlignmentOverflow) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/1, + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/std::numeric_limits::max(), + std::numeric_limits::max()); + + EXPECT_FALSE(plan.feasible); +} + +TEST(Q8taConv2dStreamPlanTest, HandlesTileCountCeilDivisionAtShaderLimit) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/std::numeric_limits::max(), + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/1, + /*scratch_budget_bytes=*/8); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.rows_per_tile, 2); + EXPECT_EQ(plan.num_tiles, 1073741824); + EXPECT_EQ(plan.scratch_bytes, 8); +} + +TEST(Q8taConv2dStreamPlanTest, RejectsRowOffsetBeyondShaderIndexRange) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/std::numeric_limits::max(), + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/1, + std::numeric_limits::max()); + + EXPECT_FALSE(plan.feasible); +} + +TEST(Q8taConv2dStreamPlanTest, CapsRowsAtGuaranteedWorkgroupCountZ) { + const auto plan = make_q8ta_conv2d_stream_plan( + /*batch=*/100000, + /*flattened_kernel_size=*/1, + /*out_height=*/1, + /*out_width=*/1, + std::numeric_limits::max()); + + EXPECT_TRUE(plan.feasible); + EXPECT_EQ(plan.num_tiles, 2); + EXPECT_EQ(plan.rows_per_tile, 50000); + EXPECT_EQ(plan.scratch_bytes, 200000); +} + +TEST(Q8taConv2dStreamPlanTest, DeviceBufferLimitMatchesActualPlan) { + constexpr int64_t kBytesPerRow = 576 * 28; + const auto rejected = make_q8ta_conv2d_stream_plan_for_device( + /*batch=*/60, + /*flattened_kernel_size=*/576, + /*out_height=*/20, + /*out_width=*/26, + /*max_buffer_bytes=*/kBytesPerRow - 1); + EXPECT_FALSE(rejected.feasible); + + const auto accepted = make_q8ta_conv2d_stream_plan_for_device( + /*batch=*/60, + /*flattened_kernel_size=*/576, + /*out_height=*/20, + /*out_width=*/26, + /*max_buffer_bytes=*/kBytesPerRow); + EXPECT_TRUE(accepted.feasible); + EXPECT_EQ(accepted.rows_per_tile, 1); + EXPECT_EQ(accepted.num_tiles, 1200); + EXPECT_EQ(accepted.scratch_bytes, kBytesPerRow); +} + +} // namespace +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index ed463b79a2f..df3dbd925f3 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -98,6 +98,16 @@ def define_common_targets(is_fbcode = False): ], ) + runtime.cxx_test( + name = "q8ta_conv2d_stream_plan_test", + srcs = ["q8ta_conv2d_stream_plan_test.cpp"], + platforms = get_platforms(), + deps = [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + ) + define_custom_op_test_binary("test_add") define_custom_op_test_binary("test_q8csw_linear") define_custom_op_test_binary("test_q8csw_conv2d") diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index 51da65da57c..56ea6e177c9 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -46,7 +46,9 @@ static TestCase create_test_case_from_config( utils::StorageType fp_storage_type, utils::GPUMemoryLayout int8_memory_layout, const std::string& impl_selector = "", - const Im2colUnsignedTestOptions* im2col_options = nullptr) { + const Im2colUnsignedTestOptions* im2col_options = nullptr, + const float input_scale_val = 0.008123f, + const DataGenType input_data_gen = DataGenType::RANDOM) { TestCase test_case; // Calculate output dimensions @@ -94,18 +96,12 @@ static TestCase create_test_case_from_config( input_dtype, fp_storage_type, fp_memory_layout, -#ifdef DEBUG_MODE - DataGenType::RANDOM -#else - DataGenType::RANDOM -#endif - ); + input_data_gen); if (debugging()) { print_valuespec_data(input_tensor, "input_tensor"); } - float input_scale_val = 0.008123; ValueSpec input_scale(input_scale_val); const int32_t input_zero_point_val = @@ -374,6 +370,123 @@ std::vector generate_quantized_conv2d_easy_cases() { static std::vector generate_im2col_unsigned_test_cases( const std::string& impl_selector); +static std::vector generate_streaming_im2col_test_cases() { + std::vector test_cases; + + Conv2dConfig full_fit_config = { + OutInChannels(4, 32), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 10}; + full_fit_config.op_name = "conv2d_q8ta_q8csw_q8to"; + full_fit_config.test_case_name = make_test_case_name( + full_fit_config, false, utils::kTexture3D, utils::kBuffer); + + Conv2dConfig streaming_fallback_config = full_fit_config; + streaming_fallback_config.channels.in = 64; + streaming_fallback_config.test_case_name = make_test_case_name( + streaming_fallback_config, false, utils::kTexture3D, utils::kBuffer); + + // Forced-fallback cases need no int8 dot-product support, so they run on + // all devices; the kAuto cases below stay gated. + test_cases.push_back(create_test_case_from_config( + full_fit_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_fallback", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + test_cases.push_back(create_test_case_from_config( + streaming_fallback_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_fallback", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + + if (!vkcompute::api::context()->adapter_ptr()->supports_int8_dot_product()) { + return test_cases; + } + + for (const utils::GPUMemoryLayout layout : + {utils::kPackedInt8_4C1W, utils::kPackedInt8_4W4C}) { + test_cases.push_back(create_test_case_from_config( + full_fit_config, + vkapi::kFloat, + utils::kTexture3D, + layout, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + } + Conv2dConfig streaming_config = full_fit_config; + streaming_config.channels.in = 64; + streaming_config.test_case_name = make_test_case_name( + streaming_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + streaming_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C1W, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + Conv2dConfig grouped_config = { + OutInChannels(16, 32), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2, + 20}; + grouped_config.op_name = "conv2d_q8ta_q8csw_q8to"; + grouped_config.test_case_name = make_test_case_name( + grouped_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + grouped_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + + Conv2dConfig output_channel_tail_config = { + OutInChannels(10, 32), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 20}; + output_channel_tail_config.op_name = "conv2d_q8ta_q8csw_q8to"; + output_channel_tail_config.test_case_name = make_test_case_name( + output_channel_tail_config, false, utils::kTexture3D, utils::kBuffer); + test_cases.push_back(create_test_case_from_config( + output_channel_tail_config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT)); + return test_cases; +} + // Generate test cases for quantized conv2d operation static std::vector generate_quantized_conv2d_test_cases() { std::vector test_cases; @@ -668,6 +781,10 @@ static std::vector generate_quantized_conv2d_test_cases() { narrow_workgroup_cases.begin(), narrow_workgroup_cases.end()); + auto streaming_cases = generate_streaming_im2col_test_cases(); + test_cases.insert( + test_cases.end(), streaming_cases.begin(), streaming_cases.end()); + return test_cases; } @@ -1066,6 +1183,99 @@ static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { return flop; } +static void execute_streaming_dynamic_shrink_test() { + Conv2dConfig config = { + OutInChannels(4, 64), + InputSize2D(30, 99), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 1, + 10}; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + TestCase test_case = create_test_case_from_config( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4W4C, + /*impl_selector=*/"im2col_auto", + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT); + for (ValueSpec& input : test_case.inputs()) { + input.ensure_data_generated(/*seed=*/0); + } + + BenchmarkGraph benchmark_graph = setup_compute_graph( + test_case, test_case.operator_name(), /*op_invocations_per_execute=*/1); + ComputeGraph& graph = *benchmark_graph.graph; + graph.prepare(); + graph.prepack(); + + const std::vector upper_input_sizes = test_case.inputs().at(0).sizes; + const std::vector shrunk_input_sizes = {1, 64, 20, 51}; + const std::vector shrunk_output_sizes = {1, 4, 20, 51}; + TestCase shrunk_case = test_case; + shrunk_case.inputs().at(0).sizes = shrunk_input_sizes; + shrunk_case.inputs().at(0).resize_data(1 * 64 * 20 * 51); + shrunk_case.outputs().at(0).sizes = shrunk_output_sizes; + shrunk_case.outputs().at(0).resize_data(1 * 4 * 20 * 51); + reference_impl(shrunk_case); + + constexpr int kRepetitions = 4; + for (int repetition = 0; repetition < kRepetitions; ++repetition) { + graph.resize_input(0, upper_input_sizes); + graph.propagate_resize(); + graph.maybe_cast_and_copy_into_staging( + graph.inputs().at(0).staging, + test_case.inputs().at(0).get_data_ptr(), + test_case.inputs().at(0).numel(), + vkapi::kFloat); + graph.execute(); + + graph.resize_input(0, shrunk_input_sizes); + graph.propagate_resize(); + if (graph.sizes_of(graph.outputs().at(0).value) != shrunk_output_sizes) { + throw std::runtime_error("streaming im2col output did not shrink"); + } + graph.maybe_cast_and_copy_into_staging( + graph.inputs().at(0).staging, + shrunk_case.inputs().at(0).get_data_ptr(), + shrunk_case.inputs().at(0).numel(), + vkapi::kFloat); + graph.execute(); + graph.maybe_cast_and_copy_from_staging( + graph.outputs().at(0).staging, + shrunk_case.outputs().at(0).get_mutable_data_ptr(), + shrunk_case.outputs().at(0).numel(), + vkapi::kFloat); + if (!shrunk_case.outputs().at(0).validate_against_reference( + shrunk_case.get_abs_tolerance(), shrunk_case.get_rel_tolerance())) { + throw std::runtime_error("streaming im2col shrink output was stale"); + } + } + + const Q8taConv2dStreamPlan upper_plan = make_q8ta_conv2d_stream_plan( + /*batch=*/10, + /*flattened_kernel_size=*/576, + /*out_height=*/30, + /*out_width=*/99, + kQ8taConv2dIm2ColScratchBudgetBytes); + if (!upper_plan.feasible || upper_plan.num_tiles != 2 || + upper_plan.rows_per_tile <= 20) { + throw std::runtime_error("dynamic shrink test did not create dead tiles"); + } + const int64_t resized_scratch_bytes = + 576 * upper_plan.rows_per_tile * utils::align_up_4(51); + if (resized_scratch_bytes > kQ8taConv2dIm2ColScratchBudgetBytes) { + throw std::runtime_error("streaming im2col scratch exceeded its cap"); + } + std::cout << "Streaming im2col dynamic shrink PASSED" << std::endl; +} + int main(int argc, char* argv[]) { const vkapi::Adapter& adapter = *vkcompute::api::context()->adapter_ptr(); const bool prefers_unsigned_dot = @@ -1079,6 +1289,8 @@ int main(int argc, char* argv[]) { std::string im2col_impl_selector; bool narrow_workgroups_only = false; + bool streaming_im2col_only = false; + bool streaming_dynamic_shrink_only = false; for (int i = 1; i < argc; ++i) { const std::string arg(argv[i]); if (arg == "--im2col-path=signed") { @@ -1089,15 +1301,21 @@ int main(int argc, char* argv[]) { im2col_impl_selector = "im2col_auto"; } else if (arg == "--narrow-workgroups-only") { narrow_workgroups_only = true; + } else if (arg == "--streaming-im2col-only") { + streaming_im2col_only = true; + } else if (arg == "--streaming-dynamic-shrink-only") { + streaming_dynamic_shrink_only = true; } else { std::cerr << "Unknown argument: " << arg << std::endl; return 2; } } - if (narrow_workgroups_only && !im2col_impl_selector.empty()) { - std::cerr - << "Narrow-workgroup and im2col path selectors are mutually exclusive" - << std::endl; + const int selected_modes = static_cast(!im2col_impl_selector.empty()) + + static_cast(narrow_workgroups_only) + + static_cast(streaming_im2col_only) + + static_cast(streaming_dynamic_shrink_only); + if (selected_modes > 1) { + std::cerr << "Test mode selectors are mutually exclusive" << std::endl; return 2; } set_debugging(false); @@ -1117,6 +1335,10 @@ int main(int argc, char* argv[]) { ReferenceComputeFunc ref_fn = reference_impl; + if (streaming_dynamic_shrink_only) { + execute_streaming_dynamic_shrink_test(); + return 0; + } #ifdef DEBUG_MODE std::function()> test_case_generator = generate_quantized_conv2d_easy_cases; @@ -1130,6 +1352,19 @@ int main(int argc, char* argv[]) { #endif if (narrow_workgroups_only) { test_case_generator = generate_narrow_workgroup_test_cases; + } else if (streaming_im2col_only) { + test_case_generator = generate_streaming_im2col_test_cases; +#ifndef DEBUG_MODE + } else if (selected_modes == 0) { + // The default run also covers the unified tile path: multi-tile + // dispatches, grouped/streaming shapes, and the fallback kernel. + test_case_generator = [base_generator = test_case_generator]() { + auto cases = base_generator(); + const auto streaming_cases = generate_streaming_im2col_test_cases(); + cases.insert(cases.end(), streaming_cases.begin(), streaming_cases.end()); + return cases; + }; +#endif } auto results = execute_test_cases( @@ -1140,5 +1375,11 @@ int main(int argc, char* argv[]) { /*benchmark_runs = */ 1, ref_fn); +#ifndef DEBUG_MODE + if (selected_modes == 0) { + execute_streaming_dynamic_shrink_test(); + } +#endif + return 0; } From 388bd233760bd8c4d98a91a7e718a8292afa1831 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Wed, 9 Sep 2026 20:37:38 -0700 Subject: [PATCH 163/190] [ET-VK][q8ta-conv] Route profitable batched Mali q8ta convolutions to im2col Batched regular q8ta convolutions default to direct convolution, leaving Mali integer-dot-product throughput unused. Route the measured profitable envelope through bounded im2col while preserving existing single-batch, legacy batched, and non-Mali decisions. The policy uses device capability, convolution geometry, scratch feasibility, and tile count instead of model-specific shape allowlists. Large batched grouped q8ta convolutions remain expensive on Mali after regular convolutions move to im2col. Route the measured profitable grouped envelope through bounded streaming im2col. Keep Adreno, PowerVR, legacy, and single-batch behavior unchanged. Require four-channel alignment within each group so packed output blocks never cross group boundaries. Authored with Codex. Differential Revision: [D119408978](https://our.internmc.facebook.com/intern/diff/D119408978/) ghstack-source-id: 427402964 Pull-Request: https://github.com/pytorch/executorch/pull/22669 --- .../runtime/graph/ops/impl/Q8taConv2d.cpp | 72 +-- .../graph/ops/impl/Q8taConv2dRoute.cpp | 147 ++++++ .../runtime/graph/ops/impl/Q8taConv2dRoute.h | 33 ++ .../test/custom_ops/impl/TestQ8taConv2d.cpp | 18 +- .../custom_ops/q8ta_conv2d_route_test.cpp | 473 ++++++++++++++++++ backends/vulkan/test/custom_ops/targets.bzl | 10 + .../test/custom_ops/test_q8ta_conv2d.cpp | 258 +++++++++- 7 files changed, 927 insertions(+), 84 deletions(-) create mode 100644 backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp create mode 100644 backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h create mode 100644 backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp index 262858b76fb..75f65d7a2a0 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2d.cpp @@ -7,6 +7,7 @@ */ #include +#include #include @@ -42,52 +43,6 @@ bool q8ta_conv2d_check_4w4c_packed_dim_info(const api::PackedDimInfo& info) { info.outer_packed_dim_block_size == 4; } -namespace { - -uint64_t q8ta_conv2d_im2col_scratch_limit(ComputeGraph& graph) { - constexpr uint64_t kMaxBatchedIm2ColScratchBytes = 32ULL * 1024ULL * 1024ULL; - const uint64_t device_scratch_limit = - graph.context()->adapter_ptr()->max_buffer_numel(); - return device_scratch_limit < kMaxBatchedIm2ColScratchBytes - ? device_scratch_limit - : kMaxBatchedIm2ColScratchBytes; -} - -bool should_use_q8ta_conv2d_im2col( - ComputeGraph& graph, - const int64_t batch, - const int64_t groups, - const int64_t in_channels_per_group, - const int64_t flattened_kernel_size, - const int64_t out_height, - const int64_t out_width) { - const bool im2col_eligible = in_channels_per_group % 4 == 0; - if (!im2col_eligible) { - return false; - } - - const int64_t spatial_out = out_height * out_width; - if (batch > 1) { - constexpr int64_t kMinFlattenedKernelSize = 1024; - constexpr int64_t kMaxSpatialOutput = 64; - const uint64_t scratch_bytes = static_cast(batch) * - static_cast(flattened_kernel_size) * - static_cast(out_height) * - static_cast(utils::align_up_4(out_width)); - return groups == 1 && flattened_kernel_size >= kMinFlattenedKernelSize && - spatial_out <= kMaxSpatialOutput && - scratch_bytes <= q8ta_conv2d_im2col_scratch_limit(graph); - } - - if (graph.device_is_mali()) { - return true; - } - - return groups == 1 && (in_channels_per_group >= 32 || spatial_out <= 4096); -} - -} // namespace - // // Workgroup size selection functions // @@ -531,27 +486,38 @@ void q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { const ValueRef output = args.at(15); const int64_t groups = graph.extract_scalar(groups_ref); + // Valid models always carry groups >= 1; fail fast on corrupt input + // instead of dividing channel counts by zero downstream (both this + // dispatcher and q8ta_conv2d_general divide by groups). + VK_CHECK_COND(groups > 0, "q8ta_conv2d requires groups >= 1"); const int64_t in_channels = graph.size_at(-3, input); const int64_t in_channels_per_group = in_channels / groups; const int64_t batch = graph.size_at(-4, input); const int64_t H_out = graph.size_at(-2, output); const int64_t W_out = graph.size_at(-1, output); - int64_t flattened_kernel_size; + const int64_t out_channels = graph.size_at(-3, output); + int64_t kernel_height; + int64_t kernel_width; { const auto kernel_size = graph.get_int_list(kernel_size_ref); - flattened_kernel_size = utils::align_up_4( - in_channels_per_group * kernel_size->at(0) * kernel_size->at(1)); + kernel_height = kernel_size->at(0); + kernel_width = kernel_size->at(1); } - const bool use_im2col = should_use_q8ta_conv2d_im2col( - graph, + const bool use_im2col = should_use_q8ta_conv2d_im2col({ + graph.device_is_mali(), + graph.can_use_int8_dot_product(), + static_cast(graph.max_buffer_numel()), batch, groups, in_channels_per_group, - flattened_kernel_size, + out_channels, + kernel_height, + kernel_width, H_out, - W_out); + W_out, + }); if (use_im2col) { q8ta_conv2d_im2col(graph, args); diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp new file mode 100644 index 00000000000..780dfcbfbf3 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.cpp @@ -0,0 +1,147 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include + +#include + +namespace vkcompute { + +namespace { + +// Computes the per-group im2col kernel columns (in-channels x kernel height +// x kernel width, aligned up to 4) for the stream-plan sizing. Returns false +// (leaving the output untouched) when any dimension is non-positive or an +// intermediate product would overflow int64; callers fail closed to the +// direct path. +bool calculate_aligned_kernel_size( + const Q8taConv2dRouteParams& params, + int64_t& aligned_kernel_size) { + if (params.in_channels_per_group <= 0 || params.kernel_height <= 0 || + params.kernel_width <= 0 || + params.in_channels_per_group > + std::numeric_limits::max() / params.kernel_height) { + return false; + } + const int64_t channel_kernel_height = + params.in_channels_per_group * params.kernel_height; + if (channel_kernel_height > + std::numeric_limits::max() / params.kernel_width) { + return false; + } + const int64_t unaligned_kernel_size = + channel_kernel_height * params.kernel_width; + if (unaligned_kernel_size > std::numeric_limits::max() - 3) { + return false; + } + aligned_kernel_size = utils::align_up_4(unaligned_kernel_size); + return true; +} + +} // namespace + +bool should_use_q8ta_conv2d_im2col(const Q8taConv2dRouteParams& params) { + if (params.batch <= 0 || params.groups <= 0 || + params.in_channels_per_group <= 0 || params.out_channels <= 0 || + params.kernel_height <= 0 || params.kernel_width <= 0 || + params.out_height <= 0 || params.out_width <= 0 || + params.out_height > + std::numeric_limits::max() / params.out_width) { + return false; + } + int64_t flattened_kernel_size; + if (!calculate_aligned_kernel_size(params, flattened_kernel_size)) { + return false; + } + const bool im2col_eligible = params.in_channels_per_group % 4 == 0; + if (!im2col_eligible) { + return false; + } + // Grouped im2col partitions output blocks per group in the PW GEMM + // (group_idx = oc_block / OC4_per_group), so each group must own whole + // packed-4 output blocks; otherwise one block straddles two groups and + // reads the wrong group's weights. + if (params.groups > 1 && + (params.out_channels % params.groups != 0 || + params.out_channels / params.groups % 4 != 0)) { + return false; + } + + const int64_t spatial_out = params.out_height * params.out_width; + if (params.batch > 1) { + constexpr int64_t kMinFlattenedKernelSize = 1024; + constexpr int64_t kMaxSpatialOutput = 64; + // Size the probe plan with the same budget the consumer uses, so + // num_tiles == 1 here means a single tile at execution too. + const Q8taConv2dStreamPlan full_plan = + make_q8ta_conv2d_stream_plan_for_device( + params.batch, + flattened_kernel_size, + params.out_height, + params.out_width, + params.max_buffer_bytes); + // Device-independent fast path: a large kernel over a tiny output makes + // the im2col materialization negligible next to the GEMM, and a single + // scratch tile means no streaming overhead, so this wins on every + // device without needing vendor-specific tuning. + const bool use_single_tile_batched_im2col = params.groups == 1 && + flattened_kernel_size >= kMinFlattenedKernelSize && + spatial_out <= kMaxSpatialOutput && full_plan.feasible && + full_plan.num_tiles == 1; + if (use_single_tile_batched_im2col) { + return true; + } + + if (!params.is_mali) { + return false; + } + + // Mali: route all eligible batched convolutions through bounded + // streaming im2col. The remaining guards are correctness bounds, not + // perf cliffs. + if (!params.supports_int8_dot_product || + // Conservative superset of the dispatch-time unsigned-path check + // (which compares the unaligned weight K): reject the aligned + // per-group K above the accumulator bound on every int8-dot path, + // failing closed for signed-path shapes near the bound as well. + flattened_kernel_size > kMaxUnsignedDotAccumulatorBytes) { + return false; + } + int64_t plan_kernel_size = flattened_kernel_size; + if (params.groups > 1) { + if (plan_kernel_size > + std::numeric_limits::max() / params.groups) { + return false; + } + plan_kernel_size *= params.groups; + } + const Q8taConv2dStreamPlan device_plan = + make_q8ta_conv2d_stream_plan_for_device( + params.batch, + plan_kernel_size, + params.out_height, + params.out_width, + params.max_buffer_bytes); + return device_plan.feasible; + } + + if (params.is_mali) { + return true; + } + + // Single-batch heuristic: im2col pays off when wide channels + // amortize the materialization over GEMM work, or when a small output + // keeps the materialized buffer cheap. Anything else stays direct. + return params.groups == 1 && + (params.in_channels_per_group >= 32 || spatial_out <= 4096); +} + +} // namespace vkcompute diff --git a/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h new file mode 100644 index 00000000000..9f60b6fd372 --- /dev/null +++ b/backends/vulkan/runtime/graph/ops/impl/Q8taConv2dRoute.h @@ -0,0 +1,33 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include + +namespace vkcompute { + +struct Q8taConv2dRouteParams final { + bool is_mali; + bool supports_int8_dot_product; + // Adapter::max_buffer_numel(), which returns maxStorageBufferRange in bytes + // (not elements). + uint64_t max_buffer_bytes; + int64_t batch; + int64_t groups; + int64_t in_channels_per_group; + int64_t out_channels; + int64_t kernel_height; + int64_t kernel_width; + int64_t out_height; + int64_t out_width; +}; + +bool should_use_q8ta_conv2d_im2col(const Q8taConv2dRouteParams& params); + +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp index 250ea51746d..e4af2c56b9d 100644 --- a/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp +++ b/backends/vulkan/test/custom_ops/impl/TestQ8taConv2d.cpp @@ -252,28 +252,34 @@ void test_q8ta_conv2d(ComputeGraph& graph, const std::vector& args) { const ValueRef dilation = args.at(idx++); const ValueRef groups = args.at(idx++); const ValueRef activation = args.at(idx++); - const ValueRef layout_int = args.at(idx++); + const ValueRef input_layout_int = args.at(idx++); + const ValueRef output_layout_int = args.at(idx++); const ValueRef impl_selector_str = args.at(idx++); const ValueRef fp_output = args.at(idx++); // Extract the layout parameter and cast to GPUMemoryLayout - int32_t layout_value = graph.extract_scalar(layout_int); - utils::GPUMemoryLayout layout = - static_cast(layout_value); + const auto input_layout = static_cast( + graph.extract_scalar(input_layout_int)); + const auto output_layout = static_cast( + graph.extract_scalar(output_layout_int)); // Extract the impl_selector string std::string impl_selector = graph.extract_string(impl_selector_str); // Create temporary packed int8 tensors for input and output TmpTensor packed_int8_input( - &graph, graph.sizes_of(fp_input), vkapi::kInt8x4, utils::kBuffer, layout); + &graph, + graph.sizes_of(fp_input), + vkapi::kInt8x4, + utils::kBuffer, + input_layout); TmpTensor packed_int8_output( &graph, graph.sizes_of(fp_output), vkapi::kInt8x4, utils::kBuffer, - layout); + output_layout); // Quantize floating point input to packed int8 add_q8ta_quantize_node( diff --git a/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp b/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp new file mode 100644 index 00000000000..63d154b170f --- /dev/null +++ b/backends/vulkan/test/custom_ops/q8ta_conv2d_route_test.cpp @@ -0,0 +1,473 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include + +#include + +namespace vkcompute { +namespace { + +constexpr uint64_t kLargeDeviceBuffer = 1ULL << 32; + +Q8taConv2dRouteParams make_mali_grouped_params() { + return { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/2, + /*in_channels_per_group=*/32, + /*out_channels=*/64, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/128, + /*out_width=*/128, + }; +} + +TEST(Q8taConv2dRouteTest, RoutesBatchedRegularMaliConvolutionToIm2Col) { + EXPECT_TRUE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + })); + + EXPECT_TRUE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/128, + /*out_channels=*/256, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/10, + /*out_width=*/13, + })); +} + +TEST(Q8taConv2dRouteTest, RoutesMeasuredSceneXGroupedConvolutionsToIm2Col) { + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(make_mali_grouped_params())); + + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 4; + params.in_channels_per_group = 32; + params.out_channels = 128; + params.kernel_height = 5; + params.kernel_width = 5; + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params = make_mali_grouped_params(); + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params = make_mali_grouped_params(); + params.groups = 3; + params.in_channels_per_group = 32; + params.out_channels = 96; + params.kernel_height = 4; + params.kernel_width = 4; + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RoutesPreviouslyOutOfEnvelopeShapesToIm2Col) { + Q8taConv2dRouteParams params = make_mali_grouped_params(); + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.is_mali = true; + params.supports_int8_dot_product = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.supports_int8_dot_product = true; + + // No group-count gate: regular and wide grouped shapes route alike, as + // long as each group owns whole packed-4 output blocks. + params.groups = 1; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 5; + params.out_channels = 80; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 2; + params.out_channels = 64; + + // No kernel squareness/size gate. + params.kernel_width = 5; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 2; + params.kernel_width = 2; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 6; + params.kernel_width = 6; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 3; + params.kernel_width = 3; + + // Packing alignment is still required. + params.in_channels_per_group = 31; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 34; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 32; + + // Output channels indivisible by groups fail closed. + params.out_channels = 63; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 64; + // No spatial window gates. + params.out_height = 63; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 128; + params.out_width = 129; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsMisalignedGroupedOutputChannels) { + // The PW GEMM derives group_idx = oc_block / OC4_per_group, so a group + // with a non-multiple-of-4 channel count would straddle packed blocks. + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 3; + params.out_channels = 66; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 5; + params.out_channels = 65; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 80; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsMisalignedGroupedSingleBatchMali) { + // The single-batch Mali branch dispatches the same grouped PW shader, so + // the packed-4 output-block requirement binds there too. + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.batch = 1; + params.groups = 2; + params.out_channels = 66; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RoutesSingleTileScratchToIm2ColOffTileToDirect) { + // The legacy full-scratch envelope (up to 32MiB) narrowed to a single + // 16MiB streaming tile. With K=1152 and 8x8 output, batch 227 needs + // 15.96MiB in one tile (routes im2col) while batch 228 needs 16.03MiB in + // two tiles (stays direct on non-Mali, though both fit the old budget). + auto make_params = [](int64_t batch) { + return Q8taConv2dRouteParams{ + /*is_mali=*/false, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + batch, + /*groups=*/1, + /*in_channels_per_group=*/128, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/8, + /*out_width=*/8, + }; + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(make_params(227))); + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(make_params(228))); +} + +TEST(Q8taConv2dRouteTest, RespectsMaliDeviceBufferLimitForGroupedShapes) { + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 4; + params.out_channels = 128; + params.kernel_height = 5; + params.kernel_width = 5; + params.out_height = 64; + params.out_width = 64; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + // One im2col row is 3200 kernel columns x 64 aligned width bytes. + constexpr uint64_t kBytesPerRow = 3200 * 64; + params.max_buffer_bytes = kBytesPerRow; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.max_buffer_bytes = kBytesPerRow - 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsOverflowingMaliGroupedGeometry) { + Q8taConv2dRouteParams params = make_mali_grouped_params(); + params.groups = 4; + params.out_channels = 128; + // max()/36 + 2 is 0 (mod 4), so it passes the alignment gate and reaches + // the grouped overflow guard: align(C*3*3) > max()/4 with groups == 4. + params.in_channels_per_group = std::numeric_limits::max() / 36 + 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params = make_mali_grouped_params(); + params.out_height = std::numeric_limits::max(); + params.out_width = 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsIneligibleBatchedMaliConvolutions) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + + // Grouped and pointwise batched shapes route alike on Mali. + params.groups = 2; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 1; + params.kernel_height = 1; + params.kernel_width = 1; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.kernel_height = 3; + params.kernel_width = 3; + + // Packing alignment is still required. + params.in_channels_per_group = 62; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 64; + + // Kernels past the unsigned-dot accumulator bound stay direct. + params.in_channels_per_group = 4096; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RoutesSmallAndLargeBatchedShapesToIm2Col) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/32, + /*out_channels=*/64, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + // No kernel-size gate. + params.in_channels_per_group = 64; + params.kernel_height = 2; + params.kernel_width = 2; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 32; + params.kernel_height = 3; + params.kernel_width = 3; + + params.supports_int8_dot_product = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.supports_int8_dot_product = true; + params.in_channels_per_group = 30; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.in_channels_per_group = 32; + + // No output-channel or spatial window gates. + params.out_channels = 63; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 64; + params.out_height = 8; + params.out_width = 16; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 8; + params.out_width = 15; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 32; + params.out_width = 32; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 1; + params.out_width = 521; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.out_height = 20; + params.out_width = 26; + + // No pointwise exclusion. + params.in_channels_per_group = 256; + params.kernel_height = 1; + params.kernel_width = 1; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RespectsMaliDeviceBufferLimit) { + EXPECT_FALSE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + /*max_buffer_bytes=*/15 * 1024, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + })); + + // One im2col row is 576 kernel columns x 28 aligned width bytes; the + // device plan is feasible down to exactly that budget. + constexpr uint64_t kBytesPerRow = 576 * 28; + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + /*max_buffer_bytes=*/kBytesPerRow, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + params.max_buffer_bytes = kBytesPerRow - 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, PreservesNonMaliBatchedHeuristic) { + Q8taConv2dRouteParams params = { + /*is_mali=*/false, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/128, + /*out_channels=*/256, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/10, + /*out_width=*/13, + }; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.out_height = 8; + params.out_width = 8; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = true; + params.supports_int8_dot_product = false; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = false; + params.max_buffer_bytes = 1024; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, LegacyBatchedRoutePrecedesMaliExtensionGates) { + EXPECT_TRUE(should_use_q8ta_conv2d_im2col({ + /*is_mali=*/true, + /*supports_int8_dot_product=*/false, + kLargeDeviceBuffer, + /*batch=*/2, + /*groups=*/1, + /*in_channels_per_group=*/1024, + /*out_channels=*/32, + /*kernel_height=*/1, + /*kernel_width=*/1, + /*out_height=*/8, + /*out_width=*/8, + })); +} + +TEST(Q8taConv2dRouteTest, PreservesSingleBatchPolicy) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/1, + /*groups=*/4, + /*in_channels_per_group=*/8, + /*out_channels=*/32, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/128, + /*out_width=*/128, + }; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); + + params.is_mali = false; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 1; + params.in_channels_per_group = 32; + EXPECT_TRUE(should_use_q8ta_conv2d_im2col(params)); +} + +TEST(Q8taConv2dRouteTest, RejectsInvalidAndOverflowingGeometry) { + Q8taConv2dRouteParams params = { + /*is_mali=*/true, + /*supports_int8_dot_product=*/true, + kLargeDeviceBuffer, + /*batch=*/60, + /*groups=*/1, + /*in_channels_per_group=*/64, + /*out_channels=*/128, + /*kernel_height=*/3, + /*kernel_width=*/3, + /*out_height=*/20, + /*out_width=*/26, + }; + + params.groups = 0; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.groups = 1; + params.batch = 0; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.batch = 60; + params.out_channels = 0; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + params.out_channels = 128; + params.out_height = std::numeric_limits::max(); + params.out_width = 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.out_height = 20; + params.out_width = 26; + params.in_channels_per_group = std::numeric_limits::max(); + params.kernel_height = 2; + params.kernel_width = 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.in_channels_per_group = 4; + params.kernel_height = std::numeric_limits::max() / 4; + params.kernel_width = 2; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); + + params.in_channels_per_group = std::numeric_limits::max() - 2; + params.kernel_height = 1; + params.kernel_width = 1; + EXPECT_FALSE(should_use_q8ta_conv2d_im2col(params)); +} + +} // namespace +} // namespace vkcompute diff --git a/backends/vulkan/test/custom_ops/targets.bzl b/backends/vulkan/test/custom_ops/targets.bzl index df3dbd925f3..1895180495a 100644 --- a/backends/vulkan/test/custom_ops/targets.bzl +++ b/backends/vulkan/test/custom_ops/targets.bzl @@ -108,6 +108,16 @@ def define_common_targets(is_fbcode = False): ], ) + runtime.cxx_test( + name = "q8ta_conv2d_route_test", + srcs = ["q8ta_conv2d_route_test.cpp"], + platforms = get_platforms(), + deps = [ + "//third-party/googletest:gtest_main", + "//executorch/backends/vulkan:vulkan_graph_runtime", + ], + ) + define_custom_op_test_binary("test_add") define_custom_op_test_binary("test_q8csw_linear") define_custom_op_test_binary("test_q8csw_conv2d") diff --git a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp index 56ea6e177c9..90fd4e1616b 100644 --- a/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp +++ b/backends/vulkan/test/custom_ops/test_q8ta_conv2d.cpp @@ -40,11 +40,12 @@ struct Im2colUnsignedTestOptions { }; // Utility function to create a test case from a Conv2dConfig -static TestCase create_test_case_from_config( +static TestCase create_test_case_from_config_with_layouts( const Conv2dConfig& config, vkapi::ScalarType input_dtype, utils::StorageType fp_storage_type, - utils::GPUMemoryLayout int8_memory_layout, + utils::GPUMemoryLayout input_int8_memory_layout, + utils::GPUMemoryLayout output_int8_memory_layout, const std::string& impl_selector = "", const Im2colUnsignedTestOptions* im2col_options = nullptr, const float input_scale_val = 0.008123f, @@ -81,7 +82,10 @@ static TestCase create_test_case_from_config( std::to_string(config.stride.h) + " p" + std::to_string(config.padding.h) + " d" + std::to_string(config.dilation.h) + " g" + std::to_string(config.groups); - std::string storage_str = repr_str(utils::kBuffer, int8_memory_layout); + std::string storage_str = repr_str(utils::kBuffer, input_int8_memory_layout); + if (input_int8_memory_layout != output_int8_memory_layout) { + storage_str += "->" + repr_str(utils::kBuffer, output_int8_memory_layout); + } std::string suffix = impl_selector.empty() ? "" : "[" + impl_selector + "]"; std::string test_name = make_test_label( prefix, dtype_str, dtype_str, shape_str, storage_str, suffix); @@ -252,8 +256,11 @@ static TestCase create_test_case_from_config( test_case.add_input_spec(activation); // Add memory layout parameter for the quantized tensors - ValueSpec layout_int(static_cast(int8_memory_layout)); - test_case.add_input_spec(layout_int); + ValueSpec input_layout_int(static_cast(input_int8_memory_layout)); + test_case.add_input_spec(input_layout_int); + + ValueSpec output_layout_int(static_cast(output_int8_memory_layout)); + test_case.add_input_spec(output_layout_int); // Add impl_selector string ValueSpec impl_selector_spec = ValueSpec::make_string(impl_selector); @@ -276,6 +283,27 @@ static TestCase create_test_case_from_config( return test_case; } +static TestCase create_test_case_from_config( + const Conv2dConfig& config, + vkapi::ScalarType input_dtype, + utils::StorageType fp_storage_type, + utils::GPUMemoryLayout int8_memory_layout, + const std::string& impl_selector = "", + const Im2colUnsignedTestOptions* im2col_options = nullptr, + const float input_scale_val = 0.008123f, + const DataGenType input_data_gen = DataGenType::RANDOM) { + return create_test_case_from_config_with_layouts( + config, + input_dtype, + fp_storage_type, + int8_memory_layout, + int8_memory_layout, + impl_selector, + im2col_options, + input_scale_val, + input_data_gen); +} + static std::vector generate_narrow_workgroup_test_cases() { std::vector test_cases; std::vector configs = { @@ -487,6 +515,94 @@ static std::vector generate_streaming_im2col_test_cases() { return test_cases; } +// SceneX route tests. The kPackedInt8_4C input + kPackedInt8_4W4C output +// layout combination is only exercised here (default generators never pair +// them); run via --scenex-regular . Zero +// tolerances are exact by construction: both pipelines accumulate in int32 +// with identical requantize, so any mismatch is a real regression, not noise. +static TestCase create_scenex_test_case( + const Conv2dConfig& source_config, + const std::string& route) { + Conv2dConfig config = source_config; + config.op_name = "conv2d_q8ta_q8csw_q8to"; + config.test_case_name = + make_test_case_name(config, false, utils::kTexture3D, utils::kBuffer); + + const std::string impl_selector = route == "auto" ? "" + : route == "direct" ? "general" + : "im2col"; + TestCase test_case = create_test_case_from_config_with_layouts( + config, + vkapi::kFloat, + utils::kTexture3D, + utils::kPackedInt8_4C, + utils::kPackedInt8_4W4C, + impl_selector, + /*im2col_options=*/nullptr, + /*input_scale_val=*/1.0f, + /*input_data_gen=*/DataGenType::RANDINT); + test_case.set_abs_tolerance(0.0f); + test_case.set_rel_tolerance(0.0f); + return test_case; +} + +static TestCase generate_scenex_regular_test_case( + const std::string& route, + const int case_index) { + const std::vector configs = { + {OutInChannels(128, 64), + InputSize2D(40, 51), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}, + {OutInChannels(256, 128), + InputSize2D(20, 26), + KernelSize(3, 3), + Stride(2, 2), + Padding(1, 1), + Dilation(1, 1), + 1, + 60}, + }; + return create_scenex_test_case(configs.at(case_index), route); +} + +static TestCase generate_scenex_grouped_test_case( + const std::string& route, + const int case_index) { + const std::vector configs = { + {OutInChannels(64, 64), + InputSize2D(128, 128), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2, + 60}, + {OutInChannels(128, 128), + InputSize2D(128, 128), + KernelSize(5, 5), + Stride(2, 2), + Padding(2, 2), + Dilation(1, 1), + 4, + 60}, + {OutInChannels(64, 64), + InputSize2D(64, 64), + KernelSize(3, 3), + Stride(1, 1), + Padding(1, 1), + Dilation(1, 1), + 2, + 60}, + }; + + return create_scenex_test_case(configs.at(case_index), route); +} + // Generate test cases for quantized conv2d operation static std::vector generate_quantized_conv2d_test_cases() { std::vector test_cases; @@ -966,6 +1082,8 @@ static void conv2d_q8ta_q8csw_q8to_reference_impl(TestCase& test_case) { const ValueSpec& activation_spec = test_case.inputs()[idx++]; const ValueSpec& layout_spec = test_case.inputs()[idx++]; (void)layout_spec; // Not used in reference implementation + const ValueSpec& output_layout_spec = test_case.inputs()[idx++]; + (void)output_layout_spec; // Not used in reference implementation const ValueSpec& impl_selector_spec = test_case.inputs()[idx++]; (void)impl_selector_spec; // Not used in reference implementation @@ -1153,6 +1271,41 @@ static void reference_impl(TestCase& test_case) { conv2d_q8ta_q8csw_q8to_reference_impl(test_case); } +// The impl selector holds one of "", "general", or "im2col"; activation and +// other string inputs use disjoint values. Overwrite it by value so a +// reordered input list fails loudly instead of mutating the wrong spec. +// Note: for route == "direct" the measured run already forces "general", so +// this reference re-executes the identical implementation and only checks +// determinism; genuine cross-implementation correctness comes from the +// auto/im2col legs. +static void scenex_direct_reference(TestCase& test_case) { + TestCase direct_case = test_case; + bool found_selector = false; + for (auto it = direct_case.inputs().rbegin(); + it != direct_case.inputs().rend(); + ++it) { + if (it->is_string() && + (it->get_string_value().empty() || + it->get_string_value() == "general" || + it->get_string_value() == "im2col")) { + it->string_data = "general"; + found_selector = true; + break; + } + } + if (!found_selector) { + throw std::runtime_error("scenex reference: impl selector input not found"); + } + execute_test_case( + direct_case, + /*warmup_runs=*/1, + /*benchmark_runs=*/1, + /*chained_dispatches=*/1, + /*write_outputs=*/true); + test_case.outputs().at(0).get_ref_float_data() = + direct_case.outputs().at(0).get_float_data(); +} + // Custom FLOP calculator for quantized conv2d operation static int64_t quantized_conv2d_flop_calculator(const TestCase& test_case) { int kernel_idx = 9; // kernel_size is at index 9 for q8ta_q8csw_q8to @@ -1276,6 +1429,14 @@ static void execute_streaming_dynamic_shrink_test() { std::cout << "Streaming im2col dynamic shrink PASSED" << std::endl; } +// Single usage string for the scenex route-test modes; argument errors +// return 2 like the other CLI errors in main. +static int print_scenex_usage(const char* mode, const char* cases) { + std::cerr << "Usage: " << mode << " <" << cases << ">" + << std::endl; + return 2; +} + int main(int argc, char* argv[]) { const vkapi::Adapter& adapter = *vkcompute::api::context()->adapter_ptr(); const bool prefers_unsigned_dot = @@ -1291,29 +1452,44 @@ int main(int argc, char* argv[]) { bool narrow_workgroups_only = false; bool streaming_im2col_only = false; bool streaming_dynamic_shrink_only = false; - for (int i = 1; i < argc; ++i) { - const std::string arg(argv[i]); - if (arg == "--im2col-path=signed") { - im2col_impl_selector = "im2col"; - } else if (arg == "--im2col-path=unsigned") { - im2col_impl_selector = "im2col_unsigned"; - } else if (arg == "--im2col-path=auto") { - im2col_impl_selector = "im2col_auto"; - } else if (arg == "--narrow-workgroups-only") { - narrow_workgroups_only = true; - } else if (arg == "--streaming-im2col-only") { - streaming_im2col_only = true; - } else if (arg == "--streaming-dynamic-shrink-only") { - streaming_dynamic_shrink_only = true; - } else { - std::cerr << "Unknown argument: " << arg << std::endl; - return 2; + const bool scenex_regular = + argc == 4 && std::string(argv[1]) == "--scenex-regular"; + const bool scenex_grouped = + argc == 4 && std::string(argv[1]) == "--scenex-grouped"; + if (argc >= 2 && std::string(argv[1]) == "--scenex-regular" && + !scenex_regular) { + return print_scenex_usage("--scenex-regular", "0|1"); + } + if (argc >= 2 && std::string(argv[1]) == "--scenex-grouped" && + !scenex_grouped) { + return print_scenex_usage("--scenex-grouped", "0|1|2"); + } + if (!scenex_regular && !scenex_grouped) { + for (int i = 1; i < argc; ++i) { + const std::string arg(argv[i]); + if (arg == "--im2col-path=signed") { + im2col_impl_selector = "im2col"; + } else if (arg == "--im2col-path=unsigned") { + im2col_impl_selector = "im2col_unsigned"; + } else if (arg == "--im2col-path=auto") { + im2col_impl_selector = "im2col_auto"; + } else if (arg == "--narrow-workgroups-only") { + narrow_workgroups_only = true; + } else if (arg == "--streaming-im2col-only") { + streaming_im2col_only = true; + } else if (arg == "--streaming-dynamic-shrink-only") { + streaming_dynamic_shrink_only = true; + } else { + std::cerr << "Unknown argument: " << arg << std::endl; + return 2; + } } } const int selected_modes = static_cast(!im2col_impl_selector.empty()) + static_cast(narrow_workgroups_only) + static_cast(streaming_im2col_only) + - static_cast(streaming_dynamic_shrink_only); + static_cast(streaming_dynamic_shrink_only) + + static_cast(scenex_regular) + static_cast(scenex_grouped); if (selected_modes > 1) { std::cerr << "Test mode selectors are mutually exclusive" << std::endl; return 2; @@ -1334,6 +1510,8 @@ int main(int argc, char* argv[]) { print_separator(); ReferenceComputeFunc ref_fn = reference_impl; + int warmup_runs = 1; + int benchmark_runs = 1; if (streaming_dynamic_shrink_only) { execute_streaming_dynamic_shrink_test(); @@ -1365,14 +1543,44 @@ int main(int argc, char* argv[]) { return cases; }; #endif + } else if (scenex_regular) { + const std::string route = argv[2]; + const std::string case_arg = argv[3]; + if ((route != "auto" && route != "direct" && route != "im2col") || + (case_arg != "0" && case_arg != "1")) { + return print_scenex_usage("--scenex-regular", "0|1"); + } + const int case_index = case_arg == "0" ? 0 : 1; + test_case_generator = [route, case_index]() { + return std::vector{ + generate_scenex_regular_test_case(route, case_index)}; + }; + ref_fn = scenex_direct_reference; + warmup_runs = 3; + benchmark_runs = 10; + } else if (scenex_grouped) { + const std::string route = argv[2]; + const std::string case_arg = argv[3]; + if ((route != "auto" && route != "direct" && route != "im2col") || + (case_arg != "0" && case_arg != "1" && case_arg != "2")) { + return print_scenex_usage("--scenex-grouped", "0|1|2"); + } + const int case_index = case_arg == "0" ? 0 : case_arg == "1" ? 1 : 2; + test_case_generator = [route, case_index]() { + return std::vector{ + generate_scenex_grouped_test_case(route, case_index)}; + }; + ref_fn = scenex_direct_reference; + warmup_runs = 3; + benchmark_runs = 10; } auto results = execute_test_cases( test_case_generator, quantized_conv2d_flop_calculator, "QuantizedConv2dQ8ToQ8To", - /*warmup_runs = */ 1, - /*benchmark_runs = */ 1, + warmup_runs, + benchmark_runs, ref_fn); #ifndef DEBUG_MODE From 3faa2e4d077123181529472a75a9e82d8cbf2542 Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Fri, 11 Sep 2026 01:01:20 +0200 Subject: [PATCH 164/190] [ET-VK] Scale reduction workers with the length of the reduction (#22348) `reduce_gwg()` sizes every texture-storage reduction from the **output** extents and then hardcodes the thread count: ```cpp constexpr uint32_t max_nthreads = 16u; constexpr uint32_t nworkers_per_group = 4u; constexpr uint32_t ngroups = 4u; ``` A global average pool has a `1x1xC/4` output, so the dispatch comes out as global `{1,1,4}` local `{4,4,1}`: four threads walk an entire HxW plane and the whole dispatch runs 16 threads, regardless of how much there is to reduce. `NWORKERS` becomes a specialization constant so the dispatch can choose it, `MAX_NTHREADS` goes 16 -> 256 (4 KiB of shared `vec4`, inside the guaranteed 16 KiB), and `reduce_nworkers()` derives the count from the reduction extent. It feeds **both** `reduce_gwg()` and the node's specialization constants, since the shader's stride and the local work group have to agree. ### Measurements MediaPipe selfie segmentation @256, where the MobileNetV3 squeeze-excitation pools land on this path. Per-dispatch GPU timings from the shader query pool; wall clock is best of 30 executions over four interleaved order-reversed rounds. | | Adreno 840 | Mali-G76 | | --- | --- | --- | | `mean2d` total | 0.767 ms (35.0% of GPU time) -> **0.115 ms (7.4%)** | | | model | 2.306 -> **1.625 ms** (-28.5%) | 30.3 -> **25.4 ms** (-17%) | Output is bit-exact against the reference on both devices (cosine 1.0000000, max abs diff 0.0), and 30/30 identical across runs on the Mali. ### Scope Covers `aten.sum.dim_IntList`, `aten.mean.dim`, `aten.amax.default` and `aten.amin.default` on texture storage, via both `add_reduce_node` and `add_reduce2d_node`. A shared-memory tree for the aggregation after the barrier was tried and **not** kept: it measured no faster here (1.625 -> 1.649 ms), because only thread 0 of each group runs that walk. It is a large win in `softmax_buffer.glsl`, where every thread runs it, and that is handled separately. ### Note This touches `reduce.glsl` and `reduce2d.glsl` in different hunks from #22326, so the two are independent and can land in either order. Fixes #22350 --- .../vulkan/runtime/graph/ops/glsl/reduce.glsl | 102 +++++++++----- .../runtime/graph/ops/glsl/reduce2d.glsl | 54 +++++--- .../vulkan/runtime/graph/ops/impl/Reduce.cpp | 130 +++++++++++++++--- 3 files changed, 211 insertions(+), 75 deletions(-) diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl index 209440cec6a..029e3b16756 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce.glsl @@ -32,12 +32,17 @@ layout(constant_id = 5) const int group_dim = 1; // A more verbose name would be NWORKERS_PER_GROUP. This describes the number of // threads that will co-operate to compute one reduction output. There may be // multiple groups computing distinct reduction outputs within one work group. -#define NWORKERS 4 +// Supplied by the dispatch so it can scale with the length of the reduction. +// A global average pool reduces a whole HxW plane into one value, and four +// workers left the GPU essentially idle for it. +layout(constant_id = 6) const int NWORKERS = 4; // Sets an upper limit on the total size of a work group based on how many // elements are allocated in the shared memory array below. Each thread in the // work group will write into its assigned element in the shared array. -#define MAX_NTHREADS 16 +// Upper bound on NWORKERS * NGROUPS, and the size of the shared array below. +// 256 vec4 is 4 KiB of shared memory, well inside the guaranteed 16 KiB. +#define MAX_NTHREADS 256 shared vec4 shared_vecs[MAX_NTHREADS]; @@ -86,19 +91,30 @@ int tid_to_smi(const ivec2 tid) { * This case is simpler because each element of a texel belongs to a separate * reduction "group", meaning we don't have to perform reduction along a texel. */ -void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_nonpacked_dim( + const ivec2 tid, + ivec3 scan_pos, + const bool in_bounds) { // shared memory index of this thread const int smi = tid_to_smi(tid); - scan_pos[reduce_dim] = 0; - vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); - - scan_pos[reduce_dim] = tid.x; - // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of - // the reduction row - for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim); - i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + // Out of bounds invocations cannot return early: barrier() below has to be + // reached by every invocation in the work group, and skipping it is undefined + // behaviour that hangs some GPUs. They still take a shared memory slot, but + // it is one that no in-bounds group aggregates over, so what they leave in it + // is never read. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim] = 0; + accum = INIT_ACCUM(load_texel(tin, scan_pos)); + + scan_pos[reduce_dim] = tid.x; + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... + // of the reduction row + for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim); + i += NWORKERS, scan_pos[reduce_dim] += NWORKERS) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } } // Write partial output to shared memory and synchronize work group shared_vecs[smi] = accum; @@ -106,7 +122,7 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { // Since the reduction row is reduced to only one element, only the "main" // thread in the group needs aggregate the partial outputs - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial outputs to obtain the overall output int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -141,7 +157,10 @@ void reduce_nonpacked_dim(const ivec2 tid, ivec3 scan_pos) { * elements in texels (which occur when the size of the packed dim is not a * multiple of 4) so that they do not influence the output of reduction. */ -void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_packed_dim( + const ivec2 tid, + ivec3 scan_pos, + const bool in_bounds) { // shared memory index of this thread const int smi = tid_to_smi(tid); @@ -151,23 +170,32 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { // handled specially if it has padding elements. const int reduce_len = safe_idx(tin_sizes, packed_dim) - nspill; - scan_pos[reduce_dim] = 0; - vec4 accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); - - // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... of - // the reduction row - scan_pos[reduce_dim] = tid.x; - for (int i = tid.x * 4; i < reduce_len; - i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); - } - // For the last texel in the dim, if there are padding elements then each - // element of the texel needs to be processed individually such that the - // padding elements are ignored - if (scan_pos[reduce_dim] == safe_idx(tin_limits, reduce_dim) - 1 && nspill > 0) { - const vec4 intex = load_texel(tin, scan_pos); - for (int i = 0; i < nspill; i++) { - accum.x = UPDATE_ACCUM(accum.x, intex[i]); + // Out of bounds invocations cannot return early: barrier() below has to be + // reached by every invocation in the work group, and skipping it is undefined + // behaviour that hangs some GPUs. They still take a shared memory slot, but + // it is one that no in-bounds group aggregates over, so what they leave in it + // is never read. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim] = 0; + accum = INIT_ACCUM(vec4(load_texel(tin, scan_pos).x)); + + // Partially accumulate over elements i, i + NWORKERS, i + 2*NWORKERS, ... + // of the reduction row + scan_pos[reduce_dim] = tid.x; + for (int i = tid.x * 4; i < reduce_len; + i += NWORKERS * 4, scan_pos[reduce_dim] += NWORKERS) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } + // For the last texel in the dim, if there are padding elements then each + // element of the texel needs to be processed individually such that the + // padding elements are ignored + if (scan_pos[reduce_dim] == safe_idx(tin_limits, reduce_dim) - 1 && + nspill > 0) { + const vec4 intex = load_texel(tin, scan_pos); + for (int i = 0; i < nspill; i++) { + accum.x = UPDATE_ACCUM(accum.x, intex[i]); + } } } // Write partial output to shared memory and synchronize work group @@ -176,7 +204,7 @@ void reduce_packed_dim(const ivec2 tid, ivec3 scan_pos) { // Since the reduction row is reduced to only one element, only the "main" // thread in the group needs aggregate the partial outputs - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial maximums to obtain the overall maximum int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -203,13 +231,13 @@ void main() { gl_LocalInvocationID[reduce_dim], gl_LocalInvocationID[group_dim]); - if (any(greaterThanEqual(scan_pos, tin_limits))) { - return; - } + const bool in_bounds = all(lessThan(scan_pos, tin_limits)); + // reduce_dim and packed_dim are specialization constants, so this branch is + // uniform across the work group and safe to take around a barrier. if (reduce_dim != packed_dim) { - reduce_nonpacked_dim(tid, scan_pos); + reduce_nonpacked_dim(tid, scan_pos, in_bounds); } else { - reduce_packed_dim(tid, scan_pos); + reduce_packed_dim(tid, scan_pos, in_bounds); } } diff --git a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl index bd55025f534..58e7c6d0b3d 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/reduce2d.glsl @@ -33,12 +33,17 @@ layout(constant_id = 6) const int group_dim = 2; // A more verbose name would be NWORKERS_PER_GROUP. This describes the number of // threads that will co-operate to compute one reduction output. There may be // multiple groups computing distinct reduction outputs within one work group. -#define NWORKERS 4 +// Supplied by the dispatch so it can scale with the length of the reduction. +// A global average pool reduces a whole HxW plane into one value, and four +// workers left the GPU essentially idle for it. +layout(constant_id = 7) const int NWORKERS = 4; // Sets an upper limit on the total size of a work group based on how many // elements are allocated in the shared memory array below. Each thread in the // work group will write into its assigned element in the shared array. -#define MAX_NTHREADS 16 +// Upper bound on NWORKERS * NGROUPS, and the size of the shared array below. +// 256 vec4 is 4 KiB of shared memory, well inside the guaranteed 16 KiB. +#define MAX_NTHREADS 256 shared vec4 shared_vecs[MAX_NTHREADS]; @@ -59,23 +64,34 @@ int tid_to_smi(const ivec2 tid) { // with the accumulator. #define POSTPROCESS(accum) ${POSTPROCESS} -void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { +void reduce_2d_non_packed_dim( + const ivec2 tid, + ivec3 scan_pos, + const bool in_bounds) { // shared memory index of this thread const int smi = tid_to_smi(tid); - scan_pos[reduce_dim1] = 0; - scan_pos[reduce_dim2] = 0; - vec4 accum = INIT_ACCUM(load_texel(tin, scan_pos)); - - // First dimension reduction - scan_pos[reduce_dim1] = tid.x; - for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim1); - i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { - - // Second dimension reduction + // Out of bounds invocations cannot return early: barrier() below has to be + // reached by every invocation in the work group, and skipping it is undefined + // behaviour that hangs some GPUs. They still take a shared memory slot, but + // it is one that no in-bounds group aggregates over, so what they leave in it + // is never read. + vec4 accum = vec4(0); + if (in_bounds) { + scan_pos[reduce_dim1] = 0; scan_pos[reduce_dim2] = 0; - for (int j = 0; j < safe_idx(tin_sizes, reduce_dim2); j++, scan_pos[reduce_dim2]++) { - accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + accum = INIT_ACCUM(load_texel(tin, scan_pos)); + + // First dimension reduction + scan_pos[reduce_dim1] = tid.x; + for (int i = tid.x; i < safe_idx(tin_sizes, reduce_dim1); + i += NWORKERS, scan_pos[reduce_dim1] += NWORKERS) { + // Second dimension reduction + scan_pos[reduce_dim2] = 0; + for (int j = 0; j < safe_idx(tin_sizes, reduce_dim2); + j++, scan_pos[reduce_dim2]++) { + accum = UPDATE_ACCUM(accum, load_texel(tin, scan_pos)); + } } } @@ -84,7 +100,7 @@ void reduce_2d_non_packed_dim(const ivec2 tid, ivec3 scan_pos) { barrier(); // Main thread aggregates results - if (tid.x == 0) { + if (in_bounds && tid.x == 0) { // Iterate over the partial outputs to obtain the overall output int group_i = tid.y * NWORKERS; accum = shared_vecs[group_i++]; @@ -121,9 +137,7 @@ void main() { gl_LocalInvocationID[reduce_dim1], gl_LocalInvocationID[group_dim]); - if (any(greaterThanEqual(scan_pos, tin_limits))) { - return; - } + const bool in_bounds = all(lessThan(scan_pos, tin_limits)); - reduce_2d_non_packed_dim(tid, scan_pos); + reduce_2d_non_packed_dim(tid, scan_pos, in_bounds); } \ No newline at end of file diff --git a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp index f684906cf57..23c8f000bed 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Reduce.cpp @@ -14,6 +14,8 @@ #include #include +#include + namespace vkcompute { using namespace utils; @@ -72,30 +74,111 @@ void resize_reduce_per_row_node( graph->virtual_resize(out, new_sizes); } -GlobalWorkGrid reduce_gwg( +// Number of threads that co-operate on one reduction output, and how many +// outputs one work group covers. +// +// The worker count used to be a flat 4 regardless of how much there was to +// reduce. A global average pool collapses a whole HxW plane, so four threads +// each walked thousands of elements while the dispatch ran 16 threads in total. +constexpr uint32_t kReduceMaxNThreads = 256u; +constexpr uint32_t kReduceNGroups = 4u; + +// The largest worker count this dispatch may ask for. Three ceilings apply and +// the smallest of them wins: +// +// - the shaders size shared_vecs at MAX_NTHREADS and every thread in the work +// group writes its own slot, so kReduceNGroups workers must fit; +// - the device bounds the invocations in one work group, and +// maxComputeWorkGroupInvocations is only guaranteed to be 128; +// - the device bounds each work group axis on its own, and the workers all sit +// on the reduction axis. +// +// Overrunning any of them aborts the dispatch in LocalWorkGroup::validate, so +// the shader capacity alone is not enough to go by. +uint32_t reduce_nworkers_cap( + ComputeGraph* graph, + const int32_t reduce_dim_whcn) { + const vkapi::Adapter* const adapter = graph->context()->adapter_ptr(); + uint32_t cap = kReduceMaxNThreads / kReduceNGroups; + cap = std::min( + cap, adapter->max_compute_workgroup_invocations() / kReduceNGroups); + cap = std::min(cap, adapter->max_compute_workgroup_size()[reduce_dim_whcn]); + return std::max(cap, 1u); +} + +uint32_t reduce_nworkers( + ComputeGraph* graph, + const ValueRef in, + const int32_t reduce_dim_whcn) { + const uint32_t cap = reduce_nworkers_cap(graph, reduce_dim_whcn); + const uint32_t extent = utils::safe_downcast( + graph->logical_limits_of(in)[reduce_dim_whcn]); + // 4 is what this used to be unconditionally; keep it as the floor so short + // reductions dispatch exactly as they did before. + uint32_t nworkers = std::min(4u, cap); + while (nworkers * 2u <= cap && nworkers < extent) { + nworkers *= 2u; + } + return nworkers; +} + +GlobalWorkGrid reduce_gwg_impl( ComputeGraph* graph, - const vkapi::ShaderInfo& shader, const std::vector& args, - const std::vector& resize_args) { - (void)shader; + const std::vector& resize_args, + const size_t reduce_dim_idx, + const size_t group_dim_idx, + const size_t nworkers_idx) { const ValueRef out = args.at(0).refs.at(0); const int32_t reduce_dim_whcn = - graph->extract_scalar(resize_args.at(1)); + graph->extract_scalar(resize_args.at(reduce_dim_idx)); const int64_t group_dim_whcn = - graph->extract_scalar(resize_args.at(2)); + graph->extract_scalar(resize_args.at(group_dim_idx)); utils::uvec3 extents = graph->logical_limits_of(out); extents[reduce_dim_whcn] = 1; - constexpr uint32_t max_nthreads = 16u; - constexpr uint32_t nworkers_per_group = 4u; - constexpr uint32_t ngroups = 4u; - VK_CHECK_COND(nworkers_per_group * ngroups <= max_nthreads); + // NWORKERS is baked into the shader as a specialization constant when the + // node is built, so it reflects the dynamic upper bound. This callback runs + // again after every resize and would see the smaller actual extent, giving a + // work group with fewer threads along the reduction dim than the shader's + // aggregation loop indexes over -- it would then fold in shared memory slots + // that no thread wrote. Read the count the node was built with instead. + // + // Launching more workers than there are elements is harmless: a worker whose + // loop body never runs contributes INIT_ACCUM, which is the identity for sum + // and mean and idempotent for amax and amin. + const uint32_t nworkers_per_group = utils::safe_downcast( + graph->extract_scalar(resize_args.at(nworkers_idx))); utils::uvec3 lwg_extents{1u, 1u, 1u}; lwg_extents[reduce_dim_whcn] = nworkers_per_group; - lwg_extents[group_dim_whcn] = ngroups; + lwg_extents[group_dim_whcn] = kReduceNGroups; return GlobalWorkGrid(extents, kTiledWorkGrid, LocalWorkGroup(lwg_extents)); } +// Resize args are {dim, reduce_dim_whcn, group_dim_whcn, nworkers}. +GlobalWorkGrid reduce_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + return reduce_gwg_impl(graph, args, resize_args, 1, 2, 3); +} + +// Resize args are {dims, reduce_dim1_whcn, reduce_dim2_whcn, group_dim_whcn, +// nworkers}, so the group dim sits one slot further along than in the 1d case. +// Sharing the 1d picker put the groups on reduce_dim2 and left the group axis +// one thread wide, so every group read tid.y == 0 and raced over the same +// shared memory slots. +GlobalWorkGrid reduce2d_gwg( + ComputeGraph* graph, + const vkapi::ShaderInfo& shader, + const std::vector& args, + const std::vector& resize_args) { + (void)shader; + return reduce_gwg_impl(graph, args, resize_args, 1, 3, 4); +} + void add_reduce_node( ComputeGraph& graph, const ValueRef in, @@ -137,6 +220,9 @@ void add_reduce_node( const ValueRef reduce_dim_whcn_ref = graph.get_or_add_value_for_int(reduce_dim); const ValueRef group_dim_whcn_ref = graph.get_or_add_value_for_int(group_dim); + const int32_t nworkers = + utils::safe_downcast(reduce_nworkers(&graph, in, reduce_dim)); + const ValueRef nworkers_ref = graph.get_or_add_value_for_int(nworkers); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, @@ -149,10 +235,12 @@ void add_reduce_node( {graph.logical_limits_ubo(in), graph.sizes_ubo(in)}, // Push Constants {}, - // Specialization Constants - {graph.packed_dim_of(out), reduce_dim, group_dim}, + // Specialization Constants. NWORKERS must match the local work group + // extent reduce_gwg picks, so the count is computed once here and passed + // to reduce_gwg through the resize args. + {graph.packed_dim_of(out), reduce_dim, group_dim, nworkers}, // Resize Args - {dim_ref, reduce_dim_whcn_ref, group_dim_whcn_ref}, + {dim_ref, reduce_dim_whcn_ref, group_dim_whcn_ref, nworkers_ref}, // Resizing Logic resize_reduce_node)); } @@ -214,11 +302,14 @@ void add_reduce2d_node( const ValueRef reduce_dim2_whcn_ref = graph.get_or_add_value_for_int(reduce_dim2); const ValueRef group_dim_whcn_ref = graph.get_or_add_value_for_int(group_dim); + const int32_t nworkers = + utils::safe_downcast(reduce_nworkers(&graph, in, reduce_dim1)); + const ValueRef nworkers_ref = graph.get_or_add_value_for_int(nworkers); graph.execute_nodes().emplace_back(new DynamicDispatchNode( graph, VK_KERNEL_FROM_STR(kernel_name), - reduce_gwg, + reduce2d_gwg, pick_required_lwg, // Inputs and Outputs {{out, vkapi::kWrite}, {in, vkapi::kRead}}, @@ -226,13 +317,16 @@ void add_reduce2d_node( {graph.logical_limits_ubo(in), graph.sizes_ubo(in)}, // Push Constants {}, - // Specialization Constants - {graph.packed_dim_of(out), reduce_dim1, reduce_dim2, group_dim}, + // Specialization Constants. NWORKERS must match the local work group + // extent reduce_gwg picks, so the count is computed once here and passed + // to reduce_gwg through the resize args. + {graph.packed_dim_of(out), reduce_dim1, reduce_dim2, group_dim, nworkers}, // Resize Args {dims_ref, reduce_dim1_whcn_ref, reduce_dim2_whcn_ref, - group_dim_whcn_ref}, + group_dim_whcn_ref, + nworkers_ref}, // Resizing Logic resize_reduce2d_node)); } From bbb9d079b0803221406f6534a5777067fb08629e Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 10 Sep 2026 12:14:11 -0700 Subject: [PATCH 165/190] [executorch][native] Add deserializer bridge (Program -> Method) ## Ulterior Motive Turn serialized native graphs into executable in-memory IR. This is stack foundation: package loading and backend execution both depend on `Method` and `Graph` objects. ## Rationale **What**: Add lazy `Program`-to-`Method` deserialization, graph-local SSA resolution, method bindings, mutation metadata, nested subgraphs, and strict wire-data validation. **Why**: `Program::load` previously verified bytes but could not materialize runtime IR. ## Details ```text verified PTG | v Program::get_method(name) | +-- cache hit --> existing Method | +-- cache miss --> build_method() | +-- GraphBuilder per graph namespace +-- resolve names to stable ValueIds +-- attach bindings and output specs +-- rebuild def-use ``` - Dynamic extents are rejected instead of silently fixed at an upper bound. - Constants contribute tensor metadata missing from graph-local tables. - Tensor-list outputs, aliases, scalar references, and `GraphArg` subgraphs are resolved. - Unknown argument tags, unresolved mutation targets, and self-aliases fail during deserialization. Valid alias chains remain supported. - Method-name lookup compares FlatBuffer bytes in place without allocating a temporary string for every candidate. Authored with Codex. Differential Revision: [D114426391](https://our.internmc.facebook.com/intern/diff/D114426391/) ghstack-source-id: 427858760 Pull-Request: https://github.com/pytorch/executorch/pull/22699 --- backends/native/runtime/Deserialize.cpp | 617 ++++++++++++++++++ backends/native/runtime/Program.cpp | 28 + backends/native/runtime/Program.h | 25 +- backends/native/runtime/targets.bzl | 5 + backends/native/test/runtime/BUCK | 8 + backends/native/test/runtime/targets.bzl | 11 + .../test/runtime/test_program_deserialize.cpp | 299 +++++++++ 7 files changed, 992 insertions(+), 1 deletion(-) create mode 100644 backends/native/runtime/Deserialize.cpp create mode 100644 backends/native/test/runtime/BUCK create mode 100644 backends/native/test/runtime/targets.bzl create mode 100644 backends/native/test/runtime/test_program_deserialize.cpp diff --git a/backends/native/runtime/Deserialize.cpp b/backends/native/runtime/Deserialize.cpp new file mode 100644 index 00000000000..75410464b7b --- /dev/null +++ b/backends/native/runtime/Deserialize.cpp @@ -0,0 +1,617 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +// The deserializer bridge: native_backend::Program (FlatBuffer) -> ptn +// in-memory IR (Method / Graph / Node / Argument / Value). In-graph references +// (SSA names) are resolved to list ValueIds; per-graph namespaces are +// resolved independently (each HOP subgraph rebuilds its own name -> id map). +// +// Everything here runs on a buffer Program::load() has already put through +// flatbuffers::Verifier, so accessors return non-null wherever the schema +// declares the field required and wherever a union discriminator matches; the +// walkers below dereference those results directly. Fields the schema leaves +// optional are still checked, because verification says nothing about whether +// they are present. Vector>::Get() computes an address rather than a +// nullable pointer, and verification bounds-checks every referenced object. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include + +namespace ptn { +namespace { + +std::string str_of(const flatbuffers::String* s) { + return s != nullptr ? s->str() : std::string(); +} + +bool nonempty(const flatbuffers::String* s) { + return s != nullptr && s->size() > 0; +} + +ScalarType map_scalar_type(fbs::ScalarType t) { + // ptn::ScalarType ids are pinned to the schema's, so the byte maps straight. + return static_cast(static_cast(t)); +} + +OpKind map_op_kind(fbs::OpKind k) { + switch (k) { + case fbs::OpKind::CALL_FUNCTION: + return OpKind::CallFunction; + case fbs::OpKind::PLACEHOLDER: + return OpKind::Placeholder; + case fbs::OpKind::OUTPUT: + return OpKind::Output; + default: + throw std::runtime_error( + "build_graph: unsupported OpKind value " + + std::to_string(static_cast(k))); + } +} + +OutputValueKind map_output_value_kind(fbs::OutputValueKind k) { + switch (k) { + case fbs::OutputValueKind::TENSOR: + return OutputValueKind::Tensor; + case fbs::OutputValueKind::TENSOR_LIST: + return OutputValueKind::TensorList; + case fbs::OutputValueKind::INT: + return OutputValueKind::Int; + case fbs::OutputValueKind::BOOL: + return OutputValueKind::Bool; + case fbs::OutputValueKind::FLOAT: + return OutputValueKind::Float; + default: + throw std::runtime_error( + "build_graph: unsupported OutputValueKind value " + + std::to_string(static_cast(k))); + } +} + +ValueRole map_input_kind(fbs::InputKind k) { + switch (k) { + case fbs::InputKind::USER_INPUT: + return ValueRole::UserInput; + case fbs::InputKind::PARAMETER: + return ValueRole::Parameter; + case fbs::InputKind::BUFFER: + return ValueRole::Buffer; + case fbs::InputKind::CONSTANT_TENSOR: + return ValueRole::ConstantTensor; + default: + throw std::runtime_error( + "build_method: unsupported InputKind value " + + std::to_string(static_cast(k))); + } +} + +OutputKind map_output_kind(fbs::OutputKind k) { + switch (k) { + case fbs::OutputKind::USER_OUTPUT: + return OutputKind::UserOutput; + case fbs::OutputKind::BUFFER_MUTATION: + return OutputKind::BufferMutation; + case fbs::OutputKind::USER_INPUT_MUTATION: + return OutputKind::UserInputMutation; + default: + throw std::runtime_error( + "build_method: unsupported OutputKind value " + + std::to_string(static_cast(k))); + } +} + +// The wire describes a dim as a min..max range, while the IR holds a concrete +// extent. Collapsing a range to its upper bound would run the graph at that +// bound and compute over elements the caller never supplied, so a dim that is +// not a single non-negative extent is refused where it enters the IR. +int64_t static_extent( + const fbs::Dim* d, + const std::string& value_name, + flatbuffers::uoffset_t i) { + if (d->min() == d->max() && d->min() >= 0) { + return d->min(); + } + throw std::runtime_error( + "build_tensor_meta: " + value_name + " dim " + std::to_string(i) + + " is not a static extent (" + std::to_string(d->min()) + ".." + + (d->max() < 0 ? std::string("inf") : std::to_string(d->max())) + + "); this runtime requires static shapes"); +} + +TensorMeta build_tensor_meta( + const fbs::TensorMeta* m, + const std::string& name) { + TensorMeta out; + if (m == nullptr) { + return out; + } + out.dtype = map_scalar_type(m->dtype()); + if (const auto* sizes = m->sizes()) { + out.sizes.reserve(sizes->size()); + for (flatbuffers::uoffset_t i = 0; i < sizes->size(); ++i) { + out.sizes.push_back(static_extent(sizes->Get(i), name, i)); + } + } + if (const auto* dord = m->dim_order()) { + out.dim_order_hint.reserve(dord->size()); + for (flatbuffers::uoffset_t i = 0; i < dord->size(); ++i) { + out.dim_order_hint.push_back(static_cast(dord->Get(i))); + } + } + return out; +} + +// value name -> tensor metadata. +using MetaTable = std::unordered_map; + +// One graph body plus the SSA-name -> value-id map used to build it. +// +// The wire format addresses values two ways: a graph body is positional, and +// the value list keeps that (a ValueId is an index), while the method-level +// side tables -- constants, mutable_buffers, output_specs -- name their targets +// by SSA string, since the serializer writes them independently of the body. +// Only the builder knows how one maps to the other, so it hands the map back +// for build_method to resolve those names against. +// +// Build-time scaffolding: nothing outside build_method sees it, and neither +// Method nor Graph stores it. Once the bindings are resolved to ids it is +// discarded, and the graph is index-addressed from then on. +struct BuiltGraph { + Graph graph; + std::unordered_map name_to_id; +}; + +// `extra_meta` supplies metadata for values the graph's own tensor_values side +// table omits: a constant placeholder's meta rides on the Method's +// NamedTensorRef binding instead, so build_method passes it in to type those +// values on creation. Subgraphs have no such bindings. +BuiltGraph build_graph(const fbs::Graph* g, const MetaTable& extra_meta = {}); + +// Builds one Graph body. The graph under construction, its SSA-name -> id map, +// and the name -> metadata table are shared by every step of the build, so they +// are members rather than threaded through each call. One builder per body: a +// subgraph gets a fresh one, which is what gives each body its own independent +// SSA namespace. +class GraphBuilder { + public: + BuiltGraph run(const fbs::Graph* g, const MetaTable& extra_meta); + + private: + // Resolve a name to its ValueId, creating the Value on first mention. + // A name in the meta side table becomes a Tensor value; otherwise a None + // value (scalar / symbolic outputs, refined once the loader models them). + ValueId id_of(const std::string& name); + + Argument convert_arg(const fbs::Argument* a); + + Graph graph_; + std::unordered_map n2i_; + MetaTable tm_; +}; + +ValueId GraphBuilder::id_of(const std::string& name) { + if (name.empty()) { + return kInvalid; + } + const auto it = n2i_.find(name); + if (it != n2i_.end()) { + return it->second; + } + const ValueId id = static_cast(graph_.values.size()); + const auto mit = tm_.find(name); + if (mit != tm_.end() && mit->second != nullptr) { + graph_.values.emplace_back(name, build_tensor_meta(mit->second, name)); + } else { + graph_.values.emplace_back(name); + } + n2i_[name] = id; + return id; +} + +Argument GraphBuilder::convert_arg(const fbs::Argument* a) { + using AV = fbs::ArgumentValue; + switch (a->value_type()) { + case AV::NONE: + case AV::NoneArg: + return NoneArg{}; + case AV::TensorArg: { + TensorArg t; + t.id = id_of(str_of(a->value_as_TensorArg()->name())); + return t; + } + case AV::IntArg: { + const auto* x = a->value_as_IntArg(); + IntArg r; + r.value = x->value(); + r.id = nonempty(x->ref()) ? id_of(x->ref()->str()) : kInvalid; + return r; + } + case AV::FloatArg: { + const auto* x = a->value_as_FloatArg(); + FloatArg r; + r.value = x->value(); + r.id = nonempty(x->ref()) ? id_of(x->ref()->str()) : kInvalid; + return r; + } + case AV::BoolArg: { + const auto* x = a->value_as_BoolArg(); + BoolArg r; + r.value = x->value(); + r.id = nonempty(x->ref()) ? id_of(x->ref()->str()) : kInvalid; + return r; + } + case AV::StringArg: { + StringArg r; + r.value = str_of(a->value_as_StringArg()->value()); + return r; + } + case AV::ScalarTypeArg: { + ScalarTypeArg r; + r.value = map_scalar_type(a->value_as_ScalarTypeArg()->value()); + return r; + } + case AV::IntListArg: { + const auto* x = a->value_as_IntListArg(); + IntListArg r; + if (const auto* vals = x->values()) { + for (flatbuffers::uoffset_t i = 0; i < vals->size(); ++i) { + r.values.push_back(vals->Get(i)); + } + } + if (const auto* refs = x->refs()) { + for (flatbuffers::uoffset_t i = 0; i < refs->size(); ++i) { + r.ids.push_back( + nonempty(refs->Get(i)) ? id_of(refs->Get(i)->str()) : kInvalid); + } + } + return r; + } + case AV::FloatListArg: { + FloatListArg r; + if (const auto* vals = a->value_as_FloatListArg()->values()) { + for (flatbuffers::uoffset_t i = 0; i < vals->size(); ++i) { + r.values.push_back(vals->Get(i)); + } + } + return r; + } + case AV::BoolListArg: { + BoolListArg r; + if (const auto* vals = a->value_as_BoolListArg()->values()) { + for (flatbuffers::uoffset_t i = 0; i < vals->size(); ++i) { + r.values.push_back(vals->Get(i)); + } + } + return r; + } + case AV::TensorListArg: { + TensorListArg r; + if (const auto* nm = a->value_as_TensorListArg()->names()) { + for (flatbuffers::uoffset_t i = 0; i < nm->size(); ++i) { + r.ids.push_back(id_of(nm->Get(i)->str())); + } + } + return r; + } + case AV::OptionalTensorListArg: { + const auto* oa = a->value_as_OptionalTensorListArg(); + const auto* nm = oa->names(); + const auto* hv = oa->has_value(); + OptionalTensorListArg r; + if (nm != nullptr) { + for (flatbuffers::uoffset_t i = 0; i < nm->size(); ++i) { + const bool present = hv != nullptr && i < hv->size() && hv->Get(i); + r.ids.push_back(present ? id_of(nm->Get(i)->str()) : kInvalid); + } + } + return r; + } + case AV::GraphArg: { + const fbs::GraphArg* ga = a->value_as_GraphArg(); + GraphArg r; + r.name = str_of(ga->name()); + r.subgraph_id = static_cast(graph_.subgraphs.size()); + graph_.subgraphs.push_back(build_graph(ga->graph()).graph); + return r; + } + default: + throw std::runtime_error( + "build_graph: unsupported ArgumentValue value " + + std::to_string(static_cast(a->value_type()))); + } +} + +BuiltGraph GraphBuilder::run(const fbs::Graph* g, const MetaTable& extra_meta) { + if (g == nullptr) { + return {}; + } + + if (const auto* tvs = g->tensor_values()) { + for (flatbuffers::uoffset_t i = 0; i < tvs->size(); ++i) { + const fbs::TensorValue* tv = tvs->Get(i); + tm_[str_of(tv->name())] = tv->meta(); + } + } + for (const auto& entry : extra_meta) { + const fbs::TensorMeta*& slot = tm_[entry.first]; + if (slot == nullptr) { + slot = entry.second; + } + } + + // Pre-create meta-carrying values in a deterministic order (nicer ids). + if (const auto* tvs = g->tensor_values()) { + for (flatbuffers::uoffset_t i = 0; i < tvs->size(); ++i) { + id_of(str_of(tvs->Get(i)->name())); + } + } + + if (const auto* nodes = g->nodes()) { + for (flatbuffers::uoffset_t i = 0; i < nodes->size(); ++i) { + const fbs::Node* nd = nodes->Get(i); + Node node; + node.name = str_of(nd->name()); + node.op_kind = map_op_kind(nd->op_kind()); + node.target = str_of(nd->target()); + + if (const auto* ins = nd->inputs()) { + for (flatbuffers::uoffset_t j = 0; j < ins->size(); ++j) { + const fbs::NamedArgument* na = ins->Get(j); + NamedArgument narg; + narg.name = str_of(na->name()); + narg.mutated = na->mutated(); + narg.arg = convert_arg(na->arg()); + node.inputs.push_back(std::move(narg)); + } + } + + if (const auto* outs = nd->outputs()) { + for (flatbuffers::uoffset_t j = 0; j < outs->size(); ++j) { + const fbs::Output* o = outs->Get(j); + Output out; + out.kind = map_output_value_kind(o->kind()); + if (o->kind() == fbs::OutputValueKind::TENSOR_LIST) { + if (const auto* nm = o->names()) { + for (flatbuffers::uoffset_t k = 0; k < nm->size(); ++k) { + out.elem_ids.push_back(id_of(nm->Get(k)->str())); + } + } + } else { + out.value_id = id_of(str_of(o->name())); + if (nonempty(o->alias_of()) && valid(out.value_id)) { + const ValueId alias_id = id_of(o->alias_of()->str()); + if (alias_id == out.value_id) { + throw std::runtime_error( + "build_graph: output '" + str_of(o->name()) + + "' cannot alias itself"); + } + graph_.values.at(static_cast(out.value_id)).alias_id = + alias_id; + } + } + node.outputs.push_back(std::move(out)); + } + } + + // A placeholder with no explicit Output still produces its named value; + // synthesize one so def-use wiring records the placeholder as producer. + if (node.op_kind == OpKind::Placeholder && node.outputs.empty() && + !node.name.empty()) { + Output out; + out.value_id = id_of(node.name); + node.outputs.push_back(out); + } + + graph_.nodes.push_back(std::move(node)); + } + } + + if (const auto* gi = g->inputs()) { + for (flatbuffers::uoffset_t i = 0; i < gi->size(); ++i) { + graph_.input_ids.push_back(id_of(gi->Get(i)->str())); + } + } + if (const auto* go = g->outputs()) { + for (flatbuffers::uoffset_t i = 0; i < go->size(); ++i) { + graph_.output_ids.push_back(id_of(go->Get(i)->str())); + } + } + + graph_.initialize_schedule(); + graph_.rebuild_def_use(); + return BuiltGraph{std::move(graph_), std::move(n2i_)}; +} + +BuiltGraph build_graph(const fbs::Graph* g, const MetaTable& extra_meta) { + return GraphBuilder().run(g, extra_meta); +} + +// Resolve a method-level binding name (namespace 2) against the top-level +// graph's SSA names, or kInvalid if the graph holds no such value. +ValueId id_of_name( + const std::unordered_map& n2i, + const std::string& name) { + const auto it = n2i.find(name); + return it != n2i.end() ? it->second : kInvalid; +} + +ValueId require_binding_value( + const std::unordered_map& n2i, + const std::string& name) { + const ValueId id = id_of_name(n2i, name); + if (!valid(id)) { + throw std::runtime_error( + "build_method: data binding '" + name + + "' does not name a graph value"); + } + return id; +} + +void stamp_role(Graph& graph, ValueId id, ValueRole role) { + if (in_bounds(id, graph.values.size())) { + graph.values.at(static_cast(id)).role = role; + } +} + +} // namespace + +Method Program::build_method(size_t index) const { + if (program_fb_ == nullptr) { + throw std::runtime_error("build_method: program is not loaded"); + } + const auto* methods = program_fb_->methods(); + if (methods == nullptr || index >= methods->size()) { + throw std::runtime_error("build_method: method index out of range"); + } + const fbs::Method* m = + methods->Get(static_cast(index)); + + Method method; + method.name = str_of(m->name()); + + MetaTable constant_meta; + if (const auto* cs = m->constants()) { + for (flatbuffers::uoffset_t i = 0; i < cs->size(); ++i) { + const fbs::NamedTensorRef* c = cs->Get(i); + constant_meta[str_of(c->name())] = c->meta(); + } + } + + BuiltGraph built = build_graph(m->graph(), constant_meta); + const std::unordered_map& n2i = built.name_to_id; + method.graph = std::move(built.graph); + Graph& graph = method.graph; + + // external-constant / buffer identity (key) -> value, for BufferMutation + // output targets (a namespace-3 fqn, not an SSA name). + std::unordered_map key_to_id; + std::unordered_set bound_ids; + + if (const auto* cs = m->constants()) { + for (flatbuffers::uoffset_t i = 0; i < cs->size(); ++i) { + const fbs::NamedTensorRef* c = cs->Get(i); + DataBinding b; + const std::string name = str_of(c->name()); + b.value_id = require_binding_value(n2i, name); + if (!bound_ids.insert(b.value_id).second) { + throw std::runtime_error( + "build_method: graph value '" + name + + "' has multiple data bindings"); + } + b.role = map_input_kind(c->kind()); + b.key = str_of(c->data_key()); + b.has_data = true; + b.mutated = c->mutated(); + stamp_role(graph, b.value_id, b.role); + if (!b.key.empty()) { + key_to_id[b.key] = b.value_id; + } + method.data_bindings.push_back(std::move(b)); + } + } + + if (const auto* mbs = m->mutable_buffers()) { + for (flatbuffers::uoffset_t i = 0; i < mbs->size(); ++i) { + const fbs::MutableBufferSpec* mb = mbs->Get(i); + DataBinding b; + const std::string name = str_of(mb->name()); + b.value_id = require_binding_value(n2i, name); + if (!bound_ids.insert(b.value_id).second) { + throw std::runtime_error( + "build_method: graph value '" + name + + "' has multiple data bindings"); + } + b.role = ValueRole::Buffer; + b.key = str_of(mb->fqn()); + b.has_data = false; + b.mutated = true; + stamp_role(graph, b.value_id, ValueRole::Buffer); + if (!b.key.empty()) { + key_to_id[b.key] = b.value_id; + } + method.data_bindings.push_back(std::move(b)); + } + } + + // Top-level graph inputs not otherwise bound are user inputs. + for (const ValueId id : graph.input_ids) { + if (in_bounds(id, graph.values.size()) && + graph.values[id].role == ValueRole::Intermediate) { + graph.values[id].role = ValueRole::UserInput; + } + } + + // output_specs are parallel to graph.outputs (same order); each classifies + // graph.output_ids[i]. The mutation target resolves to a placeholder value: + // an fqn (BufferMutation) via key_to_id, else an SSA name + // (UserInputMutation). + if (const auto* os = m->output_specs()) { + if (os->size() != graph.output_ids.size()) { + throw std::runtime_error( + "build_method: output_specs count does not match graph outputs"); + } + for (flatbuffers::uoffset_t i = 0; i < os->size(); ++i) { + const fbs::OutputSpec* o = os->Get(i); + OutputSpec spec; + spec.kind = map_output_kind(o->kind()); + const std::string target = str_of(o->target()); + if (spec.kind != OutputKind::UserOutput) { + if (spec.kind == OutputKind::BufferMutation) { + const auto it = key_to_id.find(target); + spec.target_id = it != key_to_id.end() ? it->second : kInvalid; + } else { + spec.target_id = id_of_name(n2i, target); + } + if (!valid(spec.target_id)) { + throw std::runtime_error( + "build_method: mutation target '" + target + + "' does not name a bound value"); + } + } + method.output_specs.push_back(spec); + } + } + + return method; +} + +// Defined here (rather than in Program.cpp) so it sits next to build_method and +// the deserializer helpers it drives: get_method is the public lazy entry +// point, build_method the private materializer it calls on a cache miss. +const Method& Program::get_method(const std::string& name) const { + const auto it = method_cache_.find(name); + if (it != method_cache_.end()) { + return it->second; + } + if (program_fb_ != nullptr) { + if (const auto* methods = program_fb_->methods()) { + for (flatbuffers::uoffset_t i = 0; i < methods->size(); ++i) { + const flatbuffers::String* method_name = methods->Get(i)->name(); + if (method_name != nullptr && + std::string_view(method_name->c_str(), method_name->size()) == + name) { + auto res = method_cache_.emplace(name, build_method(i)); + return res.first->second; + } + } + } + } + throw std::runtime_error( + "Program::get_method: no method named '" + name + "'"); +} + +} // namespace ptn diff --git a/backends/native/runtime/Program.cpp b/backends/native/runtime/Program.cpp index 00d8de36deb..0ee87e212dc 100644 --- a/backends/native/runtime/Program.cpp +++ b/backends/native/runtime/Program.cpp @@ -8,6 +8,8 @@ #include #include +#include +#include #include #include @@ -42,6 +44,19 @@ Program Program::load(const void* data, size_t size) { } const fbs::Program* program_fb = fbs::GetProgram(bytes.data()); + // Both accessors below are schema-required, so successful verification + // guarantees that they are non-null. + std::unordered_set method_names; + for (const fbs::Method* method : *program_fb->methods()) { + const std::string name = method->name()->str(); + if (name.empty()) { + throw std::runtime_error("native program: method name is empty"); + } + if (!method_names.insert(name).second) { + throw std::runtime_error( + "native program: duplicate method name '" + name + "'"); + } + } return Program(std::move(bytes), program_fb); } @@ -50,4 +65,17 @@ size_t Program::num_methods() const { return methods == nullptr ? 0 : methods->size(); } +std::vector Program::method_names() const { + std::vector names; + const auto* methods = program_fb_->methods(); + if (methods != nullptr) { + names.reserve(methods->size()); + for (flatbuffers::uoffset_t i = 0; i < methods->size(); ++i) { + const auto* nm = methods->Get(i)->name(); + names.push_back(nm != nullptr ? nm->str() : std::string()); + } + } + return names; +} + } // namespace ptn diff --git a/backends/native/runtime/Program.h b/backends/native/runtime/Program.h index 6db8a887f97..e035e0d4d5d 100644 --- a/backends/native/runtime/Program.h +++ b/backends/native/runtime/Program.h @@ -8,8 +8,12 @@ #include #include +#include +#include #include +#include + // Forward-declaration of the generated FlatBuffer root type, included only from // .cpp files so flatbuffers stays an implementation detail of the reader. namespace native_backend { @@ -29,6 +33,10 @@ class Program { // rather than return a null root, so accessors dereference it unchecked. std::vector bytes_; const fbs::Program* program_fb_ = nullptr; + // Lazily materialized methods, keyed by name, populated on get_method(). The + // cache is mutable so lookups work on a const Program; unordered_map keeps + // returned references stable across later insertions. Not thread-safe. + mutable std::unordered_map method_cache_; Program(std::vector bytes, const fbs::Program* program_fb) : bytes_(std::move(bytes)), program_fb_(program_fb) {} @@ -41,7 +49,8 @@ class Program { Program& operator=(const Program&) = delete; // Parse and verify serialized native-graph bytes (a *.ptg buffer). Throws - // std::runtime_error on failure. + // std::runtime_error on failure. Methods are materialized lazily (see + // get_method), not here. static Program load(const void* data, size_t size); const fbs::Program* flatbuffer() const { @@ -49,6 +58,20 @@ class Program { } size_t num_methods() const; + + // Names of the program's methods, in serialized order. + std::vector method_names() const; + + // Materialize (or return the cached) method by name. Builds the in-memory IR + // on first request and caches it; later calls return the same instance. + // Throws std::runtime_error if no method has that name. Impl in + // Deserialize.cpp. + const Method& get_method(const std::string& name) const; + + private: + // Deserialize the fb method at `index` into the in-memory IR (Graph + + // bindings). Impl in Deserialize.cpp. + Method build_method(size_t index) const; }; } // namespace ptn diff --git a/backends/native/runtime/targets.bzl b/backends/native/runtime/targets.bzl index 865f2c0242f..6cc4555fb9c 100644 --- a/backends/native/runtime/targets.bzl +++ b/backends/native/runtime/targets.bzl @@ -46,11 +46,16 @@ def define_common_targets(): runtime.cxx_library( name = "runtime", srcs = [ + "Deserialize.cpp", "Program.cpp", ], exported_headers = [ "Program.h", ], + exported_deps = [ + # Program.h publicly exposes Method (build_method), so the IR is exported. + ":method", + ], deps = [ ":native_graph_schema", ], diff --git a/backends/native/test/runtime/BUCK b/backends/native/test/runtime/BUCK new file mode 100644 index 00000000000..36909de98fe --- /dev/null +++ b/backends/native/test/runtime/BUCK @@ -0,0 +1,8 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/test/runtime/targets.bzl b/backends/native/test/runtime/targets.bzl new file mode 100644 index 00000000000..97f9110f7d9 --- /dev/null +++ b/backends/native/test/runtime/targets.bzl @@ -0,0 +1,11 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_test( + name = "program_test", + srcs = ["test_program_deserialize.cpp"], + deps = [ + "//executorch/backends/native/runtime:native_graph_schema", + "//executorch/backends/native/runtime:runtime", + ], + ) diff --git a/backends/native/test/runtime/test_program_deserialize.cpp b/backends/native/test/runtime/test_program_deserialize.cpp new file mode 100644 index 00000000000..ec3cde4f10c --- /dev/null +++ b/backends/native/test/runtime/test_program_deserialize.cpp @@ -0,0 +1,299 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include + +#include +#include + +#include + +namespace ptn { +namespace { + +namespace fbs = ::native_backend; + +flatbuffers::Offset< + flatbuffers::Vector>> +create_strings( + flatbuffers::FlatBufferBuilder& builder, + const std::vector& values) { + std::vector> strings; + strings.reserve(values.size()); + for (const std::string& value : values) { + strings.push_back(builder.CreateString(value)); + } + return builder.CreateVector(strings); +} + +flatbuffers::Offset create_graph( + flatbuffers::FlatBufferBuilder& builder, + const std::vector>& nodes = {}, + const std::vector& inputs = {}, + const std::vector& outputs = {}, + const std::vector>& tensor_values = + {}) { + return fbs::CreateGraph( + builder, + builder.CreateVector(nodes), + inputs.empty() ? 0 : create_strings(builder, inputs), + outputs.empty() ? 0 : create_strings(builder, outputs), + tensor_values.empty() ? 0 : builder.CreateVector(tensor_values)); +} + +flatbuffers::Offset create_method( + flatbuffers::FlatBufferBuilder& builder, + const std::string& name, + flatbuffers::Offset graph, + const std::vector>& output_specs = {}, + const std::vector>& constants = {}, + const std::vector>& + mutable_buffers = {}) { + return fbs::CreateMethod( + builder, + builder.CreateString(name), + graph, + constants.empty() ? 0 : builder.CreateVector(constants), + output_specs.empty() ? 0 : builder.CreateVector(output_specs), + mutable_buffers.empty() ? 0 : builder.CreateVector(mutable_buffers)); +} + +std::vector finish_program( + flatbuffers::FlatBufferBuilder& builder, + const std::vector>& methods) { + const auto program = fbs::CreateProgram( + builder, builder.CreateString("1"), builder.CreateVector(methods)); + fbs::FinishProgramBuffer(builder, program); + return { + builder.GetBufferPointer(), + builder.GetBufferPointer() + builder.GetSize()}; +} + +Program load_program(const std::vector& bytes) { + return Program::load(bytes.data(), bytes.size()); +} + +// cppcheck-suppress-begin syntaxError +TEST(ProgramTest, LoadRejectsEmptyMethodName) { + flatbuffers::FlatBufferBuilder builder; + const auto graph = create_graph(builder); + const auto bytes = + finish_program(builder, {create_method(builder, "", graph)}); + + EXPECT_THROW(load_program(bytes), std::runtime_error); +} + +TEST(ProgramTest, LoadRejectsDuplicateMethodNames) { + flatbuffers::FlatBufferBuilder builder; + const auto graph = create_graph(builder); + const auto first = create_method(builder, "forward", graph); + const auto second = create_method(builder, "forward", graph); + const auto bytes = finish_program(builder, {first, second}); + + EXPECT_THROW(load_program(bytes), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsMismatchedOutputSpecs) { + flatbuffers::FlatBufferBuilder builder; + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto first = fbs::CreateOutputSpecDirect(builder, "output"); + const auto second = fbs::CreateOutputSpecDirect(builder, "extra"); + const auto method = create_method(builder, "forward", graph, {first, second}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodPreservesAliasWhenTargetIsCreatedOnDemand) { + flatbuffers::FlatBufferBuilder builder; + const auto output = fbs::CreateOutputDirect(builder, "view", "input"); + const std::vector> outputs = {output}; + const auto node = fbs::CreateNodeDirect( + builder, + "view", + fbs::OpKind::CALL_FUNCTION, + "aten.view", + nullptr, + &outputs); + const auto graph = create_graph(builder, {node}, {"input"}, {"view"}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + + const Graph& loaded = program.get_method("forward").graph; + ASSERT_EQ(loaded.input_ids.size(), 1); + ASSERT_EQ(loaded.output_ids.size(), 1); + EXPECT_EQ(loaded.value(loaded.output_ids[0]).alias_id, loaded.input_ids[0]); +} + +TEST(ProgramTest, GetMethodRejectsSelfAlias) { + flatbuffers::FlatBufferBuilder builder; + const auto output = fbs::CreateOutputDirect(builder, "view", "view"); + const std::vector> outputs = {output}; + const auto node = fbs::CreateNodeDirect( + builder, + "view", + fbs::OpKind::CALL_FUNCTION, + "aten.view", + nullptr, + &outputs); + const auto graph = create_graph(builder, {node}, {}, {"view"}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsDynamicTensorExtent) { + flatbuffers::FlatBufferBuilder builder; + const std::vector> sizes = { + fbs::CreateDim(builder, 2, 16)}; + const auto meta = + fbs::CreateTensorMetaDirect(builder, fbs::ScalarType::FLOAT, &sizes); + const auto tensor = fbs::CreateTensorValueDirect(builder, "input", meta); + const auto graph = create_graph(builder, {}, {"input"}, {}, {tensor}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsUnknownEnumValues) { + { + flatbuffers::FlatBufferBuilder builder; + const auto node = fbs::CreateNodeDirect( + builder, "node", static_cast(127), "unknown"); + const auto graph = create_graph(builder, {node}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto output = fbs::CreateOutput( + builder, + builder.CreateString("output"), + 0, + static_cast(127)); + const std::vector> outputs = {output}; + const auto node = fbs::CreateNodeDirect( + builder, + "node", + fbs::OpKind::CALL_FUNCTION, + "unknown", + nullptr, + &outputs); + const auto graph = create_graph(builder, {node}, {}, {"output"}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto meta = fbs::CreateTensorMeta(builder); + const auto constant = fbs::CreateNamedTensorRefDirect( + builder, "input", "weight", meta, static_cast(127)); + const auto graph = create_graph(builder, {}, {"input"}); + const auto method = + create_method(builder, "forward", graph, {}, {constant}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto argument = + fbs::CreateArgument(builder, static_cast(127), 0); + const auto named_argument = + fbs::CreateNamedArgumentDirect(builder, "input", argument); + const std::vector> inputs = { + named_argument}; + const auto node = fbs::CreateNodeDirect( + builder, "node", fbs::OpKind::CALL_FUNCTION, "unknown", &inputs); + const auto graph = create_graph(builder, {node}); + const auto bytes = + finish_program(builder, {create_method(builder, "forward", graph)}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto output_spec = fbs::CreateOutputSpec( + builder, + builder.CreateString("output"), + static_cast(127)); + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto method = create_method(builder, "forward", graph, {output_spec}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } +} + +TEST(ProgramTest, GetMethodRejectsUnresolvedDataBinding) { + flatbuffers::FlatBufferBuilder builder; + const auto meta = fbs::CreateTensorMeta(builder); + const auto constant = fbs::CreateNamedTensorRefDirect( + builder, "missing", "weight", meta, fbs::InputKind::PARAMETER); + const auto graph = create_graph(builder); + const auto method = create_method(builder, "forward", graph, {}, {constant}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} + +TEST(ProgramTest, GetMethodRejectsUnresolvedMutationTarget) { + { + flatbuffers::FlatBufferBuilder builder; + const auto output_spec = fbs::CreateOutputSpecDirect( + builder, "output", fbs::OutputKind::BUFFER_MUTATION, "missing"); + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto method = create_method(builder, "forward", graph, {output_spec}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } + { + flatbuffers::FlatBufferBuilder builder; + const auto output_spec = fbs::CreateOutputSpecDirect( + builder, "output", fbs::OutputKind::USER_INPUT_MUTATION, "missing"); + const auto graph = create_graph(builder, {}, {}, {"output"}); + const auto method = create_method(builder, "forward", graph, {output_spec}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + EXPECT_THROW(program.get_method("forward"), std::runtime_error); + } +} + +TEST(ProgramTest, GetMethodRejectsDuplicateDataBinding) { + flatbuffers::FlatBufferBuilder builder; + const auto meta = fbs::CreateTensorMeta(builder); + const auto constant = fbs::CreateNamedTensorRefDirect( + builder, "state", "state", meta, fbs::InputKind::BUFFER); + const auto mutable_buffer = + fbs::CreateMutableBufferSpecDirect(builder, "state", "state"); + const auto graph = create_graph(builder, {}, {"state"}); + const auto method = create_method( + builder, "forward", graph, {}, {constant}, {mutable_buffer}); + const auto bytes = finish_program(builder, {method}); + const Program program = load_program(bytes); + + EXPECT_THROW(program.get_method("forward"), std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn From 9df9094e97d0fe3a26883bb0d280ee79cea6ea11 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 10 Sep 2026 12:14:18 -0700 Subject: [PATCH 166/190] [executorch][native] Add a safetensors index reader ## Ulterior Motive Give PTN package loading a validated constant index before archive integration. ## Rationale **What**: Add borrowed `ByteSpan` views and a `SafeTensorsReader` that preserves entry order and validates dtype, shape, byte ranges, and payload size. **Why**: Engines need trustworthy tensor metadata and byte locations before uploading weights. ## Details ```text [u64 header length][ordered JSON][tensor payload] | | v v bounds checks TensorEntry map | +-- dtype +-- shape +-- offset + byte count ``` `SafeTensorsReader` owns parsed metadata but borrows source bytes. Unsupported dtype codes, malformed JSON, overflow, inconsistent sizes, out-of-range offsets, and non-empty metadata are rejected. Authored with Codex. Differential Revision: [D118480653](https://our.internmc.facebook.com/intern/diff/D118480653/) ghstack-source-id: 427858761 Pull-Request: https://github.com/pytorch/executorch/pull/22700 --- backends/native/runtime/deserialize/BUCK | 11 + .../native/runtime/deserialize/ByteSpan.h | 18 ++ backends/native/runtime/deserialize/Json.h | 15 ++ .../runtime/deserialize/SafeTensorsReader.cpp | 244 ++++++++++++++++++ .../runtime/deserialize/SafeTensorsReader.h | 69 +++++ .../native/runtime/deserialize/targets.bzl | 33 +++ backends/native/test/runtime/deserialize/BUCK | 8 + .../test/runtime/deserialize/targets.bzl | 10 + .../deserialize/test_safetensors_reader.cpp | 83 ++++++ 9 files changed, 491 insertions(+) create mode 100644 backends/native/runtime/deserialize/BUCK create mode 100644 backends/native/runtime/deserialize/ByteSpan.h create mode 100644 backends/native/runtime/deserialize/Json.h create mode 100644 backends/native/runtime/deserialize/SafeTensorsReader.cpp create mode 100644 backends/native/runtime/deserialize/SafeTensorsReader.h create mode 100644 backends/native/runtime/deserialize/targets.bzl create mode 100644 backends/native/test/runtime/deserialize/BUCK create mode 100644 backends/native/test/runtime/deserialize/targets.bzl create mode 100644 backends/native/test/runtime/deserialize/test_safetensors_reader.cpp diff --git a/backends/native/runtime/deserialize/BUCK b/backends/native/runtime/deserialize/BUCK new file mode 100644 index 00000000000..0ab35888218 --- /dev/null +++ b/backends/native/runtime/deserialize/BUCK @@ -0,0 +1,11 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. This file can contain cell-only targets. + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/runtime/deserialize/ByteSpan.h b/backends/native/runtime/deserialize/ByteSpan.h new file mode 100644 index 00000000000..0b16f75a746 --- /dev/null +++ b/backends/native/runtime/deserialize/ByteSpan.h @@ -0,0 +1,18 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include + +namespace ptn { + +// Borrowed byte-range views. The producer defines their lifetime. +using ByteSpan = std::span; +using MutableByteSpan = std::span; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/Json.h b/backends/native/runtime/deserialize/Json.h new file mode 100644 index 00000000000..8cf1f1cd305 --- /dev/null +++ b/backends/native/runtime/deserialize/Json.h @@ -0,0 +1,15 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include + +namespace ptn { + +using Json = nlohmann::ordered_json; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/SafeTensorsReader.cpp b/backends/native/runtime/deserialize/SafeTensorsReader.cpp new file mode 100644 index 00000000000..6b89588180f --- /dev/null +++ b/backends/native/runtime/deserialize/SafeTensorsReader.cpp @@ -0,0 +1,244 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +// Reserved header member holding free-form string metadata, not a tensor. +constexpr std::string_view kMetadataKey = "__metadata__"; +constexpr size_t kHeaderLenSize = 8; + +struct DtypeCode { + std::string_view code; + ScalarType dtype; +}; + +// safetensors dtype codes, as written by safetensors.torch. Codes with no +// ScalarType counterpart (complex, 4-bit and 8-bit float variants) are absent +// and rejected by name, so an unsupported constant fails at load rather than +// being misread as another width. +constexpr std::array kDtypeCodes{{ + {"F64", kDouble}, + {"F32", kFloat}, + {"F16", kHalf}, + {"BF16", kBFloat16}, + {"I64", kLong}, + {"I32", kInt}, + {"I16", kShort}, + {"I8", kChar}, + {"U8", kByte}, + {"BOOL", kBool}, + {"U16", kUInt16}, + {"U32", kUInt32}, + {"U64", kUInt64}, +}}; + +ScalarType scalar_type_of(std::string_view code) { + const auto it = std::ranges::find(kDtypeCodes, code, &DtypeCode::code); + if (it == kDtypeCodes.end()) { + throw std::runtime_error( + "safetensors: unsupported dtype code: " + std::string(code)); + } + return it->dtype; +} + +uint64_t read_header_len(ByteSpan blob) { + static_assert( + std::endian::native == std::endian::little, + "the length prefix is little-endian; a big-endian host needs a swap"); + if (blob.size() < kHeaderLenSize) { + throw std::runtime_error( + "safetensors: blob is shorter than its length prefix"); + } + uint64_t len = 0; + std::memcpy(&len, blob.data(), kHeaderLenSize); + return len; +} + +const Json& required_member( + const Json& entry, + std::string_view key, + const std::string& name) { + const auto value = entry.find(key); + if (value == entry.end()) { + std::string message = "safetensors: entry '"; + message += name; + message += "' has no '"; + message += key; + message += "'"; + throw std::runtime_error(message); + } + return value.value(); +} + +std::vector read_sizes(const Json& shape, const std::string& name) { + if (!shape.is_array()) { + throw std::runtime_error( + "safetensors: entry '" + name + "' shape is not an array"); + } + std::vector sizes; + for (const Json& dim : shape) { + if (!dim.is_number_unsigned()) { + throw std::runtime_error( + "safetensors: entry '" + name + + "' has a non-negative integer dimension"); + } + const uint64_t value = dim.get(); + if (value > static_cast(INT64_MAX)) { + throw std::runtime_error( + "safetensors: entry '" + name + "' has an out-of-range dimension"); + } + sizes.push_back(static_cast(value)); + } + return sizes; +} + +// Element count of `sizes`, rejecting an overflowing product. A rank-0 shape is +// a scalar, whose element count is 1. +size_t numel_of(const std::vector& sizes, const std::string& name) { + size_t numel = 1; + for (const int64_t dim : sizes) { + if (dim < 0) { + throw std::runtime_error( + "safetensors: entry '" + name + "' has a negative dimension"); + } + const size_t d = static_cast(dim); + if (d != 0 && numel > SIZE_MAX / d) { + throw std::runtime_error( + "safetensors: entry '" + name + "' element count overflows"); + } + numel *= d; + } + return numel; +} + +} // namespace + +SafeTensorsReader SafeTensorsReader::open(ByteSpan blob) { + const uint64_t header_len = read_header_len(blob); + if (header_len > blob.size() - kHeaderLenSize) { + throw std::runtime_error("safetensors: header length exceeds the blob"); + } + + const std::string_view header_text( + reinterpret_cast(blob.data() + kHeaderLenSize), + static_cast(header_len)); + Json header; + try { + header = Json::parse(header_text); + } catch (const Json::exception& error) { + throw std::runtime_error( + "safetensors: invalid JSON header: " + std::string(error.what())); + } + if (!header.is_object()) { + throw std::runtime_error("safetensors: header is not a JSON object"); + } + + SafeTensorsReader out; + out.data_ = blob.subspan(kHeaderLenSize + static_cast(header_len)); + + for (auto member = header.begin(); member != header.end(); ++member) { + const std::string& name = member.key(); + if (name == kMetadataKey) { + if (!member.value().is_object() || !member.value().empty()) { + throw std::runtime_error( + "safetensors: __metadata__ must be an empty object"); + } + continue; + } + const Json& entry = member.value(); + if (!entry.is_object()) { + throw std::runtime_error( + "safetensors: entry '" + name + "' is not an object"); + } + + TensorEntry parsed; + const Json& dtype = required_member(entry, "dtype", name); + if (!dtype.is_string()) { + throw std::runtime_error( + "safetensors: entry '" + name + "' dtype is not a string"); + } + parsed.dtype = scalar_type_of(dtype.get_ref()); + parsed.sizes = read_sizes(required_member(entry, "shape", name), name); + + const Json& range = required_member(entry, "data_offsets", name); + if (!range.is_array() || range.size() != 2) { + throw std::runtime_error( + "safetensors: entry '" + name + "' data_offsets is not a pair"); + } + if (!range[0].is_number_unsigned() || !range[1].is_number_unsigned()) { + throw std::runtime_error( + "safetensors: entry '" + name + + "' data_offsets contains a non-negative integer"); + } + const uint64_t begin = range[0].get(); + const uint64_t end = range[1].get(); + if (begin > end || end > out.data_.size()) { + throw std::runtime_error( + "safetensors: entry '" + name + + "' byte range is outside the data section"); + } + parsed.offset = static_cast(begin); + parsed.nbytes = static_cast(end - begin); + + // The payload must be exactly as large as its dtype and shape imply. + // Without this, a short entry becomes an out-of-bounds read in whatever + // consumes it, sized from the metadata rather than the bytes. + const size_t numel = numel_of(parsed.sizes, name); + const size_t element_bytes = element_size(parsed.dtype); + if (element_bytes != 0 && numel > SIZE_MAX / element_bytes) { + throw std::runtime_error( + "safetensors: entry '" + name + "' byte size overflows"); + } + const size_t expected = numel * element_bytes; + if (parsed.nbytes != expected) { + throw std::runtime_error( + "safetensors: entry '" + name + "' holds " + + std::to_string(parsed.nbytes) + + " bytes but its dtype and shape need " + std::to_string(expected)); + } + + if (!out.entries_.emplace(name, std::move(parsed)).second) { + throw std::runtime_error("safetensors: duplicate entry: " + name); + } + out.names_.push_back(name); + } + + return out; +} + +const TensorEntry* SafeTensorsReader::find(const std::string& name) const { + const auto it = entries_.find(name); + return it == entries_.end() ? nullptr : &it->second; +} + +ByteSpan SafeTensorsReader::bytes(const TensorEntry& entry) const { + return data_.subspan(entry.offset, entry.nbytes); +} + +size_t SafeTensorsReader::total_bytes() const { + return std::accumulate( + entries_.begin(), + entries_.end(), + size_t{0}, + [](size_t total, const auto& entry) { + return total + entry.second.nbytes; + }); +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/SafeTensorsReader.h b/backends/native/runtime/deserialize/SafeTensorsReader.h new file mode 100644 index 00000000000..6c0ca7b44fc --- /dev/null +++ b/backends/native/runtime/deserialize/SafeTensorsReader.h @@ -0,0 +1,69 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include +#include + +#include +#include + +namespace ptn { + +// One tensor's entry in a safetensors index. +struct TensorEntry { + ScalarType dtype = kFloat; + std::vector sizes; + // Byte range within the blob's data section, not the whole blob. + size_t offset = 0; + size_t nbytes = 0; +}; + +// Reader for the safetensors format: +// +// [u64 header_len][JSON header][data section] +// +// The header maps a tensor name to its dtype, shape, and byte range within the +// data section. The reserved "__metadata__" member must be empty until the +// reader exposes metadata semantics. +// +// Tensor payloads are packed with no per-tensor padding, so an entry's absolute +// alignment within the file is arbitrary: copy through these spans rather than +// handing them to an API that requires alignment. +class SafeTensorsReader { + private: + // The data section only, i.e. the blob past its header. + ByteSpan data_; + std::unordered_map entries_; + std::vector names_; + + public: + // Parse `blob`'s index. Throws std::runtime_error if the blob is truncated, + // the header is not a JSON object, a dtype has no ScalarType, or a byte range + // is inconsistent with its dtype and shape. + // + // Borrows `blob`, which must outlive both this reader and any span from it. + static SafeTensorsReader open(ByteSpan blob); + + // Entry for `name`, or nullptr when absent. + const TensorEntry* find(const std::string& name) const; + + // Payload of an entry obtained from this reader. + ByteSpan bytes(const TensorEntry& entry) const; + + // Tensor names, in header order, excluding "__metadata__". + const std::vector& names() const { + return names_; + } + + // Total payload bytes across all entries. + size_t total_bytes() const; +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/targets.bzl b/backends/native/runtime/deserialize/targets.bzl new file mode 100644 index 00000000000..cb028f2d6cd --- /dev/null +++ b/backends/native/runtime/deserialize/targets.bzl @@ -0,0 +1,33 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + # Borrowed byte-range view shared by the package readers (a std::span alias, + # named so the borrow contract has somewhere to live). + runtime.cxx_library( + name = "byte_span", + srcs = [], + exported_headers = ["ByteSpan.h"], + visibility = ["//executorch/backends/native/..."], + ) + + # JSON representation used by package metadata readers. + runtime.cxx_library( + name = "json", + srcs = [], + exported_headers = ["Json.h"], + exported_external_deps = ["nlohmann_json"], + visibility = ["//executorch/backends/native/..."], + ) + + # safetensors index reader. + runtime.cxx_library( + name = "safetensors_reader", + srcs = ["SafeTensorsReader.cpp"], + exported_headers = ["SafeTensorsReader.h"], + exported_deps = [ + ":byte_span", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + deps = [":json"], + visibility = ["//executorch/backends/native/..."], + ) diff --git a/backends/native/test/runtime/deserialize/BUCK b/backends/native/test/runtime/deserialize/BUCK new file mode 100644 index 00000000000..36909de98fe --- /dev/null +++ b/backends/native/test/runtime/deserialize/BUCK @@ -0,0 +1,8 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/test/runtime/deserialize/targets.bzl b/backends/native/test/runtime/deserialize/targets.bzl new file mode 100644 index 00000000000..d23e793ee68 --- /dev/null +++ b/backends/native/test/runtime/deserialize/targets.bzl @@ -0,0 +1,10 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_test( + name = "safetensors_reader_test", + srcs = ["test_safetensors_reader.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:safetensors_reader", + ], + ) diff --git a/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp b/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp new file mode 100644 index 00000000000..02a5dfe0965 --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp @@ -0,0 +1,83 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +std::vector make_safetensors( + std::string_view header, + std::string_view data) { + const uint64_t header_size = header.size(); + std::vector bytes(sizeof(header_size) + header.size() + data.size()); + std::memcpy(bytes.data(), &header_size, sizeof(header_size)); + std::memcpy(bytes.data() + sizeof(header_size), header.data(), header.size()); + std::memcpy( + bytes.data() + sizeof(header_size) + header.size(), + data.data(), + data.size()); + return bytes; +} + +// cppcheck-suppress-begin syntaxError +TEST(SafeTensorsReaderTest, ReadsIndexAndPayloadsInHeaderOrder) { + const std::vector bytes = make_safetensors( + R"({"second":{"dtype":"U8","shape":[2],"data_offsets":[4,6]},"__metadata__":{},"first":{"dtype":"F32","shape":[1],"data_offsets":[0,4]}})", + "abcdXY"); + + const SafeTensorsReader reader = SafeTensorsReader::open(bytes); + + EXPECT_EQ(reader.names(), (std::vector{"second", "first"})); + ASSERT_NE(reader.find("first"), nullptr); + EXPECT_EQ(reader.find("first")->dtype, kFloat); + EXPECT_EQ(reader.find("first")->sizes, (std::vector{1})); + EXPECT_EQ(reader.find("first")->offset, 0); + EXPECT_EQ(reader.find("first")->nbytes, 4); + EXPECT_EQ(reader.bytes(*reader.find("second"))[0], 'X'); + EXPECT_EQ(reader.total_bytes(), 6); + EXPECT_EQ(reader.find("missing"), nullptr); +} + +TEST(SafeTensorsReaderTest, RejectsNonEmptyMetadata) { + EXPECT_THROW( + SafeTensorsReader::open( + make_safetensors(R"({"__metadata__":{"source":"test"}})", "")), + std::runtime_error); +} + +TEST(SafeTensorsReaderTest, RejectsInvalidMetadata) { + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors("[]", "")), std::runtime_error); + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors( + R"({"x":{"dtype":"U8","shape":[1.0],"data_offsets":[0,1]}})", "x")), + std::runtime_error); + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors( + R"({"x":{"dtype":"U8","shape":[1],"data_offsets":[0,2]}})", "x")), + std::runtime_error); +} + +TEST(SafeTensorsReaderTest, RejectsByteSizeOverflow) { + EXPECT_THROW( + SafeTensorsReader::open(make_safetensors( + R"({"x":{"dtype":"F64","shape":[2305843009213693952],"data_offsets":[0,0]}})", + "")), + std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn From 359a4244ba973505561e6dd959c35b19bab6e576 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 10 Sep 2026 12:14:23 -0700 Subject: [PATCH 167/190] [executorch][native] Add the OwnedBytes API ## Ulterior Motive Support package-sized data without forcing every caller into one loading strategy. ## Rationale **What**: Add move-only byte ownership over vectors, uninitialized heap buffers, and read-only file mappings. **Why**: Package readers need stable borrowed spans while avoiding redundant copies and zero-fill before file reads. ## Details ```text OwnedBytes | +-- vector caller-owned bytes moved in +-- HeapBuffer direct file read +-- mmap + Unmap deleter demand-paged file | +-- span() uniform read-only view ``` - Whole-file validation rejects missing paths, non-regular files, and oversized files. Virtual files that do not report their size are explicitly unsupported. - Empty mappings fall back to empty heap storage because zero-length `mmap` is invalid. - Moved-from heap and mapped instances return an empty span instead of a null span with a stale size. - Mapping errors preserve `errno` before any allocating error construction. - Tests compare complete payloads and stop at the original file-write failure. Authored with Codex. Differential Revision: [D119396096](https://our.internmc.facebook.com/intern/diff/D119396096/) ghstack-source-id: 427858769 Pull-Request: https://github.com/pytorch/executorch/pull/22701 --- .../native/runtime/deserialize/OwnedBytes.cpp | 159 ++++++++++++++++++ .../native/runtime/deserialize/OwnedBytes.h | 96 +++++++++++ .../native/runtime/deserialize/targets.bzl | 9 + .../test/runtime/deserialize/targets.bzl | 8 + .../runtime/deserialize/test_owned_bytes.cpp | 131 +++++++++++++++ 5 files changed, 403 insertions(+) create mode 100644 backends/native/runtime/deserialize/OwnedBytes.cpp create mode 100644 backends/native/runtime/deserialize/OwnedBytes.h create mode 100644 backends/native/test/runtime/deserialize/test_owned_bytes.cpp diff --git a/backends/native/runtime/deserialize/OwnedBytes.cpp b/backends/native/runtime/deserialize/OwnedBytes.cpp new file mode 100644 index 00000000000..7bb919cd172 --- /dev/null +++ b/backends/native/runtime/deserialize/OwnedBytes.cpp @@ -0,0 +1,159 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include +#include +#include + +#if !defined(_WIN32) +#include +#include +#include +#include +#endif + +namespace ptn { +namespace { + +std::string errno_suffix(int error) { + return ": " + std::error_code(error, std::system_category()).message(); +} + +} // namespace + +void OwnedBytes::Unmap::operator()(void* base) const noexcept { +#if !defined(_WIN32) + // Nothing useful to do if this fails, and it must not throw: unique_ptr calls + // this from its destructor. + ::munmap(base, size); +#endif +} + +OwnedBytes OwnedBytes::from_vector(std::vector bytes) { + return OwnedBytes(std::move(bytes)); +} + +OwnedBytes OwnedBytes::from_file(const std::string& path, bool use_mmap) { + return use_mmap ? map_file(path) : read_file(path); +} + +OwnedBytes OwnedBytes::read_file(const std::string& path) { + std::error_code error; + const std::filesystem::file_status status = + std::filesystem::status(path, error); + if (error) { + throw std::runtime_error("cannot inspect " + path + ": " + error.message()); + } + if (!std::filesystem::is_regular_file(status)) { + throw std::runtime_error("cannot read " + path + ": not a regular file"); + } + const uintmax_t file_size = std::filesystem::file_size(path, error); + if (error) { + throw std::runtime_error("cannot size " + path + ": " + error.message()); + } + if (file_size > std::numeric_limits::max() || + file_size > + static_cast(std::numeric_limits::max())) { + throw std::runtime_error("cannot read " + path + ": file is too large"); + } + + std::ifstream file(path, std::ios::binary); + if (!file) { + throw std::runtime_error("cannot open " + path); + } + const size_t size = static_cast(file_size); + if (size == 0) { + return OwnedBytes(std::vector()); + } + HeapBuffer buffer{std::make_unique_for_overwrite(size), size}; + if (!file.read( + reinterpret_cast(buffer.data.get()), + static_cast(size))) { + throw std::runtime_error("cannot read " + path); + } + return OwnedBytes(std::move(buffer)); +} + +OwnedBytes OwnedBytes::map_file(const std::string& path) { +#if defined(_WIN32) + // TODO: Implement Windows mappings with CreateFileMapping and MapViewOfFile. + throw std::runtime_error("cannot mmap " + path + ": unsupported platform"); +#else + const int fd = ::open(path.c_str(), O_RDONLY); + if (fd < 0) { + const int error = errno; + throw std::runtime_error("cannot open " + path + errno_suffix(error)); + } + + struct stat st = {}; + if (::fstat(fd, &st) < 0) { + const std::string suffix = errno_suffix(errno); + ::close(fd); + throw std::runtime_error("cannot size " + path + suffix); + } + if (!S_ISREG(st.st_mode)) { + ::close(fd); + throw std::runtime_error("cannot mmap " + path + ": not a regular file"); + } + if (st.st_size < 0) { + ::close(fd); + throw std::runtime_error("cannot mmap " + path + ": invalid file size"); + } + if (static_cast(st.st_size) > std::numeric_limits::max()) { + ::close(fd); + throw std::runtime_error("cannot mmap " + path + ": file is too large"); + } + const size_t size = static_cast(st.st_size); + if (size == 0) { + ::close(fd); + return OwnedBytes(std::vector()); + } + + // The whole file from offset 0, so the base is page-aligned and every span + // into it has the same alignment it would have in a heap buffer. MAP_SHARED + // lets other processes mapping this file share the same physical pages; the + // mapping is read-only either way. + void* base = ::mmap(nullptr, size, PROT_READ, MAP_SHARED, fd, 0); + if (base == MAP_FAILED) { + const int error = errno; + ::close(fd); + throw std::runtime_error("cannot mmap " + path + errno_suffix(error)); + } + // The mapping keeps the file alive on its own, so the descriptor is dead + // weight past this point. + ::close(fd); + return OwnedBytes(MappedFile(base, Unmap{size})); +#endif +} + +ByteSpan OwnedBytes::span() const { + if (const HeapBuffer* buffer = std::get_if(&storage_)) { + if (buffer->data == nullptr) { + return {}; + } + return ByteSpan(buffer->data.get(), buffer->size); + } + if (const MappedFile* mapped = std::get_if(&storage_)) { + if (mapped->get() == nullptr) { + return {}; + } + return ByteSpan( + static_cast(mapped->get()), mapped->get_deleter().size); + } + const std::vector& bytes = std::get>(storage_); + return ByteSpan(bytes.data(), bytes.size()); +} + +bool OwnedBytes::is_mapped() const { + return std::holds_alternative(storage_); +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/OwnedBytes.h b/backends/native/runtime/deserialize/OwnedBytes.h new file mode 100644 index 00000000000..fc376bfd68e --- /dev/null +++ b/backends/native/runtime/deserialize/OwnedBytes.h @@ -0,0 +1,96 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { + +// Owning, read-only bytes: either a heap buffer or a read-only file mapping. +// +// Hands out spans that alias the storage. Each alternative keeps its payload +// address across a move, so spans taken before a move stay valid for as long as +// the OwnedBytes lives. Copy is deleted: this holds a whole model. +// +// The mapped alternative is the one that matters for large packages: nothing is +// copied, the pages are demand-paged, and they are shared with any other +// process mapping the same file. +class OwnedBytes { + private: + // Releases a mapping. Carries the length because that is what munmap needs, + // which lets a unique_ptr supply the whole move-only lifetime — no + // hand-written destructor or move operations. + struct Unmap { + size_t size = 0; + + void operator()(void* base) const noexcept; + }; + + struct HeapBuffer { + std::unique_ptr data; + size_t size = 0; + }; + + // A read-only mapping of an entire file. + using MappedFile = std::unique_ptr; + + std::variant, HeapBuffer, MappedFile> storage_; + + explicit OwnedBytes(std::vector bytes) + : storage_(std::move(bytes)) {} + explicit OwnedBytes(HeapBuffer buffer) : storage_(std::move(buffer)) {} + explicit OwnedBytes(MappedFile mapped_file) + : storage_(std::move(mapped_file)) {} + + public: + // Empty, owning nothing. + OwnedBytes() = default; + + ~OwnedBytes() = default; + OwnedBytes(OwnedBytes&&) noexcept = default; + OwnedBytes& operator=(OwnedBytes&&) noexcept = default; + OwnedBytes(const OwnedBytes&) = delete; + OwnedBytes& operator=(const OwnedBytes&) = delete; + + // The whole payload. Valid for this OwnedBytes' lifetime. + ByteSpan span() const; + + // True when these bytes are a file mapping rather than a heap buffer. + bool is_mapped() const; + + // Take ownership of a buffer the caller already has, without copying it. + static OwnedBytes from_vector(std::vector bytes); + + // Acquire the contents of `path`. Maps it read-only by default: nothing is + // copied, pages arrive on demand, and they are shared with any other process + // mapping the same file. Pass use_mmap=false to read it into the heap + // instead, which is worth it only when the file must outlive edits to it on + // disk. A mapping sees concurrent edits, and accessing pages past a + // concurrent truncation can terminate the process with SIGBUS. + // + // `path` must report its complete size through the filesystem. Virtual files + // such as procfs entries that report zero bytes but produce data are outside + // this API's scope. + // + // Throws std::runtime_error if the file cannot be read, or cannot be mapped + // when mapping was asked for (including on a platform with no mmap). An empty + // file yields empty heap bytes either way, since mmap rejects a zero length. + static OwnedBytes from_file(const std::string& path, bool use_mmap = true); + + private: + static OwnedBytes read_file(const std::string& path); + static OwnedBytes map_file(const std::string& path); +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/targets.bzl b/backends/native/runtime/deserialize/targets.bzl index cb028f2d6cd..8290cc3f3b5 100644 --- a/backends/native/runtime/deserialize/targets.bzl +++ b/backends/native/runtime/deserialize/targets.bzl @@ -10,6 +10,15 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) + # Owning byte buffer backed by heap storage or a read-only file mapping. + runtime.cxx_library( + name = "owned_bytes", + srcs = ["OwnedBytes.cpp"], + exported_headers = ["OwnedBytes.h"], + exported_deps = [":byte_span"], + visibility = ["//executorch/backends/native/..."], + ) + # JSON representation used by package metadata readers. runtime.cxx_library( name = "json", diff --git a/backends/native/test/runtime/deserialize/targets.bzl b/backends/native/test/runtime/deserialize/targets.bzl index d23e793ee68..8989811393e 100644 --- a/backends/native/test/runtime/deserialize/targets.bzl +++ b/backends/native/test/runtime/deserialize/targets.bzl @@ -1,6 +1,14 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") def define_common_targets(): + runtime.cxx_test( + name = "owned_bytes_test", + srcs = ["test_owned_bytes.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:owned_bytes", + ], + ) + runtime.cxx_test( name = "safetensors_reader_test", srcs = ["test_safetensors_reader.cpp"], diff --git a/backends/native/test/runtime/deserialize/test_owned_bytes.cpp b/backends/native/test/runtime/deserialize/test_owned_bytes.cpp new file mode 100644 index 00000000000..43d97f765eb --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_owned_bytes.cpp @@ -0,0 +1,131 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +class OwnedBytesTest : public ::testing::Test { + private: + std::vector paths_; + + protected: + std::string temp_path(std::string_view suffix) { + const auto* info = ::testing::UnitTest::GetInstance()->current_test_info(); + const std::filesystem::path path = std::filesystem::temp_directory_path() / + (std::string("owned_bytes_") + info->name() + std::string(suffix)); + std::error_code error; + std::filesystem::remove(path, error); + paths_.push_back(path); + return path.string(); + } + + void write_file(const std::string& path, std::string_view contents) { + std::ofstream file(path, std::ios::binary | std::ios::trunc); + ASSERT_TRUE(file); + file.write(contents.data(), static_cast(contents.size())); + ASSERT_TRUE(file); + } + + void TearDown() override { + for (const std::filesystem::path& path : paths_) { + std::error_code error; + std::filesystem::remove(path, error); + } + } +}; + +static_assert(!std::is_copy_constructible_v); +static_assert(!std::is_copy_assignable_v); +static_assert(std::is_nothrow_move_constructible_v); +static_assert(std::is_nothrow_move_assignable_v); + +// cppcheck-suppress-begin syntaxError +TEST_F(OwnedBytesTest, TakesOwnershipOfVector) { + std::vector source{1, 2, 3}; + const uint8_t* data = source.data(); + + OwnedBytes bytes = OwnedBytes::from_vector(std::move(source)); + const ByteSpan span = bytes.span(); + OwnedBytes moved = std::move(bytes); + + EXPECT_FALSE(moved.is_mapped()); + EXPECT_EQ(span.data(), data); + EXPECT_EQ(moved.span().data(), data); + EXPECT_EQ( + std::vector(span.begin(), span.end()), + (std::vector{1, 2, 3})); +} + +TEST_F(OwnedBytesTest, ReadsFileIntoHeap) { + const std::string path = temp_path("_heap.bin"); + ASSERT_NO_FATAL_FAILURE(write_file(path, "abc")); + + OwnedBytes bytes = OwnedBytes::from_file(path, false); + const ByteSpan span = bytes.span(); + const OwnedBytes moved = std::move(bytes); + + EXPECT_FALSE(moved.is_mapped()); + EXPECT_TRUE(bytes.span().empty()); + EXPECT_EQ( + std::vector(span.begin(), span.end()), + (std::vector{'a', 'b', 'c'})); +} + +TEST_F(OwnedBytesTest, MapsFile) { + const std::string path = temp_path("_mapped.bin"); + ASSERT_NO_FATAL_FAILURE(write_file(path, "abc")); + +#if defined(_WIN32) + EXPECT_THROW(OwnedBytes::from_file(path), std::runtime_error); +#else + OwnedBytes bytes = OwnedBytes::from_file(path); + const ByteSpan span = bytes.span(); + const OwnedBytes moved = std::move(bytes); + EXPECT_TRUE(moved.is_mapped()); + EXPECT_TRUE(bytes.span().empty()); + EXPECT_EQ( + std::vector(span.begin(), span.end()), + (std::vector{'a', 'b', 'c'})); +#endif +} + +TEST_F(OwnedBytesTest, EmptyFileUsesHeapStorage) { + const std::string path = temp_path("_empty.bin"); + ASSERT_NO_FATAL_FAILURE(write_file(path, "")); + + const OwnedBytes bytes = OwnedBytes::from_file(path); + + EXPECT_FALSE(bytes.is_mapped()); + EXPECT_TRUE(bytes.span().empty()); +} + +TEST_F(OwnedBytesTest, RejectsInvalidPaths) { + EXPECT_THROW( + OwnedBytes::from_file(temp_path("_missing.bin"), false), + std::runtime_error); + EXPECT_THROW( + OwnedBytes::from_file( + std::filesystem::temp_directory_path().string(), false), + std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn From 0b5c907980f7ab4d0c78a584b5a0c4f1fe573491 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 10 Sep 2026 12:14:29 -0700 Subject: [PATCH 168/190] [executorch][native] Add the .ptn package reader ## Ulterior Motive Unify native graph and constant loading behind one PTN package abstraction. ## Rationale **What**: Add a `libzip` wrapper, `Package` facade, alias resolution, and `ptn_inspector`. **Why**: Runtime callers should consume one validated package instead of parsing ZIP and safetensors details themselves. ## Details ```text model.ptn | +-- program.ptg ---------> Program +-- program.safetensors -> SafeTensorsReader +-- aliases.json --------> alias key -> owner key ``` - `ZipReader` accepts stored, unencrypted members only. Transparent lookup and shared entry plumbing avoid temporary strings and duplicate map probes. - `Package` owns the archive image and extracted data. `Constant` views borrow metadata and bytes from that package. - `ptn_inspector` delegates file loading to `Package` and parses the program once in its default reporting path. - Test archives use RAII, including constructor-failure paths. Authored with Codex. Differential Revision: [D118480654](https://our.internmc.facebook.com/intern/diff/D118480654/) ghstack-source-id: 427858776 Pull-Request: https://github.com/pytorch/executorch/pull/22702 --- .../native/runtime/deserialize/Package.cpp | 170 ++++++++++++++ backends/native/runtime/deserialize/Package.h | 107 +++++++++ .../native/runtime/deserialize/ZipReader.cpp | 213 ++++++++++++++++++ .../native/runtime/deserialize/ZipReader.h | 84 +++++++ .../native/runtime/deserialize/targets.bzl | 22 ++ .../test/runtime/deserialize/targets.bzl | 9 + .../runtime/deserialize/test_zip_reader.cpp | 113 ++++++++++ 7 files changed, 718 insertions(+) create mode 100644 backends/native/runtime/deserialize/Package.cpp create mode 100644 backends/native/runtime/deserialize/Package.h create mode 100644 backends/native/runtime/deserialize/ZipReader.cpp create mode 100644 backends/native/runtime/deserialize/ZipReader.h create mode 100644 backends/native/test/runtime/deserialize/test_zip_reader.cpp diff --git a/backends/native/runtime/deserialize/Package.cpp b/backends/native/runtime/deserialize/Package.cpp new file mode 100644 index 00000000000..ba8dec00ed6 --- /dev/null +++ b/backends/native/runtime/deserialize/Package.cpp @@ -0,0 +1,170 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +// Reserved by safetensors, so it can never name a constant. +constexpr std::string_view kMetadataKey = "__metadata__"; + +std::unordered_map parse_aliases( + ByteSpan member, + const SafeTensorsReader& tensors) { + Json doc; + try { + doc = Json::parse(std::string_view( + reinterpret_cast(member.data()), member.size())); + } catch (const Json::exception& error) { + throw std::runtime_error( + "package: invalid aliases.json: " + std::string(error.what())); + } + if (!doc.is_object()) { + throw std::runtime_error("package: aliases.json is not a JSON object"); + } + + std::unordered_map aliases; + for (auto entry = doc.begin(); entry != doc.end(); ++entry) { + const std::string& key = entry.key(); + if (!entry.value().is_string()) { + throw std::runtime_error( + "package: alias '" + key + "' does not name a string owner"); + } + const std::string& owner = entry.value().get_ref(); + if (key == kMetadataKey) { + throw std::runtime_error( + "package: alias key is reserved by safetensors: " + key); + } + // An owner is always a real safetensors entry and an alias is never one, so + // resolution stays a single lookup. Enforce both rather than trusting it. + if (tensors.find(owner) == nullptr) { + std::string message = "package: alias '"; + message += key; + message += "' names owner '"; + message += owner; + message += "', which has no safetensors entry"; + throw std::runtime_error(message); + } + if (tensors.find(key) != nullptr) { + throw std::runtime_error( + "package: '" + key + "' is both a safetensors owner and an alias"); + } + if (!aliases.emplace(key, owner).second) { + throw std::runtime_error("package: duplicate alias key: " + key); + } + } + return aliases; +} + +} // namespace + +bool Package::looks_like_package(ByteSpan bytes) { + // Every zip record signature begins "PK"; a bare .ptg starts with a + // flatbuffer root offset followed by "NPTG" at offset 4, so this cannot + // collide. + return bytes.size() >= 2 && bytes[0] == 'P' && bytes[1] == 'K'; +} + +Package Package::load(std::vector bytes) { + Package out; + out.bytes_ = std::move(bytes); + const ByteSpan image{out.bytes_.data(), out.bytes_.size()}; + + out.zip_ = ZipReader::open(image); + + if (!out.zip_->member_size(kProgramEntry)) { + throw std::runtime_error( + std::string("package: missing required member ") + kProgramEntry); + } + out.program_ = out.zip_->read(kProgramEntry); + + // Absent whenever the program references no constants, which is normal for a + // graph over user inputs alone. + if (out.zip_->member_size(kSafeTensorsEntry)) { + out.tensor_bytes_ = out.zip_->read(kSafeTensorsEntry); + out.tensors_ = SafeTensorsReader::open(ByteSpan(out.tensor_bytes_)); + } + + if (out.zip_->member_size(kAliasesEntry)) { + if (!out.tensors_) { + throw std::runtime_error( + std::string("package: has ") + kAliasesEntry + " but no " + + kSafeTensorsEntry); + } + const std::vector aliases = out.zip_->read(kAliasesEntry); + out.aliases_ = parse_aliases(ByteSpan(aliases), *out.tensors_); + } + + return out; +} + +Package Package::load_file(const std::string& path) { + std::ifstream file(path, std::ios::binary | std::ios::ate); + if (!file) { + throw std::runtime_error("package: cannot open " + path); + } + const std::streamsize size = file.tellg(); + if (size < 0) { + throw std::runtime_error("package: cannot size " + path); + } + file.seekg(0, std::ios::beg); + std::vector bytes(static_cast(size)); + if (size > 0 && !file.read(reinterpret_cast(bytes.data()), size)) { + throw std::runtime_error("package: cannot read " + path); + } + return load(std::move(bytes)); +} + +std::optional Package::constant(const std::string& key) const { + if (!tensors_) { + return std::nullopt; + } + const auto alias = aliases_.find(key); + const std::string& owner = alias == aliases_.end() ? key : alias->second; + + const TensorEntry* entry = tensors_->find(owner); + if (entry == nullptr) { + return std::nullopt; + } + + Constant out; + out.dtype = entry->dtype; + out.sizes = &entry->sizes; + out.bytes = tensors_->bytes(*entry); + out.owner = owner; + return out; +} + +std::vector Package::keys() const { + std::vector out; + if (tensors_) { + out = tensors_->names(); + } + for (const auto& alias : aliases_) { + out.push_back(alias.first); + } + std::ranges::sort(out); + return out; +} + +const std::vector& Package::owner_keys() const { + static const std::vector kNone; + return tensors_ ? tensors_->names() : kNone; +} + +size_t Package::constant_bytes() const { + return tensors_ ? tensors_->total_bytes() : 0; +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/Package.h b/backends/native/runtime/deserialize/Package.h new file mode 100644 index 00000000000..3cb83f60306 --- /dev/null +++ b/backends/native/runtime/deserialize/Package.h @@ -0,0 +1,107 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace ptn { + +// Fixed member names inside a .ptn. The package survives being renamed because +// nothing depends on the file name. +constexpr const char* kProgramEntry = "program.ptg"; +constexpr const char* kSafeTensorsEntry = "program.safetensors"; +constexpr const char* kAliasesEntry = "aliases.json"; + +// One constant resolved out of a package. The sizes and bytes are borrowed from +// the Package and must not outlive it. +struct Constant { + ScalarType dtype = kFloat; + const std::vector* sizes = nullptr; + ByteSpan bytes; + // Key that actually owns these bytes. Differs from the requested key when the + // package deduplicated two byte-identical immutable constants. + std::string owner; +}; + +// A loaded .ptn package: the serialized native Program plus the constants it +// references. +// +// Owns the package image and extracted members. Copy is deleted: a package is +// model-sized. +class Package { + private: + std::vector bytes_; + std::optional zip_; + std::vector program_; + std::vector tensor_bytes_; + // Absent when the program references no constants, in which case the package + // has no safetensors member at all. + std::optional tensors_; + std::unordered_map aliases_; + + Package() = default; + + public: + ~Package() = default; + Package(Package&&) noexcept = default; + Package& operator=(Package&&) noexcept = default; + Package(const Package&) = delete; + Package& operator=(const Package&) = delete; + + // Parse a .ptn image. Takes ownership rather than copying, so a + // hundred-megabyte package is resident once. Throws std::runtime_error if the + // zip, the safetensors index, or the alias map is malformed, or if the + // required program member is missing. + static Package load(std::vector bytes); + + // Read and parse a .ptn from disk. Throws std::runtime_error if the file + // cannot be read. + static Package load_file(const std::string& path); + + // The serialized native Program flatbuffer (the program.ptg member). + ByteSpan program_bytes() const { + return ByteSpan(program_); + } + + // Zip member names present, in central-directory order. Diagnostic only. + const std::vector& member_names() const { + return zip_->names(); + } + + // Keys that own their bytes, in safetensors header order. + const std::vector& owner_keys() const; + + // Duplicate key -> owner key. + const std::unordered_map& aliases() const { + return aliases_; + } + + // Constant for `key`, resolving an alias to its owner. nullopt when the + // package holds no such constant. + std::optional constant(const std::string& key) const; + + // Every key the package resolves, owners and aliases alike, sorted. + std::vector keys() const; + + // Total bytes across owner entries, i.e. what the constants actually cost. + size_t constant_bytes() const; + + // True if `bytes` starts with the zip local-header signature, i.e. looks like + // a package rather than a bare .ptg flatbuffer. Lets a tool accept either. + static bool looks_like_package(ByteSpan bytes); +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/ZipReader.cpp b/backends/native/runtime/deserialize/ZipReader.cpp new file mode 100644 index 00000000000..c38845b040f --- /dev/null +++ b/backends/native/runtime/deserialize/ZipReader.cpp @@ -0,0 +1,213 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include + +#include + +namespace ptn { +namespace { + +struct ZipDeleter { + void operator()(zip_t* archive) const noexcept { + zip_discard(archive); + } +}; + +struct ZipFileDeleter { + void operator()(zip_file_t* file) const noexcept { + zip_fclose(file); + } +}; + +using ZipHandle = std::unique_ptr; +using ZipFileHandle = std::unique_ptr; + +[[noreturn]] void throw_zip(zip_t* archive, const std::string& operation) { + throw std::runtime_error("zip: " + operation + ": " + zip_strerror(archive)); +} + +[[noreturn]] void throw_zip_file( + zip_file_t* file, + const std::string& operation) { + throw std::runtime_error( + "zip: " + operation + ": " + zip_file_strerror(file)); +} + +ZipHandle open_path(const std::string& path) { + int error_code = 0; + ZipHandle archive( + zip_open(path.c_str(), ZIP_RDONLY | ZIP_CHECKCONS, &error_code)); + if (archive == nullptr) { + zip_error_t error; + zip_error_init_with_code(&error, error_code); + const std::string message = + "zip: cannot open " + path + ": " + zip_error_strerror(&error); + zip_error_fini(&error); + throw std::runtime_error(message); + } + return archive; +} + +ZipHandle open_memory(ByteSpan bytes) { + zip_error_t error; + zip_error_init(&error); + zip_source_t* source = + zip_source_buffer_create(bytes.data(), bytes.size(), 0, &error); + if (source == nullptr) { + const std::string message = "zip: cannot create memory source: " + + std::string(zip_error_strerror(&error)); + zip_error_fini(&error); + throw std::runtime_error(message); + } + + ZipHandle archive( + zip_open_from_source(source, ZIP_RDONLY | ZIP_CHECKCONS, &error)); + if (archive == nullptr) { + zip_source_free(source); + const std::string message = "zip: cannot open memory source: " + + std::string(zip_error_strerror(&error)); + zip_error_fini(&error); + throw std::runtime_error(message); + } + zip_error_fini(&error); + return archive; +} + +} // namespace + +struct ZipReader::Impl { + explicit Impl(ZipHandle handle) : archive(std::move(handle)) {} + + ZipHandle archive; +}; + +ZipReader::ZipReader(std::unique_ptr impl) : impl_(std::move(impl)) { + const zip_int64_t count = zip_get_num_entries(impl_->archive.get(), 0); + if (count < 0) { + throw_zip(impl_->archive.get(), "cannot enumerate members"); + } + + names_.reserve(static_cast(count)); + for (zip_uint64_t index = 0; index < static_cast(count); + ++index) { + zip_stat_t stat; + zip_stat_init(&stat); + if (zip_stat_index(impl_->archive.get(), index, ZIP_FL_UNCHANGED, &stat) != + 0) { + throw_zip(impl_->archive.get(), "cannot stat member"); + } + constexpr zip_uint64_t kRequired = ZIP_STAT_NAME | ZIP_STAT_SIZE | + ZIP_STAT_COMP_METHOD | ZIP_STAT_ENCRYPTION_METHOD; + if ((stat.valid & kRequired) != kRequired || stat.name == nullptr) { + throw std::runtime_error("zip: member metadata is incomplete"); + } + const std::string name(stat.name); + if (stat.comp_method != ZIP_CM_STORE) { + throw std::runtime_error("zip: member is compressed: " + name); + } + if (stat.encryption_method != ZIP_EM_NONE) { + throw std::runtime_error("zip: member is encrypted: " + name); + } + if (stat.size > std::numeric_limits::max()) { + throw std::runtime_error("zip: member is too large: " + name); + } + if (!entries_.emplace(name, Entry{index, static_cast(stat.size)}) + .second) { + throw std::runtime_error("zip: duplicate member name: " + name); + } + names_.push_back(name); + } +} + +ZipReader::~ZipReader() = default; +ZipReader::ZipReader(ZipReader&&) noexcept = default; +ZipReader& ZipReader::operator=(ZipReader&&) noexcept = default; + +ZipReader ZipReader::open(const std::string& path) { + return ZipReader(std::make_unique(open_path(path))); +} + +ZipReader ZipReader::open(ByteSpan archive) { + return ZipReader(std::make_unique(open_memory(archive))); +} + +std::optional ZipReader::member_size(std::string_view name) const { + const Entry* entry = find_entry(name); + return entry == nullptr ? std::nullopt : std::optional(entry->size); +} + +std::vector ZipReader::read(std::string_view name) const { + const Entry* entry = find_entry(name); + if (entry == nullptr) { + throw std::runtime_error("zip: no member named " + std::string(name)); + } + std::vector bytes(entry->size); + read_entry_into(name, *entry, 0, MutableByteSpan(bytes)); + return bytes; +} + +void ZipReader::read_into( + std::string_view name, + size_t offset, + MutableByteSpan destination) const { + const Entry* entry = find_entry(name); + if (entry == nullptr) { + throw std::runtime_error("zip: no member named " + std::string(name)); + } + read_entry_into(name, *entry, offset, destination); +} + +const ZipReader::Entry* ZipReader::find_entry(std::string_view name) const { + const auto entry = entries_.find(name); + return entry == entries_.end() ? nullptr : &entry->second; +} + +void ZipReader::read_entry_into( + std::string_view name, + const Entry& entry, + size_t offset, + MutableByteSpan destination) const { + if (offset > entry.size || destination.size() > entry.size - offset) { + throw std::runtime_error( + "zip: read range is outside member " + std::string(name)); + } + if (destination.empty()) { + return; + } + + ZipFileHandle file( + zip_fopen_index(impl_->archive.get(), entry.index, ZIP_FL_UNCHANGED)); + if (file == nullptr) { + throw_zip(impl_->archive.get(), "cannot open member " + std::string(name)); + } + if (zip_fseek(file.get(), static_cast(offset), SEEK_SET) != 0) { + throw_zip_file(file.get(), "cannot seek member " + std::string(name)); + } + + size_t written = 0; + while (written < destination.size()) { + const zip_uint64_t request = + static_cast(destination.size() - written); + const zip_int64_t count = + zip_fread(file.get(), destination.data() + written, request); + if (count < 0) { + throw_zip_file(file.get(), "cannot read member " + std::string(name)); + } + if (count == 0) { + throw std::runtime_error( + "zip: unexpected end of member " + std::string(name)); + } + written += static_cast(count); + } +} + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/ZipReader.h b/backends/native/runtime/deserialize/ZipReader.h new file mode 100644 index 00000000000..ea312ffab38 --- /dev/null +++ b/backends/native/runtime/deserialize/ZipReader.h @@ -0,0 +1,84 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +namespace ptn { + +// Read-only access to stored members of a zip archive. +class ZipReader { + private: + struct Entry { + uint64_t index = 0; + size_t size = 0; + }; + + struct StringHash { + using is_transparent = void; + + size_t operator()(std::string_view value) const noexcept { + return std::hash{}(value); + } + }; + + struct Impl; + + std::unique_ptr impl_; + std::unordered_map> entries_; + std::vector names_; + + explicit ZipReader(std::unique_ptr impl); + const Entry* find_entry(std::string_view name) const; + void read_entry_into( + std::string_view name, + const Entry& entry, + size_t offset, + MutableByteSpan destination) const; + + public: + ~ZipReader(); + ZipReader(ZipReader&&) noexcept; + ZipReader& operator=(ZipReader&&) noexcept; + ZipReader(const ZipReader&) = delete; + ZipReader& operator=(const ZipReader&) = delete; + + // Opens an archive without loading its member payloads. + static ZipReader open(const std::string& path); + + // Opens an archive over caller-owned memory. `archive` must outlive this + // reader and every read made through it. + static ZipReader open(ByteSpan archive); + + // Member size, or nullopt when the member is absent. + std::optional member_size(std::string_view name) const; + + // Copies one complete member. + std::vector read(std::string_view name) const; + + // Copies a member range directly into caller-owned storage. + void read_into( + std::string_view name, + size_t offset, + MutableByteSpan destination) const; + + const std::vector& names() const { + return names_; + } +}; + +} // namespace ptn diff --git a/backends/native/runtime/deserialize/targets.bzl b/backends/native/runtime/deserialize/targets.bzl index 8290cc3f3b5..f8028e6abf8 100644 --- a/backends/native/runtime/deserialize/targets.bzl +++ b/backends/native/runtime/deserialize/targets.bzl @@ -28,6 +28,15 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) + runtime.cxx_library( + name = "zip_reader", + srcs = ["ZipReader.cpp"], + exported_headers = ["ZipReader.h"], + exported_deps = [":byte_span"], + deps = ["fbsource//third-party/libzip:zip"], + visibility = ["//executorch/backends/native/..."], + ) + # safetensors index reader. runtime.cxx_library( name = "safetensors_reader", @@ -40,3 +49,16 @@ def define_common_targets(): deps = [":json"], visibility = ["//executorch/backends/native/..."], ) + runtime.cxx_library( + name = "package", + srcs = ["Package.cpp"], + exported_headers = ["Package.h"], + exported_deps = [ + ":byte_span", + ":safetensors_reader", + ":zip_reader", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + deps = [":json"], + visibility = ["PUBLIC"], + ) diff --git a/backends/native/test/runtime/deserialize/targets.bzl b/backends/native/test/runtime/deserialize/targets.bzl index 8989811393e..333cce5018e 100644 --- a/backends/native/test/runtime/deserialize/targets.bzl +++ b/backends/native/test/runtime/deserialize/targets.bzl @@ -16,3 +16,12 @@ def define_common_targets(): "//executorch/backends/native/runtime/deserialize:safetensors_reader", ], ) + + runtime.cxx_test( + name = "zip_reader_test", + srcs = ["test_zip_reader.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:zip_reader", + "fbsource//third-party/libzip:zip", + ], + ) diff --git a/backends/native/test/runtime/deserialize/test_zip_reader.cpp b/backends/native/test/runtime/deserialize/test_zip_reader.cpp new file mode 100644 index 00000000000..ce8362c6ea5 --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_zip_reader.cpp @@ -0,0 +1,113 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace ptn { +namespace { + +struct ZipDiscard { + void operator()(zip_t* archive) const noexcept { + zip_discard(archive); + } +}; + +class TempZip { + public: + TempZip() { + path_ = std::filesystem::temp_directory_path() / + ("ptn_zip_reader_" + std::to_string(reinterpret_cast(this)) + + ".zip"); + int error = 0; + std::unique_ptr archive( + zip_open(path_.string().c_str(), ZIP_CREATE | ZIP_TRUNCATE, &error)); + if (archive == nullptr) { + throw std::runtime_error("failed to create test zip"); + } + add(archive.get(), "program.ptg", "program"); + add(archive.get(), "program.safetensors", "0123456789"); + if (zip_close(archive.get()) != 0) { + throw std::runtime_error("failed to close test zip"); + } + archive.release(); + } + + ~TempZip() { + std::error_code error; + std::filesystem::remove(path_, error); + } + + const std::string path() const { + return path_.string(); + } + + private: + static void add(zip_t* archive, const char* name, std::string_view bytes) { + zip_source_t* source = + zip_source_buffer(archive, bytes.data(), bytes.size(), 0); + if (source == nullptr) { + throw std::runtime_error("failed to create test zip source"); + } + const zip_int64_t index = + zip_file_add(archive, name, source, ZIP_FL_ENC_UTF_8); + if (index < 0) { + zip_source_free(source); + throw std::runtime_error("failed to add test zip member"); + } + if (zip_set_file_compression(archive, index, ZIP_CM_STORE, 0) != 0) { + throw std::runtime_error("failed to store test zip member"); + } + } + + std::filesystem::path path_; +}; + +// cppcheck-suppress-begin syntaxError +TEST(ZipReaderTest, ReadsStoredMemberRanges) { + const TempZip file; + ZipReader zip = ZipReader::open(file.path()); + + EXPECT_EQ( + zip.names(), + (std::vector{"program.ptg", "program.safetensors"})); + EXPECT_EQ(zip.member_size("program.safetensors"), 10); + EXPECT_EQ(zip.member_size("missing"), std::nullopt); + + std::array bytes{}; + zip.read_into("program.safetensors", 3, MutableByteSpan(bytes)); + EXPECT_EQ(bytes, (std::array{'3', '4', '5', '6'})); + EXPECT_EQ( + zip.read("program.ptg"), + (std::vector{'p', 'r', 'o', 'g', 'r', 'a', 'm'})); +} + +TEST(ZipReaderTest, RejectsInvalidRanges) { + const TempZip file; + ZipReader zip = ZipReader::open(file.path()); + + std::array bytes{}; + EXPECT_THROW( + zip.read_into("program.safetensors", 8, MutableByteSpan(bytes)), + std::runtime_error); + EXPECT_THROW( + zip.read_into("missing", 0, MutableByteSpan(bytes)), std::runtime_error); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn From a081a3da09f4ab40320a47ba16dbd6a16fd48651 Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 10 Sep 2026 12:14:34 -0700 Subject: [PATCH 169/190] [executorch][native] Load .ptn constants on demand ## Ulterior Motive Keep large model weights off heap until a backend chooses their destination and lifetime. ## Rationale **What**: Split safetensors metadata parsing from payload access. Add metadata-only lookup, owned acquisition, direct destination loading, package identity, and checksum verification. **Why**: Eager extraction duplicates model-sized weights and prevents direct placement into backend-managed memory. ## Details ```text Package::load(path) | +-- open ZIP directory +-- read program.ptg +-- read safetensors prefix + JSON only | +-- constant_info(key) no payload read +-- acquire_constant(key) new OwnedBytes +-- load_constant_into() caller destination +-- verify_constants() streamed CRC check ``` - Package identity survives moves. Alias lookups return canonical owner keys, enabling cross-instance upload accounting. - Prefix and checksum scratch arrays are value-initialized, satisfying the blocking clang-tidy signal. - `ptn_inspector` uses the path-based `Package::load` API after the eager `load_file` API is removed. - New package-test ZIP handles use RAII on constructor failures. Authored with Codex. Differential Revision: [D119396097](https://our.internmc.facebook.com/intern/diff/D119396097/) ghstack-source-id: 427858773 Pull-Request: https://github.com/pytorch/executorch/pull/22703 --- .../native/runtime/deserialize/Package.cpp | 135 +++++++++++---- backends/native/runtime/deserialize/Package.h | 58 ++++--- .../runtime/deserialize/SafeTensorsReader.cpp | 56 +++--- .../runtime/deserialize/SafeTensorsReader.h | 20 ++- .../native/runtime/deserialize/ZipReader.cpp | 40 +++++ .../native/runtime/deserialize/ZipReader.h | 3 + .../native/runtime/deserialize/targets.bzl | 6 +- .../test/runtime/deserialize/targets.bzl | 9 + .../test/runtime/deserialize/test_package.cpp | 159 ++++++++++++++++++ .../deserialize/test_safetensors_reader.cpp | 12 +- .../runtime/deserialize/test_zip_reader.cpp | 1 + 11 files changed, 407 insertions(+), 92 deletions(-) create mode 100644 backends/native/test/runtime/deserialize/test_package.cpp diff --git a/backends/native/runtime/deserialize/Package.cpp b/backends/native/runtime/deserialize/Package.cpp index ba8dec00ed6..2b151c35b08 100644 --- a/backends/native/runtime/deserialize/Package.cpp +++ b/backends/native/runtime/deserialize/Package.cpp @@ -7,7 +7,8 @@ #include #include -#include +#include +#include #include #include @@ -18,6 +19,7 @@ namespace { // Reserved by safetensors, so it can never name a constant. constexpr std::string_view kMetadataKey = "__metadata__"; +std::atomic next_package_id{1}; std::unordered_map parse_aliases( ByteSpan member, @@ -76,57 +78,81 @@ bool Package::looks_like_package(ByteSpan bytes) { return bytes.size() >= 2 && bytes[0] == 'P' && bytes[1] == 'K'; } -Package Package::load(std::vector bytes) { +Package::Package() + : id_(next_package_id.fetch_add(1, std::memory_order_relaxed)) {} + +Package Package::load(OwnedBytes bytes) { Package out; - out.bytes_ = std::move(bytes); - const ByteSpan image{out.bytes_.data(), out.bytes_.size()}; + out.archive_bytes_ = std::move(bytes); + out.zip_ = ZipReader::open(out.archive_bytes_.span()); + out.load_metadata(); + return out; +} - out.zip_ = ZipReader::open(image); +Package Package::load(const std::string& path) { + Package out; + out.zip_ = ZipReader::open(path); + out.load_metadata(); + return out; +} + +Package& Package::operator=(Package&& other) noexcept { + if (this != &other) { + zip_.reset(); + id_ = other.id_; + archive_bytes_ = std::move(other.archive_bytes_); + zip_ = std::move(other.zip_); + program_ = std::move(other.program_); + tensors_ = std::move(other.tensors_); + tensor_data_offset_ = other.tensor_data_offset_; + aliases_ = std::move(other.aliases_); + } + return *this; +} - if (!out.zip_->member_size(kProgramEntry)) { +void Package::load_metadata() { + if (!zip_->member_size(kProgramEntry)) { throw std::runtime_error( std::string("package: missing required member ") + kProgramEntry); } - out.program_ = out.zip_->read(kProgramEntry); + program_ = zip_->read(kProgramEntry); // Absent whenever the program references no constants, which is normal for a // graph over user inputs alone. - if (out.zip_->member_size(kSafeTensorsEntry)) { - out.tensor_bytes_ = out.zip_->read(kSafeTensorsEntry); - out.tensors_ = SafeTensorsReader::open(ByteSpan(out.tensor_bytes_)); + const std::optional tensor_size = + zip_->member_size(kSafeTensorsEntry); + if (tensor_size) { + std::array prefix{}; + if (*tensor_size < prefix.size()) { + throw std::runtime_error( + "package: safetensors member is shorter than its length prefix"); + } + zip_->read_into(kSafeTensorsEntry, 0, MutableByteSpan(prefix)); + const size_t header_size = SafeTensorsReader::header_size(prefix); + if (header_size > *tensor_size - prefix.size()) { + throw std::runtime_error( + "package: safetensors header exceeds its zip member"); + } + std::vector header(header_size); + zip_->read_into(kSafeTensorsEntry, prefix.size(), MutableByteSpan(header)); + tensor_data_offset_ = prefix.size() + header_size; + tensors_ = SafeTensorsReader::open_header( + ByteSpan(header), *tensor_size - tensor_data_offset_); } - if (out.zip_->member_size(kAliasesEntry)) { - if (!out.tensors_) { + if (zip_->member_size(kAliasesEntry)) { + if (!tensors_) { throw std::runtime_error( std::string("package: has ") + kAliasesEntry + " but no " + kSafeTensorsEntry); } - const std::vector aliases = out.zip_->read(kAliasesEntry); - out.aliases_ = parse_aliases(ByteSpan(aliases), *out.tensors_); + const std::vector aliases = zip_->read(kAliasesEntry); + aliases_ = parse_aliases(ByteSpan(aliases), *tensors_); } - - return out; } -Package Package::load_file(const std::string& path) { - std::ifstream file(path, std::ios::binary | std::ios::ate); - if (!file) { - throw std::runtime_error("package: cannot open " + path); - } - const std::streamsize size = file.tellg(); - if (size < 0) { - throw std::runtime_error("package: cannot size " + path); - } - file.seekg(0, std::ios::beg); - std::vector bytes(static_cast(size)); - if (size > 0 && !file.read(reinterpret_cast(bytes.data()), size)) { - throw std::runtime_error("package: cannot read " + path); - } - return load(std::move(bytes)); -} - -std::optional Package::constant(const std::string& key) const { +std::optional Package::constant_info( + const std::string& key) const { if (!tensors_) { return std::nullopt; } @@ -138,14 +164,51 @@ std::optional Package::constant(const std::string& key) const { return std::nullopt; } - Constant out; + ConstantInfo out; + out.package_id = id_; out.dtype = entry->dtype; out.sizes = &entry->sizes; - out.bytes = tensors_->bytes(*entry); + out.nbytes = entry->nbytes; out.owner = owner; return out; } +std::optional Package::acquire_constant( + const std::string& key) const { + const std::optional info = constant_info(key); + if (!info) { + return std::nullopt; + } + std::vector bytes(info->nbytes); + load_constant_into(key, MutableByteSpan(bytes)); + return OwnedBytes::from_vector(std::move(bytes)); +} + +bool Package::load_constant_into( + const std::string& key, + MutableByteSpan destination) const { + const std::optional info = constant_info(key); + if (!info) { + return false; + } + if (destination.size() != info->nbytes) { + throw std::runtime_error( + "package: destination for '" + key + "' has " + + std::to_string(destination.size()) + " bytes; expected " + + std::to_string(info->nbytes)); + } + const TensorEntry* entry = tensors_->find(info->owner); + zip_->read_into( + kSafeTensorsEntry, tensor_data_offset_ + entry->offset, destination); + return true; +} + +void Package::verify_constants() const { + if (tensors_) { + zip_->verify(kSafeTensorsEntry); + } +} + std::vector Package::keys() const { std::vector out; if (tensors_) { diff --git a/backends/native/runtime/deserialize/Package.h b/backends/native/runtime/deserialize/Package.h index 3cb83f60306..77788e9e8d6 100644 --- a/backends/native/runtime/deserialize/Package.h +++ b/backends/native/runtime/deserialize/Package.h @@ -6,6 +6,7 @@ #pragma once +#include #include #include #include @@ -13,6 +14,7 @@ #include #include +#include #include #include #include @@ -25,12 +27,13 @@ constexpr const char* kProgramEntry = "program.ptg"; constexpr const char* kSafeTensorsEntry = "program.safetensors"; constexpr const char* kAliasesEntry = "aliases.json"; -// One constant resolved out of a package. The sizes and bytes are borrowed from +// Metadata for one constant resolved out of a package. `sizes` is borrowed from // the Package and must not outlive it. -struct Constant { +struct ConstantInfo { + uint64_t package_id = 0; ScalarType dtype = kFloat; const std::vector* sizes = nullptr; - ByteSpan bytes; + size_t nbytes = 0; // Key that actually owns these bytes. Differs from the requested key when the // package deduplicated two byte-identical immutable constants. std::string owner; @@ -39,37 +42,39 @@ struct Constant { // A loaded .ptn package: the serialized native Program plus the constants it // references. // -// Owns the package image and extracted members. Copy is deleted: a package is -// model-sized. +// Opening a file-backed package reads its directory and metadata, but leaves +// weight payloads on disk until an engine requests them. class Package { private: - std::vector bytes_; + uint64_t id_ = 0; + OwnedBytes archive_bytes_; std::optional zip_; std::vector program_; - std::vector tensor_bytes_; // Absent when the program references no constants, in which case the package // has no safetensors member at all. std::optional tensors_; + size_t tensor_data_offset_ = 0; std::unordered_map aliases_; - Package() = default; + Package(); public: ~Package() = default; Package(Package&&) noexcept = default; - Package& operator=(Package&&) noexcept = default; + Package& operator=(Package&& other) noexcept; Package(const Package&) = delete; Package& operator=(const Package&) = delete; - // Parse a .ptn image. Takes ownership rather than copying, so a - // hundred-megabyte package is resident once. Throws std::runtime_error if the - // zip, the safetensors index, or the alias map is malformed, or if the - // required program member is missing. - static Package load(std::vector bytes); + // Open and parse the .ptn at `path` without loading its weight payloads. + static Package load(const std::string& path); - // Read and parse a .ptn from disk. Throws std::runtime_error if the file - // cannot be read. - static Package load_file(const std::string& path); + // Parse a .ptn image already in hand. Takes ownership rather than copying, so + // a hundred-megabyte package is resident once. For callers that must inspect + // the bytes before deciding this is a package at all; everyone else should + // use the path overload. Throws std::runtime_error if the zip, the + // safetensors index, or the alias map is malformed, or if the required + // program member is missing. + static Package load(OwnedBytes bytes); // The serialized native Program flatbuffer (the program.ptg member). ByteSpan program_bytes() const { @@ -89,9 +94,19 @@ class Package { return aliases_; } - // Constant for `key`, resolving an alias to its owner. nullopt when the - // package holds no such constant. - std::optional constant(const std::string& key) const; + // Metadata for `key`, resolving an alias to its owner. nullopt when absent. + std::optional constant_info(const std::string& key) const; + + // Load one constant into a new owning buffer. nullopt when absent. + std::optional acquire_constant(const std::string& key) const; + + // Load one constant directly into an exact-sized destination. Returns false + // when absent and throws when the destination has the wrong size. + bool load_constant_into(const std::string& key, MutableByteSpan destination) + const; + + // Stream all weight bytes once to verify the zip member checksum. + void verify_constants() const; // Every key the package resolves, owners and aliases alike, sorted. std::vector keys() const; @@ -102,6 +117,9 @@ class Package { // True if `bytes` starts with the zip local-header signature, i.e. looks like // a package rather than a bare .ptg flatbuffer. Lets a tool accept either. static bool looks_like_package(ByteSpan bytes); + + private: + void load_metadata(); }; } // namespace ptn diff --git a/backends/native/runtime/deserialize/SafeTensorsReader.cpp b/backends/native/runtime/deserialize/SafeTensorsReader.cpp index 6b89588180f..35dc7418b1c 100644 --- a/backends/native/runtime/deserialize/SafeTensorsReader.cpp +++ b/backends/native/runtime/deserialize/SafeTensorsReader.cpp @@ -10,6 +10,7 @@ #include #include #include +#include #include #include #include @@ -21,8 +22,6 @@ namespace { // Reserved header member holding free-form string metadata, not a tensor. constexpr std::string_view kMetadataKey = "__metadata__"; -constexpr size_t kHeaderLenSize = 8; - struct DtypeCode { std::string_view code; ScalarType dtype; @@ -57,19 +56,6 @@ ScalarType scalar_type_of(std::string_view code) { return it->dtype; } -uint64_t read_header_len(ByteSpan blob) { - static_assert( - std::endian::native == std::endian::little, - "the length prefix is little-endian; a big-endian host needs a swap"); - if (blob.size() < kHeaderLenSize) { - throw std::runtime_error( - "safetensors: blob is shorter than its length prefix"); - } - uint64_t len = 0; - std::memcpy(&len, blob.data(), kHeaderLenSize); - return len; -} - const Json& required_member( const Json& entry, std::string_view key, @@ -129,15 +115,39 @@ size_t numel_of(const std::vector& sizes, const std::string& name) { } // namespace +size_t SafeTensorsReader::header_size(ByteSpan prefix) { + static_assert( + std::endian::native == std::endian::little, + "the length prefix is little-endian; a big-endian host needs a swap"); + if (prefix.size() < kLengthPrefixSize) { + throw std::runtime_error( + "safetensors: blob is shorter than its length prefix"); + } + uint64_t size = 0; + std::memcpy(&size, prefix.data(), kLengthPrefixSize); + if (size > std::numeric_limits::max()) { + throw std::runtime_error("safetensors: header is too large"); + } + return static_cast(size); +} + SafeTensorsReader SafeTensorsReader::open(ByteSpan blob) { - const uint64_t header_len = read_header_len(blob); - if (header_len > blob.size() - kHeaderLenSize) { + const size_t header_len = header_size(blob); + if (header_len > blob.size() - kLengthPrefixSize) { throw std::runtime_error("safetensors: header length exceeds the blob"); } + const ByteSpan header = + blob.subspan(kLengthPrefixSize, static_cast(header_len)); + return open_header( + header, + blob.size() - kLengthPrefixSize - static_cast(header_len)); +} +SafeTensorsReader SafeTensorsReader::open_header( + ByteSpan header_bytes, + size_t data_size) { const std::string_view header_text( - reinterpret_cast(blob.data() + kHeaderLenSize), - static_cast(header_len)); + reinterpret_cast(header_bytes.data()), header_bytes.size()); Json header; try { header = Json::parse(header_text); @@ -150,7 +160,6 @@ SafeTensorsReader SafeTensorsReader::open(ByteSpan blob) { } SafeTensorsReader out; - out.data_ = blob.subspan(kHeaderLenSize + static_cast(header_len)); for (auto member = header.begin(); member != header.end(); ++member) { const std::string& name = member.key(); @@ -188,7 +197,8 @@ SafeTensorsReader SafeTensorsReader::open(ByteSpan blob) { } const uint64_t begin = range[0].get(); const uint64_t end = range[1].get(); - if (begin > end || end > out.data_.size()) { + if (begin > end || end > data_size || + end > std::numeric_limits::max()) { throw std::runtime_error( "safetensors: entry '" + name + "' byte range is outside the data section"); @@ -227,10 +237,6 @@ const TensorEntry* SafeTensorsReader::find(const std::string& name) const { return it == entries_.end() ? nullptr : &it->second; } -ByteSpan SafeTensorsReader::bytes(const TensorEntry& entry) const { - return data_.subspan(entry.offset, entry.nbytes); -} - size_t SafeTensorsReader::total_bytes() const { return std::accumulate( entries_.begin(), diff --git a/backends/native/runtime/deserialize/SafeTensorsReader.h b/backends/native/runtime/deserialize/SafeTensorsReader.h index 6c0ca7b44fc..d54026ef805 100644 --- a/backends/native/runtime/deserialize/SafeTensorsReader.h +++ b/backends/native/runtime/deserialize/SafeTensorsReader.h @@ -6,6 +6,7 @@ #pragma once +#include #include #include #include @@ -34,29 +35,30 @@ struct TensorEntry { // reader exposes metadata semantics. // // Tensor payloads are packed with no per-tensor padding, so an entry's absolute -// alignment within the file is arbitrary: copy through these spans rather than -// handing them to an API that requires alignment. +// alignment within the file is arbitrary. Consumers must copy bytes into any +// destination that requires stronger alignment. class SafeTensorsReader { private: - // The data section only, i.e. the blob past its header. - ByteSpan data_; std::unordered_map entries_; std::vector names_; public: + static constexpr size_t kLengthPrefixSize = 8; + // Parse `blob`'s index. Throws std::runtime_error if the blob is truncated, // the header is not a JSON object, a dtype has no ScalarType, or a byte range // is inconsistent with its dtype and shape. - // - // Borrows `blob`, which must outlive both this reader and any span from it. static SafeTensorsReader open(ByteSpan blob); + // Decode the format's little-endian length prefix. + static size_t header_size(ByteSpan prefix); + + // Parse a header without loading its data section. + static SafeTensorsReader open_header(ByteSpan header, size_t data_size); + // Entry for `name`, or nullptr when absent. const TensorEntry* find(const std::string& name) const; - // Payload of an entry obtained from this reader. - ByteSpan bytes(const TensorEntry& entry) const; - // Tensor names, in header order, excluding "__metadata__". const std::vector& names() const { return names_; diff --git a/backends/native/runtime/deserialize/ZipReader.cpp b/backends/native/runtime/deserialize/ZipReader.cpp index c38845b040f..2efb11484cf 100644 --- a/backends/native/runtime/deserialize/ZipReader.cpp +++ b/backends/native/runtime/deserialize/ZipReader.cpp @@ -6,6 +6,8 @@ #include +#include +#include #include #include #include @@ -210,4 +212,42 @@ void ZipReader::read_entry_into( } } +void ZipReader::verify(std::string_view name) const { + const Entry* entry = find_entry(name); + if (entry == nullptr) { + throw std::runtime_error("zip: no member named " + std::string(name)); + } + + ZipFileHandle file( + zip_fopen_index(impl_->archive.get(), entry->index, ZIP_FL_UNCHANGED)); + if (file == nullptr) { + throw_zip(impl_->archive.get(), "cannot open member " + std::string(name)); + } + + std::array buffer{}; + size_t read = 0; + while (read < entry->size) { + const size_t request = std::min(buffer.size(), entry->size - read); + const zip_int64_t count = zip_fread(file.get(), buffer.data(), request); + if (count < 0) { + throw_zip_file(file.get(), "cannot verify member " + std::string(name)); + } + if (count == 0) { + throw std::runtime_error( + "zip: unexpected end of member " + std::string(name)); + } + read += static_cast(count); + } + + uint8_t extra = 0; + const zip_int64_t count = zip_fread(file.get(), &extra, 1); + if (count < 0) { + throw_zip_file(file.get(), "cannot verify member " + std::string(name)); + } + if (count != 0) { + throw std::runtime_error( + "zip: member is larger than its metadata: " + std::string(name)); + } +} + } // namespace ptn diff --git a/backends/native/runtime/deserialize/ZipReader.h b/backends/native/runtime/deserialize/ZipReader.h index ea312ffab38..0abc2c012b7 100644 --- a/backends/native/runtime/deserialize/ZipReader.h +++ b/backends/native/runtime/deserialize/ZipReader.h @@ -76,6 +76,9 @@ class ZipReader { size_t offset, MutableByteSpan destination) const; + // Reads a complete member and verifies its checksum without retaining it. + void verify(std::string_view name) const; + const std::vector& names() const { return names_; } diff --git a/backends/native/runtime/deserialize/targets.bzl b/backends/native/runtime/deserialize/targets.bzl index f8028e6abf8..80bb83872ab 100644 --- a/backends/native/runtime/deserialize/targets.bzl +++ b/backends/native/runtime/deserialize/targets.bzl @@ -10,7 +10,7 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) - # Owning byte buffer backed by heap storage or a read-only file mapping. + # Owning byte buffer behind a package: heap read or read-only mmap. runtime.cxx_library( name = "owned_bytes", srcs = ["OwnedBytes.cpp"], @@ -28,6 +28,8 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) + # Read-only reader for stored (uncompressed) zip archives, which is what a .ptn + # package is. runtime.cxx_library( name = "zip_reader", srcs = ["ZipReader.cpp"], @@ -49,12 +51,14 @@ def define_common_targets(): deps = [":json"], visibility = ["//executorch/backends/native/..."], ) + # The .ptn package: program flatbuffer plus its constants. runtime.cxx_library( name = "package", srcs = ["Package.cpp"], exported_headers = ["Package.h"], exported_deps = [ ":byte_span", + ":owned_bytes", ":safetensors_reader", ":zip_reader", "//executorch/backends/native/runtime/graph:scalar_type", diff --git a/backends/native/test/runtime/deserialize/targets.bzl b/backends/native/test/runtime/deserialize/targets.bzl index 333cce5018e..8a4c80f5c9d 100644 --- a/backends/native/test/runtime/deserialize/targets.bzl +++ b/backends/native/test/runtime/deserialize/targets.bzl @@ -9,6 +9,15 @@ def define_common_targets(): ], ) + runtime.cxx_test( + name = "package_test", + srcs = ["test_package.cpp"], + deps = [ + "//executorch/backends/native/runtime/deserialize:package", + "fbsource//third-party/libzip:zip", + ], + ) + runtime.cxx_test( name = "safetensors_reader_test", srcs = ["test_safetensors_reader.cpp"], diff --git a/backends/native/test/runtime/deserialize/test_package.cpp b/backends/native/test/runtime/deserialize/test_package.cpp new file mode 100644 index 00000000000..0d215409e5d --- /dev/null +++ b/backends/native/test/runtime/deserialize/test_package.cpp @@ -0,0 +1,159 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include + +namespace ptn { +namespace { + +std::vector make_safetensors( + std::string_view header, + std::string_view data) { + const uint64_t header_size = header.size(); + std::vector bytes(sizeof(header_size) + header.size() + data.size()); + std::memcpy(bytes.data(), &header_size, sizeof(header_size)); + std::memcpy(bytes.data() + sizeof(header_size), header.data(), header.size()); + std::memcpy( + bytes.data() + sizeof(header_size) + header.size(), + data.data(), + data.size()); + return bytes; +} + +ByteSpan as_bytes(std::string_view value) { + return ByteSpan(reinterpret_cast(value.data()), value.size()); +} + +struct ZipDiscard { + void operator()(zip_t* archive) const noexcept { + zip_discard(archive); + } +}; + +class TempPackage { + private: + const std::string program_ = "program"; + const std::string aliases_ = R"({"tied_weight":"weight"})"; + std::filesystem::path path_; + + public: + TempPackage() { + path_ = std::filesystem::temp_directory_path() / + ("ptn_package_" + std::to_string(reinterpret_cast(this)) + + ".ptn"); + int error = 0; + std::unique_ptr archive( + zip_open(path_.string().c_str(), ZIP_CREATE | ZIP_TRUNCATE, &error)); + if (archive == nullptr) { + throw std::runtime_error("failed to create test package"); + } + const std::vector tensors = make_safetensors( + R"({"weight":{"dtype":"U8","shape":[4],"data_offsets":[0,4]},"bias":{"dtype":"I16","shape":[1],"data_offsets":[4,6]}})", + "dataxy"); + add(archive.get(), kProgramEntry, as_bytes(program_)); + add(archive.get(), kSafeTensorsEntry, ByteSpan(tensors)); + add(archive.get(), kAliasesEntry, as_bytes(aliases_)); + if (zip_close(archive.get()) != 0) { + throw std::runtime_error("failed to close test package"); + } + archive.release(); + } + + ~TempPackage() { + std::error_code error; + std::filesystem::remove(path_, error); + } + + std::string path() const { + return path_.string(); + } + + private: + static void add(zip_t* archive, const char* name, ByteSpan bytes) { + zip_source_t* source = + zip_source_buffer(archive, bytes.data(), bytes.size(), 0); + if (source == nullptr) { + throw std::runtime_error("failed to create test package source"); + } + const zip_int64_t index = + zip_file_add(archive, name, source, ZIP_FL_ENC_UTF_8); + if (index < 0) { + zip_source_free(source); + throw std::runtime_error("failed to add test package member"); + } + if (zip_set_file_compression(archive, index, ZIP_CM_STORE, 0) != 0) { + throw std::runtime_error("failed to store test package member"); + } + } +}; + +// cppcheck-suppress-begin syntaxError +TEST(PackageTest, LoadsConstantsOnDemand) { + const TempPackage file; + const Package package = Package::load(file.path()); + + EXPECT_EQ(package.owner_keys(), (std::vector{"weight", "bias"})); + EXPECT_EQ(package.constant_bytes(), 6); + + const std::optional info = package.constant_info("tied_weight"); + ASSERT_TRUE(info); + EXPECT_NE(info->package_id, 0); + EXPECT_EQ(info->dtype, kByte); + EXPECT_EQ(*info->sizes, (std::vector{4})); + EXPECT_EQ(info->nbytes, 4); + EXPECT_EQ(info->owner, "weight"); + + std::array destination{}; + EXPECT_TRUE(package.load_constant_into("weight", destination)); + EXPECT_EQ(destination, (std::array{'d', 'a', 't', 'a'})); + + const std::optional acquired = + package.acquire_constant("tied_weight"); + ASSERT_TRUE(acquired); + EXPECT_TRUE(std::ranges::equal(acquired->span(), destination)); + EXPECT_NO_THROW(package.verify_constants()); +} + +TEST(PackageTest, ReportsMissingConstantsAndWrongDestinations) { + const TempPackage file; + const Package package = Package::load(file.path()); + + EXPECT_EQ(package.constant_info("missing"), std::nullopt); + EXPECT_EQ(package.acquire_constant("missing"), std::nullopt); + std::array destination{}; + EXPECT_FALSE(package.load_constant_into("missing", destination)); + EXPECT_THROW( + package.load_constant_into("weight", destination), std::runtime_error); +} + +TEST(PackageTest, SupportsCallerOwnedArchiveBytes) { + const TempPackage file; + Package package = Package::load(file.path()); + package = Package::load(OwnedBytes::from_file(file.path(), false)); + + const std::optional weight = package.acquire_constant("weight"); + ASSERT_TRUE(weight); + EXPECT_TRUE(std::ranges::equal( + weight->span(), (std::array{'d', 'a', 't', 'a'}))); +} +// cppcheck-suppress-end syntaxError + +} // namespace +} // namespace ptn diff --git a/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp b/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp index 02a5dfe0965..e404017b40e 100644 --- a/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp +++ b/backends/native/test/runtime/deserialize/test_safetensors_reader.cpp @@ -45,7 +45,6 @@ TEST(SafeTensorsReaderTest, ReadsIndexAndPayloadsInHeaderOrder) { EXPECT_EQ(reader.find("first")->sizes, (std::vector{1})); EXPECT_EQ(reader.find("first")->offset, 0); EXPECT_EQ(reader.find("first")->nbytes, 4); - EXPECT_EQ(reader.bytes(*reader.find("second"))[0], 'X'); EXPECT_EQ(reader.total_bytes(), 6); EXPECT_EQ(reader.find("missing"), nullptr); } @@ -57,6 +56,17 @@ TEST(SafeTensorsReaderTest, RejectsNonEmptyMetadata) { std::runtime_error); } +TEST(SafeTensorsReaderTest, ReadsHeaderWithoutPayload) { + const std::string header = + R"({"x":{"dtype":"U8","shape":[2],"data_offsets":[0,2]}})"; + const SafeTensorsReader reader = SafeTensorsReader::open_header( + ByteSpan(reinterpret_cast(header.data()), header.size()), + 2); + + ASSERT_NE(reader.find("x"), nullptr); + EXPECT_EQ(reader.find("x")->nbytes, 2); +} + TEST(SafeTensorsReaderTest, RejectsInvalidMetadata) { EXPECT_THROW( SafeTensorsReader::open(make_safetensors("[]", "")), std::runtime_error); diff --git a/backends/native/test/runtime/deserialize/test_zip_reader.cpp b/backends/native/test/runtime/deserialize/test_zip_reader.cpp index ce8362c6ea5..62cdaf4706c 100644 --- a/backends/native/test/runtime/deserialize/test_zip_reader.cpp +++ b/backends/native/test/runtime/deserialize/test_zip_reader.cpp @@ -94,6 +94,7 @@ TEST(ZipReaderTest, ReadsStoredMemberRanges) { EXPECT_EQ( zip.read("program.ptg"), (std::vector{'p', 'r', 'o', 'g', 'r', 'a', 'm'})); + EXPECT_NO_THROW(zip.verify("program.safetensors")); } TEST(ZipReaderTest, RejectsInvalidRanges) { From f07b2bea38a75697efd08d861b08c5878ed46f0f Mon Sep 17 00:00:00 2001 From: Stephen Jia Date: Thu, 10 Sep 2026 12:14:39 -0700 Subject: [PATCH 170/190] [executorch][native] Add the EngineContext / EngineExecutable interface ## Ulterior Motive Define backend-neutral execution boundaries before landing the first Vulkan implementation. ## Rationale **What**: Add pure abstract `EngineContext` and `EngineExecutable` interfaces. **Why**: Native runtime can load methods and weights but needs a stable contract for backend compilation and execution. ## Details ```text process scope compiled region scope EngineContext ----------------compile--> EngineExecutable | | +-- backend/device identity +-- input/output metadata +-- shared device state +-- set_input() +-- kernel/pipeline reuse +-- execute() +-- get_output() ``` `EngineContext` must outlive its executables. `compile()` currently accepts a whole `Method` for full delegation. Unsupported graphs fail with exceptions instead of a separate `can_run()` boolean. Polymorphic types are immovable and noncopyable; out-of-line destructors anchor vtables. Authored with Claude Code. Differential Revision: [D118480656](https://our.internmc.facebook.com/intern/diff/D118480656/) ghstack-source-id: 427858783 Pull-Request: https://github.com/pytorch/executorch/pull/22704 --- backends/native/runtime/engine/BUCK | 11 ++ backends/native/runtime/engine/Engine.cpp | 17 ++++ backends/native/runtime/engine/Engine.h | 111 +++++++++++++++++++++ backends/native/runtime/engine/targets.bzl | 19 ++++ 4 files changed, 158 insertions(+) create mode 100644 backends/native/runtime/engine/BUCK create mode 100644 backends/native/runtime/engine/Engine.cpp create mode 100644 backends/native/runtime/engine/Engine.h create mode 100644 backends/native/runtime/engine/targets.bzl diff --git a/backends/native/runtime/engine/BUCK b/backends/native/runtime/engine/BUCK new file mode 100644 index 00000000000..0ab35888218 --- /dev/null +++ b/backends/native/runtime/engine/BUCK @@ -0,0 +1,11 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. This file can contain cell-only targets. + +non_fbcode_target(_kind = define_common_targets) + +fbcode_target(_kind = define_common_targets) diff --git a/backends/native/runtime/engine/Engine.cpp b/backends/native/runtime/engine/Engine.cpp new file mode 100644 index 00000000000..53385469fc5 --- /dev/null +++ b/backends/native/runtime/engine/Engine.cpp @@ -0,0 +1,17 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include + +namespace ptn { + +// Out of line so each vtable is emitted here rather than in every translation +// unit that includes the header. +EngineExecutable::~EngineExecutable() = default; + +EngineContext::~EngineContext() = default; + +} // namespace ptn diff --git a/backends/native/runtime/engine/Engine.h b/backends/native/runtime/engine/Engine.h new file mode 100644 index 00000000000..7fc4843cd27 --- /dev/null +++ b/backends/native/runtime/engine/Engine.h @@ -0,0 +1,111 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace ptn { + +// One dependency-closed region of a Method, lowered onto a backend and ready to +// run: whatever the backend needed to compile it, plus the staging it reads +// inputs from and writes outputs to. +// +// Under full delegation -- the only mode today -- that region is the whole +// Method, and the executable's inputs and outputs are the Method's. Under +// runtime partitioning one Method yields several executables interleaved with +// other backends, and the inputs and outputs are region boundaries instead. +// +// Obtained from EngineContext::compile, never constructed directly. Not +// thread-safe and not re-entrant: one executable runs one call at a time. +// Concurrent inference means several executables (or, once the working-set +// split lands, several working sets over one compiled program). +class EngineExecutable { + protected: + EngineExecutable() = default; + + public: + EngineExecutable(const EngineExecutable&) = delete; + EngineExecutable& operator=(const EngineExecutable&) = delete; + EngineExecutable(EngineExecutable&&) = delete; + EngineExecutable& operator=(EngineExecutable&&) = delete; + virtual ~EngineExecutable(); + + // Counts and shapes of the compiled Method's user inputs / outputs, in graph + // order. Sizes are static upper bounds, so a dynamic dim reports its maximum. + virtual size_t num_inputs() const = 0; + virtual size_t num_outputs() const = 0; + virtual std::vector input_sizes(size_t i) const = 0; + virtual std::vector output_sizes(size_t i) const = 0; + virtual ScalarType input_dtype(size_t i) const = 0; + virtual ScalarType output_dtype(size_t i) const = 0; + + // Copy `numel` elements from host `data` into input i, converting from + // `src_dtype` to the input's dtype when they differ. `numel` must equal the + // input's element count. + virtual void + set_input(size_t i, const void* data, size_t numel, ScalarType src_dtype) = 0; + + // Run the compiled Method. Blocks until every output is readable, so a + // get_output right after it needs no further synchronization. + virtual void execute() = 0; + + // Copy output i back into host `data`, converting to `dst_dtype` from the + // output's dtype when they differ. Only meaningful after an execute(). + virtual void + get_output(size_t i, void* data, size_t numel, ScalarType dst_dtype) = 0; +}; + +// A compute backend, at process scope: the device context and kernel registry +// that every Method run on that device shares. One per device, constructed +// through the backend's own factory (e.g. make_vulkan_engine()) since selecting +// a backend is the caller's decision, not this interface's. +// +// Must outlive every executable it compiled. +class EngineContext { + protected: + EngineContext() = default; + + public: + EngineContext(const EngineContext&) = delete; + EngineContext& operator=(const EngineContext&) = delete; + EngineContext(EngineContext&&) = delete; + EngineContext& operator=(EngineContext&&) = delete; + virtual ~EngineContext(); + + // Backend identity ("vulkan"), and the device it selected ("SwiftShader + // Device"). Diagnostics only; nothing dispatches on either. + virtual const std::string& name() const = 0; + virtual const std::string& device_name() const = 0; + + // Lower `method` onto this backend and prepack the constants it binds, + // fetched from `package` by data_key. `method` must outlive the returned + // executable; `package` is needed only for this call. + // + // Compiles the method whole, which is full delegation -- the only mode today. + // Runtime partitioning narrows the unit to a region of a method and yields + // several executables per method; that arrives as an added entry point, not a + // change to this one. + // + // Throws std::runtime_error when the backend cannot run the method: an + // unsupported op or dtype, a binding whose constant the package does not + // hold, a constant whose byte count contradicts its TensorMeta, an unbounded + // dynamic dim, or a higher-order-op subgraph. A backend is free to reject + // anything else it cannot lower; there is no partial success. + virtual std::unique_ptr compile( + const Method& method, + const Package& package) = 0; +}; + +} // namespace ptn diff --git a/backends/native/runtime/engine/targets.bzl b/backends/native/runtime/engine/targets.bzl new file mode 100644 index 00000000000..f94d3f83d7e --- /dev/null +++ b/backends/native/runtime/engine/targets.bzl @@ -0,0 +1,19 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + # The runtime <-> compute-backend boundary: the abstract EngineContext + # (process-wide device state) and EngineExecutable (one lowered, ready-to-run + # Method). Pure std; each backend implements both in its own package. + runtime.cxx_library( + name = "engine", + srcs = ["Engine.cpp"], + exported_headers = [ + "Engine.h", + ], + exported_deps = [ + "//executorch/backends/native/runtime:method", + "//executorch/backends/native/runtime/deserialize:package", + "//executorch/backends/native/runtime/graph:scalar_type", + ], + visibility = ["//executorch/backends/native/..."], + ) From 87dab0bc3fc33ea3923c11972a637a7dd310c33b Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Fri, 11 Sep 2026 03:08:43 +0200 Subject: [PATCH 171/190] [ET-VK] Wire the missing resize functions for embedding and index_select (#22410) ### Summary `add_embedding_legacy_node()` and both `index_select` node builders pass `nullptr` as their resizing logic: ```cpp // Resize Args {}, // Resizing Logic nullptr)); ``` even though `resize_embedding_node`, `resize_index_select_channel_node` and the local `resize_fn` are already defined immediately above them and do the right thing. Under dynamic shapes the outputs keep the extents they were built with (the upper bound) instead of tracking the real input sizes, so consumers read them at the wrong size. `register_index_select()` also does not set `supports_resize`, so this patch sets it now that the op honours resize. ### Fix Pass the resize functions that already exist. Three one-line changes plus the registry flag. ### How it surfaced A TTS model on Adreno 840 produced wrong outputs after an unrelated change moved a partition boundary. The embedding output had been keeping its upper-bound extents all along; previously a CPU round-trip happened to re-establish the correct shape downstream, so the defect was masked. Once the gather stayed on the GPU the stale extents reached the consumer. Verified on a Galaxy S26 Ultra: with these wired up, the affected sub-models return correct shapes and match their CPU references (cosine >= 0.999 at sequence lengths well below the dynamic bound) where before they were wrong at every length except the bound itself. Note this class of bug is hard to see with `executor_runner`, which can only run at the dynamic upper bound -- exactly the one shape where a missing resize is harmless. --- backends/vulkan/op_registry.py | 1 + .../runtime/graph/ops/glsl/index_select.glsl | 10 +- .../runtime/graph/ops/impl/Embedding.cpp | 2 +- .../runtime/graph/ops/impl/IndexSelect.cpp | 74 +++++++++--- backends/vulkan/test/test_vulkan_delegate.py | 106 ++++++++++++++++++ 5 files changed, 175 insertions(+), 18 deletions(-) diff --git a/backends/vulkan/op_registry.py b/backends/vulkan/op_registry.py index d04d91f32e9..add4d01a78e 100644 --- a/backends/vulkan/op_registry.py +++ b/backends/vulkan/op_registry.py @@ -1412,6 +1412,7 @@ def register_index_select(): return OpFeatures( inputs_storage=utils.CHANNELS_PACKED_TEXTURE, inputs_dtypes=utils.FP_INT_BOOL_T, + supports_resize=True, ) diff --git a/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl b/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl index 4500d43b932..8b44daa9621 100644 --- a/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl +++ b/backends/vulkan/runtime/graph/ops/glsl/index_select.glsl @@ -20,7 +20,8 @@ ${layout_declare_tensor(0, "w", "t_out", DTYPE, STORAGE)} ${layout_declare_tensor(1, "r", "t_in", DTYPE, STORAGE)} ${layout_declare_tensor(2, "r", "t_idx", "int", STORAGE)} ${layout_declare_ubo(3, "ivec4", "sizes")} -${layout_declare_ubo(4, "int", "gpu_dim", "int", "stride")} +${layout_declare_ubo(4, "ivec4", "in_sizes")} +${layout_declare_ubo(5, "int", "gpu_dim")} layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in; @@ -33,6 +34,13 @@ void main() { return; } + // Selecting along the batch dim steps over the z axis in units of channel + // texels, because the batch and channel dims share that axis. The width and + // height dims each have an axis to themselves, so they step by one. The + // stride is read from in_sizes rather than baked in at build time so that it + // follows a resize that changes the channel count. + const int stride = gpu_dim == 2 ? ((in_sizes.z + 3) / 4) : 1; + const int out_idx = out_pos[gpu_dim] / stride; const int within_stride = out_pos[gpu_dim] % stride; const int in_idx = texelFetch(t_idx, ivec3(out_idx, 0, 0), 0).x; diff --git a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp index 61ba9349b45..80f7b505feb 100644 --- a/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/Embedding.cpp @@ -112,7 +112,7 @@ void add_embedding_legacy_node( // Resize Args {}, // Resizing Logic - nullptr)); + resize_embedding_node)); } void embedding(ComputeGraph& graph, const std::vector& args) { diff --git a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp index 20ab76813a7..224cc427451 100644 --- a/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp +++ b/backends/vulkan/runtime/graph/ops/impl/IndexSelect.cpp @@ -28,6 +28,35 @@ void check_index_select_args( VK_CHECK_COND(graph.packed_dim_of(out) == WHCN::kChannelsDim); } +// index_select replaces the selected dim with as many entries as the index +// tensor holds and leaves every other dim alone. +std::vector index_select_out_sizes( + ComputeGraph* graph, + const ValueRef in, + const ValueRef idx, + const DimIndex dim_idx) { + std::vector out_sizes = graph->sizes_of(in); + const int64_t ndim = static_cast(out_sizes.size()); + // dim_idx is a negative index counted from the innermost dim. + const int64_t dim = ndim + dim_idx; + VK_CHECK_COND(dim >= 0 && dim < ndim); + out_sizes.at(dim) = graph->numel_of(idx); + return out_sizes; +} + +void resize_index_select_channel_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + (void)resize_args; + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const ValueRef idx = args.at(1).refs.at(1); + + graph->virtual_resize( + out, index_select_out_sizes(graph, in, idx, kChannel4D)); +} + void add_index_select_channel_node( ComputeGraph& graph, ValueRef in, @@ -53,32 +82,43 @@ void add_index_select_channel_node( // Resize Args {}, // Resizing Logic - nullptr)); + resize_index_select_channel_node)); } struct IndexSelectParams final { int32_t gpu_dim; - int32_t stride; }; -IndexSelectParams create_index_select_params( - ComputeGraph& graph, - const int64_t dim_idx, - const ValueRef in) { +IndexSelectParams create_index_select_params(const int64_t dim_idx) { if (dim_idx == kWidth4D) { - return {0, 1}; + return {0}; } else if (dim_idx == kHeight4D) { - return {1, 1}; + return {1}; } else if (dim_idx == kBatch4D) { - const std::vector in_sizes = graph.sizes_of(in); - int64_t n_channels = dim_at(in_sizes, kChannel4D); - int64_t stride = utils::div_up_4(n_channels); - return {2, static_cast(stride)}; + // The batch axis shares the z axis with the channels, so the shader steps + // over one batch in units of channel texels. That stride is derived from + // the channel count, which a resize can change, so the shader reads it out + // of in_sizes rather than taking a value frozen at build time. + return {2}; } else { VK_THROW("Unexpected dim_idx!"); } } +void resize_index_select_node( + ComputeGraph* graph, + const std::vector& args, + const std::vector& resize_args) { + const ValueRef out = args.at(0).refs.at(0); + const ValueRef in = args.at(1).refs.at(0); + const ValueRef idx = args.at(1).refs.at(1); + + const DimIndex dim_idx = + static_cast(graph->extract_scalar(resize_args.at(0))); + + graph->virtual_resize(out, index_select_out_sizes(graph, in, idx, dim_idx)); +} + void add_index_select_node( ComputeGraph& graph, ValueRef in, @@ -87,7 +127,7 @@ void add_index_select_node( ValueRef out) { check_index_select_args(graph, in, idx, out); - IndexSelectParams params = create_index_select_params(graph, dim_idx, in); + IndexSelectParams params = create_index_select_params(dim_idx); std::string kernel_name = "index_select"; kernel_name.reserve(kShaderNameReserve); @@ -99,15 +139,17 @@ void add_index_select_node( default_pick_gwg, default_pick_lwg, {{out, vkapi::kWrite}, {{in, idx}, vkapi::kRead}}, - {graph.sizes_ubo(out), graph.create_params_buffer(params)}, + {graph.sizes_ubo(out), + graph.sizes_ubo(in), + graph.create_params_buffer(params)}, // Push Constants {}, // Specialization Constants {}, // Resize Args - {}, + {graph.get_or_add_value_for_int(dim_idx)}, // Resizing Logic - nullptr)); + resize_index_select_node)); } int64_t get_dim_idx(ComputeGraph& graph, ValueRef in, ValueRef dim_ref) { diff --git a/backends/vulkan/test/test_vulkan_delegate.py b/backends/vulkan/test/test_vulkan_delegate.py index 84d9c88fbfc..05c5084718b 100644 --- a/backends/vulkan/test/test_vulkan_delegate.py +++ b/backends/vulkan/test/test_vulkan_delegate.py @@ -1680,6 +1680,112 @@ def forward(self, x): sample_inputs, ) + def test_vulkan_backend_index_select_batch_dynamic_channels(self): + # Selecting along the batch dim walks the z axis in units of channel + # texels, so the step depends on the channel count. Vary the channels + # below the built size: a step frozen at build time reads the wrong + # texel once the count drops. + class IndexSelectModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.index = torch.tensor([1, 3, 0, 2]) + + def forward(self, x): + return torch.index_select(x, 0, self.index) + + sample_inputs = (torch.randn(size=(5, 8, 3, 4), dtype=torch.float32),) + dynamic_shapes = {"x": {1: Dim("channels", min=1, max=8)}} + test_inputs = [ + (torch.randn(5, 1, 3, 4),), + (torch.randn(5, 3, 3, 4),), + (torch.randn(5, 4, 3, 4),), + (torch.randn(5, 5, 3, 4),), + (torch.randn(5, 8, 3, 4),), + ] + + self.lower_module_and_test_output( + IndexSelectModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_index_select_width_dynamic_shapes(self): + class IndexSelectModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.index = torch.tensor([2, 0, 1]) + + def forward(self, x): + return torch.index_select(x, 2, self.index) + + sample_inputs = (torch.randn(size=(2, 3, 4, 6), dtype=torch.float32),) + dynamic_shapes = {"x": {3: Dim("width", min=1, max=6)}} + test_inputs = [ + (torch.randn(2, 3, 4, 1),), + (torch.randn(2, 3, 4, 3),), + (torch.randn(2, 3, 4, 6),), + ] + + self.lower_module_and_test_output( + IndexSelectModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_index_select_channel_dynamic_shapes(self): + # The channel path takes a separate shader and a separate resize + # callback from the one above. + class IndexSelectModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.index = torch.tensor([3, 1, 0, 1]) + + def forward(self, x): + return torch.index_select(x, 1, self.index) + + sample_inputs = (torch.randn(size=(2, 5, 4, 6), dtype=torch.float32),) + dynamic_shapes = {"x": {3: Dim("width", min=1, max=6)}} + test_inputs = [ + (torch.randn(2, 5, 4, 1),), + (torch.randn(2, 5, 4, 4),), + (torch.randn(2, 5, 4, 6),), + ] + + self.lower_module_and_test_output( + IndexSelectModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + + def test_vulkan_backend_embedding_dynamic_shapes(self): + # The output picks up the index tensor's shape, so it has to be resized + # with it rather than left at the size it was built with. + class EmbeddingModule(torch.nn.Module): + def __init__(self): + super().__init__() + self.embedding = torch.nn.Embedding(10, 8) + + def forward(self, x): + return self.embedding(x) + + sample_inputs = (torch.randint(0, 10, (2, 6), dtype=torch.int32),) + dynamic_shapes = {"x": {1: Dim("seq", min=1, max=6)}} + test_inputs = [ + (torch.randint(0, 10, (2, 1), dtype=torch.int32),), + (torch.randint(0, 10, (2, 3), dtype=torch.int32),), + (torch.randint(0, 10, (2, 6), dtype=torch.int32),), + ] + + self.lower_module_and_test_output( + EmbeddingModule(), + sample_inputs, + dynamic_shapes=dynamic_shapes, + test_inputs=test_inputs, + ) + def test_vulkan_backend_index_tensor_nonzero_axis(self): class IndexTensorModule(torch.nn.Module): def __init__(self, dim): From 028f4c03e8d27e47df460ad857298b7a3e0df68a Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Fri, 11 Sep 2026 03:33:10 +0200 Subject: [PATCH 172/190] [ET-VK] Say that dynamic activation quantization is per tensor (#22432) ### Summary `get_symmetric_quantization_config` says "per-token input quantization" in its `is_dynamic` branch, but the spec it builds is `per_tensor_affine`. The spec is correct: the kernel behind `et_vk.linear_q8ta_q8csw` takes a single input scale and zero point (`QuantizedLinear.cpp`, `input_quant_config(8, kPerTensor, ...)`). Only the comment is wrong. It is worth saying accurately, because the difference is large on encoder-style models. Mean cosine of the embedding against the fp32 eager model on `all-mpnet-base-v2` / `multi-qa-mpnet-base-dot-v1`: | config | cosine | | --- | --- | | `is_dynamic=True, weight_bits=8` | 0.910 / 0.948 | | `is_dynamic=False, weight_bits=8` | 0.9993 / 0.9993 | | torchao per-token 8da8w, same models | 0.998 / 0.998 | Dynamic int8 is not the problem; the per-tensor scale is. A per-token variant for int8 weights would be the real fix, and #22431 tracks that. This PR only makes the current behaviour discoverable, and points at `is_dynamic=False` for callers that care about fidelity. Related to #22431 ### Test plan Comment and docstring only, no functional change. `black --check` and `flake8` clean on the touched file. --- backends/vulkan/quantizer/vulkan_quantizer.py | 19 +++++++++++++++++-- 1 file changed, 17 insertions(+), 2 deletions(-) diff --git a/backends/vulkan/quantizer/vulkan_quantizer.py b/backends/vulkan/quantizer/vulkan_quantizer.py index 3d1e1eab0f2..c7cc80022f8 100644 --- a/backends/vulkan/quantizer/vulkan_quantizer.py +++ b/backends/vulkan/quantizer/vulkan_quantizer.py @@ -47,7 +47,9 @@ def get_symmetric_quantization_config( Return a QuantizationConfig for Vulkan quantizer. Args: - is_dynamic: If False, weight-only quantization. If True, dynamic quantization (activation + weight) + is_dynamic: If False, weight-only quantization. If True, dynamic + quantization (activation + weight), with the activation scale + computed per tensor at runtime weight_bits: Number of bits for weight quantization (4 or 8) act_bits: Number of bits for activation quantization (8) act_qmin: Minimum quantization value for activations (auto-calculated if None) @@ -87,7 +89,20 @@ def get_symmetric_quantization_config( act_quantization_spec = None output_activation_spec = None else: - # Dynamic quantization: per-token input quantization, no output quantization + # Dynamic quantization: a choose_qparams op computes one scale and + # zero point for the whole activation tensor at runtime, and the + # quantize/dequantize pair around the linear carries them. Per tensor, + # not per token, whatever the granularity of the surrounding graph. + # + # (The fused et_vk.linear_q8ta_q8csw kernel is a different path: it is + # matched when the input scale is a static scalar, not one chosen at + # runtime.) + # + # One scale for the whole tensor is a poor fit for transformer + # encoders, where a few outlier channels set it for everything else; on + # sentence-transformer models this costs an order of magnitude more + # accuracy than a per-token scheme. Prefer is_dynamic=False (weight + # only) when output fidelity matters. # Auto-calculate activation ranges if not provided if act_qmin is None or act_qmax is None: act_range = bits_to_range(act_bits) From 214726afce363c91174998d98469808a226e13d0 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Thu, 10 Sep 2026 20:24:09 -0700 Subject: [PATCH 173/190] Update main version to 1.6.0 (#22613) ### Summary Update the development version on `main` from `1.5.0` to `1.6.0` after cutting the `release/1.5` branch. ### Test plan - verified `version.txt` contains `1.6.0` - `git diff --check` --- version.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/version.txt b/version.txt index bc80560fad6..dc1e644a101 100644 --- a/version.txt +++ b/version.txt @@ -1 +1 @@ -1.5.0 +1.6.0 From 756cc9cc854ccd000be9c193c64248090f67e14d Mon Sep 17 00:00:00 2001 From: Andrew Grebenisan <33402477+DrJessop@users.noreply.github.com> Date: Thu, 10 Sep 2026 21:12:24 -0700 Subject: [PATCH 174/190] Support misaligned ranks for broadcasted binary ops in propagate slice (#22661) Summary: Before, we didn't support this case (ex shape(4, 4, 2) + shape(4, 3)) Now we do. Reviewed By: ethansfng Differential Revision: D119425194 Pull Request resolved: https://github.com/pytorch/executorch/pull/22661 --- backends/cadence/aot/reorder_ops.py | 17 ++++++-- .../aot/tests/test_reorder_ops_passes.py | 41 +++++++++++++++++++ 2 files changed, 55 insertions(+), 3 deletions(-) diff --git a/backends/cadence/aot/reorder_ops.py b/backends/cadence/aot/reorder_ops.py index 3df555e0fba..b105988a251 100644 --- a/backends/cadence/aot/reorder_ops.py +++ b/backends/cadence/aot/reorder_ops.py @@ -1297,17 +1297,28 @@ def _swap_binary_elementwise_slice( slice_step = get_arg(slice_node, "step", int) output_shape = op_node.meta["val"].shape + output_dim = slice_dim % len(output_shape) new_args = list(op_node.args) with graph.inserting_before(op_node): for i, inp in enumerate([lhs, rhs]): - if inp.meta["val"].shape[slice_dim] == output_shape[slice_dim]: + input_shape = inp.meta["val"].shape + # Broadcasting aligns operand dimensions to the right of the output. + input_dim = output_dim - (len(output_shape) - len(input_shape)) + if ( + input_dim >= 0 + and input_shape[input_dim] == output_shape[output_dim] + ): new_slice = graph.call_function( exir_ops.edge.aten.slice_copy.Tensor, - args=(inp, slice_dim, slice_start, slice_end, slice_step), + args=(inp, input_dim, slice_start, slice_end, slice_step), ) new_slice.meta["val"] = exir_ops.edge.aten.slice_copy.Tensor( - inp.meta["val"], slice_dim, slice_start, slice_end, slice_step + inp.meta["val"], + input_dim, + slice_start, + slice_end, + slice_step, ) new_args[i] = new_slice diff --git a/backends/cadence/aot/tests/test_reorder_ops_passes.py b/backends/cadence/aot/tests/test_reorder_ops_passes.py index 3c4443aafa1..f0083ff8a35 100644 --- a/backends/cadence/aot/tests/test_reorder_ops_passes.py +++ b/backends/cadence/aot/tests/test_reorder_ops_passes.py @@ -1425,6 +1425,47 @@ def test_swap_additional_binary_target(self) -> None: self.assertIs(sub_nodes[0].args[1], slice_nodes[0]) self.assertEqual(list(sub_nodes[0].meta["val"].shape), [2, 60, 1, 1]) + def test_swap_additional_binary_target_with_mismatched_ranks(self) -> None: + lhs_data = torch.randn(2, 3, 4) + rhs_data = torch.randn(3, 4) + builder = GraphBuilder() + lhs = builder.placeholder("lhs", lhs_data) + rhs = builder.placeholder("rhs", rhs_data) + sub = builder.call_operator( + exir_ops.edge.aten.sub.Tensor, + args=(lhs, rhs), + ) + sliced = builder.call_operator( + exir_ops.edge.aten.slice_copy.Tensor, + args=(sub, 1, 0, 2, 1), + ) + builder.output([sliced]) + gm = builder.get_graph_module() + + result = transform_and_check_numerics( + gm, + (lhs_data, rhs_data), + PropagateSlice(additional_binary_targets=[exir_ops.edge.aten.sub.Tensor]), + ) + + self.assertTrue(result.modified) + slice_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor + ) + self.assertEqual(len(slice_nodes), 2) + lhs_slice, rhs_slice = slice_nodes + self.assertIs(lhs_slice.args[0], lhs.node) + self.assertEqual(lhs_slice.args[1], 1) + self.assertEqual(list(lhs_slice.meta["val"].shape), [2, 2, 4]) + self.assertIs(rhs_slice.args[0], rhs.node) + self.assertEqual(rhs_slice.args[1], 0) + self.assertEqual(list(rhs_slice.meta["val"].shape), [2, 4]) + sub_nodes = gm.graph.find_nodes( + op="call_function", target=exir_ops.edge.aten.sub.Tensor + ) + self.assertEqual(len(sub_nodes), 1) + self.assertEqual(list(sub_nodes[0].meta["val"].shape), [2, 2, 4]) + def test_swap_broadcast_mul_slice_on_broadcast_dim(self) -> None: """[1,60,1,1] * [4,1,1,1] → [4,60,1,1] → slice(dim=0, step=2) Only the [4,1,1,1] input should be sliced.""" From dab68f8450534085a8b53e6b16bca7fa351f4ebf Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Fri, 11 Sep 2026 00:00:01 -0700 Subject: [PATCH 175/190] Declare the system frameworks the executorch SwiftPM product needs (#22711) Part of #22686 ### Problem The `executorch` SwiftPM product declared only the C++ standard library. The image processor inside it calls into Accelerate and Core Image, so a package that depends on `executorch` alone and links with `-all_load` fails to link. I reproduced it against the published prebuilt frameworks: 16 missing symbols, 8 from vImage, which lives in Accelerate, and 8 from Core Image. On iOS the Core Image half is usually hidden, because an app that imports UIKit, SwiftUI or AVFoundation already pulls Core Image in. On macOS, and on an iOS app that imports only Foundation, both halves fail. Until now each app had to name these frameworks itself, which is a detail the package should own. ### Change Declare Accelerate, Core Graphics, Core Image, Core Video and Foundation on the `executorch` product. All five are used by the image processor. For most consumers Accelerate and Core Image are the two that fix the link, because Core Graphics, Core Video and Foundation already arrive as link hints, either from the Swift files in the framework or from the consumer's own code when it is built with modules enabled, which is the default. Turn modules off in an Objective-C consumer and those hints disappear, and then Core Graphics and Core Video are needed explicitly too. Declaring all five covers that case as well and lets the product state what it actually uses. ### Note for reviewers The manifest that prebuilt-binary users get is generated from `Package.swift.template` on the `swiftpm` branch, not from this file, and it carries the same `executorch` entry with only `c++`. This change alone therefore does not reach anyone using the prebuilt frameworks. I will send the matching change against that branch as a follow up. That is why this says "Part of" rather than "Fixes": the issue should stay open until the template change lands too. ### Test plan - `swift package --disable-sandbox dump-package` parses, and both the release and debug `executorch` targets carry the five frameworks plus libc++. - Built a package that depends only on the `executorch` product with `-all_load` against the prebuilt frameworks. It links and runs with this change. Removing just the added block brings back the 16 missing symbols, so the check fails without the fix. - The Apple builds on this PR pass. The demo app jobs do not run here, because they need secrets a fork pull request cannot use, so nothing in CI currently links a real app against this product. A test that would catch this needs a consumer that links `executorch` and nothing else. Co-authored-by: PyTorch Bot --- Package.swift | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/Package.swift b/Package.swift index e1fac90ad93..356828aaa95 100644 --- a/Package.swift +++ b/Package.swift @@ -66,6 +66,13 @@ let products = deliverables([ ], ], "executorch": [ + "frameworks": [ + "Accelerate", + "CoreGraphics", + "CoreImage", + "CoreVideo", + "Foundation", + ], "libraries": [ "c++", ], From adb2ef543fcb558f9bca45bd62fff466c40e43e0 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Fri, 11 Sep 2026 09:16:07 +0100 Subject: [PATCH 176/190] Arm backend: Deduplicate constants emitted during TOSA lowering (#22638) Pool backend-generated helper constants per TOSA basic block using dtype, shape, and exact serialized bytes as the key. Implement the pool in a TosaSerializer subclass. Helper operands use the tensor returned by addConst, making reuse safe. Keep graph-owned constants unpooled because later lowering refers to their FX names. This covers model parameters, buffers, lifted constants, and CONST_SHAPE outputs, while preserving packed FP4 serialization. Deduplicate identical static CONST_SHAPE nodes in FX, where their uses can safely be rewritten to the first occurrence. Representative TOSA-FP operator comparisons: | Model | Before | After | Reduction | Bytes saved | |-------------|-------:|------:|------------:|------------:| | SD3 | 1,518 | 1,038 | 480 (31.6%) | 89,068 | | InceptionV3 | 673 | 456 | 217 (32.2%) | 42,092 | | Conformer | 488 | 365 | 123 (25.2%) | 22,852 | This change was authored with assistance from Codex. Change-Id: I4de44b034aff33cbf2888801238cc48df5a9124d Signed-off-by: Yufeng Shi --- backends/arm/_passes/__init__.py | 1 + backends/arm/_passes/arm_pass_manager.py | 2 + .../_passes/deduplicate_const_shapes_pass.py | 47 ++++++ backends/arm/operators/BUCK | 1 + backends/arm/operators/op_tosa_matmul.py | 12 +- backends/arm/operators/op_tosa_mul.py | 4 +- backends/arm/operators/op_tosa_shapes.py | 7 +- backends/arm/operators/ops_quant_utils.py | 13 +- backends/arm/process_node.py | 12 +- backends/arm/test/misc/test_process_node.py | 7 +- .../arm/test/misc/test_tosa_constant_pool.py | 144 ++++++++++++++++++ .../test_tosa_shape_node_visitors.py | 19 ++- backends/arm/test/ops/test_conv2d.py | 3 +- backends/arm/test/ops/test_conv3d.py | 3 +- backends/arm/test/ops/test_depthwise_conv.py | 3 +- .../test_deduplicate_const_shapes_pass.py | 58 +++++++ .../passes/test_fuse_duplicate_users_pass.py | 4 +- .../passes/test_remove_data_layout_noops.py | 4 +- backends/arm/test/targets.bzl | 1 + backends/arm/tosa/BUCK | 11 ++ backends/arm/tosa/backend.py | 10 +- backends/arm/tosa/constant_pool.py | 59 +++++++ 22 files changed, 395 insertions(+), 30 deletions(-) create mode 100644 backends/arm/_passes/deduplicate_const_shapes_pass.py create mode 100644 backends/arm/test/misc/test_tosa_constant_pool.py create mode 100644 backends/arm/test/passes/test_deduplicate_const_shapes_pass.py create mode 100644 backends/arm/tosa/constant_pool.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index 925faa35ca8..cdd465dd1a5 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -122,6 +122,7 @@ from .decompose_var_pass import DecomposeVarPass # noqa from .decompose_where_scalar_other_pass import DecomposeWhereScalarOtherPass # noqa from .decorate_fp32_to_int32_casting_pass import DecorateFp32toInt32CastingPass # noqa +from .deduplicate_const_shapes_pass import DeduplicateConstShapesPass # noqa from .deduplicate_get_attr_pass import DeduplicateGetAttrPass # noqa from .ensure_unique_output_nodes_pass import EnsureUniqueOutputNodesPass # noqa from .exir_to_tosa_pass import ExirToTosaPass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 0587170c23c..019093fb1b4 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -109,6 +109,7 @@ DecomposeVarPass, DecomposeWhereScalarOtherPass, DecorateFp32toInt32CastingPass, + DeduplicateConstShapesPass, DeduplicateGetAttrPass, EnsureUniqueOutputNodesPass, ExirToTosaPass, @@ -725,6 +726,7 @@ def _tosa_pipeline( # fusing generated RESCALE users can corrupt distinct quantized paths. FuseDuplicateUsersPass(), InsertRescalePass(), + DeduplicateConstShapesPass(), EnsureUniqueOutputNodesPass(), ] ) diff --git a/backends/arm/_passes/deduplicate_const_shapes_pass.py b/backends/arm/_passes/deduplicate_const_shapes_pass.py new file mode 100644 index 00000000000..61a0083c429 --- /dev/null +++ b/backends/arm/_passes/deduplicate_const_shapes_pass.py @@ -0,0 +1,47 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Set, Type + +from executorch.backends.arm._passes import ArmPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, PassResult +from torch.fx import GraphModule, Node + + +class DeduplicateConstShapesPass(ArmPass): + """Reuse the first CONST_SHAPE node with identical static values.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + + def call(self, graph_module: GraphModule) -> PassResult: + representatives: dict[tuple[int, ...], Node] = {} + modified = False + + for node in list(graph_module.graph.nodes): + if node.target != exir_ops.backend.tosa.CONST_SHAPE.default: + continue + + values = node.args[0] + if not isinstance(values, (list, tuple)) or not all( + type(value) is int for value in values + ): + continue + + key = tuple(values) + representative = representatives.get(key) + if representative is None: + representatives[key] = node + continue + + node.replace_all_uses_with(representative) + graph_module.graph.erase_node(node) + modified = True + + if modified: + graph_module.graph.lint() + graph_module.recompile() + + return PassResult(graph_module, modified) diff --git a/backends/arm/operators/BUCK b/backends/arm/operators/BUCK index 5d71cf151e5..b775d7337c5 100644 --- a/backends/arm/operators/BUCK +++ b/backends/arm/operators/BUCK @@ -33,6 +33,7 @@ fbcode_target( "fbsource//third-party/tosa_tools:tosa", ":node_visitor", ":operator_validation_utils", + "//executorch/backends/arm/tosa:constant_pool", "//executorch/backends/arm/tosa:mapping", "//executorch/backends/arm/tosa:utils", "//executorch/backends/arm/_passes:passes", diff --git a/backends/arm/operators/op_tosa_matmul.py b/backends/arm/operators/op_tosa_matmul.py index eb6a26fcc18..e46795c361c 100644 --- a/backends/arm/operators/op_tosa_matmul.py +++ b/backends/arm/operators/op_tosa_matmul.py @@ -76,8 +76,12 @@ def define_node( input_A_ZP_name = f"{output.name}_A_ZP" input_B_ZP_name = f"{output.name}_B_ZP" - tosa_graph.addConst([1], inputs[0].dtype, [input0_zp], name=input_A_ZP_name) - tosa_graph.addConst([1], inputs[1].dtype, [input1_zp], name=input_B_ZP_name) + input_A_ZP = tosa_graph.addConst( + [1], inputs[0].dtype, [input0_zp], name=input_A_ZP_name + ) + input_B_ZP = tosa_graph.addConst( + [1], inputs[1].dtype, [input1_zp], name=input_B_ZP_name + ) # Add the MATMUL to the TOSA graph. attr = ts.TosaSerializerAttribute() @@ -90,8 +94,8 @@ def define_node( [ inputs[0].name, inputs[1].name, - input_A_ZP_name, - input_B_ZP_name, + input_A_ZP.name, + input_B_ZP.name, ], [output.name], attr, diff --git a/backends/arm/operators/op_tosa_mul.py b/backends/arm/operators/op_tosa_mul.py index 65521eab88d..38ee2a6e9b0 100644 --- a/backends/arm/operators/op_tosa_mul.py +++ b/backends/arm/operators/op_tosa_mul.py @@ -47,14 +47,14 @@ def define_node( self.tosa_spec, ) - tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift") + shift = tosa_graph.addConst([1], ts.DType.INT8, 0, name=f"{output.name}_shift") attr = ts.TosaSerializerAttribute() attr.MulAttribute() self._serialize_operator( node, tosa_graph, ts.Op.MUL, - [inputs[0].name, inputs[1].name, f"{output.name}_shift"], + [inputs[0].name, inputs[1].name, shift.name], [output.name], attr, ) diff --git a/backends/arm/operators/op_tosa_shapes.py b/backends/arm/operators/op_tosa_shapes.py index b7480d78a4d..f5831f2e440 100644 --- a/backends/arm/operators/op_tosa_shapes.py +++ b/backends/arm/operators/op_tosa_shapes.py @@ -14,6 +14,7 @@ register_node_visitor, ) from executorch.backends.arm.tosa import TosaSpecification +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import TosaArg from executorch.backends.arm.tosa.utils import normalize_symint @@ -32,8 +33,10 @@ def define_node( shape_input = inputs[0].special rank = len(shape_input) vals = normalize_symint(node.meta["val"]) - tosa_graph = cast(ts.TosaSerializer, tosa_graph) - tosa_graph.addConst( + tosa_graph = cast(TosaSerializerWithConstantPool, tosa_graph) + # Downstream visitors reference this FX output by name. Pooling it with a + # serializer-generated constant could leave output.name undefined. + tosa_graph.addUnpooledConst( [ rank, ], diff --git a/backends/arm/operators/ops_quant_utils.py b/backends/arm/operators/ops_quant_utils.py index 1b5bf4caea9..fef9e3c56c1 100644 --- a/backends/arm/operators/ops_quant_utils.py +++ b/backends/arm/operators/ops_quant_utils.py @@ -25,12 +25,11 @@ def add_input_weight_zp_consts(tosa_graph, node, inputs, output_name): input_zp_name = f"{output_name}_input_zp" weight_zp_name = f"{output_name}_weight_zp" - tosa_graph.addConst([1], inputs[0].dtype, [input_zp], name=input_zp_name) - tosa_graph.addConst( - [1], - inputs[1].dtype, - weight_zp, - name=weight_zp_name, + input_zp_tensor = tosa_graph.addConst( + [1], inputs[0].dtype, [input_zp], name=input_zp_name + ) + weight_zp_tensor = tosa_graph.addConst( + [1], inputs[1].dtype, weight_zp, name=weight_zp_name ) - return input_zp_name, weight_zp_name + return input_zp_tensor.name, weight_zp_tensor.name diff --git a/backends/arm/process_node.py b/backends/arm/process_node.py index a0c2dbeb1fb..129413cabe7 100644 --- a/backends/arm/process_node.py +++ b/backends/arm/process_node.py @@ -106,12 +106,18 @@ def _add_const( tosa_arg: TosaArg, name: str, ) -> None: - """Add a constant, preserving packed FP4 storage when required.""" + """Add a graph-owned constant under its exact name. + + Parameters, buffers, and lifted constants are referenced by their FX names, + so pooling them could leave those names undefined. Preserve packed FP4 + storage when required. + + """ if _is_packed_fp4_const(values, tosa_arg): # TOSA FP4 tensors have logical FP4 shape, but constants are stored as # packed bytes (two values per byte). Add the raw bytes as INT8 first # then set TOSA dtype and shape correctly on the tensor metadata. - tosa_graph.addConst( + tosa_graph.addUnpooledConst( normalize_symint(values.shape), ts.DType.INT8, values, @@ -124,7 +130,7 @@ def _add_const( return prepared_values = _prepare_const_values_for_tosa_dtype(values, tosa_arg) - tosa_graph.addConst( + tosa_graph.addUnpooledConst( _get_const_shape(prepared_values, tosa_arg), tosa_arg.dtype, prepared_values, diff --git a/backends/arm/test/misc/test_process_node.py b/backends/arm/test/misc/test_process_node.py index 02d2a5e012b..994c1628836 100644 --- a/backends/arm/test/misc/test_process_node.py +++ b/backends/arm/test/misc/test_process_node.py @@ -10,6 +10,7 @@ import torch import tosa_serializer as ts from executorch.backends.arm.process_node import _add_const, process_placeholder +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import TosaArg, TosaSpecialDtype from executorch.backends.arm.tosa.specification import TosaSpecification from executorch.exir import to_edge @@ -39,7 +40,7 @@ def __init__(self) -> None: self.name = None self.serialized_bytes = None - def addConst(self, shape, dtype, values, name): + def addUnpooledConst(self, shape, dtype, values, name): self.shape = shape self.dtype = dtype self.values = np.asarray(values) @@ -111,7 +112,7 @@ def test_add_const_fp4_in_packed_storage() -> None: TosaArg, SimpleNamespace(dtype=ts.DType.FP4E2M1, shape=(1, 1, 8)), ) - tosa_graph = ts.TosaSerializer() + tosa_graph = TosaSerializerWithConstantPool() _add_const(tosa_graph, packed_values, tosa_arg, name="fp4_weight") @@ -140,7 +141,7 @@ def _test_add_const_fp6_in_packed_storage(dtype: int) -> None: TosaArg, SimpleNamespace(dtype=dtype, shape=(1, 1, 32)), ) - tosa_graph = ts.TosaSerializer() + tosa_graph = TosaSerializerWithConstantPool() _add_const(tosa_graph, values, tosa_arg, name="fp6_weight") diff --git a/backends/arm/test/misc/test_tosa_constant_pool.py b/backends/arm/test/misc/test_tosa_constant_pool.py new file mode 100644 index 00000000000..4a5f4502844 --- /dev/null +++ b/backends/arm/test/misc/test_tosa_constant_pool.py @@ -0,0 +1,144 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import gc +import weakref + +import numpy as np +import pytest +import tosa_serializer as ts +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool + + +def _serializer(path_prefix=""): + return TosaSerializerWithConstantPool( + path_prefix, + targetMajor=1, + targetMinor=0, + targetPatch=0, + targetDraft=False, + ) + + +@pytest.mark.parametrize("dtype", [ts.DType.INT8, ts.DType.SHAPE]) +def test_identical_constants_are_reused(dtype): + serializer = _serializer() + + first = serializer.addConst([1], dtype, [0], name="first") + second = serializer.addConst([1], dtype, [0], name="second") + + assert isinstance(serializer, ts.TosaSerializer) + assert second is first + assert first.name == "first" + block = serializer.currRegion.currBasicBlock + assert len(block.operators) == 1 + constants = block.shapes if dtype == ts.DType.SHAPE else block.tensors + assert list(constants.keys()) == ["first"] + + +@pytest.mark.parametrize("dtype", [ts.DType.INT8, ts.DType.SHAPE]) +def test_constant_pool_does_not_keep_serializer_alive(dtype): + serializer = _serializer() + serializer.addConst([1], dtype, [0], name="first") + serializer.addConst([1], dtype, [0], name="duplicate") + serializer.serialize() + serializer_ref = weakref.ref(serializer) + + del serializer + gc.collect() + + assert serializer_ref() is None + + +def test_unpooled_constants_are_not_reused(): + serializer = _serializer() + + first = serializer.addUnpooledConst([1], ts.DType.INT8, [0], name="first") + second = serializer.addUnpooledConst([1], ts.DType.INT8, [0], name="second") + + assert second is not first + assert len(serializer.currRegion.currBasicBlock.operators) == 2 + assert list(serializer.currRegion.currBasicBlock.tensors.keys()) == [ + "first", + "second", + ] + + +def test_unnamed_constant_uses_serializer_generated_name(): + serializer = _serializer() + + constant = serializer.addConst([1], ts.DType.INT8, [0]) + + assert constant.name + + +@pytest.mark.parametrize( + "first,second", + [ + (([1], ts.DType.INT8, [0]), ([1], ts.DType.INT16, [0])), + (([1], ts.DType.INT8, [0]), ([2], ts.DType.INT8, [0, 0])), + (([1], ts.DType.INT8, [0]), ([1], ts.DType.INT8, [1])), + ], +) +def test_constants_with_different_keys_remain_separate(first, second): + serializer = _serializer() + + first_const = serializer.addConst(*first, name="first") + second_const = serializer.addConst(*second, name="second") + + assert second_const is not first_const + assert len(serializer.currRegion.currBasicBlock.operators) == 2 + + +def test_float_constants_use_exact_serialized_values(): + serializer = _serializer() + + positive_zero = serializer.addConst( + [1], ts.DType.FP32, np.array([0.0]), name="positive_zero" + ) + negative_zero = serializer.addConst( + [1], ts.DType.FP32, np.array([-0.0]), name="negative_zero" + ) + repeated_negative_zero = serializer.addConst( + [1], + ts.DType.FP32, + np.array([-0.0]), + name="repeated_negative_zero", + ) + + assert negative_zero is not positive_zero + assert repeated_negative_zero is negative_zero + assert len(serializer.currRegion.currBasicBlock.operators) == 2 + + +def test_constants_are_scoped_to_basic_blocks(): + serializer = _serializer() + first = serializer.addConst([1], ts.DType.INT8, [0], name="first") + + serializer.startRegion("main") + serializer.currRegion.addBasicBlock("main") + second = serializer.addConst([1], ts.DType.INT8, [0], name="second") + + assert second is not first + assert second.name == "second" + + +def test_start_region_preserves_path_prefix(): + serializer = _serializer("artifacts") + + serializer.startRegion("other") + + assert serializer.currRegion.pathPrefix == "artifacts" + + +def test_constant_pool_serialization_is_deterministic(): + def serialize(): + serializer = _serializer() + serializer.addConst([1], ts.DType.INT8, [0], name="first") + serializer.addConst([1], ts.DType.INT8, [0], name="duplicate") + serializer.addConst([1], ts.DType.INT8, [1], name="second") + return bytes(serializer.serialize()) + + assert serialize() == serialize() diff --git a/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py index d88424599e8..3813d7a0b0f 100644 --- a/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py +++ b/backends/arm/test/misc/tosa_dialect/test_tosa_shape_node_visitors.py @@ -17,6 +17,7 @@ NodeVisitor, ) from executorch.backends.arm.test.runner_utils import TosaReferenceModelDispatch +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import TosaArg from executorch.backends.arm.tosa.specification import TosaSpecification from torch.fx import Node @@ -75,7 +76,7 @@ def _shape_spec() -> TosaSpecification: def _serializer() -> ts.TosaSerializer: - return ts.TosaSerializer( + return TosaSerializerWithConstantPool( "", targetMajor=1, targetMinor=1, @@ -289,6 +290,22 @@ def test_const_shape_node_visitor_serializes_const_operator() -> None: assert _serialized_op_codes(tosa_graph) == [ts.Op.CONST_SHAPE] +def test_const_shape_node_visitor_preserves_output_name() -> None: + visitor = get_node_visitors(_shape_spec())["tosa.CONST_SHAPE.default"] + tosa_graph = _serializer() + tosa_graph.addConst([2], ts.DType.SHAPE, [2, 3], name="helper") + + _define_node( + visitor, + SimpleNamespace(name="node", meta={"val": [2, 3]}, kwargs={}), + tosa_graph, + [SimpleNamespace(special=[2, 3])], + SimpleNamespace(name="output", shape=(2,)), + ) + + assert list(tosa_graph.currRegion.currBasicBlock.shapes) == ["helper", "output"] + + def test_dim_shape_node_visitor_serializes_operator() -> None: visitor = get_node_visitors(_shape_spec())["tosa.DIM.default"] tosa_graph = _serializer() diff --git a/backends/arm/test/ops/test_conv2d.py b/backends/arm/test/ops/test_conv2d.py index 977ffdb9a7c..5b608135406 100644 --- a/backends/arm/test/ops/test_conv2d.py +++ b/backends/arm/test/ops/test_conv2d.py @@ -566,7 +566,8 @@ def _get_dtype_count(model: torch.nn.Module): # Set nbr_conv to be the amount of groups set if necessary. nbr_convs: int = model.nbr_convs if model.groups is None else model.groups # noqa return { - "CONST": {"INT4": nbr_convs * 2}, # One for the weight, one for the zp. + # Each convolution has a distinct weight and shares the symmetric zero point. + "CONST": {"INT4": nbr_convs + 1}, "CONV2D": {"INT32": nbr_convs}, "RESCALE": {"INT8": nbr_convs}, } diff --git a/backends/arm/test/ops/test_conv3d.py b/backends/arm/test/ops/test_conv3d.py index 09cd44525c5..a43121274ce 100644 --- a/backends/arm/test/ops/test_conv3d.py +++ b/backends/arm/test/ops/test_conv3d.py @@ -585,7 +585,8 @@ def forward(self, x): def _get_dtype_count(model: torch.nn.Module): nbr_convs: int = model.nbr_convs # noqa return { - "CONST": {"INT4": nbr_convs * 2}, + # Each convolution has a distinct weight and shares the symmetric zero point. + "CONST": {"INT4": nbr_convs + 1}, "CONV3D": {"INT32": nbr_convs}, "RESCALE": {"INT8": nbr_convs}, } diff --git a/backends/arm/test/ops/test_depthwise_conv.py b/backends/arm/test/ops/test_depthwise_conv.py index a81a656017f..3dd8598dddf 100644 --- a/backends/arm/test/ops/test_depthwise_conv.py +++ b/backends/arm/test/ops/test_depthwise_conv.py @@ -264,7 +264,8 @@ def _get_dtype_count(model: torch.nn.Module): nbr_convs: int = model.nbr_convs # noqa return { - "CONST": {"INT4": nbr_convs * 2}, + # Each convolution has a distinct weight and shares the symmetric zero point. + "CONST": {"INT4": nbr_convs + 1}, "DEPTHWISE_CONV2D": {"INT32": nbr_convs}, "RESCALE": {"INT8": nbr_convs}, } diff --git a/backends/arm/test/passes/test_deduplicate_const_shapes_pass.py b/backends/arm/test/passes/test_deduplicate_const_shapes_pass.py new file mode 100644 index 00000000000..d14dbc998e2 --- /dev/null +++ b/backends/arm/test/passes/test_deduplicate_const_shapes_pass.py @@ -0,0 +1,58 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import torch +from executorch.backends.arm._passes import DeduplicateConstShapesPass +from executorch.exir.dialects._ops import ops as exir_ops +from torch.fx import Graph, GraphModule + + +def _const_shape(graph: Graph, name: str, values: list[int]): + node = graph.call_function( + exir_ops.backend.tosa.CONST_SHAPE.default, + (values,), + ) + node.name = name + node.meta["val"] = values + return node + + +def test_deduplicate_identical_const_shapes(): + graph = Graph() + first = _const_shape(graph, "first", [2, 3]) + second = _const_shape(graph, "second", [2, 3]) + graph.output((first, second)) + graph_module = GraphModule(torch.nn.Module(), graph) + + result = DeduplicateConstShapesPass()(graph_module) + + assert result is not None + assert result.modified + const_shapes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.CONST_SHAPE.default + ] + assert [node.name for node in const_shapes] == ["first"] + assert graph_module.graph.output_node().args[0] == (first, first) + + +def test_keep_const_shapes_with_different_values(): + graph = Graph() + first = _const_shape(graph, "first", [2, 3]) + second = _const_shape(graph, "second", [3, 2]) + graph.output((first, second)) + graph_module = GraphModule(torch.nn.Module(), graph) + + result = DeduplicateConstShapesPass()(graph_module) + + assert result is not None + assert not result.modified + const_shapes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.backend.tosa.CONST_SHAPE.default + ] + assert const_shapes == [first, second] diff --git a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py index 831d9d91267..aad803b98e7 100644 --- a/backends/arm/test/passes/test_fuse_duplicate_users_pass.py +++ b/backends/arm/test/passes/test_fuse_duplicate_users_pass.py @@ -8,6 +8,7 @@ import executorch.backends.arm.tosa.dialect # noqa: F401 import torch from executorch.backends.arm._passes import ( + DeduplicateConstShapesPass, EnsureUniqueOutputNodesPass, FuseDuplicateUsersPass, InsertRescalePass, @@ -241,8 +242,9 @@ def test_fuse_duplicate_users_runs_after_tosa_transformations(): for index, pass_type in enumerate(pass_types) if pass_type is RemoveNoopPass ) - assert pass_types[post_noop_index + 1 : post_noop_index + 4] == [ + assert pass_types[post_noop_index + 1 : post_noop_index + 5] == [ FuseDuplicateUsersPass, InsertRescalePass, + DeduplicateConstShapesPass, EnsureUniqueOutputNodesPass, ] diff --git a/backends/arm/test/passes/test_remove_data_layout_noops.py b/backends/arm/test/passes/test_remove_data_layout_noops.py index 87e592dcdea..7090cfbf008 100644 --- a/backends/arm/test/passes/test_remove_data_layout_noops.py +++ b/backends/arm/test/passes/test_remove_data_layout_noops.py @@ -8,6 +8,7 @@ import torch from executorch.backends.arm._passes import ( CanonicalizeViewCopyPermutePass, + DeduplicateConstShapesPass, EnsureUniqueOutputNodesPass, ExirToTosaPass, FuseDuplicateUsersPass, @@ -499,4 +500,5 @@ def test_data_layout_noop_cleanup_pipeline_order(): assert pass_types[pre_tosa_cleanup + 1] is CanonicalizeViewCopyPermutePass assert pass_types[post_tosa_cleanup + 1] is FuseDuplicateUsersPass assert pass_types[post_tosa_cleanup + 2] is InsertRescalePass - assert pass_types[post_tosa_cleanup + 3] is EnsureUniqueOutputNodesPass + assert pass_types[post_tosa_cleanup + 3] is DeduplicateConstShapesPass + assert pass_types[post_tosa_cleanup + 4] is EnsureUniqueOutputNodesPass diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 2597f74d9ac..2e5d4e1f3b3 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -66,6 +66,7 @@ def define_arm_tests(): "misc/test_external_vela_blocks.py", # "misc/test_evaluate_model.py", "misc/test_pass_pipeline_config.py", + "misc/test_tosa_constant_pool.py", "misc/tosa_dialect/test_tosa_dialect_cast_to_block_scaled.py", "misc/tosa_dialect/test_tosa_dialect_mxfp_conv2d.py", "misc/tosa_dialect/test_tosa_dialect_mxfp_linear.py", diff --git a/backends/arm/tosa/BUCK b/backends/arm/tosa/BUCK index ccf4f461a8b..7ae9919ce5c 100644 --- a/backends/arm/tosa/BUCK +++ b/backends/arm/tosa/BUCK @@ -23,6 +23,16 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "constant_pool", + srcs = [ + "constant_pool.py", + ], + deps = [ + "fbsource//third-party/tosa_tools:serializer", + ], +) + fbcode_target(_kind = runtime.python_library, name = "mapping", srcs = [ @@ -95,6 +105,7 @@ fbcode_target(_kind = runtime.python_library, ], deps = [ ":compile_spec", + ":constant_pool", "//executorch/backends/arm:constants", "//executorch/backends/arm:process_node", "//executorch/backends/arm/debug:schema", diff --git a/backends/arm/tosa/backend.py b/backends/arm/tosa/backend.py index 6ec7d078674..6efe71baec8 100644 --- a/backends/arm/tosa/backend.py +++ b/backends/arm/tosa/backend.py @@ -31,6 +31,7 @@ process_placeholder, ) from executorch.backends.arm.tosa.compile_spec import TosaCompileSpec +from executorch.backends.arm.tosa.constant_pool import TosaSerializerWithConstantPool from executorch.backends.arm.tosa.mapping import ( TOSA_CONTROL_FLOW_REGION_NAME_META, TOSA_CONTROL_FLOW_SOURCE_NODE_META, @@ -152,7 +153,7 @@ def _preprocess( # noqa: C901 artifact_path = "" version = tosa_spec.version - tosa_graph = ts.TosaSerializer( + tosa_graph = TosaSerializerWithConstantPool( artifact_path, targetMajor=version.major, targetMinor=version.minor, @@ -251,7 +252,7 @@ def _preprocess_module( # noqa: C901 graph_module: GraphModule, edge_program: ExportedProgram, compile_spec: TosaCompileSpec, - tosa_graph: ts.TosaSerializer, + tosa_graph: TosaSerializerWithConstantPool, debug_hook: DebugHook | None, submodule_name: str | None = None, containing_graph_module: GraphModule | None = None, @@ -262,9 +263,12 @@ def _preprocess_module( # noqa: C901 graph_module (GraphModule): Module to lower recursively. edge_program (ExportedProgram): Original exported program. compile_spec (TosaCompileSpec): Backend options with TOSA settings. - tosa_graph (ts.TosaSerializer): Serializer receiving operators. + tosa_graph (TosaSerializerWithConstantPool): Serializer receiving + operators. debug_hook (DebugHook | None): Optional debug instrumentation. submodule_name (str | None): Name used when visiting nested blocks. + containing_graph_module (GraphModule | None): Parent graph module for + nested control flow. Raises: RuntimeError: If an FX node with an unsupported op kind is found. diff --git a/backends/arm/tosa/constant_pool.py b/backends/arm/tosa/constant_pool.py new file mode 100644 index 00000000000..92cdc65df60 --- /dev/null +++ b/backends/arm/tosa/constant_pool.py @@ -0,0 +1,59 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Any + +import tosa_serializer as ts + + +_ConstantKey = tuple[Any, tuple[int, ...], bytes | None] + + +def _constant_key(shape, dtype, values) -> _ConstantKey: + if dtype == ts.DType.SHAPE: + if len(shape) > 1: + raise ValueError(f"CONST_SHAPE expects rank metadata, got {shape}") + rank = 0 if len(shape) == 0 else shape[0] + constant = ts.TosaSerializerShape("", rank, values) + else: + constant = ts.TosaSerializerTensor("", shape, dtype, values) + + data = None if constant.data is None else bytes(constant.data) + return constant.dtype, tuple(constant.shape), data + + +class TosaSerializerWithConstantPool(ts.TosaSerializer): + """Pool generated constants independently within each TOSA basic block.""" + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + # Native tensor wrappers retain their serializer. Cache names to avoid + # an ownership cycle that Python's garbage collector cannot release. + self._block_pools: dict[Any, dict[_ConstantKey, str]] = {} + + def addConst(self, shape, dtype, vals=None, name=""): + """Return a matching constant in the current block or add a new one.""" + block = self.currRegion.currBasicBlock + pool = self._block_pools.setdefault(block, {}) + key = _constant_key(shape, dtype, vals) + if key not in pool: + constant = super().addConst(shape, dtype, vals, name) + pool[key] = constant.name + return constant + + # Resolve the cached name to the object expected by callers. TOSA stores + # shape constants separately from tensor constants. + cached_name = pool[key] + if dtype == ts.DType.SHAPE: + constant = block.getShapeByName(cached_name) + else: + constant = block.getTensorByName(cached_name) + return constant + + def addUnpooledConst(self, shape, dtype, vals=None, name=""): + """Add a constant without pooling so its requested name remains + addressable. + """ + return super().addConst(shape, dtype, vals, name) From 3a300882e0a20149a691d90c13226e78e5e147b9 Mon Sep 17 00:00:00 2001 From: Yufeng Shi Date: Fri, 11 Sep 2026 09:19:03 +0100 Subject: [PATCH 177/190] Arm backend: Support leading full slices in index.Tensor (#22680) Generalize index.Tensor decomposition to treat leading full-slice dimensions as the TOSA GATHER batch dimension. Reshape values to [P, K, C], expand linearized indices to [P, W], and restore the PyTorch output shape without transposing the data tensor. Update support checks to accept leading None entries while rejecting interleaved indexing, and add tests for TOSA FP, TOSA INT, and VGF. Change-Id: I657dd1876c7ca2632b951bc517f8960a10242548 Signed-off-by: Yufeng Shi --- .../decompose_index_tensor_to_gather_pass.py | 255 +++++++++++------- .../operator_support/index_tensor_support.py | 96 ++++--- backends/arm/test/ops/test_index_tensor.py | 207 +++++++++++--- .../source/backends/arm-vgf/VGF_op_support.md | 2 +- 4 files changed, 387 insertions(+), 173 deletions(-) diff --git a/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py b/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py index c3e98c363dc..3413ce87238 100644 --- a/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py +++ b/backends/arm/_passes/decompose_index_tensor_to_gather_pass.py @@ -30,7 +30,7 @@ def get_index_tensor_decomposition(op): - """Return the operator overloads used to lower index.Tensor via TOSA gather. + """Return operators used to lower index.Tensor through TOSA gather. Raises: RuntimeError: If the provided operator is not supported by this pass. @@ -53,12 +53,12 @@ def get_index_tensor_decomposition(op): def _broadcast_shape( shapes: Sequence[Sequence[int]], ) -> list[int]: - """Compute the broadcasted shape (PyTorch/Numpy semantics) for a list of - shapes. + """Compute the broadcasted shape using PyTorch/NumPy semantics. Requirements: - static shape only - - shapes are right-aligned; lower-rank shapes are implicitly front-padded with 1s + - shapes are right-aligned; lower-rank shapes are implicitly + front-padded with 1s - per-axis dims must either match exactly or be 1 Raises: @@ -81,86 +81,121 @@ def _broadcast_shape( class DecomposeIndexTensorToGatherPass(ArmOpTargetedPass): - """Decompose edge.aten.index.Tensor into backend TOSA gather (+ basic - arith). + """Decompose edge.aten.index.Tensor into a TOSA gather and arithmetic. Supported subset: - y = x.index([i0, i1, ..., i{m-1}]) + y = x.index([None, ..., None, i0, i1, ..., i{m-1}]) - where each ik is a Tensor index, and m is the number of index tensors. + where each ik is a Tensor index, m is the number of index tensors, and the + optional leading None entries preserve dimensions before the indexed block. Constraints: - - `indices` list contains only Tensor indices (no None/slice/ellipsis) + - `indices` contains an optional leading run of None entries followed by + only Tensor indices - Each index tensor dtype is int32 - - Index tensor shapes are broadcastable to a common shape `S` (per index.Tensor semantics) - - Only prefix indexing is supported: the `m` tensor indices select elements - from the first `m` dimensions of `x`, so `m <= rank(x)`. + - Index tensor shapes are broadcastable to a common shape `S` (per + index.Tensor semantics) + - The `m` tensor indices select one contiguous block of dimensions after + the leading preserved dimensions. - Static shapes are required - - If `x` has more than 2^31 elements, the computed linear index may overflow int32. + - If `x` has more than 2^31 elements, the computed linear index may + overflow int32. Lowering strategy (single gather) --------------------------------- Let: + - `p` be the number of leading None entries - `S` be the broadcasted index shape - `W = prod(S)` (number of indexed positions) - - `K = prod(x.shape[:m])` (flattened size of the indexed prefix) - - `C = prod(x.shape[m:])` (flattened size of the trailing slice per index) - - `trailing = x.shape[m:]` + - `P = prod(x.shape[:p])` (flattened size of the leading preserved + dimensions) + - `K = prod(x.shape[p:p+m])` (flattened size of the indexed block) + - `C = prod(x.shape[p+m:])` (flattened size of the trailing slice per + index) + - `leading = x.shape[:p]` and `trailing = x.shape[p+m:]` Steps: 1) Compute parameters needed to lower index.Tensor - - `S`, `W`, `K`, `C`, `trailing` - - `lin_scales[i] = stride_i // C`, where `stride_i` are the contiguous-style - strides derived from `x.shape` (for dim i). - 2) Reshape x to `[1, K, C]` (`x_1kc`). - 3) Build linear indices (`lin_1w`) by scaling each flattened index and summing: - lin_1w = unsqueeze0( sum_{i=0..m-1} ( idx_flat[i] * lin_scales[i] ) ) - where: - - `m = len(indices)` - - `idx_flat[i]` is the i-th index tensor after broadcast to `S` and flatten to `[W]` - - `lin_1w` has shape `[1, W]` and is used as the `indices` input to `tosa.GATHER` + - `S`, `W`, `P`, `K`, `C`, `leading`, `trailing` + - `lin_scales` as the contiguous strides of the indexed block + `x.shape[p:p+m]`. + 2) Reshape x to `[P, K, C]` (`x_pkc`). + 3) Build linear indices by scaling and accumulating the flattened index + tensors element-wise: + For each tensor index, broadcast it to `S`, then flatten it: + + idx_broadcast[i] = broadcast_to(indices[p + i], S) + idx_flat[i] = reshape(idx_broadcast[i], [W]) + + For each j in [0, W): + + lin_w[j] = + sum_{i=0..m-1} idx_flat[i][j] * lin_scales[i] + + Equivalently, in tensor notation: + + lin_w = + sum_{i=0..m-1} idx_flat[i] * lin_scales[i] # shape [W] + + Then: + + lin_1w = unsqueeze(lin_w, 0) # [1, W] + lin_pw = lin_1w + if P > 1: + lin_pw = expand(lin_1w, [P, W]) # [P, W] 4) Single gather: - `tosa.GATHER(x=x_1kc, indices=lin_1w) -> [1,W,C]` - 5) Reshape result to `[*S, *trailing]`. + `tosa.GATHER(x=x_pkc, indices=lin_pw) -> [P,W,C]` + 5) Reshape result to `[*leading, *S, *trailing]`. - Example - ------- + Example: Consider: - x.shape = [2, 3, 4] - indices = [i0, i1] # m = 2 + x.shape = [2, 3, 4, 5] + indices = [None, i0, i1] # p = 1, m = 2 i0.shape = [2, 1] i1.shape = [1, 2] + This corresponds to ``x[:, i0, i1, :]``: the first dimension is + preserved, the next two dimensions are indexed, and the last dimension is + trailing. + 1) The index shapes broadcast to: S := [2, 2] W := prod(S) = 4 - We index the first m=2 dimensions of x, so: - K := prod(x.shape[:m]) = 2 * 3 = 6 - C := prod(x.shape[m:]) = 4 - trailing := x.shape[m:] = [4] + We preserve p=1 leading dimension and index the next m=2 dimensions: + leading := x.shape[:p] = [2] + trailing := x.shape[p+m:] = [5] + P := prod(leading) = 2 + K := prod(x.shape[p:p+m]) = 3 * 4 = 12 + C := prod(trailing) = 5 - Contiguous strides of x are [12, 4, 1], so: - lin_scales := [stride0 // C, stride1 // C] = [12//4, 4//4] = [3, 1] + The indexed block has shape [3, 4], so its contiguous strides are: + lin_scales := [4, 1] 2) Values are reshaped to: - x_1kc = view(x, [1, K, C]) = [1, 6, 4] + x_pkc = view(x, [P, K, C]) = [2, 12, 5] 3) After broadcasting and flattening the indices to length W: i0_broadcast, i1_broadcast have shape S=[2,2] i0_flat, i1_flat have shape [W]=[4] - Linear indices are computed as: - lin_w - = lin_scales * [i0_flat, i1_flat] - = 3 * i0_flat + 1 * i1_flat # shape [W] - lin_w is reshaped to [1, W] to match tosa.Gather semantics + Linear indices are computed element-wise as: + for each j in [0, W): + + lin_w[j] = 4 * i0_flat[j] + i1_flat[j] + + hence: + + lin_w = 4 * i0_flat + i1_flat # shape [W] + + lin_w is then unsqueezed to [1, W] and expanded so that + lin_pw.shape = [P, W] = [2, 4]. 4) Single Gather: - out_1wc = tosa.GATHER(values=x_1kc, indices=lin_1w) # [1, 4, 4] + out_pwc = tosa.GATHER(values=x_pkc, indices=lin_pw) # [2, 4, 5] 5) Reshape result: - out = view(out_1wc, [*S, *x.shape[m:]]) # [2, 2, 4] + out = view(out_pwc, [*leading, *S, *trailing]) # [2, 2, 2, 5] """ @@ -192,69 +227,76 @@ def _shape_to_stride( return strides @staticmethod - def _validate_tensor_indices(indices): + def _validate_and_split_indices(indices): assert ( isinstance(indices, (list, tuple)) and len(indices) > 0 ), f"index.Tensor expects non-empty indices list/tuple, got {type(indices)}." - for i, idx in enumerate(indices): - assert ( - idx is not None - ), f"index.Tensor: None indices are not supported at the moment (indices[{i}] is None)." + leading_rank = 0 + while leading_rank < len(indices) and indices[leading_rank] is None: + leading_rank += 1 + + tensor_indices = indices[leading_rank:] + assert tensor_indices, "index.Tensor expects at least one tensor index." + for i, idx in enumerate(tensor_indices, start=leading_rank): + assert idx is not None, ( + "index.Tensor supports None entries only before all tensor indices " + f"(indices[{i}] is None)." + ) assert ( idx.data.dtype == torch.int32 ), "index.Tensor requires index dtype must be int32" - def _compute_index_tensor_params(self, x, m, index_shapes): - """Compute shape/stride-derived parameters needed to lower - edge.aten.index.Tensor. + return leading_rank, tensor_indices - Derives the broadcasted index shape and the scale factors used to flatten and - acculumulate multi-dimensional indices into a single gather index, following - the S/W/K/C notation described in the class docstring. + def _compute_index_tensor_params(self, x, leading_rank, m, index_shapes): + """Compute parameters needed to lower edge.aten.index.Tensor. + + Derives the broadcasted index shape and the scale factors used to + flatten and accumulate multi-dimensional indices into a single gather + index, following the S/W/P/K/C notation described in the class + docstring. Args: - x: Values tensor being indexed. - m: Number of tensor indices (i.e., len(indices)). - index_shapes: Shapes corresponding to each tensor index. + x (ProxyValue): Values tensor being indexed. + leading_rank (int): Number of leading dimensions preserved by None + entries. + m (int): Number of tensor indices. + index_shapes (Sequence[Sequence[int]]): Shapes corresponding to + each tensor index. Returns: - (x_data, S, W, K, C, trailing, lin_scales), where: - - x_data is `x.data` (FakeTensor) - - trailing is `x.shape[m:]` as a list of ints - - lin_scales are per-dimension scale factors for linearization + tuple: `(x_data, S, W, P, K, C, leading, trailing, lin_scales)`, + where `x_data` is `x.data`, `leading` and `trailing` contain + the preserved dimensions, and `lin_scales` contains the + indexed-block strides used for linearization. """ - x_data = x.data # FakeTensor x_shape = tuple(x_data.shape) x_rank = len(x_shape) assert x_rank >= 1, f"index.Tensor expects x rank>=1, got {x_shape}." - assert ( - m <= x_rank - ), f"index.Tensor has too many indices ({m}) for x rank {x_rank}." + assert leading_rank + m <= x_rank, ( + "index.Tensor has more preserved and indexed dimensions " + f"({leading_rank + m}) than the input rank ({x_rank})." + ) # Broadcast shape S for indices, and flattened length W S = _broadcast_shape(index_shapes) W = math.prod(S) if S else 1 - # Compute gather factors K and C for leading-dims indexing - leading = list(x_shape[:m]) - trailing = list(x_shape[m:]) - K = math.prod(leading) if leading else 1 + # Compute gather factors for the preserved, indexed, and trailing blocks. + leading = list(x_shape[:leading_rank]) + indexed = list(x_shape[leading_rank : leading_rank + m]) + trailing = list(x_shape[leading_rank + m :]) + P = math.prod(leading) if leading else 1 + K = math.prod(indexed) if indexed else 1 C = math.prod(trailing) if trailing else 1 - # Strides for linearization (contiguous-style) - strides = self._shape_to_stride(x_shape) - - # Stride/C divisibility is guaranteed for contiguous strides and C=prod(trailing). - lin_scales: list[int] = [] - for i in range(m): - stride = strides[i] - lin_scales.append(stride // C) + lin_scales = self._shape_to_stride(indexed) - return x_data, S, W, K, C, trailing, lin_scales + return x_data, S, W, P, K, C, leading, trailing, lin_scales def _decompose_constant_index(self, x, indices, meta): tensor_indices = [ @@ -341,13 +383,21 @@ def call_operator(self, op, args, kwargs, meta): if constant_result is not None: return constant_result - self._validate_tensor_indices(indices) + leading_rank, indices = self._validate_and_split_indices(indices) index_shapes = [idx.data.shape for idx in indices] m = len(indices) - x_data, S, W, K, C, trailing, lin_scales = self._compute_index_tensor_params( - x, m, index_shapes - ) + ( + x_data, + S, + W, + P, + K, + C, + leading, + trailing, + lin_scales, + ) = self._compute_index_tensor_params(x, leading_rank, m, index_shapes) ( view_op, @@ -371,16 +421,16 @@ def call_operator(self, op, args, kwargs, meta): updated=True, ) - # ---- x: [1, K, C] ---- - x_1kc = super().call_operator( + # ---- x: [P, K, C] ---- + x_pkc = super().call_operator( view_op, - (x_for_gather, [1, K, C]), + (x_for_gather, [P, K, C]), {}, meta, updated=True, ) - # Build linear index [1, W] from broadcasted indices + # Build linear index [W] from broadcasted indices lin_w = None plain_meta = meta_without_qparams(meta) for i, idx in enumerate(indices): @@ -427,7 +477,7 @@ def call_operator(self, op, args, kwargs, meta): updated=True, ) - # Accumulate into lin_1w: [1, W] + # Accumulate into lin_w: [W] if lin_w is None: lin_w = idx_scaled else: @@ -441,10 +491,10 @@ def call_operator(self, op, args, kwargs, meta): if lin_w is None: raise RuntimeError( - f"[{self.__class__.__name__}] internal error: lin_1w not constructed." + f"[{self.__class__.__name__}] internal error: lin_w not constructed." ) - # Make indices shape [1, W] for tosa.GATHER + # Make indices shape [P, W] for tosa.GATHER. lin_1w = super().call_operator( unsqueeze_op, (lin_w, 0), @@ -452,22 +502,31 @@ def call_operator(self, op, args, kwargs, meta): plain_meta, updated=True, ) + lin_pw = lin_1w + if P > 1: + lin_pw = super().call_operator( + expand_op, + (lin_1w, [P, W]), + {}, + plain_meta, + updated=True, + ) # ---- backend tosa gather --- - # tosa.GATHER(x=[1,K,C], indices=[1,W]) -> [1,W,C] - gathered_1wc = super().call_operator( + # tosa.GATHER(x=[P,K,C], indices=[P,W]) -> [P,W,C] + gathered_pwc = super().call_operator( tosa_gather_op, - (x_1kc, lin_1w), + (x_pkc, lin_pw), {}, meta, updated=True, ) - # ---- output: [*S, *trailing] ---- - out_shape = list(S) + list(trailing) + # ---- output: [*leading, *S, *trailing] ---- + out_shape = list(leading) + list(S) + list(trailing) out = super().call_operator( view_op, - (gathered_1wc, out_shape), + (gathered_pwc, out_shape), {}, meta, updated=True, diff --git a/backends/arm/operator_support/index_tensor_support.py b/backends/arm/operator_support/index_tensor_support.py index 937b102ce8f..b14a30e0e94 100644 --- a/backends/arm/operator_support/index_tensor_support.py +++ b/backends/arm/operator_support/index_tensor_support.py @@ -4,12 +4,13 @@ # LICENSE file in the root directory of this source tree. """Provide TOSA support checks for ``aten.index.Tensor``. -Reject unsupported patterns such as front-positioned slice/ellipsis/None -markers and cases that exceed ``int32`` element limits. +Reject unsupported indexing layouts, zero-sized tensors, and cases that exceed +``int32`` element limits. """ import math +from typing import cast, Sequence import torch import torch.fx as fx @@ -23,6 +24,15 @@ from executorch.exir.dialects._ops import ops as exir_ops +def _has_leading_full_slices_only(indices) -> bool: + found_tensor_index = False + for index in indices: + if index is None and found_tensor_index: + return False + found_tensor_index |= index is not None + return found_tensor_index + + @register_tosa_support_check class IndexTensorSupported(SupportedTOSAOperatorCheck): """Prevent partitioning of unsupported ``index.Tensor`` usages. @@ -30,10 +40,10 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck): This support check is intended to prevent the partitioning of currently unsupported usages of the index.Tensor operator. - 1. Usages where slice, ellipsis or None are present before an indexing tensor: - t[{start}:{end}, indexTensor] - slicing - t[None, indexTensor] - unsqueeze - t[..., indexTensor] - ellipsis + 1. Usages where a slice, ellipsis, or None separates indexing tensors: + t[indexTensor, {start}:{end}, indexTensor] - slicing + t[indexTensor, None, indexTensor] - unsqueeze + t[indexTensor, ..., indexTensor] - ellipsis 2. Usages where the value tensor contains more than int32.max elements This is due to int32 TOSA limitation and the fact that we flatten out @@ -41,25 +51,21 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck): As such to avoid overflow we reject lowering of this operator if it is possible for indices to go over the int32 limit. - Extra information regarding #2: + 3. Usages where the value or an index tensor is zero-sized, because TOSA + requires every tensor dimension to be at least one. + + Extra information regarding #1: Pytorch decomposes slice and None usages before they reach aten. In the case of Slicing and Unsqueeze, Pytorch will add the relevant operation just before the index.Tensor op. In the case of Ellipsis no extra operation is added. - In all three cases Pytorch will insert "None"(s) in the index list - only if the above operations are done on a dimension BEFORE one being indexed. - - When slicing, unsqueeze and ellipsis are done on dimensions after - the ones being indexed, then they do not affect the final output - values, only the shape. Thus None is not passed to the index.Tensor op. - The purpose of None is to signify to index.Tensor that a dimension should not be indexed. - In such cases the logic behaves similar to batching along that dimension. - For the sake of simplicity we have not implemented this behavior yet - and thus have put this support check in place to prevent the partitioning - of index.Tensor ops which include None. + A leading run of None entries behaves like batching along those + dimensions and is supported. None entries after the first tensor index + remain unsupported because they interleave preserved and indexed + dimensions. Examples: #1 - Slice ----------------------------------------------------- @@ -87,12 +93,6 @@ class IndexTensorSupported(SupportedTOSAOperatorCheck): out = ...edge__ops_aten_index_Tensor(unsqueeze_res, [torch.arange(3)]) NB. - With the current implementation of flattening tensors and indices out, - supporting None (Unsqueeze) is simply a matter of ignoring the - None dimension. - This is not the case for Slice and Ellipsis operators, where - the size of the new dimension can be > 1. - Note that slice ops interleaved between indexes such as: t[1:3, torch.arange(5), 2:3, torch.arange(3).reshape(3,1)] are also possible and can result in some unintuitive behaviors @@ -108,39 +108,47 @@ def is_node_tosa_supported( """Return True if ``aten.index.Tensor`` usage fits supported patterns. Enforces the following constraints: - - No ``None`` (unsqueeze), slice, or ellipsis before an indexing tensor, - except for the U55 constant-index lowering. + - ``None`` entries may only form a leading run before all tensor indices. + - At least one tensor index is present. + - Value and index tensors must not be zero-sized. + - Boolean and byte mask indices are not supported. - The value tensor element count fits in ``int32``. """ - indices = node.args[1] - if not tosa_spec.is_U55_subset and any( - index is None for index in indices # type: ignore[union-attr] - ): + indices = cast(Sequence[fx.Node | None], node.args[1]) + if not _has_leading_full_slices_only(indices): self.reporter.report_reject( node, - ( - "None (from slice/unsqueeze/ellipsis) before an indexing tensor" - " is not supported." - ), + "Only leading None entries followed by tensor indices are supported.", ) return False - # The U55-specific check limits this to one constant tensor index. - for index in ( - index for index in indices if index is not None # type: ignore[union-attr] + if any( + get_first_fake_tensor(ensure_type(fx.Node, index)).dtype + in (torch.bool, torch.uint8) + for index in indices + if index is not None ): - index_node = ensure_type(torch.fx.Node, index) - if get_first_fake_tensor(index_node).dtype in (torch.bool, torch.uint8): - self.reporter.report_reject( - node, "Boolean and byte mask indices are not supported." - ) - return False + self.reporter.report_reject( + node, "Boolean and byte mask indices are not supported." + ) + return False - # Usage 2 guard input_node = ensure_type(torch.fx.Node, node.args[0]) input_val = get_first_fake_tensor(input_node) total_vals = math.prod(input_val.shape) + has_zero_sized_index = any( + math.prod(get_first_fake_tensor(ensure_type(fx.Node, index)).shape) == 0 + for index in indices + if index is not None + ) + if total_vals == 0 or has_zero_sized_index: + self.reporter.report_reject( + node, + "Zero-sized value or index tensors are not supported by TOSA.", + ) + return False + if total_vals > torch.iinfo(torch.int32).max: self.reporter.report_reject( node, diff --git a/backends/arm/test/ops/test_index_tensor.py b/backends/arm/test/ops/test_index_tensor.py index 20bce650513..f6b17bf5c32 100644 --- a/backends/arm/test/ops/test_index_tensor.py +++ b/backends/arm/test/ops/test_index_tensor.py @@ -44,6 +44,18 @@ def forward(self, x: torch.Tensor): return x[self.index] +class IndexTensorLeadingInt64Buffers(torch.nn.Module): + """NCHW indexing with leading full slices and int64 index buffers.""" + + def __init__(self): + super().__init__() + self.register_buffer("rows", torch.tensor([[0], [2]], dtype=torch.int64)) + self.register_buffer("columns", torch.tensor([[1, 3]], dtype=torch.int64)) + + def forward(self, x: torch.Tensor): + return x[:, :, self.rows, self.columns] + + class ConstantIndexTensor(torch.nn.Module): def __init__(self, indices: list[int]): super().__init__() @@ -100,8 +112,73 @@ def test_index_tensor_tosa_FP_int64_buffer_index(): pipeline.run() +def test_index_tensor_tosa_FP_leading_full_slice_int64_buffer_indices(): + pipeline = TosaPipelineFP[Tuple[torch.Tensor]]( + IndexTensorLeadingInt64Buffers(), + (torch.rand(1, 2, 4, 5),), + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + input_params_slice = Tuple[torch.Tensor, int, int, str, Tuple[torch.Tensor]] input_params = Tuple[torch.Tensor, Tuple[torch.Tensor]] +input_t2 = Tuple[torch.Tensor, torch.Tensor] +input_t3 = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] + + +class IndexTensorLeadingFullSlice(torch.nn.Module): + def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: + return x[:, index, :] + + +class IndexTensorLeadingFullSlicesNCHW(torch.nn.Module): + def forward( + self, x: torch.Tensor, rows: torch.Tensor, columns: torch.Tensor + ) -> torch.Tensor: + return x[:, :, rows, columns] + + +class IndexTensorLeadingFullSliceWithTrailingDim(torch.nn.Module): + def forward( + self, x: torch.Tensor, rows: torch.Tensor, columns: torch.Tensor + ) -> torch.Tensor: + return x[:, rows, columns, :] + + +leading_full_slice_test_data = { + "nchw_broadcast_indices": lambda: ( + IndexTensorLeadingFullSlicesNCHW(), + ( + torch.arange(2 * 3 * 4 * 5, dtype=torch.float32).reshape(2, 3, 4, 5), + torch.tensor([[0], [2]], dtype=torch.int32), + torch.tensor([[1, 3, 4]], dtype=torch.int32), + ), + ), + "leading_and_trailing_dims": lambda: ( + IndexTensorLeadingFullSliceWithTrailingDim(), + ( + torch.arange(2 * 4 * 5 * 3, dtype=torch.float32).reshape(2, 4, 5, 3), + torch.tensor([[0], [2]], dtype=torch.int32), + torch.tensor([[1, 3, 4]], dtype=torch.int32), + ), + ), +} + +zero_sized_test_data = { + "zero_sized_leading_dimension": ( + torch.empty(0, 3, 4), + torch.tensor([0, 2], dtype=torch.int32), + ), + "empty_index_tensor": ( + torch.rand(2, 3, 4), + torch.empty(0, dtype=torch.int32), + ), +} class IndexTensor_Ellipsis(torch.nn.Module): @@ -158,7 +235,6 @@ def forward( IndexTensor_Ellipsis.test_data_ellipsis, xfails={ # More info in index_tensor_support.py - "test_4d_ellipsis_before": "Ellipsis before index unsupported", "test_4d_ellipsis_middle": "Ellipsis before index unsupported", }, ) @@ -182,7 +258,6 @@ def test_index_tensor_tosa_FP_ellipsis(test_data: input_params): IndexTensor_Ellipsis.test_data_ellipsis, xfails={ # More info in index_tensor_support.py - "test_4d_ellipsis_before": "Ellipsis before index unsupported", "test_4d_ellipsis_middle": "Ellipsis before index unsupported", }, ) @@ -270,8 +345,6 @@ def forward( IndexTensor_Slice.test_data, xfails={ # More info in index_tensor_support.py - "test_4d_slice_before_1d_idx": "Slice before index unsupported", - "test_3d_slice_before_2d_idx": "Slice before index unsupported", "test_4d_slice_middle": "Slice before index unsupported", }, ) @@ -295,8 +368,6 @@ def test_index_tensor_tosa_FP_slice(test_data: input_params_slice): IndexTensor_Slice.test_data, xfails={ # More info in index_tensor_support.py - "test_4d_slice_before_1d_idx": "Slice before index unsupported", - "test_3d_slice_before_2d_idx": "Slice before index unsupported", "test_4d_slice_middle": "Slice before index unsupported", }, ) @@ -467,8 +538,7 @@ class IndexTensor(torch.nn.Module): ), } - # xfail - None (unsqueeze) unsupported - test_data_none: dict[input_params] = { + test_data_leading_none: dict[input_params] = { "test_3d_3_idx_with_none_before": ( torch.rand(12, 3, 7), ( @@ -484,6 +554,9 @@ class IndexTensor(torch.nn.Module): torch.randint(3, size=(12,), dtype=torch.int32), ), ), + } + + test_data_none: dict[input_params] = test_data_leading_none | { "test_3d_3_idx_with_none_around": ( torch.rand(12, 3, 7), ( @@ -573,25 +646,22 @@ def test_index_tensor_tosa_INT(test_data: input_params): IndexTensor.test_data_none, xfails={ # More info in index_tensor_support.py - "test_3d_3_idx_with_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_2_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_none_around": "None (Unsqueeze) unsupported", "test_3d_3_idx_with_none_middle": "None (Unsqueeze) unsupported", }, ) def test_index_tensor_tosa_FP_none(test_data: input_params): test_input = test_data with torch.no_grad(): - ( - TosaPipelineFP[input_params]( - IndexTensor(), - test_input, - IndexTensorTestCommon.aten_op, - IndexTensorTestCommon.exir_op, - atol=IndexTensorTestCommon.atol, - rtol=IndexTensorTestCommon.rtol, - ).run() + pipeline = TosaPipelineFP[input_params]( + IndexTensor(), + test_input, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() @common.parametrize( @@ -599,23 +669,100 @@ def test_index_tensor_tosa_FP_none(test_data: input_params): IndexTensor.test_data_none, xfails={ # More info in index_tensor_support.py - "test_3d_3_idx_with_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_2_none_before": "None (Unsqueeze) unsupported", - "test_3d_3_idx_with_none_around": "None (Unsqueeze) unsupported", "test_3d_3_idx_with_none_middle": "None (Unsqueeze) unsupported", }, ) def test_index_tensor_tosa_INT_none(test_data: input_params): test_input = test_data with torch.no_grad(): - ( - TosaPipelineINT[input_params]( - IndexTensor(), - test_input, - IndexTensorTestCommon.aten_op, - IndexTensorTestCommon.exir_op, - ).run() + pipeline = TosaPipelineINT[input_params]( + IndexTensor(), + test_input, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +def test_index_tensor_tosa_FP_leading_full_slices(test_data): + model, test_inputs = test_data() + pipeline = TosaPipelineFP[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +def test_index_tensor_tosa_INT_leading_full_slices(test_data): + model, test_inputs = test_data() + pipeline = TosaPipelineINT[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + ) + pipeline.count_tosa_ops({"GATHER": 1, "TRANSPOSE": 0}) + pipeline.run() + + +@common.parametrize("test_data", zero_sized_test_data) +def test_index_tensor_zero_sized_not_delegated_tosa_FP(test_data: input_t2): + OpNotSupportedPipeline[input_t2]( + IndexTensorLeadingFullSlice(), + test_data, + {IndexTensorTestCommon.exir_op: 1}, + ).run() + + +@common.parametrize("test_data", IndexTensor.test_data_leading_none) +@common.SkipIfNoModelConverter +def test_index_tensor_vgf_leading_full_slices(test_data: input_params): + pipeline = VgfPipeline[input_params]( + IndexTensor(), + test_data, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + quantize=False, + ) + pipeline.run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +@common.SkipIfNoModelConverter +def test_index_tensor_leading_full_slice_indexing_vgf_no_quant(test_data): + model, test_inputs = test_data() + VgfPipeline[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + atol=IndexTensorTestCommon.atol, + rtol=IndexTensorTestCommon.rtol, + quantize=False, + ).run() + + +@common.parametrize("test_data", leading_full_slice_test_data) +@common.SkipIfNoModelConverter +def test_index_tensor_leading_full_slice_indexing_vgf_quant(test_data): + model, test_inputs = test_data() + VgfPipeline[input_t3]( + model, + test_inputs, + IndexTensorTestCommon.aten_op, + IndexTensorTestCommon.exir_op, + quantize=True, + ).run() @common.parametrize("test_data", IndexTensor.test_data_int | IndexTensor.test_data_fp) diff --git a/docs/source/backends/arm-vgf/VGF_op_support.md b/docs/source/backends/arm-vgf/VGF_op_support.md index 092d9668419..f6ae2cc8780 100644 --- a/docs/source/backends/arm-vgf/VGF_op_support.md +++ b/docs/source/backends/arm-vgf/VGF_op_support.md @@ -148,7 +148,7 @@ Total supported PyTorch APIs: **155**. | `torch.t` / `torch.Tensor.t` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.tan` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.tanh` / `torch.nn.Tanh` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 | -| `torch.Tensor.__getitem__` / `tensor indexing` | FP, INT | `FP16`, `BF16`, `INT8` | 8x8 | +| `torch.Tensor.__getitem__` / `tensor indexing` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.Tensor.__getitem__` / `tensor slicing` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8` | 8x8 | | `torch.Tensor.__setitem__` / `tensor indexing assignment` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.Tensor.copy_` | FP, INT | `FP32`, `INT8` | 8x8 | From 8b7555e978115cc3766de2be6028850daf0dde4e Mon Sep 17 00:00:00 2001 From: Vaclav Novak Date: Fri, 11 Sep 2026 10:35:13 +0200 Subject: [PATCH 178/190] NXP backend: handle tests of mlperf tiny keyword spotting (#22519) ### Summary Enable and refactor MLPerf Tiny Keyword Spotting tests. ### Test plan tests can be manually run using `pytest -c /dev/null backends/nxp/tests/` cc @robert-kalmar @JakeStevens @digantdesai @rascani @roman-janik-nxp --- .../tests/generic_tests/test_aot_example.py | 86 ++++++++++- .../test_mlperf_tiny_image_classification.py | 7 +- .../test_mlperf_tiny_keyword_spotting.py | 140 ++++++++++++++++++ examples/nxp/aot_neutron_compile.py | 30 +++- .../mlperf_tiny_image_classification.py | 106 ++++--------- .../mlperf_tiny/keyword_spotting/__init__.py | 4 + .../mlperf_tiny_keyword_spotting.py | 60 ++++++++ .../models/mlperf_tiny/mlperf_tiny_model.py | 117 ++++++++++++++- 8 files changed, 453 insertions(+), 97 deletions(-) create mode 100644 backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py create mode 100644 examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py create mode 100644 examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py diff --git a/backends/nxp/tests/generic_tests/test_aot_example.py b/backends/nxp/tests/generic_tests/test_aot_example.py index 63e56ddd4bf..b75d3605d34 100644 --- a/backends/nxp/tests/generic_tests/test_aot_example.py +++ b/backends/nxp/tests/generic_tests/test_aot_example.py @@ -187,6 +187,9 @@ def test_aot_example__mlperf_tiny_ic(): """Test that the MLPerf Tiny image classification model (ResNet-8) can be lowered to Neutron backend via `aot_neutron_compile.py` and all ops are delegated.""" + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 + # Run the compilation script as a module (like run_aot_example.sh does). # The calibration data of this model is generated randomly, so no dataset download is needed. cmd = [ @@ -200,6 +203,8 @@ def test_aot_example__mlperf_tiny_ic(): "--target", "imxrt700", "--use_random_dataset", + "--num_random_samples", + str(num_random_samples), ] # Output file will be created in executorch_root @@ -216,7 +221,10 @@ def test_aot_example__mlperf_tiny_ic(): def test_aot_example__mlperf_tiny_ic__profiling(): """Test that the MLPerf Tiny image classification model (ResNet-8) can be lowered to Neutron backend via - `aot_neutron_compile.py` and all ops are delegated.""" + `aot_neutron_compile.py` and profiling works as intended.""" + + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 # Run the compilation script as a module (like run_aot_example.sh does) # Channels-last is buggy, so channels-first is used instead @@ -233,6 +241,8 @@ def test_aot_example__mlperf_tiny_ic__profiling(): "--remove-quant-io-ops", "--use_profiling", # Generate profilable model and create ETRecord "--use_random_dataset", # Avoid downloading the dataset. + "--num_random_samples", + str(num_random_samples), ] # Output files will be created in executorch_root. @@ -250,3 +260,77 @@ def test_aot_example__mlperf_tiny_ic__profiling(): with _cleanup_generated_files(pte_file, etrecord_file): result = _run_compile(cmd) _assert_profiling(result, pte_file, etrecord_file) + + +def test_aot_example__mlperf_tiny_kws(): + """Test that the MLPerf Tiny keyword spotting model (DS-CNN) can be lowered to Neutron backend via + `aot_neutron_compile.py` and all ops are delegated.""" + + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 + + # Run the compilation script as a module (like run_aot_example.sh does). + cmd = [ + sys.executable, + "-m", + "examples.nxp.aot_neutron_compile", + "--model_name", + "mlperf_tiny_keyword_spotting", + "--delegate", + "--quantize", + "--target", + "imxrt700", + "--use_random_dataset", + "--num_random_samples", + str(num_random_samples), + ] + + # Output file will be created in executorch_root + pte_file = Path( + os.path.join(EXECUTORCH_ROOT, "mlperf_tiny_keyword_spotting_nxp_delegate.pte") + ) + + with _cleanup_generated_files(pte_file): + result = _run_compile(cmd) + _assert_delegation(result, pte_file) + + +def test_aot_example__mlperf_tiny_kws__profiling(): + """Test that the MLPerf Tiny keyword spotting model (DS-CNN) can be lowered to Neutron backend via + `aot_neutron_compile.py` and profiling works as intended.""" + + # Number of random samples to generate, must be divisible by number of classes + num_random_samples = 60 + + # Run the compilation script as a module (like run_aot_example.sh does) + cmd = [ + sys.executable, + "-m", + "examples.nxp.aot_neutron_compile", + "--model_name", + "mlperf_tiny_keyword_spotting", + "--delegate", + "--quantize", + "--target", + "imxrt700", + "--remove-quant-io-ops", + "--use_profiling", # Generate profilable model and create ETRecord + "--use_random_dataset", + "--num_random_samples", + str(num_random_samples), + ] + + pte_file = Path( + os.path.join( + EXECUTORCH_ROOT, "mlperf_tiny_keyword_spotting_nxp_delegate_profile.pte" + ) + ) + etrecord_file = Path( + os.path.join( + EXECUTORCH_ROOT, "etrecord", "mlperf_tiny_keyword_spotting_etrecord.bin" + ) + ) + + with _cleanup_generated_files(pte_file, etrecord_file): + result = _run_compile(cmd) + _assert_profiling(result, pte_file, etrecord_file) diff --git a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py index 741b68aa1fd..fc0be12f500 100644 --- a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py +++ b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py @@ -67,9 +67,9 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( input_spec.dim_order = torch.channels_last quant_type_key = "QAT" if use_qat else "PTQ" - dim_order_key = "channels-last" if channels_last else "channels-first" + format_key = "channels-last" if channels_last else "channels-first" - mse = BOUNDS_MSE[quant_type_key][dim_order_key] + mse = BOUNDS_MSE[quant_type_key][format_key] comparator = NumericalStatsOutputComparator( max_mse_error=mse, use_softmax=True, is_classification_task=True ) @@ -80,7 +80,8 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( else None ) - # This model does not work in channels-last format and QAT. See more information below. + # This model does not work in channels-last format and QAT when running using portable kernels. + # See more information below. # Github issue: https://github.com/pytorch/executorch/issues/22179 # NXP internal issue ID: EIEX-1065 ref_model = ( diff --git a/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py new file mode 100644 index 00000000000..ab56849e970 --- /dev/null +++ b/backends/nxp/tests/models/test_mlperf_tiny_keyword_spotting.py @@ -0,0 +1,140 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from functools import partial + +import numpy as np +import pytest +import torch +from executorch.backends.nxp.tests.dataset_creator import ( + FromCalibrationDataDatasetCreator, +) +from executorch.backends.nxp.tests.executorch_pipeline import ModelInputSpec +from executorch.backends.nxp.tests.graph_verifier import BaseGraphVerifier +from executorch.backends.nxp.tests.model_output_comparator import ( + ClassificationAccuracyOutputComparator, + NumericalStatsOutputComparator, +) +from executorch.backends.nxp.tests.nsys_testing import ( + lower_run_compare, + lower_run_compare_ptq_qat, + ReferenceModel, +) +from executorch.backends.nxp.tests.use_qat import * # noqa F403 +from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import ( + MLPerfTinyKeywordSpotting, +) + +BOUNDS_MSE = { + "PTQ": { + "channels-last": np.inf, + "channels-first": 5.5e-7, + }, + "QAT": { + "channels-last": np.inf, + "channels-first": 3.3e-5, + }, +} + + +@pytest.fixture(autouse=True) +def reseed_model_per_test_run(): + torch.manual_seed(23) + np.random.seed(23) + + +@pytest.mark.parametrize( + "channels_last", + [ + False, + pytest.param( + True, + marks=pytest.mark.xfail( + reason="EIEX-1082, don't forget to readjust bounds when it start working", + strict=True, + ), + ), + ], +) +def test_mlperf_tiny_kws_mse_cpu_vs_npu(mocker, request, channels_last, use_qat): + # approx. 5 samples per class + num_samples = 60 + + kws = MLPerfTinyKeywordSpotting(num_samples=num_samples, use_random_dataset=True) + model = kws.get_eager_model() + dataset = kws.dataset + labels = kws.labels + + dataset_creator = FromCalibrationDataDatasetCreator( + dataset, num_examples=num_samples, idx_to_label=labels + ) + + input_spec = ModelInputSpec(kws.input_shape) + if channels_last: + model.to(memory_format=torch.channels_last) + input_spec.dim_order = torch.channels_last + + quant_type_key = "QAT" if use_qat else "PTQ" + format_key = "channels-last" if channels_last else "channels-first" + mse = BOUNDS_MSE[quant_type_key][format_key] + comparator = NumericalStatsOutputComparator( + max_mse_error=mse, is_classification_task=True + ) + model_verifier = BaseGraphVerifier(1, []) + train_fn = ( + partial(kws.train_model_fn, channels_last=channels_last) if use_qat else None + ) + + # This model does not work in channels-last format when running with portable kernels. + # See more information below. + # Github issue: https://github.com/pytorch/executorch/issues/22520 + # NXP internal issue ID: EIEX-1074 + ref_model = ( + ReferenceModel.QUANTIZED_EDGE_PYTHON + if channels_last + else ReferenceModel.QUANTIZED_EXECUTORCH_CPP + ) + + lower_run_compare( + model, + [input_spec], + model_verifier, + request, + dataset_creator=dataset_creator, + output_comparator=comparator, + mocker=mocker, + reference_model=ref_model, + use_qat=use_qat, + train_fn=train_fn, + ) + + +def test_mlperf_tiny_kws_ptq_qat_equivalence(request): + # approx. 5 samples per class + num_samples = 60 + + kws = MLPerfTinyKeywordSpotting(num_samples=num_samples, use_random_dataset=True) + + model = kws.get_eager_model() + dataset = kws.dataset + labels = kws.labels + + dataset_creator = FromCalibrationDataDatasetCreator( + dataset, num_examples=num_samples, idx_to_label=labels + ) + comparator = ClassificationAccuracyOutputComparator(class_dict=labels) + + input_spec = ModelInputSpec(kws.input_shape) + model_verifier = BaseGraphVerifier(1, []) + + lower_run_compare_ptq_qat( + model, + [input_spec], + model_verifier, + request, + train_fn=kws.train_model_fn, + dataset_creator=dataset_creator, + output_comparator=comparator, + ) diff --git a/examples/nxp/aot_neutron_compile.py b/examples/nxp/aot_neutron_compile.py index 1efdf03a0d6..b9e3298c26d 100644 --- a/examples/nxp/aot_neutron_compile.py +++ b/examples/nxp/aot_neutron_compile.py @@ -46,6 +46,9 @@ from executorch.examples.nxp.models.mlperf_tiny.image_classification.mlperf_tiny_image_classification import ( MLPerfTinyImageClassification, ) +from executorch.examples.nxp.models.mlperf_tiny.keyword_spotting.mlperf_tiny_keyword_spotting import ( + MLPerfTinyKeywordSpotting, +) from executorch.examples.nxp.models.mobilenet_v2 import MobilenetV2 from executorch.exir import ( EdgeCompileConfig, @@ -65,6 +68,7 @@ "cifar10": CifarNet, "mobilenetv2": MobilenetV2, "mlperf_tiny_image_classification": MLPerfTinyImageClassification, + "mlperf_tiny_keyword_spotting": MLPerfTinyKeywordSpotting, } FORMAT = "[%(levelname)s %(asctime)s %(filename)s:%(lineno)s] %(message)s" @@ -100,7 +104,10 @@ def _print_ops_in_edge_program(edge_program): def _get_model_info_from_name( - model_name: str, dataset_path: str | None, use_random_dataset: bool + model_name: str, + dataset_path: str | None, + use_random_dataset: bool, + num_samples: int | None, ): """Given the name of an example pytorch model and args, return the model, its class instance (can be None), example inputs and calibration inputs (can be None). @@ -119,9 +126,11 @@ def _get_model_info_from_name( ) model_cls_inst = model_cls() - elif model_cls is MLPerfTinyImageClassification: + elif model_cls in (MLPerfTinyImageClassification, MLPerfTinyKeywordSpotting): model_cls_inst = model_cls( - dataset_path=dataset_path, use_random_dataset=use_random_dataset + dataset_path=dataset_path, + use_random_dataset=use_random_dataset, + num_samples=num_samples, ) else: @@ -257,6 +266,13 @@ def _get_arg_parser(): action="store_true", help="The calibration and testing datasets will be generated randomly instead of being downloaded.", ) + parser.add_argument( + "--num_random_samples", + required=False, + default=None, + type=int, + help="Number of random samples to generate, required when `--use_random_dataset` flag is set.", + ) parser.add_argument( "-dst", "--dataset_path", @@ -286,7 +302,10 @@ def _get_arg_parser(): # 1. pick model from one of the supported lists model, example_inputs, calibration_inputs, model_cls_inst = ( _get_model_info_from_name( - args.model_name, args.dataset_path, args.use_random_dataset + args.model_name, + args.dataset_path, + args.use_random_dataset, + args.num_random_samples, ) ) model = model.eval() @@ -317,7 +336,8 @@ def _get_arg_parser(): quantizer = NeutronQuantizer(neutron_target_spec, is_qat=args.use_qat) if args.use_qat: if not isinstance( - model_cls_inst, (CifarNet, MLPerfTinyImageClassification) + model_cls_inst, + (CifarNet, MLPerfTinyImageClassification, MLPerfTinyKeywordSpotting), ): raise ValueError( f"QAT training is not supported for model '{args.model_name}'" diff --git a/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py b/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py index 60b590fa9a9..c8fdcd6226d 100644 --- a/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py +++ b/examples/nxp/models/mlperf_tiny/image_classification/mlperf_tiny_image_classification.py @@ -4,105 +4,49 @@ # LICENSE file in the root directory of this source tree. import logging -from pathlib import Path import torch -from executorch.backends.nxp.tests.calibration_dataset import ( - CalibrationDataset, - RandomCalibrationDataset, -) - from executorch.examples.models.mlperf_tiny import ResNet8 from executorch.examples.nxp.models.mlperf_tiny.mlperf_tiny_model import MLPerfTinyModel -from torch.utils.data import Dataset -from torchao.quantization.pt2e import disable_observer -from tqdm import tqdm log = logging.getLogger(__name__) -INPUT_SHAPE = (1, 3, 32, 32) -IDX_TO_LABEL = { - 0: "airplane", - 1: "automobile", - 2: "bird", - 3: "cat", - 4: "deer", - 5: "dog", - 6: "frog", - 7: "horse", - 8: "ship", - 9: "truck", -} - class MLPerfTinyImageClassification(MLPerfTinyModel): """MLPerf Tiny image classification model (ResNet-8).""" - def __init__( - self, - num_samples: int = 200, - dataset_path: Path | str | None = None, - use_random_dataset: bool = False, - ): - self._num_samples = num_samples - self._use_random_dataset = use_random_dataset - self._dataset_path = dataset_path - - super().__init__() + # ResNet-8 specific QAT training hyperparameters. + TRAIN_HYPERPARAMETERS = { + "num_epochs": 15, + "batch_size": 20, + "lr": 1e-5, + "eps": 1e-8, + "weight_decay": 1e-4, + } + + INPUT_SHAPE = (1, 3, 32, 32) + IDX_TO_LABEL = { + 0: "airplane", + 1: "automobile", + 2: "bird", + 3: "cat", + 4: "deer", + 5: "dog", + 6: "frog", + 7: "horse", + 8: "ship", + 9: "truck", + } @property def input_shape(self): - return INPUT_SHAPE + return self.INPUT_SHAPE @property def labels(self): - return IDX_TO_LABEL - - def _init_dataset(self) -> Dataset: - if self._use_random_dataset: - num_classes = len(self.labels) - sample_shape = tuple(self.input_shape)[1:] - return RandomCalibrationDataset( - self._num_samples, sample_shape, num_classes - ) - else: - if self._dataset_path is None: - raise ValueError( - "Path to dataset data cannot be empty. If you want to use random data, set `use_random_dataset = True`" - ) - return CalibrationDataset(self._dataset_path) + return self.IDX_TO_LABEL def _init_eager_model(self) -> torch.nn.Module: num_classes = len(self.labels) - return ResNet8(num_classes) - - def train_model_fn(self, model, num_epochs=15, batch_size=20, channels_last=False): - torch.manual_seed(42) - torch.use_deterministic_algorithms(True) - - optimizer = torch.optim.Adam( - params=model.parameters(), - lr=1e-5, - weight_decay=1e-4, - ) - loss_fn = torch.nn.CrossEntropyLoss() - - logging.warning("Starting training...") - - data = self.get_qat_train_inputs(batch_size=batch_size) - for nepoch in range(num_epochs): - for images, labels in tqdm(data): - if channels_last: - images = images.to(memory_format=torch.channels_last) - - optimizer.zero_grad() - outputs = model(images) - loss = loss_fn(outputs, labels) - loss.backward() - optimizer.step() - - if nepoch >= num_epochs / 3: - model.apply(disable_observer) - - return model + return ResNet8(num_classes).eval() diff --git a/examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py b/examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py new file mode 100644 index 00000000000..55dc5fccf45 --- /dev/null +++ b/examples/nxp/models/mlperf_tiny/keyword_spotting/__init__.py @@ -0,0 +1,4 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. diff --git a/examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py b/examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py new file mode 100644 index 00000000000..3bc16a5283d --- /dev/null +++ b/examples/nxp/models/mlperf_tiny/keyword_spotting/mlperf_tiny_keyword_spotting.py @@ -0,0 +1,60 @@ +# Copyright 2026 NXP +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import logging + +import torch + +from executorch.examples.models.mlperf_tiny import DSCNNKWS +from executorch.examples.nxp.models.mlperf_tiny.mlperf_tiny_model import MLPerfTinyModel + +log = logging.getLogger(__name__) + + +class MLPerfTinyKeywordSpotting(MLPerfTinyModel): + """MLPerf Tiny keyword spotting model (DS-CNN).""" + + INPUT_SHAPE = (1, 1, 49, 10) + # Because of the architecture of the model, + # non-scaled random weights tend to produce zero tensors, + # making it hard to compute numerical accuracy of the delegated model. + # Scaling the random weights makes the model produce reasonable results. + WEIGHT_INIT_SCALE = 2.0 + + IDX_TO_LABEL = { + 0: "Down", + 1: "Go", + 2: "Left", + 3: "No", + 4: "Off", + 5: "On", + 6: "Right", + 7: "Stop", + 8: "Up", + 9: "Yes", + 10: "Silence", + 11: "Unknown", + } + + @property + def input_shape(self): + return self.INPUT_SHAPE + + @property + def labels(self): + return self.IDX_TO_LABEL + + def _init_weights(self, model: torch.nn.Module): + with torch.no_grad(): + for module in model.modules(): + if isinstance(module, (torch.nn.Conv2d, torch.nn.Linear)): + module.weight *= self.WEIGHT_INIT_SCALE + + def _init_eager_model(self) -> torch.nn.Module: + num_classes = len(self.labels) + model = DSCNNKWS(num_classes) + self._init_weights(model) + + return model.eval() diff --git a/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py b/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py index d72a2ad273b..b33a8b02057 100644 --- a/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py +++ b/examples/nxp/models/mlperf_tiny/mlperf_tiny_model.py @@ -4,20 +4,56 @@ # LICENSE file in the root directory of this source tree. import itertools +import logging +import os from abc import abstractmethod +from pathlib import Path from typing import Iterator import torch +from executorch.backends.nxp.tests.calibration_dataset import ( + CalibrationDataset, + RandomCalibrationDataset, +) from executorch.examples.models import model_base from torch.utils.data import DataLoader, Dataset +from torchao.quantization.pt2e import disable_observer +from tqdm import tqdm + +log = logging.getLogger(__name__) class MLPerfTinyModel(model_base.EagerModelBase): """Base class of the MLPerf Tiny models.""" - def __init__(self): - """Create the model wrapper along with the dataset it owns.""" - self._num_workers = 4 + # Default QAT training hyperparameters. Subclasses may override them. + TRAIN_HYPERPARAMETERS = { + "num_epochs": 15, + "batch_size": 64, + "lr": 5e-6, + "eps": 1e-7, + "weight_decay": 1e-4, + } + + def __init__( + self, + dataset_path: Path | str | None = None, + use_random_dataset: bool = False, + num_samples: int | None = None, + num_workers: int = 4, + ): + """ + Create the model wrapper along with the dataset it owns. + If `use_random_dataset = True`, then `num_samples` must be set. + If `use_random_dataset = False`, then `dataset_path` must be set. + """ + self._num_workers = num_workers + self._num_samples = num_samples + self._use_random_dataset = use_random_dataset + self._dataset_path = dataset_path + + # throws ValueError if validation fails + self._validate_args() self._eager_model = self._init_eager_model() self.dataset = self._init_dataset() @@ -27,10 +63,6 @@ def _collate_fn(data: list[tuple]): data, labels = zip(*data) return torch.stack(list(data)), torch.tensor(list(labels)) - @abstractmethod - def _init_dataset(self) -> Dataset: - pass - @abstractmethod def _init_eager_model(self) -> torch.nn.Module: pass @@ -45,6 +77,26 @@ def input_shape(self): def labels(self): pass + def _validate_args(self): + valid_num_samples = isinstance(self._num_samples, int) and self._num_samples > 0 + + if self._use_random_dataset: + if not valid_num_samples: + raise ValueError( + f"Invalid number of samples to randomly generate. Got {self._num_samples}." + ) + + else: + if valid_num_samples: + raise ValueError( + "Num samples was supplied, but it is omitted because `use_random_dataset=False`." + ) + + if self._dataset_path is None or not os.path.exists(self._dataset_path): + raise ValueError( + f"Invalid dataset path for loading the data. Got {self._dataset_path}." + ) + def get_qat_train_inputs( self, batch_size: int = 5, dataset_portion: float = 0.1 ) -> Iterator[tuple[torch.Tensor]]: @@ -79,3 +131,54 @@ def get_eager_model(self): def get_example_inputs(self) -> tuple[torch.Tensor]: return (torch.randn(self.input_shape, dtype=torch.float32),) + + def train_model_fn( + self, model, num_epochs=None, batch_size=None, channels_last=False + ): + hyperparameters = self.TRAIN_HYPERPARAMETERS + num_epochs = ( + num_epochs if num_epochs is not None else hyperparameters["num_epochs"] + ) + batch_size = ( + batch_size if batch_size is not None else hyperparameters["batch_size"] + ) + + torch.manual_seed(42) + torch.use_deterministic_algorithms(True) + + optimizer = torch.optim.Adam( + params=model.parameters(), + lr=hyperparameters["lr"], + eps=hyperparameters["eps"], + weight_decay=hyperparameters["weight_decay"], + ) + loss_fn = torch.nn.CrossEntropyLoss() + + log.warning("Starting training...") + + data = self.get_qat_train_inputs(batch_size=batch_size) + for nepoch in range(num_epochs): + for samples, labels in tqdm(data): + if channels_last: + samples = samples.to(memory_format=torch.channels_last) + + optimizer.zero_grad() + outputs = model(samples) + loss = loss_fn(outputs, labels) + loss.backward() + optimizer.step() + + if nepoch >= num_epochs / 3: + model.apply(disable_observer) + + return model + + def _init_dataset(self) -> Dataset: + if self._use_random_dataset: + num_classes = len(self.labels) + sample_shape = tuple(self.input_shape)[1:] + return RandomCalibrationDataset( + self._num_samples, sample_shape, num_classes + ) + else: + return CalibrationDataset(self._dataset_path) From acd1a17b0f23615d1687f4272756f0f5786ee800 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson Date: Mon, 24 Aug 2026 08:50:10 +0200 Subject: [PATCH 179/190] Arm backend: Support isinf and isnan with FP decomposition Lower isinf and isnan to TOSA-supported comparisons for floating-point profiles. Keep quantized floating-point inputs on CPU because integer quantization cannot preserve NaN or infinity. Change-Id: Ib53653b0ebb1c5497a7db3932945a1b9c5060f0f Signed-off-by: Sebastian Larsson --- backends/arm/_passes/__init__.py | 1 + backends/arm/_passes/arm_pass_manager.py | 2 + .../arm/_passes/decompose_isinf_isnan_pass.py | 48 +++++++++++ .../tosa_profile_supported_op_lists.py | 2 + .../tosa_supported_operators.py | 10 ++- backends/arm/test/ops/test_isinf.py | 79 +++++++++++++++++++ backends/arm/test/ops/test_isnan.py | 77 ++++++++++++++++++ .../source/backends/arm-vgf/VGF_op_support.md | 4 +- 8 files changed, 220 insertions(+), 3 deletions(-) create mode 100644 backends/arm/_passes/decompose_isinf_isnan_pass.py create mode 100644 backends/arm/test/ops/test_isinf.py create mode 100644 backends/arm/test/ops/test_isnan.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index cdd465dd1a5..ff8cf87679d 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -74,6 +74,7 @@ DecomposeIndexTensorToGatherPass, ) from .decompose_int_pow_pass import DecomposeIntPowPass # noqa +from .decompose_isinf_isnan_pass import DecomposeIsInfAndIsNanPass # noqa from .decompose_large_stride_maxpool2d_pass import ( # noqa DecomposeLargeStrideMaxPool2dForU55Pass, ) diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 019093fb1b4..3145dc686b0 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -69,6 +69,7 @@ DecomposeIndexSelectToGatherPass, DecomposeIndexTensorToGatherPass, DecomposeIntPowPass, + DecomposeIsInfAndIsNanPass, DecomposeLargeStrideMaxPool2dForU55Pass, DecomposeLayerNormPass, DecomposeLeakyReLUPass, @@ -576,6 +577,7 @@ def _tosa_pipeline( RemoveGetItemPass(), FuseBatchNorm2dPass(exported_program), DecomposeBatchNormNoStatsPass(), + DecomposeIsInfAndIsNanPass(), DecomposeLogitPass(), DecomposeMaskedFillPass(), DecomposeRoundPass(), diff --git a/backends/arm/_passes/decompose_isinf_isnan_pass.py b/backends/arm/_passes/decompose_isinf_isnan_pass.py new file mode 100644 index 00000000000..7722e41ea40 --- /dev/null +++ b/backends/arm/_passes/decompose_isinf_isnan_pass.py @@ -0,0 +1,48 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from typing import Set, Type + +from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +class DecomposeIsInfAndIsNanPass(ArmOpTargetedPass): + """Decompose ``isinf`` and ``isnan`` into TOSA-supported operations.""" + + _passes_required_after: Set[Type[ExportPass]] = set() + edge_isinf = exir_ops.edge.aten.isinf.default + edge_isnan = exir_ops.edge.aten.isnan.default + target_ops = (edge_isinf, edge_isnan) + check_allowed_to_transform = True + + def call_operator(self, op, args, kwargs, meta): + if op not in self.target_ops or not self.allowed_to_transform(meta): + return super().call_operator(op, args, kwargs, meta) + + (x,) = args + abs_op = exir_ops.edge.aten.abs.default + eq_op = exir_ops.edge.aten.eq.Tensor + logical_not_op = exir_ops.edge.aten.logical_not.default + full_op = exir_ops.edge.aten.full.default + + if op is self.edge_isnan: + equal = super().call_operator(eq_op, (x, x), {}, meta, updated=True) + return super().call_operator( + logical_not_op, (equal,), {}, meta, updated=True + ) + + absolute = super().call_operator(abs_op, (x,), {}, meta, updated=True) + infinity = super().call_operator( + full_op, + (x.data.shape, float("inf")), + {"dtype": x.data.dtype}, + meta, + updated=True, + ) + return super().call_operator( + eq_op, (absolute, infinity), {}, meta, updated=True + ) diff --git a/backends/arm/operator_support/tosa_profile_supported_op_lists.py b/backends/arm/operator_support/tosa_profile_supported_op_lists.py index 78e9c617b13..a9a7c2413b0 100644 --- a/backends/arm/operator_support/tosa_profile_supported_op_lists.py +++ b/backends/arm/operator_support/tosa_profile_supported_op_lists.py @@ -186,6 +186,8 @@ exir_ops.edge.aten.expm1.default, exir_ops.edge.aten.log1p.default, exir_ops.edge.aten.log.default, + exir_ops.edge.aten.isnan.default, + exir_ops.edge.aten.isinf.default, exir_ops.edge.aten.linear.default, exir_ops.edge.aten.split_with_sizes_copy.default, exir_ops.edge.aten.split_copy.Tensor, diff --git a/backends/arm/operator_support/tosa_supported_operators.py b/backends/arm/operator_support/tosa_supported_operators.py index 2c1ae365a8b..868a4b1d23f 100644 --- a/backends/arm/operator_support/tosa_supported_operators.py +++ b/backends/arm/operator_support/tosa_supported_operators.py @@ -1153,7 +1153,7 @@ def is_node_supported( class CheckFPComparisonInputs(OperatorSupportBase): """Reject unsupported comparison inputs under the FP profile.""" - target_ops = { + comparison_ops = { exir_ops.edge.aten.eq.Tensor, exir_ops.edge.aten.eq.Scalar, exir_ops.edge.aten.ne.Tensor, @@ -1167,6 +1167,10 @@ class CheckFPComparisonInputs(OperatorSupportBase): exir_ops.edge.aten.lt.Tensor, exir_ops.edge.aten.lt.Scalar, } + target_ops = comparison_ops | { + exir_ops.edge.aten.isinf.default, + exir_ops.edge.aten.isnan.default, + } supported_dtypes = {torch.float16, torch.float32, torch.bfloat16} castable_comparison_dtypes = {torch.int8, torch.int16} @@ -1188,7 +1192,9 @@ def is_node_supported( if all(dtype in self.supported_dtypes for dtype in input_dtypes): return True - if all(dtype in self.castable_comparison_dtypes for dtype in input_dtypes): + if node.target in self.comparison_ops and all( + dtype in self.castable_comparison_dtypes for dtype in input_dtypes + ): return True unsupported_dtype = next( diff --git a/backends/arm/test/ops/test_isinf.py b/backends/arm/test/ops/test_isinf.py new file mode 100644 index 00000000000..4ed3ef6bfb7 --- /dev/null +++ b/backends/arm/test/ops/test_isinf.py @@ -0,0 +1,79 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from collections.abc import Callable + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + OpNotSupportedPipeline, + TosaPipelineFP, + VgfPipeline, +) + +aten_op = "torch.ops.aten.isinf.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_isinf_default" + + +class IsInf(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.isinf(x) + + +test_data_suite = { + "finite": lambda: torch.tensor([-1.0, 0.0, 3.14]), + "inf": lambda: torch.tensor([-float("inf"), 0.0, float("inf"), float("nan")]), + "integer": lambda: torch.tensor([-5, 0, 9], dtype=torch.int32), + "rank4": lambda: torch.tensor( + [[[[float("inf"), 0.0]]], [[[-float("inf"), float("nan")]]]] + ), +} + + +@common.parametrize( + "test_data", + {name: data for name, data in test_data_suite.items() if name != "integer"}, +) +def test_isinf_tosa_FP(test_data: Callable[[], torch.Tensor]) -> None: + TosaPipelineFP( + IsInf(), + (test_data(),), + aten_op, + exir_op, + ).run() + + +def test_isinf_tosa_FP_falls_back_for_integer() -> None: + OpNotSupportedPipeline( + IsInf(), + (test_data_suite["integer"](),), + {exir_op: 1}, + quantize=False, + ).run() + + +def test_isinf_tosa_INT_falls_back() -> None: + test_data = (test_data_suite["inf"](),) + pipeline = OpNotSupportedPipeline( + IsInf(), + test_data, + {exir_op: 1}, + quantize=True, + ) + quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0] + quantize_stage.calibration_samples = [(torch.ones_like(test_data[0]),)] + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_isinf_vgf_no_quant() -> None: + VgfPipeline( + IsInf(), + (test_data_suite["inf"](),), + aten_op, + exir_op, + quantize=False, + ).run() diff --git a/backends/arm/test/ops/test_isnan.py b/backends/arm/test/ops/test_isnan.py new file mode 100644 index 00000000000..0c6476c01ab --- /dev/null +++ b/backends/arm/test/ops/test_isnan.py @@ -0,0 +1,77 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from collections.abc import Callable + +import torch + +from executorch.backends.arm.test import common +from executorch.backends.arm.test.tester.test_pipeline import ( + OpNotSupportedPipeline, + TosaPipelineFP, + VgfPipeline, +) + +aten_op = "torch.ops.aten.isnan.default" +exir_op = "executorch_exir_dialects_edge__ops_aten_isnan_default" + + +class IsNan(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.isnan(x) + + +test_data_suite = { + "finite": lambda: torch.tensor([-1.0, 0.0, 3.14]), + "nan": lambda: torch.tensor([float("nan"), 0.0, float("inf")]), + "integer": lambda: torch.tensor([-5, 0, 9], dtype=torch.int32), + "rank4": lambda: torch.tensor([[[[float("nan"), 0.0]]], [[[float("inf"), -3.0]]]]), +} + + +@common.parametrize( + "test_data", + {name: data for name, data in test_data_suite.items() if name != "integer"}, +) +def test_isnan_tosa_FP(test_data: Callable[[], torch.Tensor]) -> None: + TosaPipelineFP( + IsNan(), + (test_data(),), + aten_op, + exir_op, + ).run() + + +def test_isnan_tosa_FP_falls_back_for_integer() -> None: + OpNotSupportedPipeline( + IsNan(), + (test_data_suite["integer"](),), + {exir_op: 1}, + quantize=False, + ).run() + + +def test_isnan_tosa_INT_falls_back() -> None: + test_data = (test_data_suite["nan"](),) + pipeline = OpNotSupportedPipeline( + IsNan(), + test_data, + {exir_op: 1}, + quantize=True, + ) + quantize_stage = pipeline._stages[pipeline.find_pos("quantize")].args[0] + quantize_stage.calibration_samples = [(torch.ones_like(test_data[0]),)] + pipeline.run() + + +@common.SkipIfNoModelConverter +def test_isnan_vgf_no_quant() -> None: + VgfPipeline( + IsNan(), + (test_data_suite["nan"](),), + aten_op, + exir_op, + quantize=False, + ).run() diff --git a/docs/source/backends/arm-vgf/VGF_op_support.md b/docs/source/backends/arm-vgf/VGF_op_support.md index f6ae2cc8780..98a3623244f 100644 --- a/docs/source/backends/arm-vgf/VGF_op_support.md +++ b/docs/source/backends/arm-vgf/VGF_op_support.md @@ -6,7 +6,7 @@ This page lists VGF-supported PyTorch APIs and the dtype and quantization modes `8x8` means 8-bit activations and 8-bit weights. `16x8` means 16-bit activations and 8-bit weights. `8x4` means 8-bit activations and 4-bit weights. -Total supported PyTorch APIs: **155**. +Total supported PyTorch APIs: **157**. | PyTorch API | Support profile | DType | Quantization mode | | --- | --- | --- | --- | @@ -69,6 +69,8 @@ Total supported PyTorch APIs: **155**. | `torch.gt` / `>` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.index_put_` | INT | `INT8` | 8x8 | | `torch.index_select` | FP, INT | `FP32`, `FP16`, `BF16`, `INT8`, `BOOL` | 8x8 | +| `torch.isinf` | FP | `FP32` | - | +| `torch.isnan` | FP | `FP32` | - | | `torch.layer_norm` | FP, INT | `FP32`, `INT8` | 8x8 | | `torch.le` / `<=` | FP, INT | `FP32`, `INT8`, `INT16` | 8x8, 16x8 | | `torch.linspace` | FP, INT | `FP32`, `INT8` | 8x8 | From a9d3407e47aeed37874f45d91331a35833efb336 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson Date: Tue, 8 Sep 2026 13:10:51 +0200 Subject: [PATCH 180/190] Arm backend: Register isinf and isnan tests with Buck2 Change-Id: I3eaa54c7b794f583214bb252015e38960ae9be45 Signed-off-by: Sebastian Larsson --- backends/arm/test/targets.bzl | 2 ++ 1 file changed, 2 insertions(+) diff --git a/backends/arm/test/targets.bzl b/backends/arm/test/targets.bzl index 2e5d4e1f3b3..a38743526cc 100644 --- a/backends/arm/test/targets.bzl +++ b/backends/arm/test/targets.bzl @@ -19,6 +19,8 @@ def define_arm_tests(): "ops/test_avg_pool2d.py", "ops/test_cat.py", "ops/test_conv2d.py", + "ops/test_isinf.py", + "ops/test_isnan.py", "ops/test_linear.py", "ops/test_log10.py", "ops/test_max_pool1d.py", From 5ac62a4962a04e6b6943944a5a5a2a43749240b2 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson Date: Thu, 10 Sep 2026 15:10:52 +0200 Subject: [PATCH 181/190] Arm backend: Update SD3.5 T5 expectations for isinf delegation Expect no CPU isinf calls after FP and VGF partitioning. Account for four fewer delegate calls now that isinf is lowered inside delegates. Keep the TOSA INT fallback expectations unchanged. Signed-off-by: Sebastian Larsson Change-Id: Ib52c3a9e8bf55d9b8b1ef600a7e7d0fc694e34a7 --- .../test_T5EncoderModel_sd35_large.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py b/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py index 2d850676d0f..c6e9d268248 100644 --- a/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py +++ b/backends/arm/test/models/stable_diffusion_3_5_large/test_T5EncoderModel_sd35_large.py @@ -29,10 +29,9 @@ class TestT5EncoderModel: ops_after_partitioner_FP = { "executorch_exir_dialects_edge__ops_aten_clamp_Tensor": 4, - "executorch_exir_dialects_edge__ops_aten_isinf_default": 4, "executorch_exir_dialects_edge__ops_aten_where_self": 1, "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2, - "torch.ops.higher_order.executorch_call_delegate": 10, + "torch.ops.higher_order.executorch_call_delegate": 6, } ops_after_partitioner_INT = { @@ -48,9 +47,8 @@ class TestT5EncoderModel: ops_after_partitioner_vgf_quantize = { "executorch_exir_dialects_edge__ops_aten_clamp_Tensor": 4, - "executorch_exir_dialects_edge__ops_aten_isinf_default": 4, "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 1, - "torch.ops.higher_order.executorch_call_delegate": 9, + "torch.ops.higher_order.executorch_call_delegate": 5, } ops_after_partitioner_vgf_no_quantize = ops_after_partitioner_vgf_quantize From dfeac513bf583a96484aac28bd4c9b979bad8da3 Mon Sep 17 00:00:00 2001 From: Sebastian Larsson <38941629+Sebastian-Larsson@users.noreply.github.com> Date: Fri, 11 Sep 2026 13:06:13 +0200 Subject: [PATCH 182/190] Arm backend: Add static public API manifest for 1.5 (#22731) Signed-off-by: Sebastian Larsson --- .../api_manifest_1_5.toml | 287 ++++++++++++++++++ 1 file changed, 287 insertions(+) create mode 100644 backends/arm/public_api_manifests/api_manifest_1_5.toml diff --git a/backends/arm/public_api_manifests/api_manifest_1_5.toml b/backends/arm/public_api_manifests/api_manifest_1_5.toml new file mode 100644 index 00000000000..dc87ac9ee96 --- /dev/null +++ b/backends/arm/public_api_manifests/api_manifest_1_5.toml @@ -0,0 +1,287 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +# +# This file is generated by +# backends/arm/scripts/public_api_manifest/generate_public_api_manifest.py + +[python] + +[python.EthosUBackend] +kind = "class" +signature = "EthosUBackend()" + +[python.EthosUBackend.preprocess] +kind = "function" +signature = "EthosUBackend.preprocess(edge_program: torch.export.exported_program.ExportedProgram, compile_specs: List[executorch.exir.backend.compile_spec_schema.CompileSpec]) -> executorch.exir.backend.backend_details.PreprocessResult" + +[python.EthosUCompileSpec] +kind = "class" +signature = "EthosUCompileSpec(target: str, system_config: str | None = None, memory_mode: str | None = None, extra_flags: list[str] | None = None, config_ini: str | None = 'Arm/vela.ini', external_block_placements: executorch.backends.arm.ethosu.compile_spec.VelaExternalBlockPlacements | None = None)" + +[python.EthosUCompileSpec.DebugMode] +kind = "enum" +signature = "EthosUCompileSpec.DebugMode(*values)" + +[python.EthosUCompileSpec.__eq__] +kind = "function" +signature = "EthosUCompileSpec.__eq__(self, other)" + +[python.EthosUCompileSpec.__repr__] +kind = "function" +signature = "EthosUCompileSpec.__repr__(self)" + +[python.EthosUCompileSpec.dump_debug_info] +kind = "function" +signature = "EthosUCompileSpec.dump_debug_info(self, debug_mode: executorch.backends.arm.common.arm_compile_spec.ArmCompileSpec.DebugMode | None)" + +[python.EthosUCompileSpec.dump_intermediate_artifacts_to] +kind = "function" +signature = "EthosUCompileSpec.dump_intermediate_artifacts_to(self, output_path: str | None)" + +[python.EthosUCompileSpec.set_pass_pipeline_config] +kind = "function" +signature = "EthosUCompileSpec.set_pass_pipeline_config(self, config: executorch.backends.arm.common.pipeline_config.ArmPassPipelineConfig) -> None" + +[python.EthosUPartitioner] +kind = "class" +signature = "EthosUPartitioner(compile_spec: executorch.backends.arm.ethosu.compile_spec.EthosUCompileSpec, additional_checks: Optional[Sequence[torch.fx.passes.operator_support.OperatorSupportBase]] = None) -> None" + +[python.EthosUPartitioner.ops_to_not_decompose] +kind = "function" +signature = "EthosUPartitioner.ops_to_not_decompose(self, ep: torch.export.exported_program.ExportedProgram) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.node.Node], bool]]]" + +[python.EthosUPartitioner.partition] +kind = "function" +signature = "EthosUPartitioner.partition(self, exported_program: torch.export.exported_program.ExportedProgram) -> executorch.exir.backend.partitioner.PartitionResult" + +[python.EthosUPartitioner.register_custom_partition_op] +kind = "function" +signature = "EthosUPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" + +[python.EthosUPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "EthosUPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + +[python.EthosUQuantizer] +kind = "class" +signature = "EthosUQuantizer(compile_spec: 'EthosUCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" + +[python.EthosUQuantizer.annotate] +kind = "function" +signature = "EthosUQuantizer.annotate(self, model: 'GraphModule') -> 'GraphModule'" + +[python.EthosUQuantizer.set_global] +kind = "function" +signature = "EthosUQuantizer.set_global(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.set_io] +kind = "function" +signature = "EthosUQuantizer.set_io(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.set_module_name] +kind = "function" +signature = "EthosUQuantizer.set_module_name(self, module_name: 'str', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.set_module_type] +kind = "function" +signature = "EthosUQuantizer.set_module_type(self, module_type: 'Callable', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.EthosUQuantizer.transform_for_annotation] +kind = "function" +signature = "EthosUQuantizer.transform_for_annotation(self, model: 'GraphModule') -> 'GraphModule'" + +[python.EthosUQuantizer.validate] +kind = "function" +signature = "EthosUQuantizer.validate(self, model: 'GraphModule') -> 'None'" + +[python.VelaExternalBlockPlacements] +kind = "class" +signature = "VelaExternalBlockPlacements(cmd_data: str | None = None, weight_data: str | None = None) -> None" + +[python.VelaExternalBlockPlacements.__delattr__] +kind = "function" +signature = "VelaExternalBlockPlacements.__delattr__(self, name)" + +[python.VelaExternalBlockPlacements.__eq__] +kind = "function" +signature = "VelaExternalBlockPlacements.__eq__(self, other)" + +[python.VelaExternalBlockPlacements.__hash__] +kind = "function" +signature = "VelaExternalBlockPlacements.__hash__(self)" + +[python.VelaExternalBlockPlacements.__post_init__] +kind = "function" +signature = "VelaExternalBlockPlacements.__post_init__(self) -> None" + +[python.VelaExternalBlockPlacements.__repr__] +kind = "function" +signature = "VelaExternalBlockPlacements.__repr__(self)" + +[python.VelaExternalBlockPlacements.__setattr__] +kind = "function" +signature = "VelaExternalBlockPlacements.__setattr__(self, name, value)" + +[python.VelaExternalBlockPlacements.to_block_placements] +kind = "function" +signature = "VelaExternalBlockPlacements.to_block_placements(self) -> dict[str, str]" + +[python.VgfBackend] +kind = "class" +signature = "VgfBackend()" + +[python.VgfBackend.preprocess] +kind = "function" +signature = "VgfBackend.preprocess(edge_program: torch.export.exported_program.ExportedProgram, compile_specs: List[executorch.exir.backend.compile_spec_schema.CompileSpec]) -> executorch.exir.backend.backend_details.PreprocessResult" + +[python.VgfCompileSpec] +kind = "class" +signature = "VgfCompileSpec(tosa_spec: executorch.backends.arm.tosa.specification.TosaSpecification | str | None = None, compiler_flags: list[str] | None = None)" + +[python.VgfCompileSpec.DebugMode] +kind = "enum" +signature = "VgfCompileSpec.DebugMode(*values)" + +[python.VgfCompileSpec.__eq__] +kind = "function" +signature = "VgfCompileSpec.__eq__(self, other)" + +[python.VgfCompileSpec.__repr__] +kind = "function" +signature = "VgfCompileSpec.__repr__(self)" + +[python.VgfCompileSpec.dump_debug_info] +kind = "function" +signature = "VgfCompileSpec.dump_debug_info(self, debug_mode: executorch.backends.arm.common.arm_compile_spec.ArmCompileSpec.DebugMode | None)" + +[python.VgfCompileSpec.dump_intermediate_artifacts_to] +kind = "function" +signature = "VgfCompileSpec.dump_intermediate_artifacts_to(self, output_path: str | None)" + +[python.VgfCompileSpec.set_pass_pipeline_config] +kind = "function" +signature = "VgfCompileSpec.set_pass_pipeline_config(self, config: executorch.backends.arm.common.pipeline_config.ArmPassPipelineConfig) -> None" + +[python.VgfCompileSpec.validate_environment] +kind = "function" +signature = "VgfCompileSpec.validate_environment(self, build_dir: str | None = None, *, require_runtime_build: bool = False) -> 'VgfEnvironmentReport'" + +[python.VgfPartitioner] +kind = "class" +signature = "VgfPartitioner(compile_spec: executorch.backends.arm.vgf.compile_spec.VgfCompileSpec, additional_checks: Optional[Sequence[torch.fx.passes.operator_support.OperatorSupportBase]] = None) -> None" + +[python.VgfPartitioner.ops_to_not_decompose] +kind = "function" +signature = "VgfPartitioner.ops_to_not_decompose(self, ep: torch.export.exported_program.ExportedProgram) -> Tuple[List[torch._ops.OpOverload], Optional[Callable[[torch.fx.node.Node], bool]]]" + +[python.VgfPartitioner.partition] +kind = "function" +signature = "VgfPartitioner.partition(self, exported_program: torch.export.exported_program.ExportedProgram) -> executorch.exir.backend.partitioner.PartitionResult" + +[python.VgfPartitioner.register_custom_partition_op] +kind = "function" +signature = "VgfPartitioner.register_custom_partition_op(self, op: torch._ops.OpOverload) -> None" + +[python.VgfPartitioner.transform_for_pre_decomposition] +kind = "function" +signature = "VgfPartitioner.transform_for_pre_decomposition(self, exported_program: torch.export.exported_program.ExportedProgram) -> torch.export.exported_program.ExportedProgram" + +[python.VgfQuantizer] +kind = "class" +signature = "VgfQuantizer(compile_spec: 'VgfCompileSpec', use_composable_quantizer: 'bool' = True) -> 'None'" + +[python.VgfQuantizer.annotate] +kind = "function" +signature = "VgfQuantizer.annotate(self, model: 'GraphModule') -> 'GraphModule'" + +[python.VgfQuantizer.set_global] +kind = "function" +signature = "VgfQuantizer.set_global(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.set_io] +kind = "function" +signature = "VgfQuantizer.set_io(self, quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.set_module_name] +kind = "function" +signature = "VgfQuantizer.set_module_name(self, module_name: 'str', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.set_module_type] +kind = "function" +signature = "VgfQuantizer.set_module_type(self, module_type: 'Callable', quantization_config: 'Optional[QuantizationConfig]') -> 'TOSAQuantizer'" + +[python.VgfQuantizer.transform_for_annotation] +kind = "function" +signature = "VgfQuantizer.transform_for_annotation(self, model: 'GraphModule') -> 'GraphModule'" + +[python.VgfQuantizer.validate] +kind = "function" +signature = "VgfQuantizer.validate(self, model: 'GraphModule') -> 'None'" + +[python.get_symmetric_a16w8_quantization_config] +kind = "function" +signature = "get_symmetric_a16w8_quantization_config(is_per_channel: 'bool' = True, is_qat: 'bool' = False, is_dynamic: 'bool' = False, weight_qmin: 'int' = -127, weight_qmax: 'int' = 127, epsilon: 'float' = 0.000244140625) -> 'QuantizationConfig'" + +[python.get_symmetric_quantization_config] +kind = "function" +signature = "get_symmetric_quantization_config(is_per_channel: 'bool' = True, is_qat: 'bool' = False, is_dynamic: 'bool' = False, act_qmin: 'int' = -128, act_qmax: 'int' = 127, weight_qmin: 'int' = -127, weight_qmax: 'int' = 127, eps: 'float' = 1.52587890625e-05) -> 'QuantizationConfig'" + +[cmake] + +[cmake.arm_runner_add_minimal_executable] +kind = "function" +signature = "arm_runner_add_minimal_executable(*, TARGET, SOURCE, OPS_PREFIX, COMPILE_DEFINITIONS=())" + +[cmake.arm_runner_add_standalone_executorch] +kind = "macro" +signature = "arm_runner_add_standalone_executorch()" + +[cmake.arm_runner_configure_ethos_u_platform] +kind = "function" +signature = "arm_runner_configure_ethos_u_platform(*, SDK_PATH, SYSTEM_CONFIG, MEMORY_MODE)" + +[cmake.arm_runner_configure_linker_script] +kind = "function" +signature = "arm_runner_configure_linker_script(*, TARGET, SYSTEM_CONFIG, OUTPUT_NAME=None)" + +[cmake.arm_runner_configure_model] +kind = "function" +signature = "arm_runner_configure_model(*, TARGET, PTE_FILE=None, MODEL_PTE_ADDR=None, MODEL_PTE_SIZE=None, PUBLIC=False)" + +[cmake.arm_runner_configure_runtime_output] +kind = "function" +signature = "arm_runner_configure_runtime_output(TARGET_NAME, FALLBACK_DIR)" + +[cmake.arm_runner_create_default_selected_ops_libs] +kind = "function" +signature = "arm_runner_create_default_selected_ops_libs(*, PREFIX, SUFFIX=None, OP_LIST=None, OPS_FROM_MODEL=None, DTYPE_SELECTIVE_BUILD=None, OUT_LIBS=None, DEPS=())" + +[cmake.arm_runner_create_selected_ops_lib] +kind = "function" +signature = "arm_runner_create_selected_ops_lib(*, LIB_NAME, FUNCTIONS_YAML=None, CUSTOM_OPS_YAML=None, OP_LIST=None, OPS_FROM_MODEL=None, DTYPE_SELECTIVE_BUILD=None, KERNEL_LIBS=(), DEPS=(), INCLUDE_ALL_OPS=False, PRIM_OPS=False)" + +[cmake.arm_runner_define_cache_options] +kind = "function" +signature = "arm_runner_define_cache_options(*, METHOD_ALLOCATOR_SIZE=None)" + +[cmake.arm_runner_link_minimal_specs] +kind = "function" +signature = "arm_runner_link_minimal_specs(TARGET_NAME)" + +[cmake.arm_runner_link_registration_libraries] +kind = "function" +signature = "arm_runner_link_registration_libraries(*, TARGET, SCOPE=None, BASE_LIBS=(), REGISTRATION_LIBS=(), NORMAL_LIBS=(), SUPPRESS_LIBS=())" + +[cmake.arm_runner_require_baremetal_targets] +kind = "function" +signature = "arm_runner_require_baremetal_targets()" + +[cmake.arm_runner_require_python] +kind = "macro" +signature = "arm_runner_require_python()" + +[cmake.arm_runner_validate_model_source] +kind = "function" +signature = "arm_runner_validate_model_source(*, ALLOW_SEMIHOSTING=False)" From 473d150b5b957698ffd28b79e96495722ba69880 Mon Sep 17 00:00:00 2001 From: Youngsik Yang Date: Tue, 8 Sep 2026 03:16:42 +0900 Subject: [PATCH 183/190] Arm backend: Stop re-creating Ethos-U driver objects per inference on Linux **Problem** On the Cortex-A path, execute() created the driver objects (the network and the DMA buffers, including a full copy of the weights) on every inference. **Fix** Create them once in platform_init() and keep them for the lifetime of the loaded method. Measured on the Corstone-1000 FVP: 53-77% less CPU time in execute() per inference. Signed-off-by: Youngsik Yang --- backends/arm/runtime/EthosUBackend.cpp | 2 +- .../arm/runtime/EthosUBackend_Cortex_A.cpp | 123 ++++++++++++------ .../arm/runtime/EthosUBackend_Cortex_M.cpp | 3 +- backends/arm/runtime/EthosUBackend_Internal.h | 3 +- 4 files changed, 86 insertions(+), 45 deletions(-) diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index 079da95391d..abcdae40671 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -143,7 +143,7 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { return read_status; } - handle->platform_state = platform_init(compile_specs, allocator); + handle->platform_state = platform_init(compile_specs, allocator, handle); // Return the same buffer we were passed - this data will be // executed directly diff --git a/backends/arm/runtime/EthosUBackend_Cortex_A.cpp b/backends/arm/runtime/EthosUBackend_Cortex_A.cpp index 41c1bca97bf..2a7f083dff4 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_A.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_A.cpp @@ -49,6 +49,11 @@ struct LinuxDriverOptions { struct PlatformState { LinuxDriverOptions options; + std::shared_ptr network; + std::shared_ptr constant_buffer; + std::shared_ptr intermediate_buffer; + std::vector> ifm_buffers; + std::vector> ofm_buffers; }; namespace { @@ -181,35 +186,9 @@ Error invoke_linux_driver( const std::vector& output_ptrs, const std::vector& input_copy_sizes, const std::vector& output_copy_sizes, - const LinuxDriverOptions& options) { - if (handles.outputs == nullptr) { - ET_LOG(Error, "Ethos-U backend missing output metadata"); - return Error::InvalidProgram; - } + const PlatformState& state) { + const LinuxDriverOptions& options = state.options; try { - EthosU::Device& device = get_linux_device_cache().get(options.device_path); - auto network = std::make_shared( - device, - reinterpret_cast(handles.cmd_data), - handles.cmd_data_size); - - std::shared_ptr constant_buffer = - std::make_shared(); - if (handles.weight_data_size > 0) { - auto constant_buffers = device.createBuffers({handles.weight_data_size}); - constant_buffer = constant_buffers.front(); - constant_buffer->write( - const_cast(handles.weight_data), handles.weight_data_size); - } - - std::shared_ptr intermediate_buffer = - std::make_shared(); - if (handles.scratch_data_size > 0) { - auto scratch_buffers = device.createBuffers({handles.scratch_data_size}); - intermediate_buffer = scratch_buffers.front(); - } - - std::vector> ifm_buffers; if (handles.inputs != nullptr && handles.inputs->count > 0) { if (input_copy_sizes.size() != static_cast(handles.inputs->count)) { @@ -228,7 +207,6 @@ Error invoke_linux_driver( input_copy_sizes.size()); return Error::InvalidState; } - ifm_buffers = device.createBuffers(input_copy_sizes); for (int i = 0; i < handles.inputs->count; ++i) { const size_t copy_size = input_copy_sizes[i]; if (copy_size == 0) { @@ -239,7 +217,7 @@ Error invoke_linux_driver( ET_LOG(Error, "Missing input buffer for index %d", i); return Error::InvalidState; } - ifm_buffers[i]->write(const_cast(src), copy_size); + state.ifm_buffers[i]->write(const_cast(src), copy_size); } } @@ -260,16 +238,14 @@ Error invoke_linux_driver( output_copy_sizes.size()); return Error::InvalidState; } - auto ofm_buffers = device.createBuffers(output_copy_sizes); - auto inference = std::make_unique( - network, - ifm_buffers.begin(), - ifm_buffers.end(), - ofm_buffers.begin(), - ofm_buffers.end(), - intermediate_buffer, - constant_buffer, + state.network, + state.ifm_buffers.begin(), + state.ifm_buffers.end(), + state.ofm_buffers.begin(), + state.ofm_buffers.end(), + state.intermediate_buffer, + state.constant_buffer, options.pmu_events, options.enable_cycle_counter); @@ -311,7 +287,7 @@ Error invoke_linux_driver( ET_LOG(Error, "Missing output buffer for index %d", i); return Error::InvalidState; } - ofm_buffers[i]->read(dst, copy_size); + state.ofm_buffers[i]->read(dst, copy_size); } } catch (const std::exception& e) { ET_LOG(Error, "Ethos-U Linux driver invocation failed: %s", e.what()); @@ -320,19 +296,82 @@ Error invoke_linux_driver( return Error::Ok; } + +// Get the byte size of an IO tensor from its Vela descriptor. +size_t vela_io_bytes(const VelaIO& io) { + size_t count = 1; + for (int i = 0; i < shapeDim; i++) { + count *= static_cast(io.shape[i]); + } + return count * static_cast(io.elem_size); +} + +// Created once in platform_init(), reused by every invoke_linux_driver(). +Error create_driver_objects(const VelaHandles& handles, PlatformState* state) { + if (handles.outputs == nullptr) { + ET_LOG(Error, "Ethos-U backend missing output metadata"); + return Error::InvalidProgram; + } + const LinuxDriverOptions& options = state->options; + try { + EthosU::Device& device = get_linux_device_cache().get(options.device_path); + state->network = std::make_shared( + device, + reinterpret_cast(handles.cmd_data), + handles.cmd_data_size); + + state->constant_buffer = std::make_shared(); + if (handles.weight_data_size > 0) { + auto constant_buffers = device.createBuffers({handles.weight_data_size}); + state->constant_buffer = constant_buffers.front(); + state->constant_buffer->write( + const_cast(handles.weight_data), handles.weight_data_size); + } + + state->intermediate_buffer = std::make_shared(); + if (handles.scratch_data_size > 0) { + auto scratch_buffers = device.createBuffers({handles.scratch_data_size}); + state->intermediate_buffer = scratch_buffers.front(); + } + + if (handles.inputs != nullptr && handles.inputs->count > 0) { + std::vector ifm_sizes; + for (int i = 0; i < handles.inputs->count; ++i) { + ifm_sizes.push_back(vela_io_bytes(handles.inputs->io[i])); + } + state->ifm_buffers = device.createBuffers(ifm_sizes); + } + + std::vector ofm_sizes; + for (int i = 0; i < handles.outputs->count; ++i) { + ofm_sizes.push_back(vela_io_bytes(handles.outputs->io[i])); + } + state->ofm_buffers = device.createBuffers(ofm_sizes); + } catch (const std::exception& e) { + ET_LOG(Error, "Ethos-U Linux driver setup failed: %s", e.what()); + return Error::InvalidState; + } + + return Error::Ok; +} } // namespace // Used by EthosUBackend.cpp through EthosUBackend_Internal.h. // cppcheck-suppress unusedFunction PlatformState* platform_init( ArrayRef specs, - MemoryAllocator* allocator) { + MemoryAllocator* allocator, + const ExecutionHandle* handle) { (void)allocator; PlatformState* state = new (std::nothrow) PlatformState(); if (state == nullptr) { return nullptr; } state->options = parse_linux_options(specs); + if (create_driver_objects(handle->handles, state) != Error::Ok) { + delete state; + return nullptr; + } return state; } @@ -401,7 +440,7 @@ Error platform_execute( linux_output_ptrs, input_copy_sizes, output_io_bytes, - state->options); + *state); if (status != Error::Ok) { return status; } diff --git a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp index eab91382247..0f591c18b50 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp @@ -55,7 +55,8 @@ struct PlatformState {}; PlatformState* platform_init( executorch::runtime::ArrayRef /*specs*/, - executorch::runtime::MemoryAllocator* /*allocator*/) { + executorch::runtime::MemoryAllocator* /*allocator*/, + const ExecutionHandle* /*handle*/) { return nullptr; } diff --git a/backends/arm/runtime/EthosUBackend_Internal.h b/backends/arm/runtime/EthosUBackend_Internal.h index 5b3cc58858f..51f1629b6a7 100644 --- a/backends/arm/runtime/EthosUBackend_Internal.h +++ b/backends/arm/runtime/EthosUBackend_Internal.h @@ -88,7 +88,8 @@ extern size_t ethosu_fast_scratch_size; PlatformState* platform_init( executorch::runtime::ArrayRef specs, - executorch::runtime::MemoryAllocator* allocator); + executorch::runtime::MemoryAllocator* allocator, + const ExecutionHandle* handle); void platform_destroy(PlatformState* state); From 013de8294265159fb7b8b21f16f9fc7c1c156a75 Mon Sep 17 00:00:00 2001 From: Youngsik Yang Date: Fri, 11 Sep 2026 02:06:55 +0900 Subject: [PATCH 184/190] Arm backend: Properly handle platform init failures Properly handle platform setup failures by propagating platform initialization failures and releasing the execution handle. Signed-off-by: Youngsik Yang --- backends/arm/runtime/EthosUBackend.cpp | 7 ++++++- backends/arm/runtime/EthosUBackend_Cortex_A.cpp | 14 ++++++++------ backends/arm/runtime/EthosUBackend_Cortex_M.cpp | 6 +++--- backends/arm/runtime/EthosUBackend_Internal.h | 4 ++-- 4 files changed, 19 insertions(+), 12 deletions(-) diff --git a/backends/arm/runtime/EthosUBackend.cpp b/backends/arm/runtime/EthosUBackend.cpp index abcdae40671..185623d8504 100644 --- a/backends/arm/runtime/EthosUBackend.cpp +++ b/backends/arm/runtime/EthosUBackend.cpp @@ -143,7 +143,12 @@ class EthosUBackend final : public ::executorch::runtime::BackendInterface { return read_status; } - handle->platform_state = platform_init(compile_specs, allocator, handle); + const Error platform_status = + platform_init(compile_specs, allocator, handle); + if (platform_status != Error::Ok) { + delete handle; + return platform_status; + } // Return the same buffer we were passed - this data will be // executed directly diff --git a/backends/arm/runtime/EthosUBackend_Cortex_A.cpp b/backends/arm/runtime/EthosUBackend_Cortex_A.cpp index 2a7f083dff4..ef5009c25fb 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_A.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_A.cpp @@ -358,21 +358,23 @@ Error create_driver_objects(const VelaHandles& handles, PlatformState* state) { // Used by EthosUBackend.cpp through EthosUBackend_Internal.h. // cppcheck-suppress unusedFunction -PlatformState* platform_init( +Error platform_init( ArrayRef specs, MemoryAllocator* allocator, - const ExecutionHandle* handle) { + ExecutionHandle* handle) { (void)allocator; PlatformState* state = new (std::nothrow) PlatformState(); if (state == nullptr) { - return nullptr; + return Error::MemoryAllocationFailed; } state->options = parse_linux_options(specs); - if (create_driver_objects(handle->handles, state) != Error::Ok) { + const Error status = create_driver_objects(handle->handles, state); + if (status != Error::Ok) { delete state; - return nullptr; + return status; } - return state; + handle->platform_state = state; + return Error::Ok; } // Used by EthosUBackend.cpp through EthosUBackend_Internal.h. diff --git a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp index 0f591c18b50..0697b783a36 100644 --- a/backends/arm/runtime/EthosUBackend_Cortex_M.cpp +++ b/backends/arm/runtime/EthosUBackend_Cortex_M.cpp @@ -53,11 +53,11 @@ namespace arm { struct PlatformState {}; -PlatformState* platform_init( +executorch::runtime::Error platform_init( executorch::runtime::ArrayRef /*specs*/, executorch::runtime::MemoryAllocator* /*allocator*/, - const ExecutionHandle* /*handle*/) { - return nullptr; + ExecutionHandle* /*handle*/) { + return executorch::runtime::Error::Ok; } void platform_destroy(PlatformState* /*state*/) {} diff --git a/backends/arm/runtime/EthosUBackend_Internal.h b/backends/arm/runtime/EthosUBackend_Internal.h index 51f1629b6a7..a62926676b3 100644 --- a/backends/arm/runtime/EthosUBackend_Internal.h +++ b/backends/arm/runtime/EthosUBackend_Internal.h @@ -86,10 +86,10 @@ extern unsigned char* ethosu_fast_scratch; extern size_t ethosu_fast_scratch_size; } -PlatformState* platform_init( +executorch::runtime::Error platform_init( executorch::runtime::ArrayRef specs, executorch::runtime::MemoryAllocator* allocator, - const ExecutionHandle* handle); + ExecutionHandle* handle); void platform_destroy(PlatformState* state); From 6ab65d0fbcad829101ce78b61b628e214ada9a80 Mon Sep 17 00:00:00 2001 From: Per Held Date: Fri, 11 Sep 2026 09:51:02 +0200 Subject: [PATCH 185/190] Arm backend: Silence successful VGF coverage output Capture VGF coverage checker output during pre-push and print it only when the check fails. Keep strict attribution diagnostics available for direct checker use while making successful pushes concise. Signed-off-by: Per Held Assisted-by: Codex Change-Id: I4246d3f4ef07e4ec3dd3ace3292fc9b098a97931 --- backends/arm/scripts/docgen/generate_vgf_op_support.py | 2 +- backends/arm/scripts/pre-push | 5 ++++- backends/arm/test/misc/test_docgen_op_support.py | 4 ++++ 3 files changed, 9 insertions(+), 2 deletions(-) diff --git a/backends/arm/scripts/docgen/generate_vgf_op_support.py b/backends/arm/scripts/docgen/generate_vgf_op_support.py index b2068d1d157..1f84966fe9b 100644 --- a/backends/arm/scripts/docgen/generate_vgf_op_support.py +++ b/backends/arm/scripts/docgen/generate_vgf_op_support.py @@ -2669,7 +2669,7 @@ def run_check(repo_root: Path, *, strict_ast: bool = False) -> int: # noqa: C90 print(f"| `{op}` | {profile} | {classification} | {sat} | {test_cell} |") print() - if unresolved: + if strict_ast and unresolved: _print_unresolved(unresolved) if diagnostics: print("AST normalisation diagnostics:") diff --git a/backends/arm/scripts/pre-push b/backends/arm/scripts/pre-push index be12e14fccf..0cdbb85e17f 100755 --- a/backends/arm/scripts/pre-push +++ b/backends/arm/scripts/pre-push @@ -74,6 +74,8 @@ run_docgen_check() { } run_vgf_op_support_checks() { + local coverage_output + echo -e "${INFO} Generating VGF operator support documentation" if ! python "$VGF_OP_SUPPORT_SCRIPT"; then @@ -90,7 +92,8 @@ run_vgf_op_support_checks() { echo -e "${INFO} Checking VGF operator support coverage" - if ! python "$VGF_OP_SUPPORT_SCRIPT" --check; then + if ! coverage_output=$(python "$VGF_OP_SUPPORT_SCRIPT" --check 2>&1); then + echo "$coverage_output" >&2 echo -e "${ERROR} VGF operator support coverage check failed" >&2 FAILED=1 else diff --git a/backends/arm/test/misc/test_docgen_op_support.py b/backends/arm/test/misc/test_docgen_op_support.py index 4956b15e06b..775eb97174d 100644 --- a/backends/arm/test/misc/test_docgen_op_support.py +++ b/backends/arm/test/misc/test_docgen_op_support.py @@ -411,6 +411,7 @@ def test_run_check_reports_missing_profile( def test_run_check_strict_ast_fails_on_unresolved_attribution( monkeypatch: pytest.MonkeyPatch, + capsys: pytest.CaptureFixture[str], ) -> None: unresolved = [ docgen.UnresolvedPipelineEvidence( @@ -431,7 +432,10 @@ def test_run_check_strict_ast_fails_on_unresolved_attribution( monkeypatch.setattr(docgen, "_collect_backend_supported_ops", lambda _root: {}) assert docgen.run_check(Path("/repo"), strict_ast=False) == 0 + assert "Unresolved VgfPipeline attribution" not in capsys.readouterr().out + assert docgen.run_check(Path("/repo"), strict_ast=True) == 1 + assert "Unresolved VgfPipeline attribution" in capsys.readouterr().out def test_main_writes_requested_markdown_and_html( From 19ab68266157daa85067d74582cdfd8fe8909980 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Fri, 11 Sep 2026 14:12:09 +0100 Subject: [PATCH 186/190] Arm backend: Enable Vulkan BF16 shader support in the VGF backend when supported by the device. (#22727) Enable Vulkan BF16 shader support in the VGF backend when supported by the device. This fixes validation errors for BF16 workloads by enabling VK_KHR_shader_bfloat16 and shaderBFloat16Type. cc @SS-JIA @manuelcandales @digantdesai @cbilgin @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Elena Zhelezina --- backends/arm/runtime/VGFBackend.cpp | 39 ++++++++++++++++++++++++++++- 1 file changed, 38 insertions(+), 1 deletion(-) diff --git a/backends/arm/runtime/VGFBackend.cpp b/backends/arm/runtime/VGFBackend.cpp index e29bbe7de9f..3365cd259cf 100644 --- a/backends/arm/runtime/VGFBackend.cpp +++ b/backends/arm/runtime/VGFBackend.cpp @@ -955,9 +955,13 @@ VkResult vkml_allocate_basics( }; // Query features + VkPhysicalDeviceShaderBfloat16FeaturesKHR available_bfloat16{ + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_BFLOAT16_FEATURES_KHR, + .pNext = nullptr, + }; VkPhysicalDeviceVulkan12Features available_12 = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_2_FEATURES, - .pNext = NULL, + .pNext = &available_bfloat16, }; VkPhysicalDeviceVulkan11Features available_11 = { .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_VULKAN_1_1_FEATURES, @@ -995,6 +999,11 @@ VkResult vkml_allocate_basics( } // Select features + VkPhysicalDeviceShaderBfloat16FeaturesKHR features_bfloat16{ + .sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_BFLOAT16_FEATURES_KHR, + .pNext = nullptr, + .shaderBFloat16Type = VK_FALSE, + }; VkPhysicalDeviceShaderReplicatedCompositesFeaturesEXT features_c{ VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_REPLICATED_COMPOSITES_FEATURES_EXT, nullptr}; @@ -1062,6 +1071,34 @@ VkResult vkml_allocate_basics( vector requested_exts; + const bool bfloat16_extension_available = std::any_of( + available.begin(), available.end(), [](const auto& ext_avail) { + return std::strcmp( + VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME, + ext_avail.extensionName) == 0; + }); + const bool bfloat16_feature_available = + available_bfloat16.shaderBFloat16Type == VK_TRUE; + + if (bfloat16_extension_available && bfloat16_feature_available) { + requested_exts.push_back(VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + features_bfloat16.shaderBFloat16Type = VK_TRUE; + features_c.pNext = &features_bfloat16; + ET_LOG( + Info, + "Enabled %s with shaderBFloat16Type", + VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + } else if (!bfloat16_extension_available) { + ET_LOG( + Info, + "VGF BF16 shaders are unavailable: Vulkan device does not expose %s", + VK_KHR_SHADER_BFLOAT16_EXTENSION_NAME); + } else { + ET_LOG( + Info, + "VGF BF16 shaders are unavailable: shaderBFloat16Type is not supported"); + } + #if defined(VK_ARM_data_graph_neural_accelerator_statistics) const bool neural_statistics_extension_available = std::any_of( available.begin(), available.end(), [](const auto& ext_avail) { From f18ce146148ff072974b2ecd6c89ec9bb8c93067 Mon Sep 17 00:00:00 2001 From: Elena Zhelezina Date: Fri, 11 Sep 2026 14:12:47 +0100 Subject: [PATCH 187/190] Arm backend: Fix generate_vgf_op_support to work with custom ops. (#22733) In automatic generation of supported ops we don't take into account registry of custom ops. This PR fixes this. cc @digantdesai @freddan80 @per @zingo @oscarandersson8218 @mansnils @Sebastian-Larsson @robell @rascani Signed-off-by: Elena Zhelezina --- .../scripts/docgen/generate_vgf_op_support.py | 39 +++++++++++++++++++ .../arm/test/misc/test_docgen_op_support.py | 27 +++++++++++++ 2 files changed, 66 insertions(+) diff --git a/backends/arm/scripts/docgen/generate_vgf_op_support.py b/backends/arm/scripts/docgen/generate_vgf_op_support.py index 1f84966fe9b..a4df1ba55d3 100644 --- a/backends/arm/scripts/docgen/generate_vgf_op_support.py +++ b/backends/arm/scripts/docgen/generate_vgf_op_support.py @@ -54,6 +54,10 @@ BACKEND_PIPELINE_CLASS_NAMES = frozenset({"VgfPipeline"}) BACKEND_PIPELINE_LABEL = "VgfPipeline" BACKEND_TOSA_SPEC = "TOSA-1.0+FP+INT+int4+int16" +BACKEND_PROFILE_TOSA_SPECS = { + "FP": "TOSA-1.0+FP", + "INT": "TOSA-1.0+INT", +} GENERATOR_PATH = Path("backends/arm/scripts/docgen/generate_vgf_op_support.py") GENERATOR_COMMAND = f"python {GENERATOR_PATH}" @@ -2197,6 +2201,37 @@ def _profiles_for_checker( return profiles +def _collect_backend_custom_partition_ops( + backend_tosa_spec: TosaSpecificationLike, +) -> dict[str, set[object]]: + """Collect VGF custom partition ops for each enabled support profile. + + Instantiate the partitioner with a single-profile compile spec so custom + registrations that are conditional on the compile spec are attributed only + to the profiles for which they are actually registered. + + """ + from executorch.backends.arm.vgf import VgfCompileSpec, VgfPartitioner + + enabled_profiles = { + "FP": backend_tosa_spec.support_float(), + "INT": backend_tosa_spec.support_integer(), + } + custom_ops_by_profile: dict[str, set[object]] = {} + + for profile, enabled in enabled_profiles.items(): + if not enabled: + continue + partitioner = VgfPartitioner( + VgfCompileSpec(BACKEND_PROFILE_TOSA_SPECS[profile]) + ) + custom_ops_by_profile[profile] = set( + getattr(partitioner, "_custom_partition_ops", ()) + ) + + return custom_ops_by_profile + + def _collect_backend_supported_ops( # noqa: C901 repo_root: Path, ) -> dict[str, SupportedOperatorEvidence]: @@ -2248,6 +2283,10 @@ def add(target: object, profile: str, evidence: str) -> None: for profile in _profiles_for_checker(checker, tosa_spec): add(target, profile, checker_evidence) + for profile, targets in _collect_backend_custom_partition_ops(tosa_spec).items(): + for target in targets: + add(target, profile, "VgfPartitioner.register_custom_partition_op") + # Lowering visitors are not the source of partitioner support, but they are # useful evidence when the exported op name matches a registered visitor # target directly. diff --git a/backends/arm/test/misc/test_docgen_op_support.py b/backends/arm/test/misc/test_docgen_op_support.py index 775eb97174d..3622cfc5228 100644 --- a/backends/arm/test/misc/test_docgen_op_support.py +++ b/backends/arm/test/misc/test_docgen_op_support.py @@ -364,6 +364,33 @@ def test_matching_evidence_accepts_stage_equivalent_alias() -> None: assert records[0].asserted_op == alias +def test_collect_backend_custom_partition_ops_discovers_fp_and_int_profiles() -> None: + from executorch.backends.arm.tosa import TosaSpecification + + tosa_spec = TosaSpecification.create_from_string(docgen.BACKEND_TOSA_SPEC) + custom_ops = docgen._collect_backend_custom_partition_ops(tosa_spec) + canonical_by_profile = { + profile: { + docgen._canonical_pytorch_op_from_target(target) for target in targets + } + for profile, targets in custom_ops.items() + } + + expected = "torch.ops.aten.grid_sampler_2d.default" + assert expected in canonical_by_profile["FP"] + assert expected in canonical_by_profile["INT"] + + +def test_collect_backend_supported_ops_includes_vgf_custom_partition_ops() -> None: + repo_root = Path(__file__).resolve().parents[4] + + expected = docgen._collect_backend_supported_ops(repo_root) + row = expected["torch.ops.aten.grid_sampler_2d.default"] + + assert row.support_profiles == {"FP", "INT"} + assert "VgfPartitioner.register_custom_partition_op" in row.evidence + + def test_run_check_reports_missing_profile( monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str] ) -> None: From cf1e8b7d29f6f3a17830c599fef47af22483fb3d Mon Sep 17 00:00:00 2001 From: Sebastian Larsson <38941629+Sebastian-Larsson@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:32:17 +0200 Subject: [PATCH 188/190] Arm backend: Support boolean sums (#22345) Boolean sums produce int64 outputs, so the TOSA partitioner rejects them before lowering. Cast boolean inputs to int32 and run the sum with an explicit int32 dtype. Cast the result back to the original output dtype to preserve model semantics. Signed-off-by: Sebastian Larsson --- backends/arm/_passes/__init__.py | 1 + backends/arm/_passes/arm_pass_manager.py | 3 + backends/arm/_passes/convert_bool_sum_pass.py | 81 ++++++++++++++++ backends/arm/test/ops/test_sum.py | 11 +++ .../test/passes/test_convert_bool_sum_pass.py | 96 +++++++++++++++++++ 5 files changed, 192 insertions(+) create mode 100644 backends/arm/_passes/convert_bool_sum_pass.py create mode 100644 backends/arm/test/passes/test_convert_bool_sum_pass.py diff --git a/backends/arm/_passes/__init__.py b/backends/arm/_passes/__init__.py index ff8cf87679d..7d93aa38621 100644 --- a/backends/arm/_passes/__init__.py +++ b/backends/arm/_passes/__init__.py @@ -20,6 +20,7 @@ from .cast_int_comparison_inputs_pass import CastIntComparisonInputsPass # noqa from .cast_to_int32_pass import CastToInt32Pass # noqa from .constant_folding_pass import ConstantFoldingPass # noqa +from .convert_bool_sum_pass import ConvertBoolSumPass # noqa from .convert_elu_params import ConvertELUParamsPass # noqa from .convert_expand_copy_to_repeat import ConvertExpandCopyToRepeatPass # noqa from .convert_full_like_to_full_pass import ConvertFullLikeToFullPass # noqa diff --git a/backends/arm/_passes/arm_pass_manager.py b/backends/arm/_passes/arm_pass_manager.py index 3145dc686b0..b74860954ba 100644 --- a/backends/arm/_passes/arm_pass_manager.py +++ b/backends/arm/_passes/arm_pass_manager.py @@ -22,6 +22,7 @@ ComputeConstantOpsAOTPass, ConstantFoldingPass, ControlFlowConstInlinePass, + ConvertBoolSumPass, ConvertEluFamilyToEluPass, ConvertELUParamsPass, ConvertExpandCopyToRepeatPass, @@ -595,6 +596,7 @@ def _tosa_pipeline( DecomposeExpm1Pass(), DecomposeIntPowPass(), DecomposeLog1pPass(), + ConvertBoolSumPass(), PromoteBoolOperandsPass(), DecomposeSinhPass(), DecomposeSignPass(), @@ -769,6 +771,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule): DecomposeDynamicFullPass(tfa_pass=True), ConvertInt64ConstOpsToInt32Pass(tfa_pass=True), ConvertInt64OutputOpsToInt32Pass(tfa_pass=True), + ConvertBoolSumPass(tfa_pass=True), InsertInt32CastsAfterInt64PlaceholdersPass(tfa_pass=True), FoldScalarMulIntoConvPass(tfa_pass=True), DecomposeEmbeddingPass(tfa_pass=True), diff --git a/backends/arm/_passes/convert_bool_sum_pass.py b/backends/arm/_passes/convert_bool_sum_pass.py new file mode 100644 index 00000000000..0b77281272c --- /dev/null +++ b/backends/arm/_passes/convert_bool_sum_pass.py @@ -0,0 +1,81 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import math +from typing import Set, Type + +import torch + +from executorch.backends.arm._passes.arm_pass import ArmOpTargetedPass +from executorch.backends.arm._passes.decompose_sum_pass import DecomposeSumPass +from executorch.backends.arm.tosa.specification import ( + get_context_shape_env, + get_context_spec, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass, NodeMetadata + + +class ConvertBoolSumPass(ArmOpTargetedPass): + """Compute boolean sums with an int32 accumulator.""" + + _passes_required_after: Set[Type[ExportPass]] = {DecomposeSumPass} + _MAX_EXACT_SUM = torch.iinfo(torch.int32).max + target_ops = ( + torch.ops.aten.sum.dim_IntList, + exir_ops.edge.aten.sum.dim_IntList, + ) + check_allowed_to_transform = True + + def call_operator(self, op, args, kwargs, meta): + tosa_spec = get_context_spec() + if ( + op not in self.target_ops + or args[0].data.dtype != torch.bool + or not self.allowed_to_transform(meta) + or not tosa_spec.support_integer() + or tosa_spec.support_float() + ): + return super().call_operator(op, args, kwargs, meta) + + dims = args[1] + if not dims: + dims = range(args[0].data.dim()) + reduced_elements = math.prod(args[0].data.shape[dim] for dim in dims) + if isinstance(reduced_elements, torch.SymInt): + reduced_elements = ( + get_context_shape_env().bound_sympy(reduced_elements.node.expr).upper + ) + if reduced_elements > self._MAX_EXACT_SUM: + return super().call_operator(op, args, kwargs, meta) + + cast_op = ( + exir_ops.edge.dim_order_ops._to_dim_order_copy.default + if op == exir_ops.edge.aten.sum.dim_IntList + else torch.ops.dim_order_ops._to_dim_order_copy.default + ) + accumulator_input = super().call_operator( + cast_op, + (args[0],), + {"dtype": torch.int32}, + NodeMetadata(args[0].node.meta), + updated=True, + ) + sum_kwargs = dict(kwargs) + sum_kwargs["dtype"] = torch.int32 + accumulator_sum = super().call_operator( + op, + (accumulator_input, *args[1:]), + sum_kwargs, + meta, + updated=True, + ) + return super().call_operator( + cast_op, + (accumulator_sum,), + {"dtype": meta["val"].dtype}, + meta, + updated=True, + ) diff --git a/backends/arm/test/ops/test_sum.py b/backends/arm/test/ops/test_sum.py index 8b4cfe46075..ee9ecaed05f 100644 --- a/backends/arm/test/ops/test_sum.py +++ b/backends/arm/test/ops/test_sum.py @@ -8,6 +8,7 @@ import pytest import torch + from executorch.backends.arm.test import common from executorch.backends.arm.test.tester.test_pipeline import ( @@ -81,6 +82,16 @@ def test_sum_dim_intlist_scalar_input_tosa_FP_not_delegated(): pipeline.run() +def test_sum_bool_tosa_INT() -> None: + pipeline = TosaPipelineINT( + Sum(), + (torch.ones(1, dtype=torch.bool), [], False), + aten_op, + exir_op=[], + ) + pipeline.run() + + @common.parametrize( "test_data", Sum.test_parameters | Sum.test_parameters_bf16 | Sum.test_parameters_fp16, diff --git a/backends/arm/test/passes/test_convert_bool_sum_pass.py b/backends/arm/test/passes/test_convert_bool_sum_pass.py new file mode 100644 index 00000000000..1b5842a9ed7 --- /dev/null +++ b/backends/arm/test/passes/test_convert_bool_sum_pass.py @@ -0,0 +1,96 @@ +# Copyright 2026 Arm Limited and/or its affiliates. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import pytest + +import torch + +from executorch.backends.arm._passes import ConvertBoolSumPass +from executorch.backends.arm.tosa.specification import ( + TosaLoweringContext, + TosaSpecification, +) +from torch._export.utils import _get_shape_env_from_gm +from torch._subclasses import FakeTensorMode +from torch.fx import Graph, GraphModule + + +class BoolSum(torch.nn.Module): + def forward(self, x: torch.Tensor, dim: int, keepdim: bool): + return x.sum(dim=dim, keepdim=keepdim) + + +def _run_pass(graph_module: GraphModule, tosa_spec: str) -> GraphModule: + with TosaLoweringContext( + TosaSpecification.create_from_string(tosa_spec), + _get_shape_env_from_gm(graph_module), + ): + return ConvertBoolSumPass().call(graph_module).graph_module + + +def _operator_targets(graph_module: GraphModule) -> list[torch._ops.OpOverload]: + return [ + node.target for node in graph_module.graph.nodes if node.op == "call_function" + ] + + +def test_sum_bool_skips_inexact_accumulator() -> None: + graph = Graph() + with FakeTensorMode(): + fake_input = torch.empty(torch.iinfo(torch.int32).max + 1, dtype=torch.bool) + fake_output = torch.empty((), dtype=torch.int64) + x = graph.placeholder("x") + x.meta["val"] = fake_input + output = graph.call_function(torch.ops.aten.sum.dim_IntList, (x, [0], False)) + output.meta["val"] = fake_output + graph.output(output) + + result = _run_pass(GraphModule(torch.nn.Module(), graph), "TOSA-1.0+INT") + + assert _operator_targets(result) == [torch.ops.aten.sum.dim_IntList] + + +@pytest.mark.parametrize( + "max_dynamic_dim,expect_transform", + [ + (1024, True), + (torch.iinfo(torch.int32).max + 1, False), + ], +) +def test_sum_bool_dynamic_shape(max_dynamic_dim: int, expect_transform: bool) -> None: + dynamic_dim = torch.export.Dim("dynamic_dim", min=1, max=max_dynamic_dim) + exported_program = torch.export.export( + BoolSum(), + (torch.ones(2, dtype=torch.bool), 0, False), + dynamic_shapes=({0: dynamic_dim}, None, None), + ) + + result = _run_pass(exported_program.graph_module, "TOSA-1.0+INT") + + if not expect_transform: + assert _operator_targets(result) == [torch.ops.aten.sum.dim_IntList] + return + + sum_node = next( + node + for node in result.graph.nodes + if node.target == torch.ops.aten.sum.dim_IntList + ) + (output,) = result(torch.ones(5, dtype=torch.bool), 0, False) + assert sum_node.kwargs["dtype"] == torch.int32 + assert output.dtype == torch.int64 + assert output == 5 + + +@pytest.mark.parametrize("tosa_spec", ["TOSA-1.0+FP", "TOSA-1.0+FP+INT"]) +def test_sum_bool_skips_float_profiles(tosa_spec: str) -> None: + exported_program = torch.export.export( + BoolSum(), + (torch.ones(2, dtype=torch.bool), 0, False), + ) + + result = _run_pass(exported_program.graph_module, tosa_spec) + + assert _operator_targets(result) == [torch.ops.aten.sum.dim_IntList] From a3ae32961d2907140265cf8c6ff33db34181aecf Mon Sep 17 00:00:00 2001 From: Mateusz Sluszniak <56299341+msluszniak@users.noreply.github.com> Date: Fri, 11 Sep 2026 15:33:37 +0200 Subject: [PATCH 189/190] Add a portable _fft_r2c kernel (#22055) ### Summary Addresses the first half of #21950. `_fft_r2c.out` only had an optimized (pocketfft-backed) kernel. A program that leaves it undelegated builds fine and writes a `.pte`, but a runtime carrying only the portable kernels cannot load it. The failure surfaces as `0x14 OperatorMissing` from `load_method`, long after the export reported success. `_fft_r2c.out` and `_fft_c2r.out` are two of only three ops in `optimized.yaml` with no counterpart in the portable `functions.yaml`: ``` portable ops: 206 optimized ops: 23 in optimized but NOT in portable (3): _fft_c2r.out _fft_r2c.out linear.out ``` The third, `linear.out`, ExecuTorch decomposes anyway, so the FFT pair is the real gap. ### Approach This is a direct O(n^2) evaluation of the transform sum, not a fast Fourier transform, in keeping with the portable library's role as the dependency-free reference; `kernels/optimized` keeps the asymptotically faster pocketfft path for anyone who links it. Audio front-ends transforming a few hundred points per frame, which is where this op tends to show up, are the intended case. Two details worth calling out: - **Multi-dimensional transforms** run the real transform along the last requested dimension and complex transforms along the rest, matching pocketfft's multi-axis `r2c`. A complex pass has to read a whole line before overwriting it; lines up to 128 elements (2 KB for double) use a stack buffer and longer ones ask the runtime for temporary memory. The single-dimension case, which is what `torch.fft.rfft` lowers to, needs no line buffer at all. - **The quarter-turn twiddle factors are returned exactly** rather than through `cos`/`sin`, so a real input's Nyquist bin comes out with a zero imaginary part instead of rounding noise around 1e-16. That is what pocketfft produces and what the existing tests expect. ### Test plan Registers the shared `op_fft_r2c_test` for `portable` alongside `aten` and `optimized`. All five cases pass against both kernel libraries: ``` $ ./cmake-out/kernels/test/portable_kernels_test --gtest_filter='OpFftR2c*' [ PASSED ] 5 tests. $ ./cmake-out/kernels/test/optimized_kernels_test --gtest_filter='OpFftR2c*' [ PASSED ] 5 tests. ``` End to end on the model from the issue (`torch.fft.rfft(x).abs().pow(2)` over `[8, 512]`, exported with no partitioner so the op stays on the CPU), run through `executor_runner` built with portable kernels only: | | result | | --- | --- | | before | exit 134 (abort on load) | | after | exit 0, output `[262144., 1.2e-27, 2.4e-27, ...]` | which is the correct `abs(rfft(ones))**2`: `512**2` in bin 0 and zero elsewhere. Accuracy against numpy, since the tests above use exactly representable values: | case | max abs error | | --- | --- | | `rfft`, length 5 (odd, no Nyquist bin, no quarter-turn twiddles) | 2.5e-15 | | `rfft`, 4x512 | 1.8e-13 (1.4e-15 relative to the largest output) | | `rfft2`, 6x8 | 1.8e-14 | I wanted to add the length-5 case as a regular test, but `tensors_are_close()` falls through to a bitwise `memcmp` for complex dtypes, so a complex comparison cannot have a tolerance and only exactly representable expected values can pass. That is why the existing cases all use small integers. Happy to fix that separately if it would be welcome. ### Not addressed here The issue also asks for lowering to fail when an op is left undelegated and no kernel exists, rather than emitting a `.pte` that dies at load. That one is a change to the AOT/runtime contract, since the export side does not know which kernel library the runtime will link (portable, optimized, or a selective build), so it needs its own discussion rather than riding along here. cc @larryliu0820 @manuelcandales @JakeStevens Co-authored-by: Jacob Stevens --- kernels/portable/cpu/op_fft_r2c.cpp | 307 ++++++++++++++++++ kernels/portable/functions.yaml | 5 + kernels/test/CMakeLists.txt | 1 + kernels/test/targets.bzl | 2 +- .../kernels/portable/op_registration_util.bzl | 7 + 5 files changed, 321 insertions(+), 1 deletion(-) create mode 100644 kernels/portable/cpu/op_fft_r2c.cpp diff --git a/kernels/portable/cpu/op_fft_r2c.cpp b/kernels/portable/cpu/op_fft_r2c.cpp new file mode 100644 index 00000000000..ddd5593ef5c --- /dev/null +++ b/kernels/portable/cpu/op_fft_r2c.cpp @@ -0,0 +1,307 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +#include +#include +#include +#include + +namespace torch::executor::native { + +namespace { + +constexpr double kTwoPi = 6.283185307179586476925286766559; + +// A complex-to-complex pass has to read a whole line before it can overwrite +// it. Lines up to this length go through a stack buffer, which keeps the buffer +// at 2 KB for double; longer ones ask the runtime for temporary memory. Only +// multi-dimensional transforms reach this path at all: the single-dimension +// case, which is what torch.fft.rfft lowers to, needs no line buffer. +constexpr size_t kStackLineLimit = 128; + +// Mirrors ATen's fft_norm_mode (ATen/native/SpectralOpsUtils.h), which is how +// the normalization argument is encoded. +enum class fft_norm_mode { + none, // No normalization + by_root_n, // Divide by sqrt(signal_size) + by_n, // Divide by signal_size +}; + +template +std::optional compute_fct( + KernelRuntimeContext& ctx, + const Tensor& t, + IntArrayRef dim, + int64_t normalization) { + constexpr auto one = static_cast(1); + const auto mode = static_cast(normalization); + if (mode == fft_norm_mode::none) { + return one; + } + int64_t n = 1; + for (auto idx : dim) { + n *= t.sizes()[idx]; + } + switch (mode) { + case fft_norm_mode::none: + return one; + case fft_norm_mode::by_n: + return one / static_cast(n); + case fft_norm_mode::by_root_n: + return one / std::sqrt(static_cast(n)); + } + ET_KERNEL_CHECK_MSG( + ctx, + false, + InvalidArgument, + std::nullopt, + "Unsupported normalization type: %" PRId64, + normalization); +} + +// cos and sin of -2*pi*idx/n. +// +// The quarter turns are returned exactly rather than through cos/sin, so that a +// real input's Nyquist bin comes out with a zero imaginary part instead of +// rounding noise on the order of 1e-16. Reducing idx modulo n first also keeps +// the angle inside one period, which matters for accuracy once k * j is large. +void twiddle(size_t idx, size_t n, double& cos_out, double& sin_out) { + idx %= n; + if (idx == 0) { + cos_out = 1.0; + sin_out = 0.0; + } else if (2 * idx == n) { + cos_out = -1.0; + sin_out = 0.0; + } else if (4 * idx == n) { + cos_out = 0.0; + sin_out = -1.0; + } else if (4 * idx == 3 * n) { + cos_out = 0.0; + sin_out = 1.0; + } else { + const double angle = + -kTwoPi * static_cast(idx) / static_cast(n); + cos_out = std::cos(angle); + sin_out = std::sin(angle); + } +} + +// Offset of the start of the line_index'th line along `axis`, for a tensor with +// the given sizes and strides. Lines are enumerated over every dimension except +// `axis`, so the same index names the same logical line in two tensors that +// agree on all dimensions but that one. +size_t line_offset( + size_t line_index, + ArrayRef sizes, + ArrayRef strides, + size_t axis) { + size_t offset = 0; + for (size_t d = sizes.size(); d-- > 0;) { + if (d == axis) { + continue; + } + const size_t size = static_cast(sizes[d]); + offset += (line_index % size) * static_cast(strides[d]); + line_index /= size; + } + return offset; +} + +// Forward real-to-complex DFT along `axis`, writing the onesided output. +// The normalization factor is folded in here: every later pass is linear, so +// applying it once at the front scales the whole transform. +template +void dft_r2c_axis(const Tensor& in, Tensor& out, size_t axis, T fct) { + using C = executorch::runtime::etensor::complex; + const T* const in_data = in.const_data_ptr(); + C* const out_data = out.mutable_data_ptr(); + + const size_t n = static_cast(in.size(axis)); + const size_t n_out = static_cast(out.size(axis)); + const size_t in_stride = static_cast(in.strides()[axis]); + const size_t out_stride = static_cast(out.strides()[axis]); + const size_t num_lines = n == 0 ? 0 : static_cast(in.numel()) / n; + + for (size_t line = 0; line < num_lines; ++line) { + const size_t in_off = line_offset(line, in.sizes(), in.strides(), axis); + const size_t out_off = line_offset(line, out.sizes(), out.strides(), axis); + for (size_t k = 0; k < n_out; ++k) { + double real = 0; + double imag = 0; + for (size_t j = 0; j < n; ++j) { + double c = 0; + double s = 0; + twiddle(k * j, n, c, s); + const double x = static_cast(in_data[in_off + j * in_stride]); + real += x * c; + imag += x * s; + } + out_data[out_off + k * out_stride] = + C{static_cast(real * static_cast(fct)), + static_cast(imag * static_cast(fct))}; + } + } +} + +// In-place forward complex-to-complex DFT along `axis`. `scratch` must hold at +// least out.size(axis) elements. +template +void dft_c2c_axis_(Tensor& out, size_t axis, void* scratch) { + using C = executorch::runtime::etensor::complex; + C* const out_data = out.mutable_data_ptr(); + C* const line_buf = static_cast(scratch); + + const size_t n = static_cast(out.size(axis)); + const size_t stride = static_cast(out.strides()[axis]); + const size_t num_lines = n == 0 ? 0 : static_cast(out.numel()) / n; + + for (size_t line = 0; line < num_lines; ++line) { + const size_t off = line_offset(line, out.sizes(), out.strides(), axis); + for (size_t j = 0; j < n; ++j) { + line_buf[j] = out_data[off + j * stride]; + } + for (size_t k = 0; k < n; ++k) { + double real = 0; + double imag = 0; + for (size_t j = 0; j < n; ++j) { + double c = 0; + double s = 0; + twiddle(k * j, n, c, s); + const double xr = static_cast(line_buf[j].real_); + const double xi = static_cast(line_buf[j].imag_); + real += xr * c - xi * s; + imag += xr * s + xi * c; + } + out_data[off + k * stride] = + C{static_cast(real), static_cast(imag)}; + } + } +} + +} // namespace + +// Reference discrete Fourier transform. +// +// This is a direct O(n^2) evaluation of the transform sum, not a fast Fourier +// transform. kernels/optimized provides a pocketfft-backed _fft_r2c.out that is +// asymptotically faster; this exists so that a graph containing _fft_r2c can be +// run by a build that only has the portable kernels, rather than failing to +// load with OperatorMissing. Audio front-ends that transform a few hundred +// points per frame are the intended case. +Tensor& _fft_r2c_out( + KernelRuntimeContext& ctx, + const Tensor& in, + IntArrayRef dim, + int64_t normalization, + bool onesided, + Tensor& out) { + auto in_sizes = in.sizes(); + ET_KERNEL_CHECK( + ctx, + static_cast(in.dim()) <= kTensorDimensionLimit, + InvalidArgument, + out); + ET_KERNEL_CHECK(ctx, !dim.empty(), InvalidArgument, out); + ET_KERNEL_CHECK( + ctx, tensors_have_same_dim_order(in, out), InvalidArgument, out); + + ET_KERNEL_CHECK_MSG( + ctx, + onesided, + InvalidArgument, + out, + "onesided=False is not supported yet in _fft_r2c"); + + ET_KERNEL_CHECK_MSG( + ctx, + out.scalar_type() == executorch::runtime::toComplexType(in.scalar_type()), + InvalidArgument, + out, + "the output type for _fft_r2c must be the Complex type corresponding to the input type"); + + for (auto d : dim) { + ET_KERNEL_CHECK_MSG( + ctx, + d >= 0 && d < in.dim(), + InvalidArgument, + out, + "dims must be in bounds (got %" PRId64 ")", + d); + } + + std::array out_sizes_storage; + executorch::runtime::Span out_sizes( + out_sizes_storage.data(), in_sizes.size()); + std::copy(in_sizes.begin(), in_sizes.end(), out_sizes.begin()); + out_sizes[dim.back()] = out_sizes[dim.back()] / 2 + 1; + + ET_KERNEL_CHECK_MSG( + ctx, + resize_tensor( + out, + executorch::runtime::ArrayRef( + out_sizes.data(), out_sizes.size())) == Error::Ok, + InvalidArgument, + out, + "Failed to resize output tensor (last dim %d).", + out_sizes[dim.back()]); + + // NOTE: as of this writing, upstream PyTorch only supports float/double, so + // we follow suit. + ET_SWITCH_FLOAT_TYPES(in.scalar_type(), ctx, "_fft_r2c.out", CTYPE_IN, [&] { + auto fct = compute_fct(ctx, in, dim, normalization); + if (!fct) { + // Check failed, just bail out of the lambda. + return; + } + + // The real transform runs along the last requested dimension, which is the + // one that is halved; the remaining dimensions are complex transforms of + // the result, matching pocketfft's multi-axis r2c. + const size_t real_axis = static_cast(dim.back()); + dft_r2c_axis(in, out, real_axis, *fct); + + if (dim.size() == 1) { + return; + } + + using Complex = executorch::runtime::etensor::complex; + size_t max_line = 0; + for (size_t i = 0; i + 1 < dim.size(); ++i) { + max_line = std::max(max_line, static_cast(out.size(dim[i]))); + } + + std::array stack_buf; + void* line_buf = stack_buf.data(); + if (max_line > kStackLineLimit) { + Result scratch = ctx.allocate_temp(max_line * sizeof(Complex)); + ET_KERNEL_CHECK_MSG( + ctx, + scratch.ok(), + MemoryAllocationFailed, + , + "_fft_r2c needs %zu bytes of temporary memory to transform a " + "dimension of length %zu, but no temp allocator is available", + max_line * sizeof(Complex), + max_line); + line_buf = scratch.get(); + } + + for (size_t i = 0; i + 1 < dim.size(); ++i) { + dft_c2c_axis_(out, static_cast(dim[i]), line_buf); + } + }); + + return out; +} + +} // namespace torch::executor::native diff --git a/kernels/portable/functions.yaml b/kernels/portable/functions.yaml index ecf62ee3606..61f32677b99 100644 --- a/kernels/portable/functions.yaml +++ b/kernels/portable/functions.yaml @@ -32,6 +32,11 @@ - arg_meta: null kernel_name: torch::executor::_conj_physical_out +- op: _fft_r2c.out + kernels: + - arg_meta: null + kernel_name: torch::executor::_fft_r2c_out + - op: _log_softmax.out kernels: - arg_meta: null diff --git a/kernels/test/CMakeLists.txt b/kernels/test/CMakeLists.txt index a8e703cb5aa..fe70f09f00d 100644 --- a/kernels/test/CMakeLists.txt +++ b/kernels/test/CMakeLists.txt @@ -207,6 +207,7 @@ set(all_test_sources "op_exp_test.cpp" "op_expand_copy_test.cpp" "op_expm1_test.cpp" + "op_fft_r2c_test.cpp" "op_fill_test.cpp" "op_flip_test.cpp" "op_floor_divide_test.cpp" diff --git a/kernels/test/targets.bzl b/kernels/test/targets.bzl index 9084dd2b16d..3d36071260b 100644 --- a/kernels/test/targets.bzl +++ b/kernels/test/targets.bzl @@ -252,7 +252,7 @@ def define_common_targets(): _common_op_test("op_expand_copy_test", ["aten", "portable"]) _common_op_test("op_expm1_test", ["aten", "portable"]) _common_op_test("op_fft_c2r_test", ["aten", "optimized"]) - _common_op_test("op_fft_r2c_test", ["aten", "optimized"]) + _common_op_test("op_fft_r2c_test", ["aten", "portable", "optimized"]) _common_op_test("op_fill_test", ["aten", "portable"]) _common_op_test("op_flip_test", ["aten", "portable"]) _common_op_test("op_floor_divide_test", ["aten", "portable"]) diff --git a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl index 0038df45dc1..c76c1af67ba 100644 --- a/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl +++ b/shim_et/xplat/executorch/kernels/portable/op_registration_util.bzl @@ -573,6 +573,13 @@ ATEN_OPS = ( "//executorch/kernels/portable/cpu/pattern:pattern", ], ), + op_target( + name = "op_fft_r2c", + deps = [ + "//executorch/runtime/core/exec_aten/util:scalar_type_util", + "//executorch/runtime/core/exec_aten/util:tensor_util", + ], + ), op_target( name = "op_fill", deps = [ From 82c3a8244d05e4f8a251c7cd5a1a208fcf90956c Mon Sep 17 00:00:00 2001 From: Jake Stevens Date: Fri, 11 Sep 2026 08:33:24 -0700 Subject: [PATCH 190/190] NXP backend: remove #22179 workaround for QAT channels-last segfault What: test_mlperf_tiny_classification_mse_cpu_vs_npu used the Python edge reference instead of the portable-kernel C++ reference for the QAT plus channels-last variant, because that configuration segfaulted in the portable kernels. Why: the crash is fixed on current main. The reporter's stack predates two out-of-bounds fixes in the portable dequantize path that this model exercises at runtime (16 per-channel dequantize ops): #21517 fixed an out-of-bounds traversal for non-contiguous (channels-last) inputs, and #21773 fixed misreading int32 zero points as int64. The int32 zero points are QAT-only: QAT emits int32 bias zero points while PTQ emits int64, matching the issue's QAT-only signature. Verification: exported the exact failing configuration (QAT, channels-last, 15-epoch training, dataset calibration, NXP edge passes) and ran it with the portable-kernel executor_runner. It runs cleanly and its outputs bit-match the eager quantized reference. The quantized_kernels_test suite passes 74/74, including the regression tests from both fixes. lintrunner reports no issues on the touched file. The NXP SDK-gated test itself was not run here (no SDK); NXP CI will exercise it. Fixes https://github.com/pytorch/executorch/issues/22179 Authored with AI assistance (Muse Code). --- .../models/test_mlperf_tiny_image_classification.py | 12 +----------- 1 file changed, 1 insertion(+), 11 deletions(-) diff --git a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py index fc0be12f500..6dabb505268 100644 --- a/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py +++ b/backends/nxp/tests/models/test_mlperf_tiny_image_classification.py @@ -80,16 +80,6 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( else None ) - # This model does not work in channels-last format and QAT when running using portable kernels. - # See more information below. - # Github issue: https://github.com/pytorch/executorch/issues/22179 - # NXP internal issue ID: EIEX-1065 - ref_model = ( - ReferenceModel.QUANTIZED_EDGE_PYTHON - if channels_last and use_qat - else ReferenceModel.QUANTIZED_EXECUTORCH_CPP - ) - lower_run_compare( model, [input_spec], @@ -97,7 +87,7 @@ def test_mlperf_tiny_classification_mse_cpu_vs_npu( request, dataset_creator=dataset_creator, output_comparator=comparator, - reference_model=ref_model, + reference_model=ReferenceModel.QUANTIZED_EXECUTORCH_CPP, mocker=mocker, use_qat=use_qat, train_fn=train_fn,