[release/1.6] Fix OpenVINO scalars, precision, dynamic outputs and reductions

.ci/scripts/test_openvino.sh

@@ -12,5 +12,8 @@ source "$(dirname "${BASH_SOURCE[0]}")/utils.sh"
 
 source openvino/setupvars.sh
 cd backends/openvino/tests
-python test_runner.py --test_type ops
-python test_runner.py --test_type models
+# Run the model tests even when an op test fails, and still fail the job.
+status=0
+python test_runner.py --test_type ops || status=1
+python test_runner.py --test_type models || status=1
+exit "${status}"

backends/openvino/README.md

@@ -48,6 +48,63 @@ print(_get_registered_backend_names())
 # Should include 'OpenvinoBackend'
 ```
 
+## Precision
+
+When the device is `CPU`, a program computes in float32 by default. Without this default,
+OpenVINO would compute in bfloat16 on x86 CPUs with bfloat16 units, which keeps only two to
+three significant digits, and in float16 on ARM CPUs, which keeps about three.
+Both are faster. To trade accuracy for speed, ask for the faster precision when
+lowering:
+
+```python
+from executorch.backends.openvino.partitioner import OpenvinoPartitioner
+from executorch.exir.backend.compile_spec_schema import CompileSpec
+
+OpenvinoPartitioner(
+    [CompileSpec("device", b"CPU"), CompileSpec("inference_precision", b"bf16")]
+)
+```
+
+On an ARM CPU, ask for `f16` instead: there float32 is about 1.8x slower than
+float16, as measured on an Orin Nano.
+
+The value is OpenVINO's `INFERENCE_PRECISION_HINT`: `f32`, `f16` or `bf16`. An
+empty value means the default, and any other value fails at lowering. A CPU
+without float16 support computes an `f16` program in float32. OpenVINO has no
+bfloat16 on ARM CPUs, so there `bf16` keeps the ARM default, float16 where the
+CPU supports it. Which x86 CPUs compute in bfloat16 is up to OpenVINO; see its
+[Precision Control](https://docs.openvino.ai/2026/openvino-workflow/running-inference/optimize-inference/precision-control.html)
+page. Other devices, such as GPUs, keep OpenVINO's default unless the option is
+set.
+
+The default also applies to a program exported in float16: on a CPU it computes
+in float32 unless you lower it with `inference_precision`.
+
+On x86 CPUs, OpenVINO applies the precision when the runtime loads a program.
+So a program lowered without `inference_precision` computes in float32, even
+one exported with an ExecuTorch release that did not have this default. On an
+x86 CPU with bfloat16 units, float32 is slower and can use more memory; export it
+again with `inference_precision` set to `bf16` to get the earlier speed. Other
+x86 CPUs already computed such a program in float32. On ARM CPUs, OpenVINO applies the
+precision when the program is lowered, so a program exported with such an
+earlier release keeps computing in float16 until it is exported again.
+
+OpenVINO's CPU plugin computes 64-bit integers in 32 bits. An integer result
+outside the int32 range wraps around without an error. This applies to every
+int64 tensor, including the int64 result of `cumsum` and `prod` over smaller
+integer types. In OpenVINO 2026.4.1, the CPU plugin also computes some int32 and
+int64 ops in float32. On x86 CPUs these are `sum`, `amax`, `amin`, `max` and
+`min`. On ARM CPUs they are element-wise ops such as `add`, `maximum` and
+`clamp`, and the same reductions along any dimension but the last. Once a value
+or a running total is above 16,777,216 (2^24) or below -16,777,216, such a
+result can be off, even inside the int32 range. For example, on x86 `sum` over a
+tensor that holds 16,777,216 and 1 gives 16,777,216, and so does `amax` over one
+that holds 16,777,217.
+
+In OpenVINO 2026.4.1, the CPU plugin's `min(dim)` and `max(dim)` can return a
+later index on ties than PyTorch, which returns the first one. This happens with
+float32 as well as integer dtypes such as int32 and uint8.
+
 ## Directory Structure
 
 ```

