[ET-VK][ops] Add eq.Scalar operator

Pull Request resolved: #20383

Adds Vulkan support for `aten.eq.Scalar`. This is the second of two ops needed to collapse the Llama4-mini TISO en_US backbone export to a single Vulkan partition after `bitwise_or`: the discrete-speech mask compares the int token-id tensor against scalar constants via `aten.eq.Scalar`, which previously had no Vulkan implementation and forced a CPU fallback that split the delegated graph.

Implemented by extending the existing tensor-scalar binary-op path with a comparison-output variant: `binary_scalar_buffer.glsl` / `binary_scalar_texture.glsl` gain an `IS_COMPARISON_OP` code path that writes a `uint8` (bool) output while leaving the existing arithmetic path unchanged; `binary_scalar_buffer.yaml` / `binary_scalar_texture.yaml` add compatible dtype/scalar pairs only (`half`/`float`, `float`/`float`, `int32`/`int32`) so unsupported mixed pairs fail instead of silently truncating scalar values; `BinaryScalarOp.cpp` adds an `eq_tensor_scalar` dispatch and registers `aten.eq.Scalar`; `op_registry.py` registers `aten.eq.Scalar` features with FP/INT tensor input and bool output.

The generated op-test graph builders now preserve `at::Scalar` tags when adding graph scalars, so integer, boolean, and floating scalar literals exercise the correct graph scalar type instead of all being converted through `double`.

The int64 token tensor is serialized to int32 via the existing `downcast_64_bit` path, so the dispatch resolves to the int32 shader variant; no dtype-conversion pass is added.

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

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

Author: ssjia

backends/vulkan/op_registry.py

@@ -327,6 +327,17 @@ def register_pow_tensor_scalar():
     )
 
 
+@update_features(exir_ops.edge.aten.eq.Scalar)
+def register_eq_scalar():
+    return OpFeatures(
+        inputs_storage=utils.ANY_STORAGE,
+        inputs_dtypes=utils.FP_INT_T,
+        outputs_dtypes=utils.BOOL_T,
+        supports_resize=True,
+        supports_highdim=True,
+    )
+
+
 # =============================================================================
 # ToCopy.cpp
 # =============================================================================

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

@@ -6,15 +6,27 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// Binary comparison ops write a bool/uint8 output dtype, which differs from
+// the input dtype. IS_COMPARISON_OP is set explicitly per shader variant in the
+// .yaml.
+
 #version 450 core
 
 ${define_required_extensions(STORAGE, DTYPE)}
+${define_explicit_type_extensions(SCALAR_VALUE_TYPE)}
+$if IS_COMPARISON_OP:
+  ${define_required_extensions(STORAGE, "uint8")}
 
 #define PRECISION ${PRECISION}
 
 #define NAME ${VARIANT_NAME}
 
 #define T ${buffer_scalar_type(DTYPE)}
+#define SCALAR_T ${buffer_scalar_type(SCALAR_VALUE_TYPE)}
+$if IS_COMPARISON_OP:
+  #define OUT_T ${buffer_scalar_type("uint8")}
+$else:
+  #define OUT_T ${buffer_scalar_type(DTYPE)}
 
 #define op(X, Y) ${OPERATOR}
 
@@ -24,25 +36,30 @@ layout(std430) buffer;
 
 #include "indexing.glslh"
 
-${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_ubo(B, "BufferMetadata", "outp")}
 ${layout_declare_ubo(B, "BufferMetadata", "inp")}
 
 layout(push_constant) uniform restrict Block {
-  float scalar_value;
+  SCALAR_T scalar_value;
 };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
-#include "binary_op_defs.glslh"
+$if not IS_COMPARISON_OP:
+  #include "binary_op_defs.glslh"
 
 void main() {
   const uint out_bufi = gl_GlobalInvocationID.x;
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }
 
-  t_out[out_bufi] = T(op(t_in[out_bufi], T(scalar_value)));
+  t_out[out_bufi] = OUT_T(op(t_in[out_bufi], T(scalar_value)));
 }

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

