a few fixes & improvements

microsoft · May 18, 2024 · 2daa05c · 2daa05c
1 parent 8596dbf
commit 2daa05c
Show file tree

Hide file tree

Showing 7 changed files with 2,054 additions and 34 deletions.
diff --git a/ark/codegen.cpp b/ark/codegen.cpp
@@ -143,8 +143,7 @@ CodeGenerator::Impl::Impl(const Json &plan,
             }
         }
         for (auto &rg : pg["ResourceGroups"]) {
-            body_ss << this->resource_group(rg, plan.at("TaskInfos"),
-                                            proc_range);
+            body_ss << resource_group(rg, plan.at("TaskInfos"), proc_range);
         }
         unsynced.push_back(proc_range);
         pg_idx++;

diff --git a/ark/include/kernels/kernel_template.in b/ark/include/kernels/kernel_template.in
@@ -15,14 +15,16 @@ __forceinline__ __device__ void task_seq(char *_buf) {
     constexpr size_t SramBytesPerWarp = SlotSramBytes / SlotNumWarps;
     size_t p = ((blockIdx.x + gridDim.x - ProcCurrent) % gridDim.x) / ProcStep;
     size_t k = threadIdx.x / SlotNumThreads;
+    if constexpr (ARK_WARPS_PER_BLOCK > SlotNumWarps) {
+      if (k >= NumSlots) return;
+    }
     size_t task_id_base = TaskBegin + p * TaskStep * TaskGranularity;
     for (size_t t = k; ; t += NumSlots) {
       size_t task_id = task_id_base + TaskStep *
         (t % TaskGranularity + t / TaskGranularity * TaskGranularity * NumProcs);
       if (task_id >= TaskEnd) break;
       task(_buf, task_id, SramBytesPerWarp);
     }
-    __syncthreads();
   }
 }