[ET-VK][binary] Select binary ops with a specialization constant

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

@@ -0,0 +1,110 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ * All rights reserved.
+ *
+ * This source code is licensed under the BSD-style license found in the
+ * LICENSE file in the root directory of this source tree.
+ */
+
+/*
+ * Element-wise binary ops, selected by the op_type specialization constant.
+ *
+ * The including shader declares op_type and defines:
+ * - IN_T and OUT_T: operand and result types; vectors if BINARY_OP_VEC4 is
+ *   defined, otherwise scalars
+ * - ALPHA_T: type of the alpha that add and sub scale Y by
+ * - BINARY_OP_IS_FLOAT if the operands are floating point
+ * - BINARY_OP_IS_COMPARISON (OUT_T holds booleans) or BINARY_OP_IS_BITWISE to
+ *   compile in those ops instead of the arithmetic ones
+ */
+
+#ifndef BINARY_OP_GLSLH
+#define BINARY_OP_GLSLH
+
+// Values of op_type. Must match BinaryOpType in BinaryOp.cpp.
+#define BINARY_OP_ADD 0
+#define BINARY_OP_SUB 1
+#define BINARY_OP_MUL 2
+#define BINARY_OP_DIV 3
+#define BINARY_OP_FLOOR_DIVIDE 4
+#define BINARY_OP_POW 5
+#define BINARY_OP_MINIMUM 6
+#define BINARY_OP_SWIGLU 7
+#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_BITWISE_AND 13
+#define BINARY_OP_BITWISE_OR 14
+
+#ifdef BINARY_OP_VEC4
+#define binary_eq(X, Y) equal(X, Y)
+#define binary_lt(X, Y) lessThan(X, Y)
+#define binary_le(X, Y) lessThanEqual(X, Y)
+#define binary_gt(X, Y) greaterThan(X, Y)
+#define binary_ge(X, Y) greaterThanEqual(X, Y)
+#else
+#define binary_eq(X, Y) ((X) == (Y))
+#define binary_lt(X, Y) ((X) < (Y))
+#define binary_le(X, Y) ((X) <= (Y))
+#define binary_gt(X, Y) ((X) > (Y))
+#define binary_ge(X, Y) ((X) >= (Y))
+#endif
+
+#if defined(BINARY_OP_IS_FLOAT) && !defined(BINARY_OP_IS_COMPARISON)
+float swiglu(float gate, float up) {
+  return (gate * (1.0 / (1.0 + exp(-gate)))) * up;
+}
+
+vec4 swiglu(vec4 gate, vec4 up) {
+  return (gate * (1.0 / (1.0 + exp(-gate)))) * up;
+}
+#endif
+
+OUT_T apply_binary_op(const IN_T X, const IN_T Y, const ALPHA_T A) {
+  switch (op_type) {
+#if defined(BINARY_OP_IS_COMPARISON)
+    case BINARY_OP_LT:
+      return OUT_T(binary_lt(X, Y));
+    case BINARY_OP_LE:
+      return OUT_T(binary_le(X, Y));
+    case BINARY_OP_GT:
+      return OUT_T(binary_gt(X, Y));
+    case BINARY_OP_GE:
+      return OUT_T(binary_ge(X, Y));
+    default: // BINARY_OP_EQ
+#ifdef BINARY_OP_IS_FLOAT
+      return OUT_T(binary_lt(abs(X - Y), IN_T(1e-5)));
+#else
+      return OUT_T(binary_eq(X, Y));
+#endif
+#elif defined(BINARY_OP_IS_BITWISE)
+    case BINARY_OP_BITWISE_OR:
+      return OUT_T(X | Y);
+    default: // BINARY_OP_BITWISE_AND
+      return OUT_T(X & Y);
+#else
+    case BINARY_OP_SUB:
+      return OUT_T(X - A * Y);
+    case BINARY_OP_MUL:
+      return OUT_T(X * Y);
+    case BINARY_OP_DIV:
+      return OUT_T(X / Y);
+    case BINARY_OP_FLOOR_DIVIDE:
+      return OUT_T(floor(X / Y));
+    case BINARY_OP_POW:
+      return OUT_T(pow(X, Y));
+    case BINARY_OP_MINIMUM:
+      return OUT_T(min(X, Y));
+#ifdef BINARY_OP_IS_FLOAT
+    case BINARY_OP_SWIGLU:
+      return OUT_T(swiglu(X, Y));
+#endif
+    default: // BINARY_OP_ADD
+      return OUT_T(X + A * Y);
+#endif
+  }
+}
+
+#endif // BINARY_OP_GLSLH

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

@@ -6,8 +6,7 @@
  * LICENSE file in the root directory of this source tree.
  */
 
-// Binary comparison ops require that the output is boolean and not the same as input.
-$IS_COMPARISON_OP = (any([name in VARIANT_NAME for name in ["binary_eq",  "binary_lt", "binary_le", "binary_gt", "binary_ge"]]))
+$IS_FLOAT = DTYPE in ["half", "float"]
 
 #version 450 core
 
@@ -19,15 +18,21 @@ $if IS_COMPARISON_OP:
 
 #define NAME ${VARIANT_NAME}
 
-#define VEC4_T ${texel_load_type(DTYPE, "buffer")}
+#define IN_T ${buffer_scalar_type(DTYPE)}
+// Comparison ops output booleans, stored as uint8.
 $if IS_COMPARISON_OP:
   #define T ${buffer_scalar_type("uint8")}
-  #define VEC4_OUT_T ${texel_load_type("uint8", "buffer")}
 $else:
-  #define T ${buffer_scalar_type(DTYPE)}
-  #define VEC4_OUT_T VEC4_T
+  #define T IN_T
 
-#define op(X, Y, A) ${OPERATOR}
+#define OUT_T T
+#define ALPHA_T IN_T
+$if IS_FLOAT:
+  #define BINARY_OP_IS_FLOAT
+$if IS_COMPARISON_OP:
+  #define BINARY_OP_IS_COMPARISON
+$elif DTYPE == "uint8":
+  #define BINARY_OP_IS_BITWISE
 
 ${define_active_storage_type("buffer")}
 
@@ -61,34 +66,33 @@ ${layout_declare_spec_const(C, "int", "other_layout", "CONTIG_LAYOUT_INT")}
 // Unused by buffer path, declared for spec constant ID consistency with texture
 ${layout_declare_spec_const(C, "int", "in_broadcast_packed_dim", "0")}
 ${layout_declare_spec_const(C, "int", "other_broadcast_packed_dim", "0")}
+${layout_declare_spec_const(C, "int", "op_type", "0")}
+// Cleared when all three tensors share sizes and layout, so they share buffer
+// indices and no index conversion is needed.
+${layout_declare_spec_const(C, "int", "should_broadcast", "1")}
 
