[ET-VK][ops] Consolidate binary scalar shaders behind op_type

Pull Request resolved: #24756

## Ulterior Motive

Make buffer storage behave more like texture storage, so one shader can service
both storage types. Earlier commits in this stack gave unary and binary
tensor-tensor ops one shader each, selected by an `op_type` specialization
constant. This commit does the same for tensor-scalar ops.

## Rationale

**What**: The tensor-scalar ops (`pow.Tensor_Scalar`, `mul.Scalar`, and `eq`,
`ne`, `lt`, `le`, `gt`, `ge` `.Scalar`) share one arithmetic and one comparison
shader variant family. The op is selected at dispatch through an `op_type`
specialization constant.

**Why**: Each op had its own variant family, which multiplied the number of
compiled shaders by the number of ops for no change in the code each one runs.

## Details

- Shader variants
  - **Before**: one family per op, `{op}_scalar_{storage}_{dtype}_{scalar}`,
    with the op inlined through an `OPERATOR` yaml parameter: 8 ops × 2
    storages × 4 dtype pairs = 64 shaders.
  - **After**: `binary_scalar_op` (`pow`, `mul`) and `binary_scalar_cmp`
    (comparisons, which write uint8): 2 × 2 × 4 = 16 shaders. A host `dev`
    build of `etvk_ops_bench` shrinks by 514 KiB.
- Op selection
  - `apply_binary_scalar_op(op_type, X, Y)` in `binary_scalar.glsl` switches on
    the spec constant, which the compiler folds. Each op computes what its
    `OPERATOR` expression did, in the same promoted type.
  - `op_type` takes `BinaryOpType` values, which gains `kNe`. Only the
    tensor-scalar shader implements `kNe`.
- Dispatch
  - `add_binary_scalar_op_node` takes a `BinaryOpType` instead of an op name
    and picks the comparison family for every op other than `mul` and `pow`.
- `test_vulkan_shader_names.py` checks that every name the tensor-scalar
  dispatch can build is generated.

Authored with Claude Code.
ghstack-source-id: 467556317
@exported-using-ghexport

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

Author: SS-JIA

backends/vulkan/runtime/graph/ops/glsl/binary/binary_scalar.glsl

@@ -6,9 +6,9 @@
  * 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.
+// The op is selected by the op_type specialization constant. Comparison ops
+// write a bool/uint8 output dtype, which differs from the input dtype, so they
+// get their own variants with IS_COMPARISON_OP set.
 
 #version 450 core
 
@@ -33,9 +33,9 @@ $if IS_COMPARISON_OP:
 #define OUT_VEC4_T ${texel_load_type(OUT_DTYPE, STORAGE)}
 #define OUT_COMPUTE_VEC4_T ${texel_type(OUT_DTYPE)}
 
-#define op(X, Y) ${OPERATOR}
-
-$if not IS_COMPARISON_OP and DTYPE in ("float", "half"):
+$if IS_COMPARISON_OP:
+  #define BINARY_SCALAR_IS_COMPARISON
+$elif DTYPE in ("float", "half"):
   #define ROUND_HALF_OUTPUT
 
 ${define_active_storage_type(STORAGE)}
@@ -63,12 +63,50 @@ layout(local_size_x_id = 0, local_size_y_id = 1, local_size_z_id = 2) in;
 layout(constant_id = 3) const bool input_is_half = false;
 layout(constant_id = 4) const int out_layout = CONTIG_LAYOUT_INT;
 const int out_packed_dim = get_packed_dim(out_layout);
+layout(constant_id = 5) const int op_type = 0;
 
 #include "dispatch.glslh"
 
 $if not IS_COMPARISON_OP:
   #include "binary_op_defs.glslh"
 