backends/openvino/_passes/__init__.py

@@ -4,6 +4,12 @@
 # except in compliance with the License. See the license file found in the
 # LICENSE file in the root directory of this source tree.
 
+from .cast_reduction_input_pass import CastReductionInputPass
 from .decompose_floor_divide_pass import DecomposeFloorDividePass
+from .fold_constant_cumsum_pass import FoldConstantCumsumPass
 
-__all__ = ["DecomposeFloorDividePass"]
+__all__ = [
+    "CastReductionInputPass",
+    "DecomposeFloorDividePass",
+    "FoldConstantCumsumPass",
+]

backends/openvino/_passes/cast_reduction_input_pass.py

@@ -0,0 +1,45 @@
+# 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.exir.dialects._ops import ops as exir_ops
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+
+REDUCTION_OPS = {
+    exir_ops.edge.aten.cumsum.default,
+    exir_ops.edge.aten.prod.default,
+    exir_ops.edge.aten.prod.dim_int,
+}
+
+
+class CastReductionInputPass(PassBase):
+    """Cast the input of cumsum and prod to the dtype of their result.
+
+    PyTorch reduces bool and integer tensors in int64, but OpenVINO keeps a
+    narrower type for these two, though not for sum. A cumsum over a bool mask
+    is computed in 8 bits and wraps after 255, and a cumsum or prod over int32
+    returns int32.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        modified = False
+        for node in graph_module.graph.nodes:
+            if node.op != "call_function" or node.target not in REDUCTION_OPS:
+                continue
+            source = node.args[0]
+            dtype = node.meta["val"].dtype
+            if source.meta["val"].dtype == dtype:
+                continue
+            with graph_module.graph.inserting_before(node):
+                cast = graph_module.graph.call_function(
+                    exir_ops.edge.aten._to_copy.default, (source,), {"dtype": dtype}
+                )
+            cast.meta["val"] = source.meta["val"].to(dtype)
+            node.replace_input_with(source, cast)
+            modified = True
+        if modified:
+            graph_module.recompile()
+        return PassResult(graph_module, modified)

backends/openvino/_passes/fold_constant_cumsum_pass.py

@@ -0,0 +1,161 @@
+# 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 logging
+import operator
+
+import torch
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.passes.constant_prop_pass import _nodes_written_in_place
+from torch.fx.passes.infra.pass_base import PassBase, PassResult
+from torch.utils import _pytree as pytree
+
+
+def _reads_only_shape(node: torch.fx.Node) -> bool:
+    """full_like with a constant fill value reads only the shape and dtype of
+    its input. With a static shape its result is a constant, and OpenVINO
+    folds it into one too.
+    """
+    if node.target != exir_ops.edge.aten.full_like.default:
+        return False
+    if any(
+        isinstance(a, torch.fx.Node) for a in (*node.args[1:], *node.kwargs.values())
+    ):
+        return False
+    val = node.meta.get("val")
+    return isinstance(val, torch.Tensor) and all(isinstance(d, int) for d in val.shape)
+
+
+def _inputs(node: torch.fx.Node) -> list:
+    return [] if _reads_only_shape(node) else node.all_input_nodes
+
+
+def _on_cpu(value):
+    if isinstance(value, torch.Tensor):
+        return value.cpu()
+    if isinstance(value, torch.device):
+        return torch.device("cpu")
+    return value
+
+
+def _constant_nodes(graph_module: torch.fx.GraphModule) -> set:
+    """Nodes that depend only on constants. Graph order puts inputs first."""
+    constants = set()
+    written = _nodes_written_in_place(graph_module)
+    for node in graph_module.graph.nodes:
+        # Only tensors: a higher-order op such as cond reads its subgraphs,
+        # which may hold a random op, through get_attr nodes too.
+        if node.op == "get_attr":
+            value = operator.attrgetter(node.target)(graph_module)
+            if isinstance(value, torch.Tensor) and node not in written:
+                constants.add(node)
+        elif node.op == "call_function" and (
+            _reads_only_shape(node)
+            or (
+                torch.Tag.nondeterministic_seeded
+                not in getattr(node.target, "tags", ())
+                and all(n in constants for n in node.all_input_nodes)
+            )
+        ):
+            constants.add(node)
+    return constants
+
+
+def _evaluate(
+    graph_module: torch.fx.GraphModule, target: torch.fx.Node
+) -> torch.Tensor:
+    """Value of target, computing only what it reads, on the CPU."""
+    needed = set()
+    stack: list = [target]
+    while stack:
+        node = stack.pop()
+        if node not in needed:
+            needed.add(node)
+            stack.extend(_inputs(node))
+    order = [n for n in graph_module.graph.nodes if n in needed]
+    last_use = {i: n for n in order for i in _inputs(n)}
+    values = {}
+    for node in order:
+        if node.op == "get_attr":
+            values[node] = operator.attrgetter(node.target)(graph_module)
+            continue
+        if _reads_only_shape(node):
+            like = node.args[0].meta["val"]
+            args = (
+                torch.empty(like.shape, dtype=like.dtype, device="cpu"),
+                *node.args[1:],
+            )
+            kwargs = dict(node.kwargs)
+        else:
+            args, kwargs = torch.fx.node.map_arg(
+                (node.args, node.kwargs), lambda n: values[n]
+            )
+        # The values are the same on any device, and lowering for OpenVINO
+        # must not need the GPU a constant was made for.
+        args, kwargs = pytree.tree_map(_on_cpu, (args, kwargs))
+        if kwargs.get("pin_memory"):
+            kwargs = {**kwargs, "pin_memory": False}
+        values[node] = node.target(*args, **kwargs)
+        for i in _inputs(node):
+            if last_use[i] is node:
+                del values[i]
+    return values[target]
+
+
+class FoldConstantCumsumPass(PassBase):
+    """Replace an integer cumsum that reads only constants with its result.
+
+    With an integer cumsum over a constant, OpenVINO's CPU plugin (2025.1 to
+    2026.4) gives wrong results from the second inference on: masks built from
+    it keep only what earlier calls allowed. pi0 builds its attention masks
+    this way. The OpenVINO nightly no longer does this, so the pass can go once
+    the oldest supported OpenVINO release has the fix. Only these cumsums are
+    folded; constants that feed no such cumsum, such as compressed weights, are
+    left as they are. Runs on an unlifted module, where constants are get_attr
+    nodes; buffers the graph writes to are not constants, and full_like of an
+    input with a static shape is one.
+    """
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        graph = graph_module.graph
+        constants = _constant_nodes(graph_module)
+        cumsums = [
+            n
+            for n in graph.nodes
+            if n in constants
+            and n.target == exir_ops.edge.aten.cumsum.default
+            and not n.meta["val"].dtype.is_floating_point
+            and not n.meta["val"].dtype.is_complex
+        ]
+        modified = False
+        # One at a time, so a cumsum that cannot be evaluated does not keep
+        # the others from being folded.
+        for node in cumsums:
+            try:
+                value = _evaluate(graph_module, node)
+            except RuntimeError as e:
+                # Some ops have modes the CPU lacks.
+                logging.warning(
+                    "OpenVINO: could not fold %s, so on a CPU some OpenVINO "
+                    "releases give wrong results for it from the second "
+                    "inference on: %s",
+                    node.name,
+                    e,
+                )
+                continue
+            name = f"_folded_cumsum_{node.name}"
+            while hasattr(graph_module, name):
+                name += "_"
+            graph_module.register_buffer(name, value)
+            with graph.inserting_before(node):
+                folded = graph.get_attr(name)
+            folded.meta = node.meta.copy()
+            node.replace_all_uses_with(folded)
+            modified = True
+        if modified:
+            graph.eliminate_dead_code()
+            graph_module.recompile()
+        return PassResult(graph_module, modified)