[UR][CUDA][HIP] Fix large kernel arg buffer overflow (#17538)

CUDA and HIP have a fixed maximum size for kernel arguments, this patch
ensures that passing arguments large than this maximum size is handled
properly by the adapters.

This code really needs a refactoring to avoid duplication and allow
larger argument sizes (supported by newer CUDA versions), but this will
be in follow up work.

This patch is a quick fix to avoid potential buffer overflows.

---------

Co-authored-by: aarongreig <[email protected]>

Author: npmiller

unified-runtime/source/adapters/cuda/kernel.cpp

@@ -362,24 +362,28 @@ urKernelSetArgPointer(ur_kernel_handle_t hKernel, uint32_t argIndex,
                       const ur_kernel_arg_pointer_properties_t *pProperties,
                       const void *pArgValue) {
   std::ignore = pProperties;
-  // setKernelArg is expecting a pointer to our argument
-  hKernel->setKernelArg(argIndex, sizeof(pArgValue), &pArgValue);
+  try {
+    // setKernelArg is expecting a pointer to our argument
+    hKernel->setKernelArg(argIndex, sizeof(pArgValue), &pArgValue);
+  } catch (ur_result_t Err) {
+    return Err;
+  }
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
                      const ur_kernel_arg_mem_obj_properties_t *Properties,
                      ur_mem_handle_t hArgValue) {
-  // Below sets kernel arg when zero-sized buffers are handled.
-  // In such case the corresponding memory is null.
-  if (hArgValue == nullptr) {
-    hKernel->setKernelArg(argIndex, 0, nullptr);
-    return UR_RESULT_SUCCESS;
-  }
-
   ur_result_t Result = UR_RESULT_SUCCESS;
   try {
+    // Below sets kernel arg when zero-sized buffers are handled.
+    // In such case the corresponding memory is null.
+    if (hArgValue == nullptr) {
+      hKernel->setKernelArg(argIndex, 0, nullptr);
+      return UR_RESULT_SUCCESS;
+    }
+
     auto Device = hKernel->getProgram()->getDevice();
     ur_mem_flags_t MemAccess =
         Properties ? Properties->memoryAccess

unified-runtime/source/adapters/cuda/kernel.hpp

@@ -117,6 +117,9 @@ struct ur_kernel_handle_t_ {
 
       // Copy new argument to storage if it hasn't been added before.
       if (ParamSizes[Index] == 0) {
+        if ((InsertPos + Size) > MaxParamBytes) {
+          throw UR_RESULT_ERROR_OUT_OF_RESOURCES;
+        }
         ParamSizes[Index] = Size;
         std::memcpy(&Storage[InsertPos], Arg, Size);
         ArgPointers[Index] = &Storage[InsertPos];

unified-runtime/source/adapters/hip/kernel.cpp

@@ -295,24 +295,28 @@ urKernelGetSubGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
 UR_APIEXPORT ur_result_t UR_APICALL urKernelSetArgPointer(
     ur_kernel_handle_t hKernel, uint32_t argIndex,
     const ur_kernel_arg_pointer_properties_t *, const void *pArgValue) {
-  // setKernelArg is expecting a pointer to our argument
-  hKernel->setKernelArg(argIndex, sizeof(pArgValue), &pArgValue);
+  try {
+    // setKernelArg is expecting a pointer to our argument
+    hKernel->setKernelArg(argIndex, sizeof(pArgValue), &pArgValue);
+  } catch (ur_result_t Err) {
+    return Err;
+  }
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
                      const ur_kernel_arg_mem_obj_properties_t *Properties,
                      ur_mem_handle_t hArgValue) {
-  // Below sets kernel arg when zero-sized buffers are handled.
-  // In such case the corresponding memory is null.
-  if (hArgValue == nullptr) {
-    hKernel->setKernelArg(argIndex, 0, nullptr);
-    return UR_RESULT_SUCCESS;
-  }
-
   ur_result_t Result = UR_RESULT_SUCCESS;
   try {
+    // Below sets kernel arg when zero-sized buffers are handled.
+    // In such case the corresponding memory is null.
+    if (hArgValue == nullptr) {
+      hKernel->setKernelArg(argIndex, 0, nullptr);
+      return UR_RESULT_SUCCESS;
+    }
+
     auto Device = hKernel->getProgram()->getDevice();
     hKernel->Args.addMemObjArg(argIndex, hArgValue,
                                Properties ? Properties->memoryAccess : 0);

unified-runtime/source/adapters/hip/kernel.hpp

@@ -110,6 +110,9 @@ struct ur_kernel_handle_t_ {
 
       // Copy new argument to storage if it hasn't been added before.
       if (ParamSizes[Index] == 0) {
+        if ((InsertPos + Size) > MAX_PARAM_BYTES) {
+          throw UR_RESULT_ERROR_OUT_OF_RESOURCES;
+        }
         ParamSizes[Index] = Size;
         std::memcpy(&Storage[InsertPos], Arg, Size);
         ArgPointers[Index] = &Storage[InsertPos];

unified-runtime/test/adapters/cuda/kernel_tests.cpp

@@ -160,6 +160,27 @@ TEST_P(cudaKernelTest, URKernelArgumentSimple) {
   ASSERT_EQ(storedValue, number);
 }
 
+TEST_P(cudaKernelTest, URKernelArgumentLarge) {
+  uur::raii::Program program = nullptr;
+  auto Length = std::strlen(ptxSource);
+  ASSERT_SUCCESS(urProgramCreateWithBinary(context, 1, &device, &Length,
+                                           (const uint8_t **)(&ptxSource),
+                                           nullptr, program.ptr()));
+  ASSERT_NE(program, nullptr);
+  ASSERT_SUCCESS(urProgramBuild(context, program, nullptr));
+
+  uur::raii::Kernel kernel = nullptr;
+  ASSERT_SUCCESS(urKernelCreate(program, "_Z8myKernelPi", kernel.ptr()));
+  ASSERT_NE(kernel, nullptr);
+
+  // The CUDA adapter can't do proper argument validation so any kernel will
+  // work for this test.
+  std::array<uint8_t, 4004> data;
+  data.fill(0);
+  ASSERT_EQ_RESULT(urKernelSetArgValue(kernel, 0, 4004, nullptr, data.data()),
+                   UR_RESULT_ERROR_OUT_OF_RESOURCES);
+}
+
 TEST_P(cudaKernelTest, URKernelArgumentSetTwice) {
   uur::raii::Program program = nullptr;
   auto Length = std::strlen(ptxSource);

unified-runtime/test/adapters/hip/CMakeLists.txt

@@ -4,14 +4,15 @@
 # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 add_adapter_test(hip
-    FIXTURE DEVICES
+    FIXTURE KERNELS
     SOURCES
         fixtures.h
         urContextGetNativeHandle.cpp
         urDeviceGetNativeHandle.cpp
         urEventGetNativeHandle.cpp
         test_context.cpp
         test_event.cpp
+        kernel_tests.cpp
     ENVIRONMENT
         "UR_ADAPTERS_FORCE_LOAD=\"$<TARGET_FILE:ur_adapter_hip>\""
 )

unified-runtime/test/adapters/hip/kernel_tests.cpp

@@ -0,0 +1,21 @@
+// Copyright (C) 2025 Intel Corporation
+// Part of the Unified-Runtime Project, under the Apache License v2.0 with LLVM
+// Exceptions. See LICENSE.TXT
+//
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include "kernel.hpp"
+#include "uur/fixtures.h"
+#include "uur/raii.h"
+
+using hipKernelTest = uur::urKernelExecutionTest;
+UUR_INSTANTIATE_DEVICE_TEST_SUITE(hipKernelTest);
+
+TEST_P(hipKernelTest, URKernelArgumentLarge) {
+  // The HIP adapter can't do proper argument validation so any kernel will
+  // work for this test.
+  std::array<uint8_t, 4004> data;
+  data.fill(0);
+  ASSERT_EQ_RESULT(urKernelSetArgValue(kernel, 0, 4004, nullptr, data.data()),
+                   UR_RESULT_ERROR_OUT_OF_RESOURCES);
+}