[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` generate variants for the `half`/`float`, `float`/`float`, `int32`/`int32`, and `int32`/`float` tensor/scalar dtype pairs. Mixed tensor/scalar dtypes are computed in a promoted type rather than narrowing the scalar to the tensor dtype: a new `get_higher_precision_dtype` helper in `gen_vulkan_spv.py` picks the higher-precision of the two dtypes, and `binary_op_defs.glslh` evaluates operands in that promoted `COMPUTE_T` before narrowing to the output dtype (so e.g. an `int32` tensor compared against a `float` scalar compares in `float`). `BinaryScalarOp.cpp` adds an `eq_tensor_scalar` dispatch, registers `aten.eq.Scalar`, and coerces the scalar extract dtype to match a generated variant; `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: 397529325
@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/gen_vulkan_spv.py

@@ -255,6 +255,61 @@ def texel_load_component_type(dtype: str, storage_type: str) -> str:
         return texel_component_type(dtype)
 
 
+def get_higher_precision_dtype(dtype_a: str, dtype_b: str) -> str:
+    """Return the higher-precision of two dtypes, intended for picking the type
+    in which to perform a mixed-dtype computation (e.g. a tensor of one dtype
+    combined with a scalar of another). The result is the operand dtype that can
+    represent the other without loss for the value ranges these shaders handle.
+
+    Ranking rule (highest first):
+      - The float family ALWAYS outranks the int family. So any float dtype
+        beats any int dtype (e.g. int32 + float -> float), because integer
+        values in these shaders are small enough to be exactly representable in
+        the float compute type.
+      - Within the float family: double > float > half.
+      - Within the int family: rank by bit width
+        (int64/uint64 > int32/uint32 > int16/uint16 > int8/uint8/bool).
+      - Signed-vs-unsigned tie at the same bit width resolves to the SIGNED
+        dtype (matches PyTorch's same-width promotion preference and keeps the
+        compute type able to hold negative operands).
+
+    Dtype aliases are normalized first: int<->int32, uint<->uint32, and bool is
+    treated as the lowest-ranked 8-bit integer.
+    """
+    alias_map = {"int": "int32", "uint": "uint32"}
+    a = alias_map.get(dtype_a, dtype_a)
+    b = alias_map.get(dtype_b, dtype_b)
+
+    if a == b:
+        return a
+
+    # (family_rank, bit_width, signed_rank). Higher tuple wins. family_rank 1 for
+    # the float family ensures it always outranks the int family (family_rank 0).
+    # signed_rank breaks same-width int ties in favor of signed.
+    float_ranks = {"half": 16, "float": 32, "double": 64}
+    int_ranks = {
+        "bool": 8,
+        "uint8": 8,
+        "int8": 8,
+        "uint16": 16,
+        "int16": 16,
+        "uint32": 32,
+        "int32": 32,
+        "uint64": 64,
+        "int64": 64,
+    }
+
+    def rank(dtype: str) -> tuple:
+        if dtype in float_ranks:
+            return (1, float_ranks[dtype], 0)
+        if dtype in int_ranks:
+            signed_rank = 0 if dtype in ("bool",) or dtype.startswith("uint") else 1
+            return (0, int_ranks[dtype], signed_rank)
+        raise AssertionError(f"Cannot rank precision of dtype: {dtype}")
+
+    return a if rank(a) >= rank(b) else b
+
+
 def get_access_qualifier(access_type: Optional[str]) -> str:
     if access_type is None:
         return ""
@@ -497,6 +552,7 @@ def define_required_extensions(storage_type: str, dtypes: Union[str, List[str]])
     "texel_component_type": texel_component_type,
     "texel_load_type": texel_load_type,
     "texel_load_component_type": texel_load_component_type,
+    "get_higher_precision_dtype": get_higher_precision_dtype,
     "layout_declare_buffer": layout_declare_buffer,
     "layout_declare_image": layout_declare_image,
     "layout_declare_sampler": layout_declare_sampler,

backends/vulkan/runtime/graph/ops/glsl/binary_op_defs.glslh

@@ -19,38 +19,44 @@
 // - Uses standard pow() for x > 0
 //
 
+// Operands are evaluated in the promoted compute type COMPUTE_T (and its vector
+// form COMPUTE_VEC4_T), which the including shader sets from
+// get_higher_precision_dtype(DTYPE, SCALAR_VALUE_TYPE) so mixed tensor/scalar
+// dtype combinations compute without loss before narrowing back to the output
+// dtype.
+
 // Scalar overload
-T power_of(T x, T y) {
+COMPUTE_T power_of(COMPUTE_T x, COMPUTE_T y) {
   if (x == 0.0) {
     // Handle 0^y: 0^0 = 1, 0^y = 0 for y > 0
-    return (y == 0.0) ? T(1.0) : T(0.0);
+    return (y == 0.0) ? COMPUTE_T(1.0) : COMPUTE_T(0.0);
   }
 
   // Use absolute value to avoid undefined behavior
-  float result = pow(abs(x), y);
+  float result = pow(abs(float(x)), float(y));
 
   // For negative bases with odd integer exponents, preserve the negative sign
   if (x < 0.0) {
-    float int_y = round(y);
-    if (abs(y - int_y) < 1e-5 && int(int_y) % 2 == 1) {
+    float int_y = round(float(y));
+    if (abs(float(y) - int_y) < 1e-5 && int(int_y) % 2 == 1) {
       result = -result;
     }
   }
 
-  return T(result);
+  return COMPUTE_T(result);
 }
 
-#ifdef VEC4_T
+#ifdef COMPUTE_VEC4_T
 
 // Vector overload
-VEC4_T power_of(VEC4_T x, VEC4_T y) {
-  VEC4_T result;
+COMPUTE_VEC4_T power_of(COMPUTE_VEC4_T x, COMPUTE_VEC4_T y) {
+  COMPUTE_VEC4_T result;
   for (int i = 0; i < 4; i++) {
     result[i] = power_of(x[i], y[i]);
   }
   return result;
 }
 
-#endif // VEC4_T
+#endif // COMPUTE_VEC4_T
 
 #endif // BINARY_OP_DEFS_GLSLH

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

@@ -6,15 +6,32 @@
  * 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
 
+$PROMOTED_DTYPE = get_higher_precision_dtype(DTYPE, SCALAR_VALUE_TYPE)
+
 ${define_required_extensions(STORAGE, DTYPE)}
+${define_required_extensions(STORAGE, PROMOTED_DTYPE)}
+${define_explicit_type_extensions(SCALAR_VALUE_TYPE)}
+${define_explicit_type_extensions(PROMOTED_DTYPE)}
+$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)}
+#define COMPUTE_T ${buffer_scalar_type(PROMOTED_DTYPE)}
+$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 +41,31 @@ 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(COMPUTE_T(t_in[out_bufi]), COMPUTE_T(scalar_value)));
 }

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

