[ET-VK][ops] Add bitwise_or / logical_or operators

Pull Request resolved: #20382

Adds Vulkan support for `aten.bitwise_or.Tensor` and `aten.logical_or.default`, mirroring the existing `bitwise_and` / `logical_and` implementation. This is the first of two ops needed to collapse the Llama4-mini TISO en_US backbone export to a single Vulkan partition: the discrete-speech mask OR-s several bool tensors via `bitwise_or`, which previously had no Vulkan implementation and forced a CPU fallback that split the delegated graph.

Implementation mirrors `bitwise_and`: a `X | Y` uint8 shader variant in `binary_op_buffer.yaml` / `binary_op_texture.yaml`, a `DEFINE_BINARY_OP_FN(bitwise_or)` dispatch with `VK_REGISTER_OP` for both `aten.bitwise_or.Tensor` and `aten.logical_or.default` in `BinaryOp.cpp`, and `register_bitwise_or` / `register_logical_or` `OpFeatures` (bool inputs) in `op_registry.py`.

This change was authored with Claude.
ghstack-source-id: 397529322
@exported-using-ghexport

Differential Revision: [D108457794](https://our.internmc.facebook.com/intern/diff/D108457794/)

Author: ssjia

backends/vulkan/op_registry.py

@@ -294,28 +294,16 @@ def register_comparison_ops():
 # =============================================================================
 
 
-@update_features(exir_ops.edge.aten.bitwise_and.Tensor)
-def register_bitwise_and():
-    return OpFeatures(
-        inputs_storage=utils.ANY_STORAGE,
-        inputs_dtypes=utils.BOOL_T,
-        supports_resize=True,
-        supports_highdim=True,
-    )
-
-
-@update_features(exir_ops.edge.aten.bitwise_not.default)
-def register_bitwise_not():
-    return OpFeatures(
-        inputs_storage=utils.ANY_STORAGE,
-        inputs_dtypes=utils.BOOL_T,
-        supports_resize=True,
-        supports_highdim=True,
-    )
-
-
-@update_features(exir_ops.edge.aten.logical_and.default)
-def register_logical_and():
+@update_features(
+    [
+        exir_ops.edge.aten.bitwise_and.Tensor,
+        exir_ops.edge.aten.bitwise_or.Tensor,
+        exir_ops.edge.aten.bitwise_not.default,
+        exir_ops.edge.aten.logical_and.default,
+        exir_ops.edge.aten.logical_or.default,
+    ]
+)
+def register_bool_binary_ops():
     return OpFeatures(
         inputs_storage=utils.ANY_STORAGE,
         inputs_dtypes=utils.BOOL_T,

backends/vulkan/runtime/graph/ops/glsl/binary_op_buffer.yaml

@@ -52,3 +52,8 @@ binary_op_buffer:
       generate_variant_forall:
         DTYPE:
           - VALUE: uint8
+    - NAME: binary_bitwise_or_buffer
+      OPERATOR: X | Y
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: uint8

backends/vulkan/runtime/graph/ops/glsl/binary_op_texture.yaml

@@ -54,3 +54,8 @@ binary_op_texture:
       generate_variant_forall:
         DTYPE:
           - VALUE: uint8
+    - NAME: binary_bitwise_or_texture3d
+      OPERATOR: X | Y
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: uint8

backends/vulkan/runtime/graph/ops/impl/BinaryOp.cpp

@@ -143,6 +143,7 @@ DEFINE_BINARY_OP_FN(le);
 DEFINE_BINARY_OP_FN(gt);
 DEFINE_BINARY_OP_FN(ge);
 DEFINE_BINARY_OP_FN(bitwise_and);
+DEFINE_BINARY_OP_FN(bitwise_or);
 
 REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.add.Tensor, add);
@@ -159,6 +160,8 @@ REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.ge.Tensor, ge);
   VK_REGISTER_OP(aten.bitwise_and.Tensor, bitwise_and);
   VK_REGISTER_OP(aten.logical_and.default, bitwise_and);
+  VK_REGISTER_OP(aten.bitwise_or.Tensor, bitwise_or);
+  VK_REGISTER_OP(aten.logical_or.default, bitwise_or);
 }
 
 } // namespace vkcompute

backends/vulkan/test/op_tests/cases.py

@@ -2141,14 +2141,18 @@ def get_where_inputs():
     return test_suite
 
 
-@register_test_suite("aten.bitwise_and.Tensor")
-def get_bitwise_and_inputs():
+@register_test_suite(
+    ["aten.bitwise_and.Tensor", "aten.bitwise_or.Tensor", "aten.logical_or.default"]
+)
+def get_bitwise_binary_inputs():
     test_suite = VkTestSuite(
         [
             ((M1, M2), (M1, M2)),
             ((S, S1, S2), (S, S1, S2)),
             ((XS, S, S1, S2), (XS, S, S1, S2)),
             ((1, M1), (1, M1)),
+            ((1, M2), (M1, M2)),
+            ((XS, 1, S1, 1), (1, S, 1, S2)),
         ]
     )
     test_suite.layouts = [
@@ -2160,7 +2164,6 @@ def get_bitwise_and_inputs():
         "utils::kTexture3D",
     ]
     test_suite.dtypes = ["at::kBool"]
-    test_suite.data_gen = "make_seq_tensor"
     return test_suite