Skip to content

Commit 86896d4

Browse files
committed
[Vulkan] Allow assert-only SymBool chains in the partitioner
Summary: ## Ulterior Motive Restore the "Test WebGPU Native (Dawn)" workflow on `main`, which has failed on every commit since "[Vulkan] Keep unrepresentable and symbolic scalars out of the delegate (pytorch#23247)". ## Rationale **What**: The Vulkan partitioner no longer applies its SymBool/SymFloat rejection to assert ops (`_assert_scalar`, `sym_constrain_range_for_size`) or to nodes whose only users are assert ops. **Why**: pytorch#23247 rejects every node that produces or consumes a SymBool. Models that guard an `.item()` value with `torch._check` produce `u <= N` / `u >= 0` comparisons feeding `_assert_scalar`. Rejecting them cascades through the new "symbolic scalar has an unsupported consumer" check to the `_local_scalar_dense` producer, splitting the graph into multiple delegates. `test_rope_hf.py` then fails with `AssertionError: expected one delegate, got 2`. Asserts are stripped by `RemoveAssertsTransform` during preprocessing, so these SymBool values never need a runtime representation. ## Details - Assert-only SymBool chains - **Before**: Comparisons and asserts rejected; `.item()` producer rejected as a consequence; graph split into multiple delegates. - **After**: Delegated, then removed during preprocessing; graph stays in one delegate. - SymBool/SymFloat values feeding non-assert ops - **Before**: Rejected. - **After**: Unchanged (still rejected). Test Plan: ## Observed Results | Validation | Device | Result | |---|---|---| | Vulkan partition count for `DynamicHfRope` (`test_rope_hf.py`), parent commit | macOS (Apple Silicon) host, partitioner only | 2 partitions; skip logs match the CI failure | | Vulkan partition count for `DynamicHfRope`, this commit | macOS (Apple Silicon) host, partitioner only | 1 partition, no skipped nodes | | `test_rope_hf.py` end-to-end ("Test WebGPU Native (Dawn)") | — | Not run: needs the Dawn/SwiftShader runtime; relying on CI | | `backends/vulkan/test/test_vulkan_dynamic.py` | — | Not run: needs the Vulkan runtime; relying on CI | | Lint | — | Not run: lint tooling not installed locally; relying on CI | ## Reproduction Guide ### Setup ```bash git clone https://github.com/pytorch/executorch.git cd executorch ./install_executorch.sh ``` ### Count Vulkan partitions for the dynamic HF RoPE model ```bash cat > /tmp/count_partitions.py <<'PY' import torch from executorch.exir import to_edge from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.webgpu.test.ops import test_rope_hf as t inputs = t._dynamic_inputs() with torch._dynamo.config.patch(capture_scalar_outputs=True): ep = torch.export.export(t.DynamicHfRope().eval(), inputs) res = VulkanPartitioner().partition(to_edge(ep).exported_program()) tags = { n.meta["delegation_tag"] for n in res.tagged_exported_program.graph.nodes if n.meta.get("delegation_tag") } print("partitions:", len(tags)) PY python /tmp/count_partitions.py 2>&1 | grep -E "Vulkan Partitioner|partitions:" ``` Run once on this commit (expect `partitions: 1`) and once on its parent (`git checkout HEAD~1`, expect `partitions: 2` plus "symbolic float or bool values are not supported" skip logs). Reviewers: This PR was authored with Claude Code.
1 parent 1f517e1 commit 86896d4

1 file changed

Lines changed: 10 additions & 1 deletion

File tree

‎backends/vulkan/partitioner/vulkan_partitioner.py‎

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414

1515
import torch
1616

17+
from executorch.backends.vulkan._passes.remove_asserts import RemoveAssertsTransform
1718
from executorch.backends.vulkan.op_registry import (
1819
get_op_features,
1920
has_impl,
@@ -229,7 +230,15 @@ def _check_node_support(self, node: torch.fx.Node) -> bool: # noqa: C901
229230
self.log_skip(node, "fusable pattern requires 64-bit tensor downcasting")
230231
return False
231232

232-
if any(
233+
# Asserts and the symbolic comparisons feeding them are removed during
234+
# preprocessing, so their SymBool values never need to be represented.
235+
is_assert_chain = node.target in RemoveAssertsTransform.assert_ops or (
236+
len(node.users) > 0
237+
and all(
238+
user.target in RemoveAssertsTransform.assert_ops for user in node.users
239+
)
240+
)
241+
if not is_assert_chain and any(
233242
isinstance(arg.meta.get("val"), (torch.SymFloat, torch.SymBool))
234243
for arg in [node, *node.all_input_nodes]
235244
):

0 commit comments

Comments
 (0)