@@ -7,14 +7,20 @@
 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]
+       - parameter_values: [int32, float]
   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,34 @@
  * 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
 
+$PROMOTED_DTYPE = get_higher_precision_dtype(DTYPE, SCALAR_VALUE_TYPE)
+
 ${define_required_extensions(STORAGE, DTYPE)}
+${define_required_extensions(STORAGE, PROMOTED_DTYPE)}
+${define_explicit_type_extensions(SCALAR_VALUE_TYPE)}
+${define_explicit_type_extensions(PROMOTED_DTYPE)}
+$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)}
+#define COMPUTE_VEC4_T ${texel_load_type(PROMOTED_DTYPE, STORAGE)}
+#define COMPUTE_T ${texel_load_component_type(PROMOTED_DTYPE, STORAGE)}
+$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 +43,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 +70,8 @@ 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(COMPUTE_VEC4_T(in_texel), COMPUTE_VEC4_T(scalar_value)));
 
   imageStore(t_out, pos, out_texel);
 }

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

@@ -7,14 +7,20 @@
 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]
+       - parameter_values: [int32, float]
   shader_variants:
     - NAME: pow_scalar_texture3d
+    - NAME: eq_scalar_texture3d
+      OPERATOR: equal(X, Y)
+      IS_COMPARISON_OP: true