@@ -7,14 +7,19 @@
 binary_scalar_buffer:
   parameter_names_with_default_values:
     OPERATOR: power_of(X, Y)
-    NDIM: 3
+    IS_COMPARISON_OP: false
     DTYPE: float
-    PACKING: C_packed
+    SCALAR_VALUE_TYPE: float
     STORAGE: buffer
   generate_variant_forall:
-    DTYPE:
-      - VALUE: half
-      - VALUE: float
-      - VALUE: int32
+    combination:
+      parameter_names: [DTYPE, SCALAR_VALUE_TYPE]
+      combos:
+       - parameter_values: [half, float]
+       - parameter_values: [float, float]
+       - parameter_values: [int32, int32]
   shader_variants:
     - NAME: pow_scalar_buffer
+    - NAME: eq_scalar_buffer
+      OPERATOR: X == Y
+      IS_COMPARISON_OP: true

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

@@ -6,16 +6,28 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+// Binary comparison ops write a bool/uint8 output dtype, which differs from
+// the input dtype. IS_COMPARISON_OP is set explicitly per shader variant in the
+// .yaml.
+
 #version 450 core
 
 ${define_required_extensions(STORAGE, DTYPE)}
+${define_explicit_type_extensions(SCALAR_VALUE_TYPE)}
+$if IS_COMPARISON_OP:
+  ${define_required_extensions(STORAGE, "uint8")}
 
 #define PRECISION ${PRECISION}
 
 #define NAME ${VARIANT_NAME}
 
 #define VEC4_T ${texel_load_type(DTYPE, STORAGE)}
 #define T ${texel_load_component_type(DTYPE, STORAGE)}
+#define SCALAR_T ${buffer_scalar_type(SCALAR_VALUE_TYPE)}
+$if IS_COMPARISON_OP:
+  #define VEC4_OUT_T ${texel_load_type("uint8", STORAGE)}
+$else:
+  #define VEC4_OUT_T VEC4_T
 
 #define op(X, Y) ${OPERATOR}
 
@@ -25,19 +37,24 @@ layout(std430) buffer;
 
 #include "indexing.glslh"
 
-${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_ubo(B, "TextureMetadata", "outp")}
 ${layout_declare_ubo(B, "TextureMetadata", "inp")}
 
 layout(push_constant) uniform restrict Block {
-  float scalar_value;
+  SCALAR_T scalar_value;
 };
 
 layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 
-#include "binary_op_defs.glslh"
+$if not IS_COMPARISON_OP:
+  #include "binary_op_defs.glslh"
 
 void main() {
   const ivec3 pos = ivec3(gl_GlobalInvocationID);
@@ -47,7 +64,7 @@ void main() {
   }
 
   VEC4_T in_texel = texelFetch(t_in, pos, 0);
-  VEC4_T out_texel = VEC4_T(op(in_texel, VEC4_T(scalar_value)));
+  VEC4_OUT_T out_texel = VEC4_OUT_T(op(in_texel, VEC4_T(scalar_value)));
 
   imageStore(t_out, pos, out_texel);
 }

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

@@ -7,14 +7,19 @@
 binary_scalar_texture:
   parameter_names_with_default_values:
     OPERATOR: power_of(X, Y)
-    NDIM: 3
+    IS_COMPARISON_OP: false
     DTYPE: float
-    PACKING: C_packed
+    SCALAR_VALUE_TYPE: float
     STORAGE: texture3d
   generate_variant_forall:
-    DTYPE:
-      - VALUE: half
-      - VALUE: float
-      - VALUE: int32
+    combination:
+      parameter_names: [DTYPE, SCALAR_VALUE_TYPE]
+      combos:
+       - parameter_values: [half, float]
+       - parameter_values: [float, float]
+       - parameter_values: [int32, int32]
   shader_variants:
     - NAME: pow_scalar_texture3d
+    - NAME: eq_scalar_texture3d
+      OPERATOR: equal(X, Y)
+      IS_COMPARISON_OP: true

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

