Address comments #3

Signed-off-by: Harsh Menon <harsh@nod-labs.com>
iree-org · Oct 30, 2024 · e09dff2 · e09dff2
1 parent 9d6bb0f
commit e09dff2
Showing 1 changed file with 6 additions and 3 deletions.
diff --git a/iree/turbine/kernel/wave/utils.py b/iree/turbine/kernel/wave/utils.py
@@ -289,11 +289,14 @@ def find_mma_in_slice(node: CustomOp) -> Optional[MMA]:
                 return prev_mma
         return None
 
+    # Look in the backward slices of both the LHS and RHS to find
+    # mmas. If found, add reshapes if necessary.
     for mma in mma_nodes:
         custom_mma = get_custom(mma)
-        prev_mma = find_mma_in_slice(custom_mma.lhs) or find_mma_in_slice(
-            custom_mma.rhs
-        )
+        prev_mma = find_mma_in_slice(custom_mma.lhs)
+        if prev_mma:
+            add_reshape_if_needed(custom_mma, prev_mma)
+        prev_mma = find_mma_in_slice(custom_mma.rhs)
         if prev_mma:
             add_reshape_if_needed(custom_mma, prev_mma)