Update base for Update on "[ET-VK] Contiguous buffer implementation for embedding"

Title says it all!

Also moved around some structs throughout the `glslh` files for better code organization.

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

[ghstack-poisoned]

Author: ssjia

backends/vulkan/utils.py

@@ -259,6 +259,32 @@ def tensor_node_is_bool(node: torch.fx.Node) -> bool:
     return False
 
 
+def ndim_of(node: Any) -> Optional[int]:
+    """
+    Returns the number of dimensions of the tensor produced by the given node
+    """
+    if not is_single_tensor_node(node):
+        return None
+
+    return node.meta["val"].ndim
+
+
+def is_unsqueezed_vector(node: torch.fx.Node) -> bool:
+    """
+    Returns True if the node's tensor has all dimensions equal to 1 except for the last dimension.
+    """
+    if not is_single_tensor_node(node):
+        return False
+
+    tensor = node.meta["val"]
+    assert isinstance(tensor, FakeTensor)
+
+    if len(tensor.shape) < 1:
+        return False
+    # All dims except last are 1, last can be any size
+    return all(dim == 1 for dim in tensor.shape[:-1])
+
+
 def op_contains_bool_tensor(node: torch.fx.Node) -> bool:
     """
     Returns true if the operator used to compute the given node contains a bool tensor
@@ -267,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
 
@@ -582,6 +609,16 @@ def make_intersect(self, other: "TensorRepSet") -> "TensorRepSet":
             self.valid_texture_layouts & other.valid_texture_layouts,
         )
 
+    def make_union(self, other: "TensorRepSet") -> "TensorRepSet":
+        """
+        Merge this TensorRepSet with another TensorRepSet, returning a new TensorRepSet
+        with the union of the two.
+        """
+        return TensorRepSet(
+            self.valid_buffer_layouts | other.valid_buffer_layouts,
+            self.valid_texture_layouts | other.valid_texture_layouts,
+        )
+
     def is_compatible(self, storage: TensorRepr) -> bool:
         """
         Check if this TensorRepr is compatible with the given TensorRepSet.
@@ -1240,6 +1277,26 @@ def is_in_8bit_range(tensor: torch.Tensor) -> bool:
 ##
 
 
+def normalize_dims(dims: Union[int, List[int]], ndim: int) -> Union[int, List[int]]:
+    """
+    Normalize dimension indices to be non-negative and within [0, ndim).
+    Accepts a single int or a list of ints.
+    """
+    if isinstance(dims, int):
+        if dims < 0:
+            dims += ndim
+
+        return dims
+
+    normalized = []
+    for d in dims:
+        if d < 0:
+            d += ndim
+        normalized.append(d)
+
+    return normalized
+
+
 def nchw_dim_to_whcn_dim(nchw_dim: int, ndim: int) -> int:
     # Handle negative indices for nchw_dim
     if nchw_dim < 0: