Update on "[ET-VK] Introduce specialized implementation for per-row reduction"

Title says it all!

This diff also adds support for argmin and argmax, but only for per-row reduction.

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

[ghstack-poisoned]

Author: ssjia

backends/vulkan/op_registry.py

@@ -422,6 +422,7 @@ def register_softmax_op():
 
 def get_dims_reduced(node: torch.fx.Node) -> Union[int, List[int]]:
     ndim = utils.ndim_of(node.args[0])
+    assert ndim is not None
     dims_reduced = None
     if len(node.args) >= 1:
         dims_reduced = node.args[1]
@@ -438,6 +439,7 @@ def get_dims_reduced(node: torch.fx.Node) -> Union[int, List[int]]:
     if isinstance(dims_reduced, (list, tuple)) and len(dims_reduced) == 1:
         dims_reduced = dims_reduced[0]
 
+    assert isinstance(dims_reduced, (int, list, tuple))
     return utils.normalize_dims(dims_reduced, ndim)
 
 
@@ -456,6 +458,7 @@ def is_reduce_node_supported_by_per_row_impl(node: torch.fx.Node) -> bool:
     special case implementation.
     """
     input_ndim = utils.ndim_of(node.args[0])
+    assert input_ndim is not None
     dims_reduced = get_dims_reduced(node)
 
     return dims_reduced == input_ndim - 1
@@ -505,7 +508,9 @@ def pick_storage_for_reduce(node: torch.fx.Node):
 
     # For 2D reductions, the packed dimension cannot be one of the reduced dims
     if isinstance(dim_list, (list, tuple)) and len(dim_list) == 2:
+        # pyre-ignore[6]
         reduce_dim1_whcn = utils.nchw_dim_to_whcn_dim(dim_list[0], ndim)
+        # pyre-ignore[6]
         reduce_dim2_whcn = utils.nchw_dim_to_whcn_dim(dim_list[1], ndim)
 
         possible_packed_dims = {0, 1, 2}
@@ -569,6 +574,7 @@ def register_2d_pool_op():
 def register_convolution_op():
     def check_conv_node(node: torch.fx.Node) -> bool:
         x = node.args[0]
+        assert isinstance(x, torch.fx.Node)
         x_shape = x.meta["val"].size()
         # 4-D input implies 2D convolution
         if len(x_shape) == 4:

backends/vulkan/utils.py

@@ -293,6 +293,7 @@ def op_contains_bool_tensor(node: torch.fx.Node) -> bool:
         return True
 
     for arg_node in node.args:
+        # pyre-ignore[6]
         if is_tensor_node(arg_node) and tensor_node_is_bool(arg_node):
             return True
 
@@ -756,6 +757,7 @@ def make_filtered_tensor_repset(
 CONTIGUOUS_BUFFER = TensorRepSet({VkMemoryLayout.TENSOR_WIDTH_PACKED}, set())
 
 WIDTH_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_WIDTH_PACKED})
+HEIGHT_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_HEIGHT_PACKED})
 CHANNELS_PACKED_TEXTURE = TensorRepSet(set(), {VkMemoryLayout.TENSOR_CHANNELS_PACKED})
 
 ANY_TEXTURE = TensorRepSet(set(), all_memory_layouts)