[llvm] [offload] Add CUDA_ROOT to path for unit tests (PR #214282)

[Nick Sarnie via llvm-commits]

https://github.com/sarnex created https://github.com/llvm/llvm-project/pull/214282

We need this to fix the test added in https://github.com/llvm/llvm-project/pull/212860. Right now it errors saying it can't find `ptxas`.
We already have code doing this for the e2e tests, but we also need it for the unit tests.
We had a similar fix for AMDGPU in https://github.com/llvm/llvm-project/pull/213149.
Locally reproduced the issue and verified the fix.

>From 9f4174bbb68bd9c7604d198cc677044a3e8db813 Mon Sep 17 00:00:00 2001
From: Nick Sarnie <nick.sarnie at intel.com>
Date: Wed, 5 Aug 2026 08:55:12 -0700
Subject: [PATCH] [offload] Add CUDA_ROOT to path for unit tests

Signed-off-by: Nick Sarnie <nick.sarnie at intel.com>

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
 offload/test/unit/lit.cfg.py      | 19 +++++++++++++------
 offload/test/unit/lit.site.cfg.in |  1 +
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/offload/test/unit/lit.cfg.py b/offload/test/unit/lit.cfg.py
index 41d64da9ae9c2..1dc1fa6bec49c 100644
--- a/offload/test/unit/lit.cfg.py
+++ b/offload/test/unit/lit.cfg.py
@@ -13,16 +13,23 @@
 # suffixes: A list of file extensions to treat as test files.
 config.suffixes = []
 
+def prepend_executable_path(path):
+    old_path = config.environment.get("PATH")
+    config.environment["PATH"] = (
+        f"{path}{os.path.pathsep}{old_path}" if old_path else path
+    )
+
 # Add the tools bin dir to the search path so the JIT tests can
 # find it. Prefer the tools from the configured LLVM install or
 # bootstrapping build.
 if config.bin_llvm_tools_dir:
-    old_path = config.environment.get("PATH")
-    config.environment["PATH"] = (
-        f"{config.bin_llvm_tools_dir}{os.path.pathsep}{old_path}"
-        if old_path
-        else config.bin_llvm_tools_dir
-    )
+    prepend_executable_path(config.bin_llvm_tools_dir)
+
+# The CUDA plugin invokes 'ptxas' via the PATH when JIT-compiling PTX for
+# NVPTX devices. Add the CUDA bin dir to the search path so the unit tests can
+# find it, mirroring what the main lit.cfg does for the lit tests.
+if config.cuda_path:
+    prepend_executable_path(f"{config.cuda_path}{os.path.sep}bin")
 
 # test_source_root: The root path where tests are located.
 # test_exec_root: The root path where tests should be run.
diff --git a/offload/test/unit/lit.site.cfg.in b/offload/test/unit/lit.site.cfg.in
index 7c1becaefbee1..821503a92327c 100644
--- a/offload/test/unit/lit.site.cfg.in
+++ b/offload/test/unit/lit.site.cfg.in
@@ -4,6 +4,7 @@ config.library_dir = "@LIBOMPTARGET_LIBRARY_DIR@"
 config.unittest_dir = "@OFFLOAD_UNITTEST_DIR@"
 config.llvm_build_mode = lit_config.substitute("@LLVM_BUILD_MODE@")
 config.bin_llvm_tools_dir = "@LLVM_TOOLS_BINARY_DIR@"
+config.cuda_path = "@CUDA_ROOT@"
 
 # Let the main config do the real work.
 lit_config.load_config(config, "@CMAKE_CURRENT_SOURCE_DIR@/unit/lit.cfg.py")