[ET-VK][ez] Clean up organization of supported_ops

As title. Group supported ops by features instead of op category. This will make it easier to mark that an op has increased its feature set. This also allows the registration code to be simplified a lot.

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

[ghstack-poisoned]

Author: SS-JIA

backends/vulkan/partitioner/supported_ops.py

@@ -47,17 +47,16 @@ def __contains__(self, op):
     operator.getitem,
 ]
 
-BINARY_OPS = [
+SUPPORTS_DYNAMIC_SHAPE = [
+    # Binary broadcasting operators
     exir_ops.edge.aten.add.Tensor,
     exir_ops.edge.aten.sub.Tensor,
     exir_ops.edge.aten.minimum.default,
     exir_ops.edge.aten.mul.Tensor,
     exir_ops.edge.aten.div.Tensor,
     exir_ops.edge.aten.div.Tensor_mode,
     exir_ops.edge.aten.pow.Tensor_Tensor,
-]
-
-UNARY_OPS = [
+    # Unary elementwise operators
     exir_ops.edge.aten.abs.default,
     exir_ops.edge.aten.clamp.default,
     exir_ops.edge.aten.cos.default,
@@ -71,60 +70,48 @@ def __contains__(self, op):
     exir_ops.edge.aten.sin.default,
     exir_ops.edge.aten.sqrt.default,
     exir_ops.edge.aten.tanh.default,
-]
-
-MATMUL_OPS = [
+    # Matrix Multiplication Operators
     exir_ops.edge.aten.bmm.default,
     exir_ops.edge.aten.mm.default,
     exir_ops.edge.aten.addmm.default,
     exir_ops.edge.aten.linear.default,
-]
-
-POOLING_OPS = [
+    # Reduction operators
+    exir_ops.edge.aten._log_softmax.default,
+    exir_ops.edge.aten._softmax.default,
+    # 2D Pooling ops
     exir_ops.edge.aten.avg_pool2d.default,
     exir_ops.edge.aten.max_pool2d_with_indices.default,
-]
-
-CONVOLUTION_OPS = [
+    # Convolution ops
     exir_ops.edge.aten.convolution.default,
     exir_ops.edge.et_vk.conv_with_clamp.default,
+    # Custom ops
+    "llama::sdpa_with_kv_cache",
 ]
 
-REDUCTION_OPS = [
+NO_DYNAMIC_SHAPE = [
+    # Reduction operators
     exir_ops.edge.aten.mean.dim,
     exir_ops.edge.aten.sum.dim_IntList,
-    exir_ops.edge.aten._log_softmax.default,
-    exir_ops.edge.aten._softmax.default,
-]
-
-NORMALIZATION_OPS = [
+    # Normalization operators
     exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
     exir_ops.edge.aten.native_layer_norm.default,
-]
-
-SHAPE_MANIPULATION_OPS = [
+    # Shape Manipulation operators
     exir_ops.edge.aten.squeeze_copy.dims,
     exir_ops.edge.aten.unsqueeze_copy.default,
     exir_ops.edge.aten.view_copy.default,
     exir_ops.edge.aten.permute_copy.default,
     exir_ops.edge.aten.t_copy.default,
-]
-
-INDEXING_OPS = [
+    # Indexing and lookup operators
     exir_ops.edge.aten.embedding.default,
     exir_ops.edge.aten.index_select.default,
     exir_ops.edge.aten.select_copy.int,
     exir_ops.edge.aten.slice_copy.Tensor,
-]
-
-ORCHESTRATION_OPS = [
+    # Tensor combination operators
     exir_ops.edge.aten.cat.default,
     exir_ops.edge.aten.split_with_sizes_copy.default,
     exir_ops.edge.aten.split.Tensor,
     exir_ops.edge.aten.repeat.default,
-]
-
-CREATION_OPS = [
+    # Tensor creation operators
     exir_ops.edge.aten.arange.start_step,
     exir_ops.edge.aten.clone.default,
     exir_ops.edge.aten.constant_pad_nd.default,
@@ -139,46 +126,20 @@ def __contains__(self, op):
 ]
 
 
-def register_prim_ops(ops: OpList):
-    for op in PRIM_OPS:
-        ops[op].supports_texture = True
-        ops[op].supports_buffer = True
-        ops[op].supports_dynamic_shape = True
+def enumerate_supported_ops():
+    ops = OpList()
 
+    # Register in order of least to most capabilities
 
-def register_no_dynamic_shape_ops(ops: OpList):
-    for op in [
-        *REDUCTION_OPS,
-        *NORMALIZATION_OPS,
-        *SHAPE_MANIPULATION_OPS,
-        *INDEXING_OPS,
-        *ORCHESTRATION_OPS,
-        *CREATION_OPS,
-    ]:
+    for op in NO_DYNAMIC_SHAPE:
         ops[op].supports_dynamic_shape = False
 
-
-def register_dynamic_shape_ops(ops: OpList):
-    for op in [
-        *BINARY_OPS,
-        *UNARY_OPS,
-        *MATMUL_OPS,
-        *POOLING_OPS,
-        *CONVOLUTION_OPS,
-    ]:
+    for op in SUPPORTS_DYNAMIC_SHAPE:
         ops[op].supports_dynamic_shape = True
 
-
-def register_custom_ops(ops: OpList):
-    for op in CUSTOM_OPS:
-        ops[op].supports_dynamic_shape = True
+    for op in PRIM_OPS:
         ops[op].supports_texture = True
+        ops[op].supports_buffer = True
+        ops[op].supports_dynamic_shape = True
 
-
-def enumerate_supported_ops():
-    ops = OpList()
-    register_prim_ops(ops)
-    register_no_dynamic_shape_ops(ops)
-    register_dynamic_shape_ops(ops)
-    register_custom_ops(ops)
     return ops