Address comments

Signed-off-by: Harsh Menon <harsh@nod-labs.com>

iree/turbine/kernel/wave/decompose_reduce_ops.py

@@ -66,7 +66,9 @@ def determine_shuffle_config(
     offset = subs_idxc(offset)
     offset_table = [offset.subs({THREAD_0: i}) for i in range(subgroup_size)]
     unique_offsets = list(dict.fromkeys(offset_table))
-    cluster_size = len(set(offset_table))
+    # The cluster size represents the number of unique threads that are participating in a shuffle.
+    # We can obtain this information by just computing the number of unique entries in the offset table.
+    cluster_size = len(unique_offsets)
     thread_ids = []
     for thread_offset in unique_offsets:
         thread_ids.append(offset_table.index(thread_offset))

iree/turbine/kernel/wave/index_sequence_analysis.py

@@ -220,6 +220,11 @@ def set_vector_shapes(
             or node in mma_slices[mma][MMA_RHS]
         ):
             # Ensure that the operators indexing dims are present in the anchor.
+            # For example, say we have a write node with indexing dimensions [B, M, N]
+            # and there are two potential options for anchors: an MMA with
+            # indexing dimensions [B, M, K1, K2] and another with indexing dimensions
+            # [B, M, N, K2], we want to pick the second one otherwise the index
+            # that is set from the anchor will not be accurate.
             if not set(custom.indexing_dims).issubset(mma.indexing_dims):
                 continue
             custom.anchor = mma