-$if OPERATOR == "swiglu(X, Y)":
-  float swiglu(float gate, float up) {
-    return (gate * (1.0 / (1.0 + exp(-gate)))) * up;
-  }
+#include "binary_op.glslh"
 
 void main() {
   const uint out_bufi = linear_idx_from_gid();
   if (out_of_bounds(out_bufi, outp)) {
     return;
   }
 
-  // Simple case; no broadcasting
-  if (are_equal(inp, otherp)) {
-    t_out[out_bufi] = T(op(t_in[out_bufi], t_other[out_bufi], T(alpha)));
-    return;
-  }
-
-  TensorIndex outp_tidx = linear_idx_to_tensor_idx(outp, out_bufi);
+  uint inp_bufi = out_bufi;
+  uint other_bufi = out_bufi;
+  if (should_broadcast != 0) {
+    TensorIndex outp_tidx = linear_idx_to_tensor_idx(outp, out_bufi);
 
-  TensorIndex inp_tidx = outp_tidx;
-  clamp_tensor_idx(inp, inp_tidx);
+    TensorIndex inp_tidx = outp_tidx;
+    clamp_tensor_idx(inp, inp_tidx);
 
-  TensorIndex other_tidx = outp_tidx;
-  clamp_tensor_idx(otherp, other_tidx);
+    TensorIndex other_tidx = outp_tidx;
+    clamp_tensor_idx(otherp, other_tidx);
 
-  uint inp_bufi = tensor_idx_to_linear_idx(inp, inp_tidx);
-  uint other_bufi = tensor_idx_to_linear_idx(otherp, other_tidx);
+    inp_bufi = tensor_idx_to_linear_idx(inp, inp_tidx);
+    other_bufi = tensor_idx_to_linear_idx(otherp, other_tidx);
+  }
 
-  t_out[out_bufi] = T(op(t_in[inp_bufi], t_other[other_bufi], T(alpha)));
+  t_out[out_bufi] = apply_binary_op(t_in[inp_bufi], t_other[other_bufi], IN_T(alpha));
 }

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

@@ -4,68 +4,27 @@
 # 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_op_buffer:
   parameter_names_with_default_values:
-    OPERATOR: X + A * Y
     NDIM: 3
     DTYPE: float
     PACKING: C_packed
     STORAGE: buffer
-  generate_variant_forall:
-    DTYPE:
-      - VALUE: half
-      - VALUE: float
-      - VALUE: int32
+    IS_COMPARISON_OP: 0
   shader_variants:
-    - NAME: binary_add_buffer
-    - NAME: binary_sub_buffer
-      OPERATOR: X - A * Y
-    - NAME: binary_mul_buffer
-      OPERATOR: X * Y
-    - NAME: binary_swiglu_buffer
-      OPERATOR: swiglu(X, Y)
+    # Bitwise ops only exist for uint8; uint8 has no arithmetic ops.
+    - NAME: binary_op_buffer
       generate_variant_forall:
         DTYPE:
           - VALUE: half
           - VALUE: float
-    - NAME: binary_div_buffer
-      OPERATOR: X / Y
-    - NAME: binary_pow_buffer
-      OPERATOR: pow(X, Y)
-    - NAME: binary_floor_divide_buffer
-      OPERATOR: floor(X / Y)
-    - NAME: binary_minimum_buffer
-      OPERATOR: min(X, Y)
-    # Named to match the dispatcher, which builds "binary_" + op + storage +
-    # dtype. A variant called binary_eq_int32_buffer generates shaders nothing
-    # ever asks for, and aten.eq on an int32 tensor aborts at dispatch with
-    # "Could not find ShaderInfo with name binary_eq_buffer_int32".
-    - NAME: binary_eq_buffer
-      OPERATOR: X == Y
-      generate_variant_forall:
-        DTYPE:
           - VALUE: int32
-    - NAME: binary_eq_buffer
-      OPERATOR: abs(X - Y) < 1e-5
+          - VALUE: uint8
+    - NAME: binary_cmp_buffer
+      IS_COMPARISON_OP: 1
       generate_variant_forall:
         DTYPE:
           - VALUE: half
           - VALUE: float
-    - NAME: binary_lt_buffer
-      OPERATOR: X < Y
-    - NAME: binary_le_buffer
-      OPERATOR: X <= Y
-    - NAME: binary_gt_buffer
-      OPERATOR: X > Y
-    - NAME: binary_ge_buffer
-      OPERATOR: X >= Y
-    - NAME: binary_bitwise_and_buffer
-      OPERATOR: X & Y
-      generate_variant_forall:
-        DTYPE:
-          - VALUE: uint8
-    - NAME: binary_bitwise_or_buffer
-      OPERATOR: X | Y
-      generate_variant_forall:
-        DTYPE:
-          - VALUE: uint8
+          - VALUE: int32