Repository navigation
Enabled pi05 support - #24024
Closed
suryasidd wants to merge 3 commits into
Closed
Enabled pi05 support#24024suryasidd wants to merge 3 commits into
suryasidd wants to merge 3 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24024
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit 158116e with merge base 4db778e ( This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
Missing replacement-node metadata breaks chained rewrites, and automated regression coverage is absent.
2 open issues
What changed in this PR
Adds OpenVINO boolean-operation rewrites to support pi05 in Physical AI Studio through ExecuTorch.
Changes:
- Casts boolean reduction inputs to integers and replaces boolean multiplication with bitwise AND.
- Registers the new pass in OpenVINO preprocessing.
| File | Description |
|---|---|
| backends/openvino/preprocess.py | Runs the boolean rewrite pass. |
| backends/openvino/_passes/decompose_boolean_reduction_pass.py | Implements boolean reduction and multiplication rewrites. |
| backends/openvino/_passes/__init__.py | Exports the new pass. |
Review effort: Balanced
Comment on lines
+86
to
+91
| if node.target in SUM_OPS: | ||
| self._rewrite_reduction(graph, node, "sum") | ||
| elif node.target in CUMSUM_OPS: | ||
| self._rewrite_reduction(graph, node, "cumsum") | ||
| elif node.target in MUL_OPS: | ||
| self._rewrite_bool_mul(graph, node) |
shoumikhin
approved these changes
Oct 27, 2026
shoumikhin
left a comment
shoumikhin
left a comment
Contributor
There was a problem hiding this comment.
Thanks, this makes pi0.5 lower through OpenVINO, and the cases it targets match eager in my runs. Approving; the comments below are suggestions, not blockers.
| for pass_cls in [ | ||
| DimOrderOpsRevertPass, | ||
| DecomposeFloorDividePass, | ||
| DecomposeBooleanReductionPass, |
Contributor
There was a problem hiding this comment.
I think the root cause is one step earlier. When preprocess hands the graph to OpenVINO, the input placeholders have meta["val"] but no meta["tensor_meta"], and OpenVINO 2026.4 reads an input's dtype only from tensor_meta. So a bool input gets an unknown type and OpenVINO skips its own bool handling. Filling it in before openvino_compile fixes bool sum and mul without this pass, and also bool add and prod and int32 sum, which are still wrong with this PR:
import torch
from torch.fx.passes.shape_prop import _extract_tensor_metadata
for node in edge_program.graph.nodes:
if node.op == "placeholder" and isinstance(node.meta.get("val"), torch.Tensor):
node.meta["tensor_meta"] = _extract_tensor_metadata(node.meta["val"])With this plus the cumsum change below, pi0 and pi0.5 give the same outputs as this PR. SmolVLA also gets through lowering, which it does not with this PR. OpenVINO main already falls back to meta["val"] (openvinotoolkit/openvino#38078), so the sum and mul rewrites would not be needed after the next OpenVINO update. Happy for this to be a follow-up.
| if input_node is None or _node_dtype(input_node) != torch.bool: | ||
| return | ||
|
|
||
| # Pin the output dtype: OpenVINO keeps int32 for cumsum(int32) while |
Contributor
There was a problem hiding this comment.
The cumsum part fixes a real bug: without it, cumsum over a bool input runs and silently returns bools instead of counts. The comment names cumsum(int32), but the pass only rewrites bool inputs, so a cumsum over an int32 input still returns int32 where the program expects int64. Setting only the dtype keyword on any integer or bool cumsum fixes both, and the extra int32 cast is not needed.
Contributor
There was a problem hiding this comment.
Correction to my suggestion here: setting only dtype on cumsum is not safe either. pi0 and pi0.5 build their attention masks with a cumsum over a constant. When that cumsum is int64, OpenVINO's CPU plugin goes wrong on a reused program. The first call is right. After that, each call gets the AND of its mask with every earlier mask, so the mask never grows back. It reproduces in plain OpenVINO on 2025.4.1, 2026.0.0 and 2026.4.1, but not on the current nightly. The PR's int32 cast hits it too. Running constant_prop_pass before lowering avoided it. This was checked on a small model with the same mask code, not on full pi0.5.
| return None | ||
|
|
||
|
|
||
| class DecomposeBooleanReductionPass(ExportPass): |
Contributor
There was a problem hiding this comment.
Could you add a small op test under backends/openvino/tests/ops/? This pass runs on every OpenVINO compile, but no existing test gives it a bool input, so its rewrites never run in CI. A bool-input cumsum returns the wrong dtype on main today, so one case for it, and one for a bool-input sum, would guard both paths.
shoumikhin
requested changes
Oct 27, 2026
shoumikhin
left a comment
shoumikhin
left a comment
Contributor
There was a problem hiding this comment.
Sorry, changing my approval, see my cumsum reply. With current OpenVINO releases, the rewritten cumsum gives wrong pi0-style masks on a reused program, while main is right on every call. The bug is in OpenVINO, but this change exposes it. Folding the constant cumsum before lowering avoided it. A test that runs one loaded program twice, with more valid tokens the second time, would catch it.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 29, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws inside cond and float cumsums
are not folded, that a cumsum over a split constant is, and that a
folded result does not replace a buffer.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 29, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are. The fold computes on the CPU, so lowering
does not need the GPU a constant was created on.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws (direct and inside cond) and
float cumsums are not folded, that cumsums over a split constant and
over a CUDA constant without a GPU are, and that a folded result does
not replace a buffer.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 29, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are. The fold computes on the CPU, so lowering
does not need the GPU, or the pinned memory, a constant was created
with.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws (direct and inside cond) and
float cumsums are not folded, that cumsums over a split constant and
over CUDA and pinned constants without a GPU are, and that a folded
result does not replace a buffer.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 30, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are. The fold computes on the CPU, so lowering
does not need the GPU, or the pinned memory, a constant was created
with. If the CPU cannot compute such a constant, the cumsum is left to
OpenVINO as before.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws (direct and inside cond) and
float cumsums are not folded, that cumsums over a split constant and
over CUDA and pinned constants without a GPU are, and that a folded
result does not replace a buffer.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 30, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are. The fold computes on the CPU, so lowering
does not need the GPU, or the pinned memory, a constant was created
with. If the CPU cannot compute such a constant, the cumsum is left to
OpenVINO as before.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws (direct and inside cond) and
float cumsums are not folded, that cumsums over a split constant and
over CUDA and pinned constants without a GPU are, and that a folded
result does not replace a buffer.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 30, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are. The fold computes on the CPU, so lowering
does not need the GPU, or the pinned memory, a constant was created
with. If the CPU cannot compute such a constant, the cumsum is left to
OpenVINO as before.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws (direct and inside cond) and
float cumsums are not folded, that cumsums over a split constant and
over CUDA and pinned constants without a GPU are, and that a folded
result does not replace a buffer.
shoumikhin
added a commit
to shoumikhin/executorch
that referenced
this pull request
Oct 30, 2026
PyTorch computes sum, cumsum and prod of a bool or integer tensor in
int64. OpenVINO's PyTorch frontend keeps a narrower type for some of
them. A cumsum over a bool mask is computed in 8 bits, so the count
wraps after 255. A cumsum or prod over int32 returns int32. When such a
result is a delegate output, the runtime rejects it at execute:
m = torch.rand(2, 7) > 0.5
torch.cumsum(m, 1) # CALL_DELEGATE execute failed: 0x1
Inside a graph the wrong values flow on silently. pi0 and pi0.5
compute their position ids as cumsum(pad_masks) - 1 over more than
255 tokens, so every position from 255 on was wrong and the policies
gave wrong actions.
Cast the input of these reductions to their result dtype before the
graph goes to OpenVINO. The open pytorch#24024 proposes a similar cast for
bool sum and cumsum.
The cast exposes an OpenVINO CPU bug, seen on 2025.4.1, 2026.0.0 and
2026.4.1 but not on the nightly: with an int64 cumsum over a constant,
a loaded program gives wrong results from the second inference on.
pi0 builds its attention masks from such a cumsum, and each call's mask
kept only what earlier calls allowed. So the backend now replaces an
int64 cumsum that reads only constants with its result before
converting. Buffers the graph writes to, random ops and higher-order ops
are not treated as constants, and other constants, such as compressed
weights, are left as they are. The fold computes on the CPU, so lowering
does not need the GPU, or the pinned memory, a constant was created
with. If the CPU cannot compute such a constant, the cumsum is left to
OpenVINO as before.
The new tests run cumsum over bool and int32, prod over int32, and
pi0-style attention masks and positions over 816 tokens through one
loaded program twice, with more valid tokens the second time. Pass tests
check that mutated buffers, random draws (direct and inside cond) and
float cumsums are not folded, that cumsums over a split constant and
over CUDA and pinned constants without a GPU are, and that a folded
result does not replace a buffer.
Collaborator
Author
|
Changes addressed in #24294. Closing this. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Enabled the necessary support required for pi05 model integration in Physical AI studio with Executorch backend support.
Test plan
Tested this with Physical AI sample application.