[Openmp-commits] [llvm] [openmp] [OpenMPOpt] Ask the runtime how many of a block's threads can be workers (PR #217080)

[Larry Meadows via Openmp-commits]

================
@@ -4534,26 +4529,23 @@ struct AAKernelInfoFunction : AAKernelInfo {
     IsWorker->setDebugLoc(DLoc);
     CondBrInst::Create(IsWorker, IsWorkerCheckBB, UserCodeEntryBB, InitBB);
 
+    // How many of the block's threads can be worker threads is a property of
+    // the launch geometry, which the runtime knows and the block size alone
+    // does not determine: a target may host the main thread in a whole warp
+    // above the workers, or in a single thread above them. Ask the runtime
+    // rather than subtracting a warp here, or the workers in between are left
+    // in neither group, waiting for no parallel region while it is handed
+    // iterations.
     Module &M = *Kernel->getParent();
-    FunctionCallee BlockHwSizeFn =
+    FunctionCallee MaxTeamThreadsFn =
         OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
-            M, OMPRTL___kmpc_get_hardware_num_threads_in_block);
-    FunctionCallee WarpSizeFn =
-        OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
-            M, OMPRTL___kmpc_get_warp_size);
-    CallInst *BlockHwSize =
-        CallInst::Create(BlockHwSizeFn, "block.hw_size", IsWorkerCheckBB);
-    OMPInfoCache.setCallingConvention(BlockHwSizeFn, BlockHwSize);
-    BlockHwSize->setDebugLoc(DLoc);
-    CallInst *WarpSize =
-        CallInst::Create(WarpSizeFn, "warp.size", IsWorkerCheckBB);
-    OMPInfoCache.setCallingConvention(WarpSizeFn, WarpSize);
-    WarpSize->setDebugLoc(DLoc);
-    Instruction *BlockSize = BinaryOperator::CreateSub(
-        BlockHwSize, WarpSize, "block.size", IsWorkerCheckBB);
-    BlockSize->setDebugLoc(DLoc);
+            M, OMPRTL___kmpc_get_max_team_threads);
+    CallInst *MaxTeamThreads =
+        CallInst::Create(MaxTeamThreadsFn, "max_team_threads", IsWorkerCheckBB);
+    OMPInfoCache.setCallingConvention(MaxTeamThreadsFn, MaxTeamThreads);
+    MaxTeamThreads->setDebugLoc(DLoc);
----------------
lfmeadow wrote:

The premise is right — `AAFoldRuntimeCall` folds `__kmpc_get_hardware_num_threads_in_block` to `omp_target_thread_limit`, and `__kmpc_get_max_team_threads` is not registered — but there is no folding opportunity being lost, because the bound was never a constant to begin with.

`__kmpc_get_warp_size` is not one of the folded runtime calls, so in the old code the subtraction always had an opaque operand. The in-tree expectations on `main` show it:

```llvm
%block.hw_size = call i32 @__kmpc_get_hardware_num_threads_in_block()
%warp.size = call i32 @__kmpc_get_warp_size()
%block.size = sub i32 %block.hw_size, %warp.size
%thread_is_main_or_worker = icmp slt i32 %0, %block.size
```

Even where `omp_target_thread_limit` is present and `%block.hw_size` folds, `%block.size` stays dynamic and the comparison never becomes constant. One opaque call in place of two is if anything less work. Folding elsewhere is unaffected, since `AAFoldRuntimeCall` works per call site and user-code calls to the old getter still fold.

Teaching the folder to treat the two equivalently would be actively wrong: `omp_target_thread_limit` is the block size, and the whole point of this patch is that the worker bound is *not* the block size — the threads hosting the main thread are excluded, by a convention the runtime owns. Folding the new getter to the thread limit would reintroduce the bug the patch fixes, which downstream shows up as workers stranded between the two bounds.


https://github.com/llvm/llvm-project/pull/217080