add abs, sigmoid, tanh to ET-VK (#2605)

Summary:
Pull Request resolved: #2605

tsia

bypass-github-pytorch-ci-checks

Reviewed By: jorgep31415

Differential Revision: D55169458

fbshipit-source-id: fa980d73c9544f5be7e44141d99456008f343b73

Author: copyrightly

backends/vulkan/partitioner/vulkan_partitioner.py

@@ -35,10 +35,13 @@ def is_node_supported(self, submodules, node: torch.fx.Node) -> bool:
             exir_ops.edge.aten.div.Tensor,
             exir_ops.edge.aten.div.Tensor_mode,
             exir_ops.edge.aten.pow.Tensor_Tensor,
-            # Activation operators
+            # Unary operators
+            exir_ops.edge.aten.abs.default,
             exir_ops.edge.aten.clamp.default,
             exir_ops.edge.aten.hardtanh.default,
             exir_ops.edge.aten.relu.default,
+            exir_ops.edge.aten.sigmoid.default,
+            exir_ops.edge.aten.tanh.default,
             # Matrix multiplication operators
             exir_ops.edge.aten.mm.default,
             # Pooling operators

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

@@ -10,5 +10,11 @@ unary_op:
       - VALUE: float
         SUFFIX: float
   shader_variants:
+    - NAME: abs
+      OPERATOR: abs(X)
     - NAME: clamp
       OPERATOR: clamp(X, A, B)
+    - NAME: sigmoid
+      OPERATOR: 1 / (1 + exp(-1 * X))
+    - NAME: tanh
+      OPERATOR: tanh(clamp(X, -15.0, 15.0))

backends/vulkan/runtime/graph/ops/impl/Clamp.cpp renamed to backends/vulkan/runtime/graph/ops/impl/UnaryOp.cpp

@@ -19,7 +19,10 @@ namespace at {
 namespace native {
 namespace vulkan {
 
-void resize_clamp_node(
+constexpr float kDummyFloat = -1.0f;
+const std::string kClampShaderName = "clamp";
+
+void resize_unary_op_node(
     ComputeGraph* graph,
     const std::vector<ArgGroup>& args,
     const std::vector<ValueRef>& extra_args) {
@@ -30,20 +33,21 @@ void resize_clamp_node(
   out.virtual_resize(self.sizes());
 }
 
-void add_clamp_node(
+void add_unary_op_node(
     ComputeGraph& graph,
     const ValueRef in,
     const float min,
     const float max,
-    const ValueRef out) {
+    const ValueRef out,
+    const std::string& op_name) {
   ValueRef arg = prepack_if_tensor_ref(graph, in);
 
   vTensor& t_out = graph.get_val(out).toTensor();
   api::utils::uvec3 global_size = t_out.virtual_extents();
   api::utils::uvec3 local_size = adaptive_work_group_size(global_size);
 
   std::stringstream kernel_name;
-  kernel_name << "clamp";
+  kernel_name << op_name;
   apply_dtype_suffix(kernel_name, t_out);
 
   graph.execute_nodes().emplace_back(new ExecuteNode(
@@ -58,7 +62,7 @@ void add_clamp_node(
        graph.create_params_buffer(min),
        graph.create_params_buffer(max)},
       // Resizing
-      resize_clamp_node));
+      resize_unary_op_node));
 }
 
 float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) {
@@ -69,30 +73,48 @@ float get_val_or_inf(ComputeGraph& graph, const ValueRef& val, bool max) {
              : -std::numeric_limits<float>::infinity();
 }
 
+#define DEFINE_ACTIVATION_FN(op_name)                                    \
+  void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) { \
+    return add_unary_op_node(                                            \
+        graph, args[0], kDummyFloat, kDummyFloat, args[1], #op_name);    \
+  }
+
 #define DEFINE_CLAMP_FN(op_name)                                         \
   void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) { \
-    return add_clamp_node(                                               \
+    return add_unary_op_node(                                            \
         graph,                                                           \
         args[0],                                                         \
         get_val_or_inf(graph, args[1], /*max =*/false),                  \
         get_val_or_inf(graph, args[2], /*max =*/true),                   \
-        args[3]);                                                        \
+        args[3],                                                         \
+        kClampShaderName);                                               \
   }
 
-#define DEFINE_RELU_FN(op_name)                                              \
-  void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) {     \
-    return add_clamp_node(                                                   \
-        graph, args[0], 0, std::numeric_limits<float>::infinity(), args[1]); \
+#define DEFINE_RELU_FN(op_name)                                          \
+  void op_name(ComputeGraph& graph, const std::vector<ValueRef>& args) { \
+    return add_unary_op_node(                                            \
+        graph,                                                           \
+        args[0],                                                         \
+        0,                                                               \
+        std::numeric_limits<float>::infinity(),                          \
+        args[1],                                                         \
+        kClampShaderName);                                               \
   }
 
+DEFINE_ACTIVATION_FN(abs);
+DEFINE_ACTIVATION_FN(sigmoid);
+DEFINE_ACTIVATION_FN(tanh);
 DEFINE_CLAMP_FN(clamp);
 DEFINE_CLAMP_FN(hardtanh);
 DEFINE_RELU_FN(relu);
 
 REGISTER_OPERATORS {
+  VK_REGISTER_OP(aten.abs.default, abs);
   VK_REGISTER_OP(aten.clamp.default, clamp);
   VK_REGISTER_OP(aten.hardtanh.default, hardtanh);
   VK_REGISTER_OP(aten.relu.default, relu);
+  VK_REGISTER_OP(aten.sigmoid.default, sigmoid);
+  VK_REGISTER_OP(aten.tanh.default, tanh);
 }
 
 } // namespace vulkan

backends/vulkan/test/test_vulkan_delegate.py

@@ -383,6 +383,36 @@ def forward(self, x):
             first_output_only=True,
         )
 
+    def test_vulkan_backend_abs(self):
+        class AbsModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.abs(x)
+
+        self.lower_clamp_module_and_test_output(AbsModule())
+
+    def test_vulkan_backend_sigmoid(self):
+        class SigmoidModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.sigmoid(x)
+
+        self.lower_clamp_module_and_test_output(SigmoidModule())
+
+    def test_vulkan_backend_tanh(self):
+        class TanhModule(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return torch.tanh(x)
+
+        self.lower_clamp_module_and_test_output(TanhModule())
+
     def test_vulkan_backend_partial(self):
         class SimpleModel(torch.nn.Module):
             def __init__(self):