[llvm] [Offload] Implement better level zero dispatch (PR #218367)

[Łukasz Plewa via llvm-commits]

================
@@ -52,9 +54,62 @@ Error L0QueueTy::dispatchLaunchKernel(ze_kernel_handle_t Kernel,
                                       ze_event_handle_t *WaitEvents) {
   // Unlock KEnv lock after launching the kernel.
   llvm::scope_exit UnlockGuard([&KEnv]() { KEnv.Lock.unlock(); });
-  return CmdList->appendLaunchKernelWithArgs(
-      Kernel, &KEnv.GroupCounts, &KEnv.GroupSizes, KEnv.ArgPtrs, SignalEvent,
-      NumWaitEvents, WaitEvents, KEnv.IsCooperative);
+
+  bool AppendLaunchKernelWithArgsAvailable =
+      Device.getL0Context().LaunchKernelWithArguments.available();
+
+  if (AppendLaunchKernelWithArgsAvailable &&
+      Device.getL0Context().AppendLaunchKernelSupported.load(
+          std::memory_order_acquire)) {
+    auto Err = CmdList->appendLaunchKernelWithArgs(
+        Kernel, &KEnv.GroupCounts, &KEnv.GroupSizes, KEnv.ArgPtrs, SignalEvent,
+        NumWaitEvents, WaitEvents, KEnv.IsCooperative);
+
+    if (!Err)
+      return Plugin::success();
+
+    // Check if Err is ErrorCode::UNSUPPORTED, if so consume it
+    Err = llvm::handleErrors(
+        std::move(Err),
+        [&](std::unique_ptr<error::OffloadError> E) -> llvm::Error {
+          if (E->convertToErrorCode() ==
+              error::make_error_code(error::ErrorCode::UNSUPPORTED))
+            return llvm::Error::success(); // Swallow error
+          return llvm::Error(std::move(E));
+        });
+
+    if (Err) {
+      // Err is still here, so it was not ErrorCode::UNSUPPORTED
+      return Err;
+    }
+
+    // No Err - it was ErrorCode::UNSUPPORTED, continue into fallback
+  }
+
+  // Submit kernel using older set of APIs - zeKernelSetArgumentValue
+  auto &GroupSizes = KEnv.GroupSizes;
+  auto Res = zeKernelSetGroupSize(Kernel, GroupSizes.groupSizeX,
+                                  GroupSizes.groupSizeY, GroupSizes.groupSizeZ);
+  if (Res != ZE_RESULT_SUCCESS)
+    return error::createOffloadError(ErrorCode::UNKNOWN,
+                                     "Could not set group size!");
+
+  auto &KernelProperties = KEnv.KernelPR;
+
+  for (uint32_t KernelArg = 0; KernelArg < KernelProperties.NumKernelArgs;
+       KernelArg++) {
+    uint32_t ArgSize = KernelProperties.ArgSizes[KernelArg];
----------------
lplewa wrote:

AI thinks that this array might be null - and we have null deference risk

> Scenario A (most direct): driver lacks zeCommandListAppendLaunchKernelWithArguments >entirely.
>- L0Queue.cpp:61 — AppendLaunchKernelWithArgsAvailable is false, so we skip the >modern path and fall straight into the legacy branch at line 89.
>- If the same driver also lacks zexKernelGetArgumentSize, L0Kernel.cpp:46 never allocates >KernelPR.ArgSizes — it stays as a default-constructed null unique_ptr<uint32_t[]>.
>- Line 101 does KernelProperties.ArgSizes[KernelArg] → null dereference on any kernel >with NumKernelArgs > 0.
>
>Scenario B: driver reports LaunchKernelWithArguments as available but returns ZE_RESULT_ERROR_UNSUPPORTED_FEATURE at runtime (the case the sticky AppendLaunchKernelSupported flag exists to handle). Same fallback → same null deref, if KernelGetArgumentSize was never loaded.
>
>Is that combination realistic? The PR's own compatibility table doesn't list zexKernelGetArgumentSize — it's an older Intel-specific extension. On a modern Intel driver the two tend to coexist, so on typical Intel hardware this may not fire in practice. But:
>
>- The old code had the guard baked in: zeCommandListAppendLaunchKernelWithArgumentsFallbackAvailable() returned false unless zexKernelGetArgumentSize_ptr != nullptr, so canCall<> would report the API as unavailable and the caller got a clean UNSUPPORTED error. The new code dropped that coupling — the two dispatchers are now independent.
>- Any non-Intel L0 driver, or a very old Intel driver, that has one but not the other would hit this. The whole point of the fallback existing is to handle drivers that don't have the modern API — those are exactly the drivers most likely to also be missing the zex* helper.
>
>So it's not a "purely theoretical UB" — it's a real regression in the safety net the removed code was providing. Cheap fix: gate the fallback on Context.KernelGetArgumentSize.available() (return UNSUPPORTED if false), or equivalently on KernelProperties.ArgSizes != nullptr.

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