[release/1.6] Fix OpenVINO bool operations and masked reads

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

@@ -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 @@
 
         # 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
         )

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))

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)
+        )