LWP100-1464: Fix fusing of tail loops

client/include/DataParallelGEMMSolution.hpp

@@ -549,7 +549,7 @@ namespace rocRoller
                     // predicates
                     // unrollK size match predicates
 
-                    if(params->unrollX <= 1 && params->unrollY <= 1 && !params->streamK)
+                    if(params->tailLoops && !params->streamK)
                     {
                         auto unrollKPredicate = (aSizeExps[1] % macKExp == zero);
                         setComment(unrollKPredicate, "K must be a multiple of macK.");

lib/include/rocRoller/KernelGraph/ControlToCoordinateMapper.hpp

@@ -70,6 +70,11 @@ namespace rocRoller::KernelGraph
             return a.id < b.id;
         }
 
+        bool inline operator==(TypeAndSubDimension const& a, TypeAndSubDimension const& b)
+        {
+            return a.id == b.id && a.subdimension == b.subdimension;
+        }
+
         struct TypeAndNaryArgument
         {
             std::string  id;
@@ -89,6 +94,11 @@ namespace rocRoller::KernelGraph
             return a.id < b.id;
         }
 
+        bool inline operator==(TypeAndNaryArgument const& a, TypeAndNaryArgument const& b)
+        {
+            return a.id == b.id && a.argument == b.argument;
+        }
+
         enum class ComputeIndexArgument : int
         {
             TARGET = 0,
@@ -117,6 +127,11 @@ namespace rocRoller::KernelGraph
             return a.argument < b.argument;
         }
 
+        bool inline operator==(ComputeIndex const& a, ComputeIndex const& b)
+        {
+            return a.argument == b.argument && a.index == b.index;
+        }
+
         using ConnectionSpec = std::variant<std::monostate,
                                             JustNaryArgument,
                                             ComputeIndex,

lib/include/rocRoller/KernelGraph/Transforms/FuseLoops.hpp

@@ -40,8 +40,9 @@ namespace rocRoller
         class FuseLoops : public GraphTransform
         {
         public:
-            KernelGraph apply(KernelGraph const& original) override;
-            std::string name() const override
+            KernelGraph                          apply(KernelGraph const& original) override;
+            virtual std::vector<GraphConstraint> postConstraints() const override;
+            std::string                          name() const override
             {
                 return "FuseLoops";
             }

lib/source/KernelGraph/Transformations/RemoveDuplicates.cpp

@@ -358,6 +358,27 @@ namespace rocRoller
                     }
                 }
 
+                // If there exists a reindexing "chain" A -> B -> C, then processing
+                // B -> C before A -> B causes the old references to A to be left
+                // dangling. To avoid this, we collapse the chain so A -> C and B -> C,
+                // which matches the behaviour we would see if A -> B was processed first.
+                for(auto [oldTag, newTag] : expressionReindexer.coordinates)
+                {
+                    while(expressionReindexer.coordinates.count(newTag) > 0 && newTag != oldTag)
+                    {
+                        auto nextTag = expressionReindexer.coordinates.at(newTag);
+                        if(nextTag != newTag)
+                        {
+                            newTag = expressionReindexer.coordinates.at(newTag);
+                        }
+                        else
+                        {
+                            break;
+                        }
+                    }
+                    expressionReindexer.coordinates[oldTag] = newTag;
+                }
+
                 // Update tile references
                 auto kernel = *graph.control.roots().begin();
                 reindexExpressions(graph, kernel, expressionReindexer);

lib/source/KernelGraph/Transformations/UnrollLoops.cpp

@@ -622,22 +622,6 @@ namespace rocRoller
                 }
             }
 
-            {
-                auto containingForLoops = graph.control.nodesContaining(loop).filter(
-                    graph.control.isElemType<ForLoopOp>());
-                for(auto containingLoop : containingForLoops)
-                {
-                    if(getUnrollAmount(graph, containingLoop, m_params) > 1)
-                    {
-                        Log::debug("Not adding tail loop for {} because it is contained by {} "
-                                   "which is also unrolled.",
-                                   loop,
-                                   containingLoop);
-                        return std::nullopt;
-                    }
-                }
-            }
-
             auto loopSizeType        = resultVariableType(loopSize);
             auto amount              = Expression::literal(unrollAmount, loopSizeType);
             auto loopSizeRoundedDown = (loopSize / amount) * amount;

scripts/lib/rrperf/rrsuites.py

@@ -430,6 +430,26 @@ def tail_loop_reproducer():
     )
 
 
+def jammed_tail_loop():
+    yield mkGEMM(
+        M=256,
+        N=256,
+        K=24,
+        trans_A="T",
+        trans_B="N",
+        wave_m=32,
+        wave_n=32,
+        wave_k=2,
+        wave_b=1,
+        mac_m=64,
+        mac_n=64,
+        mac_k=8,
+        unroll_x=2,
+        unroll_y=2,
+        # unroll_k=2 implicit due to prefetchInFlight=2
+    )
+
+
 def guidepost_1():
     yield mkGEMM(
         HGEMM_7680x8448x8448,

test/unit/GEMMFusion.cpp

@@ -77,7 +77,7 @@ namespace GEMMDriverTest
             AssertFatal(M % gemm.macM == 0, "MacroTile size mismatch (M)");
             AssertFatal(N % gemm.macN == 0, "MacroTile size mismatch (N)");
 
-            if(gemm.unrollK > 0)
+            if(gemm.unrollK > 0 && !gemm.tailLoops)
             {
                 AssertFatal(K % (gemm.macK * gemm.unrollK) == 0,
                             "MacroTile size mismatch (K unroll)");