+// Values of op_type. Must match BinaryOpType in BinaryOp.h.
+#define BINARY_OP_MUL 2
+#define BINARY_OP_POW 5
+#define BINARY_OP_EQ 8
+#define BINARY_OP_LT 9
+#define BINARY_OP_LE 10
+#define BINARY_OP_GT 11
+#define BINARY_OP_GE 12
+#define BINARY_OP_NE 15
+
+OUT_COMPUTE_VEC4_T apply_binary_scalar_op(
+    const int op_type,
+    const COMPUTE_VEC4_T X,
+    const COMPUTE_VEC4_T Y) {
+  switch (op_type) {
+#ifdef BINARY_SCALAR_IS_COMPARISON
+    case BINARY_OP_NE:
+      return OUT_COMPUTE_VEC4_T(notEqual(X, Y));
+    case BINARY_OP_LT:
+      return OUT_COMPUTE_VEC4_T(lessThan(X, Y));
+    case BINARY_OP_LE:
+      return OUT_COMPUTE_VEC4_T(lessThanEqual(X, Y));
+    case BINARY_OP_GT:
+      return OUT_COMPUTE_VEC4_T(greaterThan(X, Y));
+    case BINARY_OP_GE:
+      return OUT_COMPUTE_VEC4_T(greaterThanEqual(X, Y));
+    default: // BINARY_OP_EQ
+      return OUT_COMPUTE_VEC4_T(equal(X, Y));
+#else
+    case BINARY_OP_POW:
+      return OUT_COMPUTE_VEC4_T(power_of(X, Y));
+    default: // BINARY_OP_MUL
+      return OUT_COMPUTE_VEC4_T(X * Y);
+#endif
+  }
+}
+
 void main() {
 #ifdef USING_BUFFER
   const uint out_idx = linear_idx_from_gid();
@@ -84,8 +122,8 @@ void main() {
   const VEC4_T in_texel = texelFetch(t_in, out_idx, 0);
 #endif
 
-  OUT_COMPUTE_VEC4_T out_texel = OUT_COMPUTE_VEC4_T(
-      op(COMPUTE_VEC4_T(in_texel), COMPUTE_VEC4_T(scalar_value)));
+  OUT_COMPUTE_VEC4_T out_texel = apply_binary_scalar_op(
+      op_type, COMPUTE_VEC4_T(in_texel), COMPUTE_VEC4_T(scalar_value));
 
 #ifdef ROUND_HALF_OUTPUT
   if (input_is_half) {

backends/vulkan/runtime/graph/ops/glsl/binary/binary_scalar.yaml

@@ -4,9 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+# The op is selected at dispatch through the op_type specialization constant.
 binary_scalar:
   parameter_names_with_default_values:
-    OPERATOR: power_of(X, Y)
     IS_COMPARISON_OP: false
     DTYPE: float
     SCALAR_VALUE_TYPE: float
@@ -23,24 +23,6 @@ binary_scalar:
        - parameter_values: [int32, int32]
        - parameter_values: [int32, float]
   shader_variants:
-    - NAME: pow_scalar
-    - NAME: mul_scalar
-      OPERATOR: X * Y
-    - NAME: eq_scalar
-      OPERATOR: equal(X, Y)
-      IS_COMPARISON_OP: true
-    - NAME: ne_scalar
-      OPERATOR: notEqual(X, Y)
-      IS_COMPARISON_OP: true
-    - NAME: lt_scalar
-      OPERATOR: lessThan(X, Y)
-      IS_COMPARISON_OP: true
-    - NAME: le_scalar
-      OPERATOR: lessThanEqual(X, Y)
-      IS_COMPARISON_OP: true
-    - NAME: gt_scalar
-      OPERATOR: greaterThan(X, Y)
-      IS_COMPARISON_OP: true
-    - NAME: ge_scalar
-      OPERATOR: greaterThanEqual(X, Y)
+    - NAME: binary_scalar_op
+    - NAME: binary_scalar_cmp
       IS_COMPARISON_OP: true

backends/vulkan/runtime/graph/ops/impl/binary/BinaryOp.h

@@ -12,8 +12,8 @@
 
 namespace vkcompute {
 
-// Values of the op_type specialization constant of the binary op shaders. Must
-// match binary_op.glslh.
+// Values of the op_type specialization constant of the binary op and
+// tensor-scalar shaders. Must match binary_op.glslh and binary_scalar.glsl.
 enum class BinaryOpType : int32_t {
   kAdd = 0,
   kSub = 1,
@@ -30,6 +30,8 @@ enum class BinaryOpType : int32_t {
   kGe = 12,
   kBitwiseAnd = 13,
   kBitwiseOr = 14,
+  // Implemented by the tensor-scalar shader only.
+  kNe = 15,
 };
 
 } // namespace vkcompute

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

@@ -10,6 +10,7 @@
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Common.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
+#include <executorch/backends/vulkan/runtime/graph/ops/impl/binary/BinaryOp.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/ScalarUtils.h>
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
@@ -75,7 +76,7 @@ void add_binary_scalar_op_node(
     const ValueRef in,
     const ValueRef scalar,
     const ValueRef out,
-    const std::string& op_name) {
+    const BinaryOpType op_type) {
   ValueRef arg = prepack_standard_like(graph, in, out, true);
   VK_CHECK_COND(check_same_packed_dim(graph, arg, out));
   if (!graph.is_buffer_storage(out)) {
@@ -97,7 +98,10 @@ void add_binary_scalar_op_node(
     VK_THROW("Unsupported tensor-scalar op scalar dtype: ", scalar_dtype);
   }
 
-  std::string kernel_name = op_name + "_scalar";
+  const bool is_comparison =
+      op_type != BinaryOpType::kMul && op_type != BinaryOpType::kPow;
+  std::string kernel_name(
+      is_comparison ? "binary_scalar_cmp" : "binary_scalar_op");
   kernel_name.reserve(kShaderNameReserve);
   add_storage_type_suffix(kernel_name, graph.storage_type_of(out));
   add_dtype_suffix(kernel_name, tensor_dtype);
@@ -118,44 +122,28 @@ void add_binary_scalar_op_node(
       push_constants,
       // Specialization Constants
       {int32_t(graph.original_dtype_of(out) == vkapi::kHalf),
-       graph.hashed_layout_of(out)},
+       graph.hashed_layout_of(out),
+       static_cast<int32_t>(op_type)},
       // Resize Args
       {},
       // Resizing Logic
       resize_binary_scalar_op_node));
 }
 
-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 mul_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "mul");
-}
-
-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");
-}
-
-void ne_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "ne");
-}
-
-void lt_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "lt");
-}
-
-void le_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "le");
-}
-
-void gt_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "gt");
-}
+#define DEFINE_BINARY_SCALAR_OP_FN(op_name, op_type)                     \
+  void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) { \
+    return add_binary_scalar_op_node(                                    \
+        graph, args[0], args[1], args[2], BinaryOpType::op_type);        \
+  }
 
