skip op in partitioning if there are bool input tensors

Summary: Vulkan backend does not support bool tensors

Differential Revision: D69273733

Author: nathanaelsee

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -98,6 +98,10 @@ def op_node_is_compatible(
                 and utils.is_tensor_node(arg)
                 and i not in features.skip_limits_check
             ):
+                # Check for bool inputs
+                if utils.is_tensor_node(arg) and utils.tensor_node_is_bool(arg):
+                    return False, "contains bool tensor"
+
                 # Check for high dimensional tensors
                 if utils.is_tensor_node(arg) and utils.tensor_node_is_high_dim(arg):
                     return False, "contains high dim tensor"

backends/vulkan/utils.py

@@ -80,6 +80,19 @@ def is_tensor_node(node: torch.fx.Node) -> bool:
     return False
 
 
+def tensor_node_is_bool(node: torch.fx.Node) -> bool:
+    """
+    Returns true if a given node contains a tensor with bool dtype
+    """
+    if isinstance(node.meta["val"], FakeTensor):
+        return node.meta["val"].dtype == torch.bool
+    if isinstance(node.meta["val"], list) or isinstance(node.meta["val"], tuple):
+        for fake_tensor in node.meta["val"]:
+            if isinstance(fake_tensor, FakeTensor):
+                if fake_tensor.dtype == torch.bool:
+                    return True
+    return False
+
 ##
 ## Memory Layout, Storage Type Determination
 ##