[ET-VK][ez] Implement rsqrt (#6472)

Pull Request resolved: #6456

TSIA. This op is used in Llama model architecture.
ghstack-source-id: 249709740
@exported-using-ghexport

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

Co-authored-by: Stephen Jia <[email protected]>

Author: pytorchbot

backends/vulkan/partitioner/supported_ops.py

@@ -76,6 +76,7 @@ def __contains__(self, op):
     exir_ops.edge.aten.sigmoid.default,
     exir_ops.edge.aten.sin.default,
     exir_ops.edge.aten.sqrt.default,
+    exir_ops.edge.aten.rsqrt.default,
     exir_ops.edge.aten.tanh.default,
     exir_ops.edge.aten._to_copy.default,
     # Matrix Multiplication

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

@@ -32,6 +32,8 @@ unary_op:
       OPERATOR: sin(X)
     - NAME: sqrt
       OPERATOR: sqrt(X)
+    - NAME: rsqrt
+      OPERATOR: (1 / sqrt(X))
     - NAME: tanh
       OPERATOR: tanh(clamp(X, -15.0, 15.0))
     - NAME: hardshrink

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

@@ -129,6 +129,7 @@ DEFINE_ACTIVATION_FN(neg);
 DEFINE_ACTIVATION_FN(sigmoid);
 DEFINE_ACTIVATION_FN(sin);
 DEFINE_ACTIVATION_FN(sqrt);
+DEFINE_ACTIVATION_FN(rsqrt);
 DEFINE_ACTIVATION_FN(tanh);
 DEFINE_CLAMP_FN(clamp);
 DEFINE_CLAMP_FN(hardtanh);
@@ -149,6 +150,7 @@ REGISTER_OPERATORS {
   VK_REGISTER_OP(aten.sigmoid.default, sigmoid);
   VK_REGISTER_OP(aten.sin.default, sin);
   VK_REGISTER_OP(aten.sqrt.default, sqrt);
+  VK_REGISTER_OP(aten.rsqrt.default, rsqrt);
   VK_REGISTER_OP(aten.tanh.default, tanh);
   VK_REGISTER_OP(aten.hardshrink.default, hardshrink);
   VK_REGISTER_OP(aten.hardswish.default, hardswish);

backends/vulkan/test/op_tests/cases.py

@@ -988,6 +988,7 @@ def get_softmax_inputs():
 @register_test_suite(
     [
         "aten.sqrt.default",
+        "aten.rsqrt.default",
         "aten.exp.default",
         "aten.hardshrink.default",
         "aten.sin.default",