Create ReLU6 Vulkan Compute Shader & Add Op To Backend Delegate

Summary: This diff adds support for the ReLU6 activation function in the Vulkan backend of Executorch. It introduces a Vulkan compute shader for ReLU6 and integrates the new op into the Vulkan delegate for execution. This enables models using ReLU6 to run efficiently on Vulkan targets, improving compatibility and expanding the delegate’s supported op set. The shader supports both buffer and texture data formats, enabling flexible deployment across different memory storage modes and ensuring compatibility with a wide range of models and execution environments.

Differential Revision: D78124607

Author: leafs1

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

@@ -104,6 +104,11 @@ float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) {
         kClampShaderName);                                               \
   }
 
+#define DEFINE_RELU6_FN(op_name)                                         \
+  void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) { \
+    return add_unary_op_node(graph, args[0], 0, 6, args[1], kClampShaderName);   \
+  }
+
 #define DEFINE_HARDSHRINK_FN(op_name)                                    \
   void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) { \
     return add_unary_op_node(                                            \
@@ -146,6 +151,7 @@ DEFINE_ACTIVATION_FN(tanh);
 DEFINE_CLAMP_FN(clamp);
 DEFINE_CLAMP_FN(hardtanh);
 DEFINE_RELU_FN(relu);
+DEFINE_RELU6_FN(relu6);
 DEFINE_HARDSHRINK_FN(hardshrink);
 DEFINE_ACTIVATION_FN(hardswish);
 DEFINE_ACTIVATION_FN(hardsigmoid);
@@ -161,6 +167,7 @@ REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.hardtanh.default, hardtanh);
   VK_REGISTER_OP(aten.neg.default, neg);
   VK_REGISTER_OP(aten.relu.default, relu);
+  VK_REGISTER_OP(aten.relu6.default, relu6);
   VK_REGISTER_OP(aten.sigmoid.default, sigmoid);
   VK_REGISTER_OP(aten.sin.default, sin);
   VK_REGISTER_OP(aten.sqrt.default, sqrt);

backends/vulkan/test/op_tests/cases.py

@@ -1511,6 +1511,7 @@ def get_var_inputs():
         "aten.leaky_relu.default",
         "aten.round.default",
         "aten.tan.default",
+        "aten.relu6.default",
     ]
 )
 def get_unary_ops_inputs():