From c7fad4469b5c10e35f0ec5331e3b509e45b466b2 Mon Sep 17 00:00:00 2001 From: Jacob Szwejbka Date: Tue, 3 Nov 2026 00:02:47 -0800 Subject: [PATCH] [release/1.6] Fix OpenVINO bool operations and masked reads Cherry-pick the bool partitioning and tensor-metadata fixes from #24396. Keep the matching bool tests in their own module and leave the core OpenVINO follow-ups in the preceding release pick. Prepared with OpenAI Codex. --- backends/openvino/partitioner.py | 26 +++++++++ backends/openvino/preprocess.py | 8 +++ .../openvino/tests/ops/test_bool_min_max.py | 56 +++++++++++++++++++ backends/openvino/tests/ops/test_index_put.py | 20 +++++++ 4 files changed, 110 insertions(+) create mode 100644 backends/openvino/tests/ops/test_bool_min_max.py diff --git a/backends/openvino/partitioner.py b/backends/openvino/partitioner.py index 0d407e33f6e..5b01995daa5 100644 --- a/backends/openvino/partitioner.py +++ b/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: diff --git a/backends/openvino/preprocess.py b/backends/openvino/preprocess.py index 7a3f490b258..55bad961782 100644 --- a/backends/openvino/preprocess.py +++ b/backends/openvino/preprocess.py @@ -24,6 +24,7 @@ TorchFXPythonDecoder, ) from openvino.preprocess import PrePostProcessor # type: ignore[import-untyped] +from torch.fx.passes.shape_prop import _extract_tensor_metadata @final @@ -76,6 +77,13 @@ def preprocess( # The converter chooses some lowerings (e.g. boolean-mask indexing) from # the input types, so they must be known before conversion starts. + # 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) decoder = TorchFXPythonDecoder( edge_program.module(), input_shapes=input_shapes, input_types=input_types ) diff --git a/backends/openvino/tests/ops/test_bool_min_max.py b/backends/openvino/tests/ops/test_bool_min_max.py new file mode 100644 index 00000000000..5c349cc8c37 --- /dev/null +++ b/backends/openvino/tests/ops/test_bool_min_max.py @@ -0,0 +1,56 @@ +# mypy: disable-error-code=import-not-found + +import torch +from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( + BaseOpenvinoOpTest, +) + + +class TestBoolMinMax(BaseOpenvinoOpTest): + """OpenVINO converts min, max, clamp and sign to ops that take no bool tensors.""" + + def test_bool_ops_run_on_the_cpu(self): + class BoolMinMax(torch.nn.Module): + def forward(self, x): + m = x > 0 + n = m.flip(1) + return ( + *m.min(1), + *m.max(1), + m.amin(1), + m.amax(1), + m.min(), + m.max(), + torch.minimum(m, n), + torch.maximum(m, n), + torch.sign(m), + m.clamp(min=0.5), + ) + + # One row all true and one all false, so min and max differ by row and + # their indices must pick the first of the ties. + x = torch.randn(3, 7) + x[0] = x[0].abs() + 0.1 + x[1] = -x[1].abs() - 0.1 + self.execute_layer_test(BoolMinMax(), (x,), expected_partitions=2) + + def test_bool_min_of_a_buffer(self): + class BufferMin(torch.nn.Module): + def __init__(self): + super().__init__() + self.register_buffer("w", torch.tensor([1.0, -2.0, 3.0])) + + def forward(self, x): + return x + (self.w > 0).min().float() + + # The comparison reads only the buffer, so its partition has no inputs. + self.execute_layer_test(BufferMin(), (torch.randn(3),), expected_partitions=2) + + def test_maximum_of_bool_and_float_stays_with_openvino(self): + class Maximum(torch.nn.Module): + def forward(self, mask, x): + return torch.maximum(mask, x) + + # Computed in float, which OpenVINO converts. + x = torch.randn(2, 7) + self.execute_layer_test(Maximum(), (x > 0, x)) diff --git a/backends/openvino/tests/ops/test_index_put.py b/backends/openvino/tests/ops/test_index_put.py index 238ffbc3206..0dd57cac923 100644 --- a/backends/openvino/tests/ops/test_index_put.py +++ b/backends/openvino/tests/ops/test_index_put.py @@ -1,7 +1,12 @@ +# mypy: disable-error-code=import-not-found + import torch from executorch.backends.openvino.tests.ops.base_openvino_op_test import ( BaseOpenvinoOpTest, ) +from openvino.frontend.pytorch.torchdynamo.op_support import ( # type: ignore[import-untyped] + OperatorSupport, +) class TestIndexPutOperator(BaseOpenvinoOpTest): @@ -26,3 +31,18 @@ def test_index_put_bool_mask(self): torch.randn(2, 8).to(dtype), ) self.execute_layer_test(module, sample_input) + + def test_index_after_index_put(self): + if "torch.ops.aten.index_put.default" not in OperatorSupport([])._support_dict: + self.skipTest("this OpenVINO release does not delegate index_put") + + class WriteThenRead(torch.nn.Module): + def forward(self, x, mask): + x = x.clone() + x[mask] = x[mask] * 2 + # The read needs the rank of what the write returns. + return x + x[mask].sum() + + self.execute_layer_test( + WriteThenRead(), (torch.randn(2, 8), torch.rand(2, 8) > 0.5) + )