make blocksync compatible with warp specialization

NVIDIA · Dec 20, 2024 · 0e213d6 · 0e213d6
1 parent 120cd8a
commit 0e213d6
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/csrc/codegen.cpp b/csrc/codegen.cpp
@@ -3325,10 +3325,16 @@ class CudaKernelGenerator : private kir::ConstIrVisitor {
     // Use a custom synchronization method if enabled
     if (getNvFuserEnv("USE_BLOCK_SYNC_ATOMIC")) {
       indent() << "block_sync::sync();\n";
-    } else if (isAligned()) {
-      indent() << "__syncthreads();\n";
     } else {
-      indent() << "__barrier_sync(0);\n";
+      ArgumentBuilder sync_call_template_parms;
+      sync_call_template_parms.arg(isAligned());
+
+      ArgumentBuilder sync_call_args;
+      sync_call_args.arg(genComputeBlockDim());
+
+      auto sync_call =
+          genCall("block_sync::sync", sync_call_template_parms, sync_call_args);
+      indent() << sync_call << ";\n";
     }
   }