OpenVINO: fix bool min/max and masked read lowering, and follow-ups to #24294

.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

@@ -82,16 +82,28 @@ 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. That
-is slower and can use more memory. Export it again with `inference_precision`
-set to `bf16` to get the earlier speed. On ARM CPUs, OpenVINO applies the
+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.
+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/fold_constant_cumsum_pass.py

@@ -64,10 +64,12 @@ def _constant_nodes(graph_module: torch.fx.GraphModule) -> set:
     return constants
 
 
-def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict:
-    """Values of targets, computing only what they read, on the CPU."""
+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(targets)
+    stack: list = [target]
     while stack:
         node = stack.pop()
         if node not in needed:
@@ -82,7 +84,10 @@ def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict:
             continue
         if _reads_only_shape(node):
             like = node.args[0].meta["val"]
-            args = (torch.empty(like.shape, dtype=like.dtype), *node.args[1:])
+            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(
@@ -95,15 +100,15 @@ def _evaluate(graph_module: torch.fx.GraphModule, targets: list) -> dict:
             kwargs = {**kwargs, "pin_memory": False}
         values[node] = node.target(*args, **kwargs)
         for i in _inputs(node):
-            if last_use[i] is node and i not in targets:
+            if last_use[i] is node:
                 del values[i]
-    return values
+    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.4 to
+    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
@@ -130,7 +135,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         # the others from being folded.
         for node in cumsums:
             try:
-                value = _evaluate(graph_module, [node])[node]
+                value = _evaluate(graph_module, node)
             except RuntimeError as e:
                 # Some ops have modes the CPU lacks.
                 logging.warning(

backends/openvino/partitioner.py

@@ -33,11 +33,33 @@ def __init__(self):
         self.op_types = {}
 
 
+def _reads_only_bool(node: torch.fx.Node) -> bool:
+    """Whether OpenVINO would compute the node in bool. It promotes mixed tensor
+    inputs, so maximum(bool, float) is computed in float and converts fine."""
+    vals = [n.meta.get("val") for n in node.all_input_nodes]
+    tensors = [val for val in vals if isinstance(val, torch.Tensor)]
+    return bool(tensors) and all(t.dtype == torch.bool for t in tensors)
+
+
 class OpenvinoOperatorsSupport(OperatorSupportBase):
     extended_support_dict = {
         "torch.ops.dim_order_ops._clone_dim_order.default": None,
         "torch.ops.dim_order_ops._to_dim_order_copy.default": None,
     }
+    # OpenVINO converts these to ReduceMin, ReduceMax, Minimum, Maximum or Sign,
+    # none of which take bool tensors.
+    bool_unsupported_ops = {
+        "aten.amax.default",
+        "aten.amin.default",
+        "aten.clamp.default",
+        "aten.max.default",
+        "aten.max.dim",
+        "aten.maximum.default",
+        "aten.min.default",
+        "aten.min.dim",
+        "aten.minimum.default",
+        "aten.sign.default",
+    }
 
     def __init__(
         self,
@@ -87,6 +109,10 @@ def is_node_supported(self, _, node: torch.fx.Node) -> bool:
         if op_type == "getitem":
             return True
 
+        if op_type in self.bool_unsupported_ops and _reads_only_bool(node):
+            print("Op not supported for bool tensors: ", "torch.ops." + str(op_type))
+            return False
+
         if "torch.ops." + str(op_type) in supported_ops:
             return True
         else:

backends/openvino/preprocess.py

@@ -27,13 +27,42 @@
 from openvino.frontend.pytorch.fx_decoder import (  # type: ignore[import-untyped]
     TorchFXPythonDecoder,
 )
+from openvino.frontend.pytorch.utils import (  # type: ignore[import-untyped]
+    pt_to_ov_type_map,
+)
 from openvino.preprocess import PrePostProcessor  # type: ignore[import-untyped]
+from torch.fx.passes.shape_prop import _extract_tensor_metadata
 
 # The values of the inference_precision compile spec. The runtime
 # (OpenvinoBackend.cpp) accepts the same ones.
 _PRECISIONS = ("f32", "f16", "bf16")
 
 
+def _match_partition_inputs(
+    model, module: torch.fx.GraphModule, input_shapes: list, input_types: list
+) -> None:
+    """Make the converted model take the inputs the runtime passes."""
+    # The converter takes an unused first input for forward's self and removes
+    # it, but the runtime passes every partition input. The cumsum fold can
+    # leave the first input unused.
+    parameters = list(model.get_parameters())
+    placeholders = [n for n in module.graph.nodes if n.op == "placeholder"]
+    if len(parameters) == len(input_types) - 1 and not placeholders[0].users:
+        for parameter in parameters:
+            model.remove_parameter(parameter)
+        first = opset13.parameter(
+            input_shapes[0], pt_to_ov_type_map[str(input_types[0])]
+        )
+        model.add_parameters([first, *parameters])
+    # The runtime binds inputs by position and refuses a wrong count, so any
+    # other mismatch fails here instead of on every call.
+    if len(model.get_parameters()) != len(input_types):
+        raise RuntimeError(
+            f"OpenVINO backend: expected {len(input_types)} parameters in the "
+            f"converted model, got {len(model.get_parameters())}"
+        )
+
+
 @final
 class OpenvinoBackend(BackendDetails):
 
@@ -96,15 +125,24 @@ def preprocess(
                         f"{type(val).__name__}; only tensors and ints are supported"
                     )
 
-        # The converter chooses some lowerings (e.g. boolean-mask indexing) from
-        # the input types, so they must be known before conversion starts.
         module = edge_program.module()
         FoldConstantCumsumPass()(module)
+        # The converter reads ranks only from tensor_meta, which edge graphs
+        # leave empty, and some ops need one, like aten.index after
+        # aten.index_put.
+        for node in module.graph.nodes:
+            val = node.meta.get("val")
+            if isinstance(val, torch.Tensor) and not node.meta.get("tensor_meta"):
+                node.meta["tensor_meta"] = _extract_tensor_metadata(val)
+        # The converter chooses some lowerings (e.g. boolean-mask indexing) from
+        # the input types, so they must be known before conversion starts.
         decoder = TorchFXPythonDecoder(
             module, input_shapes=input_shapes, input_types=input_types
         )
         frontend = FrontEndManager().load_by_framework("pytorch")
-        ppp = PrePostProcessor(frontend.convert(frontend.load(decoder)))
+        model = frontend.convert(frontend.load(decoder))
+        _match_partition_inputs(model, module, input_shapes, input_types)
+        ppp = PrePostProcessor(model)
         # The converter takes an int input, like a SymInt size, as a scalar, but
         # the runtime passes it as an int64 tensor of shape [1].
         for i in int_inputs:

backends/openvino/runtime/OpenvinoApi.h

@@ -22,6 +22,7 @@ typedef struct ov_core ov_core_t;
 typedef struct ov_compiled_model ov_compiled_model_t;
 typedef struct ov_infer_request ov_infer_request_t;
 typedef struct ov_tensor ov_tensor_t;
+typedef struct ov_output_const_port ov_output_const_port_t;
 
 // Value types reproduced from openvino/c/ov_shape.h and ov_common.h.
 // These are stable C ABI — pinned via version constraint in pyproject.toml.
@@ -30,6 +31,17 @@ typedef struct {
   int64_t* dims;
 } ov_shape_t;
 
+// Reproduced from openvino/c/ov_dimension.h and ov_partial_shape.h.
+typedef struct {
+  int64_t min;
+  int64_t max;
+} ov_dimension_t;
+
+typedef struct {
+  ov_dimension_t rank;
+  ov_dimension_t* dims;
+} ov_partial_shape_t;
+
 typedef struct {
   char** devices;
   size_t size;
@@ -92,6 +104,13 @@ using ov_compiled_model_inputs_size_fn =
 using ov_compiled_model_outputs_size_fn =
     ov_status_e (*)(const ov_compiled_model_t*, size_t*);
 using ov_compiled_model_free_fn = void (*)(ov_compiled_model_t*);
+using ov_compiled_model_output_by_index_fn = ov_status_e (*)(
+    const ov_compiled_model_t*,
+    size_t,
+    ov_output_const_port_t**);
+using ov_port_get_partial_shape_fn =
+    ov_status_e (*)(const ov_output_const_port_t*, ov_partial_shape_t*);
+using ov_output_const_port_free_fn = void (*)(ov_output_const_port_t*);
 using ov_infer_request_set_input_tensor_by_index_fn =
     ov_status_e (*)(ov_infer_request_t*, size_t, const ov_tensor_t*);
 using ov_infer_request_set_output_tensor_by_index_fn =
@@ -111,6 +130,7 @@ using ov_tensor_data_fn = ov_status_e (*)(const ov_tensor_t*, void**);
 using ov_shape_create_fn =
     ov_status_e (*)(int64_t, const int64_t*, ov_shape_t*);
 using ov_shape_free_fn = ov_status_e (*)(ov_shape_t*);
+using ov_partial_shape_free_fn = void (*)(ov_partial_shape_t*);
 
 // Defined in OpenvinoBackend.cpp, where the platform's loader API is included.
 struct DlCloser {
@@ -132,6 +152,9 @@ struct OpenvinoFunctions {
   ov_compiled_model_inputs_size_fn compiled_model_inputs_size = nullptr;
   ov_compiled_model_outputs_size_fn compiled_model_outputs_size = nullptr;
   ov_compiled_model_free_fn compiled_model_free = nullptr;
+  ov_compiled_model_output_by_index_fn compiled_model_output_by_index = nullptr;
+  ov_port_get_partial_shape_fn port_get_partial_shape = nullptr;
+  ov_output_const_port_free_fn output_const_port_free = nullptr;
   ov_infer_request_set_input_tensor_by_index_fn
       infer_request_set_input_tensor_by_index = nullptr;
   ov_infer_request_set_output_tensor_by_index_fn
@@ -148,6 +171,7 @@ struct OpenvinoFunctions {
   ov_tensor_data_fn tensor_data = nullptr;
   ov_shape_create_fn shape_create = nullptr;
   ov_shape_free_fn shape_free = nullptr;
+  ov_partial_shape_free_fn partial_shape_free = nullptr;
 };
 
 } // namespace openvino