Binary Comparison Ops

Summary:
This change introduces the `binary_eq`, `binary_lt`, `binary_le`, `binary_gt`, `binary_ge` operators.
1. Introduced the operators in the binary_op.yaml file
2. we now store the shader variant base name to better handle special cases in the glsl of binary op
3. binary ops assumed that the output of the operation is the same of the input data. For `binary_eq` this is not true. This is now handled both in the shader and the in the graph creation.
4. added test case

Reviewed By: SS-JIA

Differential Revision: D76049244

Author: andreanicastro

backends/vulkan/op_registry.py

@@ -259,6 +259,7 @@ def register_ephemeral_op(features: OpFeatures):
         exir_ops.edge.aten.div.Tensor,
         exir_ops.edge.aten.div.Tensor_mode,
         exir_ops.edge.aten.pow.Tensor_Tensor,
+        exir_ops.edge.aten.eq.Tensor,
     ]
 )
 def register_binary_op(features: OpFeatures):

backends/vulkan/runtime/gen_vulkan_spv.py

@@ -728,9 +728,19 @@ def parseTemplateYaml(self, yaml_file: str) -> None:
                 )
 
                 for variant in params_dict["shader_variants"]:
+                    default_iterated_params_names = set(
+                        default_iterated_params.keys()
+                        if default_iterated_params is not None
+                        else {}
+                    )
                     variant_params_names = set(variant.keys())
+
+                    print(
+                        f"default_iterated_params_names: {default_iterated_params_names}"
+                    )
                     invalid_keys = (
                         variant_params_names
+                        - default_iterated_params_names
                         - params_names
                         - {"generate_variant_forall"}
                     )
@@ -758,6 +768,7 @@ def parseTemplateYaml(self, yaml_file: str) -> None:
                                     variant_name = f"{variant_name}_{param_value[1]}"
 
                             default_params_copy["NAME"] = variant_name
