Merge pull request #8 from rohany/fix-cannon

tests: fix cannonMM implementation to scale k with gx
tensor-compiler · Jun 6, 2021 · 7f780f3 · 7f780f3
2 parents ba8addd + 07e0b21
commit 7f780f3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/test/tests-distributed.cpp b/test/tests-distributed.cpp
@@ -189,14 +189,14 @@ TEST(distributed, cannonMM) {
       //  version of the code does.
       // .distributeOnto({i, j}, {in, jn}, {il, jl}, a(i, j))
       .distribute({i, j}, {in, jn}, {il, jl}, grid)
-      .divide(k, ko, ki, 2)
+      .divide(k, ko, ki, gx)
       .reorder({ko, il, jl})
       .stagger(ko, {in, jn}, kos)
       .pushCommUnder(b(i, k), kos)
       .pushCommUnder(c(k, j), kos)
       .pushCommUnder(a(i, j), in)
       // This can be enabled on Sapling where we have an OpenMP + OpenBLAS build.
-      // .swapLeafKernel(il, gemm)
+       .swapLeafKernel(il, gemm)
       ;
 
   auto lowered = lower(stmt, "computeLegion", false, true);