@@ -16,8 +16,42 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
 
+#include <executorch/backends/vulkan/runtime/utils/VecUtils.h>
+
+#include <vector>
+
 namespace vkcompute {
 
+namespace {
+
+vkapi::ScalarType scalar_value_dtype(
+    ComputeGraph& graph,
+    const ValueRef scalar) {
+  // Use a 32-bit integer push constant for bool scalars. Push constants in this
+  // file otherwise use 32-bit lanes, and T(scalar_value) preserves bool as 0/1.
+  if (graph.val_is_bool(scalar) || graph.val_is_symint(scalar)) {
+    return vkapi::kInt;
+  }
+  return graph.dtype_of(scalar);
+}
+
+int32_t extract_int32_scalar(ComputeGraph& graph, const ValueRef scalar) {
+  if (graph.val_is_int(scalar)) {
+    return utils::safe_downcast<int32_t>(graph.get_int(scalar));
+  }
+  if (graph.val_is_bool(scalar)) {
+    return graph.get_bool(scalar) ? 1 : 0;
+  }
+  if (graph.val_is_symint(scalar)) {
+    return graph.read_symint(scalar);
+  }
+  VK_THROW(
+      "Expected int, bool, or SymInt scalar, got: ",
+      graph.get_val_type(scalar));
+}
+
+} // namespace
+
 void resize_binary_scalar_op_node(
     ComputeGraph* graph,
     const std::vector<ArgGroup>& args,
@@ -39,14 +73,32 @@ void add_binary_scalar_op_node(
     const std::string& op_name) {
   ValueRef arg = prepack_standard_like(graph, in, out, true);
 
-  // Extract scalar value
-  float scalar_val = graph.extract_scalar<float>(scalar);
-
-  // Pick shader
+  const vkapi::ScalarType scalar_dtype = scalar_value_dtype(graph, scalar);
+  std::vector<PushConstantDataInfo> push_constants;
+  if (scalar_dtype == vkapi::kInt) {
+    const int32_t scalar_val = extract_int32_scalar(graph, scalar);
+    push_constants.emplace_back(&scalar_val, sizeof(scalar_val));
+  } else if (scalar_dtype == vkapi::kFloat) {
+    const float scalar_val = graph.extract_scalar<float>(scalar);
+    push_constants.emplace_back(&scalar_val, sizeof(scalar_val));
+  } else {
+    VK_THROW("Unsupported tensor-scalar op scalar dtype: ", scalar_dtype);
+  }
+
+  // Pick shader -- note that currently, only a few shader variants are
+  // generated for "compatible" tensor dtype / scalar dtype pairs. In particular
+  // float/half tensor + float scalar, and int32 tensor + int32 scalar. This
+  // decision is to prevent combinations like int tensor + float scalar, which
+  // would currently produce wrong values with the way the shader is currently
+  // written. The current combinations should be sufficient for practical
+  // applications, but in the future if "mixed precision" combinations are
+  // needed, then the shaders would need to be updated with more robust dtype
+  // handling.
   std::string kernel_name = op_name + "_scalar";
   kernel_name.reserve(kShaderNameReserve);
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
   add_dtype_suffix(kernel_name, graph.dtype_of(in));
+  add_dtype_suffix(kernel_name, scalar_dtype);
 
   vkapi::ParamsBindList param_ubos = {graph.meta_ubo(out), graph.meta_ubo(in)};
 
@@ -60,7 +112,7 @@ void add_binary_scalar_op_node(
       // Shader params buffers
       param_ubos,
       // Push Constants
-      {PushConstantDataInfo(&scalar_val, sizeof(scalar_val))},
+      push_constants,
       // Specialization Constants
       {},
       // Resize Args
@@ -73,8 +125,13 @@ void pow_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
   return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "pow");
 }
 
+void eq_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
+  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "eq");
+}
+
 REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.pow.Tensor_Scalar, pow_tensor_scalar);