+                            default_params_copy["VARIANT_NAME"] = variant["NAME"]
 
                             self.shader_template_params[template_name].append(
                                 default_params_copy

backends/vulkan/runtime/graph/ops/glsl/binary_op.glsl

@@ -10,17 +10,35 @@
 
 #define PRECISION ${PRECISION}
 
+// Binary comparison ops require that the output is boolean and not the same as input.
+$IS_COMPARISON_OP = ([name in VARIANT_NAME for name in ["binary_eq",  "binary_lt", "binary_le", "binary_gt", "binary_ge"]])
+
+#define NAME ${VARIANT_NAME}
+
 #define VEC4_T ${texel_type(DTYPE)}
-#define T ${buffer_scalar_type(DTYPE)}
+$if IS_COMPARISON_OP:
+  #define T ${buffer_scalar_type("uint8")}
+  #define VEC4_OUT_T ${texel_type("uint8")}
+$else:
+  #define T ${buffer_scalar_type(DTYPE)}
+  #define VEC4_OUT_T VEC4_T
 
 #define op(X, Y, A) ${OPERATOR}
 
 ${define_active_storage_type(STORAGE)}
 ${define_required_extensions(DTYPE)}
 
+
+$if IS_COMPARISON_OP:
+  ${define_required_extensions("uint8")}
+
 layout(std430) buffer;
 
-${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
+$if IS_COMPARISON_OP:
+  ${layout_declare_tensor(B, "w", "t_out", "uint8", STORAGE)}
+$else:
+  ${layout_declare_tensor(B, "w", "t_out", DTYPE, STORAGE)}
+
 ${layout_declare_tensor(B, "r", "t_in", DTYPE, STORAGE)}
 ${layout_declare_tensor(B, "r", "t_other", DTYPE, STORAGE)}
 
@@ -121,7 +139,7 @@ void main() {
   write_texel_lpos(
     t_out,
     lpos,
-    VEC4_T(op(in_texel, other_texel, alpha)),
+    VEC4_OUT_T(op(in_texel, other_texel, alpha)),
     out_axis_map);
 }
 

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

@@ -32,3 +32,84 @@ binary_op:
       OPERATOR: floor(X / Y)
     - NAME: binary_minimum
       OPERATOR: min(X, Y)
+    - NAME: binary_eq_int32
+      OPERATOR: X == Y
+      DTYPE: int32
+    - NAME: binary_eq_buffer
+      OPERATOR: abs(X - Y) < 1e-5
+      STORAGE: buffer
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+    - NAME: binary_eq_texture3d
+      OPERATOR: all(lessThanEqual(abs(X - Y), VEC4_T(1e-5)))
+      STORAGE: texture3d
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+    - NAME: binary_lt_buffer
+      OPERATOR: X < Y
+      STORAGE: buffer
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_lt_texture3d
+      OPERATOR: all(lessThan(X, Y))
+      STORAGE: texture3d
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_le_buffer
+      OPERATOR: X <= Y
+      STORAGE: buffer
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_le_texture3d
+      OPERATOR: all(lessThanEqual(X, Y))
+      STORAGE: texture3d
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_gt_buffer
+      OPERATOR: X > Y
+      STORAGE: buffer
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_gt_texture3d
+      OPERATOR: all(greaterThan(X, Y))
+      STORAGE: texture3d
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_ge_buffer
+      OPERATOR: X >= Y
+      STORAGE: buffer
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32
+    - NAME: binary_ge_texture3d
+      OPERATOR: all(greaterThanEqual(X, Y))
+      STORAGE: texture3d
+      generate_variant_forall:
+        DTYPE:
+          - VALUE: half
+          - VALUE: float
+          - VALUE: int32

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

@@ -77,7 +77,7 @@ void add_binary_op_texture_node(
   kernel_name.reserve(kShaderNameReserve);
   kernel_name += op_name;
   add_storage_type_suffix(kernel_name, *t_out);
-  add_dtype_suffix(kernel_name, *t_out);
+  add_dtype_suffix(kernel_name, graph.dtype_of(in1));
 
   graph.execute_nodes().emplace_back(new DynamicDispatchNode(
       graph,
@@ -121,7 +121,8 @@ void add_binary_op_buffer_node(
   kernel_name.reserve(kShaderNameReserve);
   kernel_name += op_name;
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
-  add_dtype_suffix(kernel_name, graph.dtype_of(out));
+
+  add_dtype_suffix(kernel_name, graph.dtype_of(in1));
 
   graph.execute_nodes().emplace_back(new DynamicDispatchNode(
       graph,
@@ -189,6 +190,11 @@ DEFINE_BINARY_OP_FN(mul);
 DEFINE_BINARY_OP_FN(div);
 DEFINE_BINARY_OP_FN(pow);
 DEFINE_BINARY_OP_FN(minimum);
+DEFINE_BINARY_OP_FN(eq);
+DEFINE_BINARY_OP_FN(lt);
+DEFINE_BINARY_OP_FN(le);
+DEFINE_BINARY_OP_FN(gt);
+DEFINE_BINARY_OP_FN(ge);
 
 REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.add.Tensor, add);
@@ -198,6 +204,11 @@ REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.div.Tensor_mode, floor_divide);
   VK_REGISTER_OP(aten.pow.Tensor_Tensor, pow);
   VK_REGISTER_OP(aten.minimum.default, minimum);
+  VK_REGISTER_OP(aten.eq.Tensor, eq);
+  VK_REGISTER_OP(aten.lt.Tensor, lt);
+  VK_REGISTER_OP(aten.le.Tensor, le);
+  VK_REGISTER_OP(aten.gt.Tensor, gt);
+  VK_REGISTER_OP(aten.ge.Tensor, ge);
 }
 
 } // namespace vkcompute

backends/vulkan/test/op_tests/cases.py

@@ -63,6 +63,42 @@ def get_binary_elementwise_inputs():
         "utils::kBuffer",
         "utils::kTexture3D",
     ]
+
+    return test_suite
+
+
+# Eq requires a different test generator so it was split from the other test case.
+@register_test_suite(
+    [
+        "aten.eq.Tensor",
+        "aten.gt.Tensor",
+        "aten.lt.Tensor",
+        "aten.ge.Tensor",
+        "aten.le.Tensor",
+    ]
+)
+def get_binary_elementwise_compare_inputs():
+    test_suite = VkTestSuite(
+        [
+            ((M1, M2), (M1, M2)),
+            ((M1, M2), (M1, 1), 2.0),
+            ((M1, M2), (1, M2)),
+            ((S, S1, S2), (S, S1, S2)),
+            ((S, S1, S2), (S, S1, 1), 2.0),
+            ((S, S1, S2), (S, 1, S2), 2.0),
+            ((XS, S, S1, S2), (XS, S, 1, 1), 2.0),
+            ((3, 64, 1), (1, 64, 1)),
+        ]
+    )
+    test_suite.layouts = [
+        "utils::kWidthPacked",
+        "utils::kChannelsPacked",
+    ]
+    test_suite.storage_types = [
+        "utils::kBuffer",
+        "utils::kTexture3D",
+    ]
+    test_suite.data_gen = "make_casted_randint_tensor"
     return test_suite
 
 

backends/vulkan/test/op_tests/utils/gen_benchmark_vk.py

@@ -196,6 +196,15 @@ def generate_benchmark_fixture(self) -> str:
   }}
 }}
 
+at::Tensor make_casted_randint_tensor(
+    std::vector<int64_t> sizes,
+    at::ScalarType dtype = at::kFloat,
+    int low = 0,
+    int high = 10) {{
+
+  return at::randint(high, sizes, at::device(at::kCPU).dtype(dtype));
+}}
+
 at::Tensor make_rand_tensor(
     std::vector<int64_t> sizes,
     at::ScalarType dtype = at::kFloat,

backends/vulkan/test/op_tests/utils/gen_correctness_base.py

@@ -283,6 +283,15 @@ def generate_suite_cpp(self) -> str:
 
 {preamble}
 
+at::Tensor make_casted_randint_tensor(
+    std::vector<int64_t> sizes,
+    at::ScalarType dtype = at::kFloat,
+    int low = 0,
+    int high = 10) {{
+
+  return at::randint(high, sizes, at::device(at::kCPU).dtype(dtype));
+}}
+
 at::Tensor make_rand_tensor(
     std::vector<int64_t> sizes,
     at::ScalarType dtype = at::kFloat,