-void ge_tensor_scalar(ComputeGraph& graph, const std::vector<ValueRef>& args) {
-  return add_binary_scalar_op_node(graph, args[0], args[1], args[2], "ge");
-}
+DEFINE_BINARY_SCALAR_OP_FN(pow_tensor_scalar, kPow);
+DEFINE_BINARY_SCALAR_OP_FN(mul_tensor_scalar, kMul);
+DEFINE_BINARY_SCALAR_OP_FN(eq_tensor_scalar, kEq);
+DEFINE_BINARY_SCALAR_OP_FN(ne_tensor_scalar, kNe);
+DEFINE_BINARY_SCALAR_OP_FN(lt_tensor_scalar, kLt);
+DEFINE_BINARY_SCALAR_OP_FN(le_tensor_scalar, kLe);
+DEFINE_BINARY_SCALAR_OP_FN(gt_tensor_scalar, kGt);
+DEFINE_BINARY_SCALAR_OP_FN(ge_tensor_scalar, kGe);
 
 REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.pow.Tensor_Scalar, pow_tensor_scalar);

backends/vulkan/test/test_vulkan_shader_names.py

@@ -44,6 +44,8 @@
 
 # The yaml template whose variants `add_binary_op_node` dispatches into.
 _BINARY_TEMPLATE = os.path.join("binary", "binary")
+# The yaml template whose variants `add_binary_scalar_op_node` dispatches into.
+_BINARY_SCALAR_TEMPLATE = os.path.join("binary", "binary_scalar")
 
 
 def _function_body(text: str, name: str) -> str:
@@ -138,6 +140,21 @@ def test_binary_variants_end_in_a_storage_and_dtype_suffix(self) -> None:
             "see add_binary_op_node in BinaryOp.cpp",
         )
 
+    def test_binary_scalar_shaders_cover_every_dispatch(self) -> None:
+        """Every name `add_binary_scalar_op_node` can build must be generated.
+
+        It appends a storage suffix, the tensor's dtype suffix and the scalar's
+        dtype suffix. The scalar is read as float for float and half tensors,
+        and as int32 or float for int32 tensors.
+        """
+        dtype_pairs = ("_half_float", "_float_float", "_int32_int32", "_int32_float")
+        for prefix in ("binary_scalar_op", "binary_scalar_cmp"):
+            for storage in ("_buffer", "_texture3d"):
+                for dtypes in dtype_pairs:
+                    expected = prefix + storage + dtypes
+                    with self.subTest(expected=expected):
+                        self.assertIn(expected, self.names[_BINARY_SCALAR_TEMPLATE])
+
 
 if __name__ == "__main__":
     unittest.main()