+  VK_REGISTER_OP(aten.eq.Scalar, eq_tensor_scalar);
 }
 
 } // namespace vkcompute

backends/vulkan/test/op_tests/cases.py

@@ -2217,3 +2217,32 @@ def get_pow_tensor_scalar_inputs():
     ]
     test_suite.dtypes = ["at::kFloat"]
     return test_suite
+
+
+@register_test_suite("aten.eq.Scalar")
+def get_eq_scalar_inputs():
+    # Scalars are chosen to fall within the make_seq_tensor range (1..numel),
+    # so each case exercises a genuine mix of equal / not-equal elements rather
+    # than a trivially all-false comparison.
+    test_suite = VkTestSuite(
+        [
+            ((M1,), 5),
+            ((M2, M1), 100),
+            ((S1, M1, M2), 1000),
+            ((S1, S2, S2, M2), 2000),
+            ((S, S1, S2), 50),
+            ((M1, M2), 700),
+            ((S1, S2), 20),
+        ]
+    )
+    test_suite.storage_types = [
+        "utils::kBuffer",
+        "utils::kTexture3D",
+    ]
+    test_suite.layouts = [
+        "utils::kWidthPacked",
+        "utils::kChannelsPacked",
+    ]
+    test_suite.dtypes = ["at::kInt"]
+    test_suite.data_gen = "make_seq_tensor"
+    return test_suite

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

@@ -198,6 +198,19 @@ def generate_benchmark_fixture(self) -> str:
   }}
 }}
 
+ValueRef add_scalar_to_graph(ComputeGraph& graph, const at::Scalar& scalar) {{
+  if (scalar.isBoolean()) {{
+    return graph.add_scalar<bool>(scalar.toBool());
+  }}
+  if (scalar.isIntegral(/*includeBool=*/false)) {{
+    return graph.add_scalar<int64_t>(scalar.toLong());
+  }}
+  if (scalar.isFloatingPoint()) {{
+    return graph.add_scalar<double>(scalar.toDouble());
+  }}
+  VK_THROW("Unsupported at::Scalar!");
+}}
+
 at::Tensor make_casted_randint_tensor(
     std::vector<int64_t> sizes,
     at::ScalarType dtype = at::kFloat,

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

@@ -476,8 +476,8 @@ def create_value_for(  # noqa: C901
             ret_str += f"from_at_scalartype({ref.src_cpp_name}.scalar_type()), "
             ret_str += f"{ref.src_cpp_name}.const_data_ptr()); \n"
         elif ref.src_cpp_type == AT_SCALAR:
-            # TODO(ssjia): generalize this to work with all scalar types
-            ret_str += f"add_scalar<double>({ref.src_cpp_name}.toDouble()); \n"
+            ret_str = f"{cpp_type} {ref.name} = "
+            ret_str += f"add_scalar_to_graph(*{self.graph}, {ref.src_cpp_name}); \n"
         elif ref.src_cpp_type == AT_INT_ARRAY_REF:
             ret_str += f"add_scalar_list({ref.src_cpp_name}.vec()); \n"
         elif ref.src_cpp_type == BOOL:

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

@@ -129,6 +129,19 @@ def gen_parameterization(self) -> str:
   }
 }
 
+ValueRef add_scalar_to_graph(ComputeGraph& graph, const at::Scalar& scalar) {
+  if (scalar.isBoolean()) {
+    return graph.add_scalar<bool>(scalar.toBool());
+  }
+  if (scalar.isIntegral(/*includeBool=*/false)) {
+    return graph.add_scalar<int64_t>(scalar.toLong());
+  }
+  if (scalar.isFloatingPoint()) {
+    return graph.add_scalar<double>(scalar.toDouble());
+  }
+  VK_THROW("Unsupported at::Scalar!");
+}
+
 #ifdef USE_VULKAN_FP16_INFERENCE
 bool check_close(at::Tensor& t1, at::Tensor& t2, float rtol=1e-2, float atol=1e-2) {
 #else