Update on "[ET-VK] Allow specifying multiple storage types/memory layouts for an operator + register group norm operator"

## Changes

* Handle cases where an operator needs to specify a separate storage type / memory layout for each individual output.

## Motivation

Required for the group norm operator.

## Future Work

Currently, the `tag_memory_meta_pass` graph pass assumes that all tensors participating in a computation (aside from weights) will have the same storage type and memory layout. As more operators are being added, there are more exceptions to this rule.

The pass may need an update in the near future to make it possible to specify required storage types and memory layouts on a more granular level.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/_passes/tag_memory_meta_pass.py

@@ -91,7 +91,7 @@ def __init__(
         self.default_layout: VkMemoryLayout = default_memory_layout
         self.texture_limits = texture_limits
 
-    def propose_node_storage(
+    def propose_node_storage(  # noqa: C901
         self,
         node: torch.fx.Node,
     ) -> Optional[VkStorageType]:

backends/vulkan/op_registry.py

@@ -660,7 +660,7 @@ def register_ported_ops_with_prepacking(features: OpFeatures):
         exir_ops.edge.aten.native_group_norm.default,
     ]
 )
-def register_ported_ops_with_prepacking(features: OpFeatures):
+def register_native_group_norm(features: OpFeatures):
     features.texture_impl = TextureImplFeatures(
         valid_packed_dims={PackedDim.CHANNELS},
     )

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

@@ -12,7 +12,6 @@
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/Staging.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/KernelUtils.h>
-#include <executorch/backends/vulkan/runtime/graph/ops/impl/utils/TensorUtils.h>
 
 #include <executorch/backends/vulkan/runtime/graph/ops/utils/ShaderNameUtils.h>
 

backends/vulkan/test/op_tests/cases.py

@@ -27,10 +27,8 @@
 test_suites = {}
 
 
-def register_test_suite(aten_op, skip=True):
+def register_test_suite(aten_op):
     def test_suite_decorator(fn: Callable) -> Callable:
-        if skip:
-            return fn
         if isinstance(aten_op, str):
             test_suites[aten_op] = fn()
         elif isinstance(aten_op, list):
@@ -143,7 +141,7 @@ def get_linear_inputs():
     inputs_list += [((M, K), (N, K), (N)) for M, K, N in MKN_list]
     inputs_list += [((3, M, K), (N, K), None) for M, K, N in MKN_list]
     inputs_list += [((3, M, K), (N, K), (N)) for M, K, N in MKN_list]
-    inputs_list += [((3, 6, K), (N, K), (N)) for _, K, N in MKN_list]
+    inputs_list += [((3, 6, K), (N, K), (N)) for M, K, N in MKN_list]
 
     test_suite = VkTestSuite(inputs_list)
     test_suite.dtypes = ["at::kFloat"]
@@ -648,7 +646,7 @@ def get_native_layer_norm_inputs():
     return test_suite
 
 
-@register_test_suite("aten.native_group_norm.default", skip=False)
+@register_test_suite("aten.native_group_norm.default")
 def get_native_group_norm_inputs():
     test_suite = VkTestSuite(
         [
@@ -662,24 +660,18 @@ def get_native_group_norm_inputs():
         ]
     )
     test_suite.layouts = [
-        # "utils::kWidthPacked",
-        # "utils::kHeightPacked",
         "utils::kChannelsPacked",
     ]
     test_suite.storage_types = [
-        # "utils::kBuffer",
-        "utils::kTexture3D"
+        "utils::kTexture3D",
+    ]
+    test_suite.dtypes = [
+        "at::kFloat",
     ]
-    test_suite.dtypes = ["at::kFloat"]
     test_suite.arg_storage_types = {
         "out": [None, "utils::kBuffer", "utils::kBuffer"],
     }
 
-    # test_suite.arg_data_gen_fn = {
-    #     "weight": "make_zeros_tensor",
-    #     "bias": "make_rand_tensor",
-    # }
-
     test_suite.prepacked_args = ["weight", "bias"]
     test_suite.requires_prepack = True
 
@@ -798,7 +790,6 @@ def get_permute_inputs():
     ]
     test_suite.dtypes = [
         "at::kFloat",
-        "at::kInt",
     ]
     return test_suite