More updates for 3.1 (#958)

* Updates for 3.1

* Minor change

* doc link fix

* Minor updates

examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage.h

@@ -641,6 +641,11 @@ public:
 
     }
 
+    // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
+    cutlass::arch::cp_async_fence();
+    cutlass::arch::cp_async_wait<0>();
+    __syncthreads();
+
     // 2nd Gemm
 
     /// Iterator to load a warp-scoped tile of A1 operand from intermediate accumulator tile
@@ -871,7 +876,10 @@ public:
 
     }
 
-
+    // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
+    cutlass::arch::cp_async_fence();
+    cutlass::arch::cp_async_wait<0>();
+    __syncthreads();
 
   }
 };

examples/13_two_tensor_op_fusion/threadblock/b2b_mma_multistage_smem_accumulator.h

@@ -664,6 +664,11 @@ public:
 
     }
 
+    // Insert fence and wait for all outstanding cp.async operations to commit.
+    cutlass::arch::cp_async_fence();
+    cutlass::arch::cp_async_wait<0>();
+    __syncthreads();
+
     /// Epilogue for the first Implicit Gemm
     Epilogue0 epilogue0;
 
@@ -855,7 +860,10 @@ public:
 
     }
 
-
+    // Commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
+    cutlass::arch::cp_async_fence();
+    cutlass::arch::cp_async_wait<0>();
+    __syncthreads();
 
   }
 };

examples/45_dual_gemm/threadblock/dual_mma_multistage.h

@@ -759,13 +759,10 @@ public:
       accum1 = plus_accum(accum1, tmp_accum1); 
     }
  
-    if (SharedMemoryClear == SharedMemoryClearOption::kZfill) {
-      // commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
-      cutlass::arch::cp_async_fence();
-      cutlass::arch::cp_async_wait<0>();
-      __syncthreads();
-    }
-
+    // commit and drain all pending and predicated cp.async pnz from the GEMM mainloop
+    cutlass::arch::cp_async_fence();
+    cutlass::arch::cp_async_wait<0>();
+    __syncthreads();
   }
 };
 

examples/47_ampere_gemm_universal_streamk/ampere_gemm_universal_streamk_broadcast.cu

@@ -461,11 +461,6 @@ Result run(std::string description, Options &options)
     std::cout << "  GFLOPs: " << result.gflops << std::endl;
   }
 
-  // TODO: uncomment when results match
-  //if (!result.passed) {
-  //  exit(-1);
-  //}
-
   return result;
 }