[CUDA][HIP] Fix for command-buffer local argument update

After setting kernel arguments during update, we
need to reset the amount of local memory used.

Author: Ewan Crawford

source/adapters/cuda/command_buffer.cpp

@@ -1396,14 +1396,22 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferUpdateKernelLaunchExp(
 
   CUDA_KERNEL_NODE_PARAMS &Params = KernelCommandHandle->Params;
 
+  const auto LocalSize = KernelCommandHandle->Kernel->getLocalSize();
+  if (LocalSize != 0) {
+    // Clean the local size, otherwise calling updateKernelArguments() in
+    // future updates with local arguments will incorrectly increase the
+    // size further.
+    KernelCommandHandle->Kernel->clearLocalSize();
+  }
+
   Params.func = CuFunc;
-  Params.gridDimX = BlocksPerGrid[0];
-  Params.gridDimY = BlocksPerGrid[1];
-  Params.gridDimZ = BlocksPerGrid[2];
-  Params.blockDimX = ThreadsPerBlock[0];
-  Params.blockDimY = ThreadsPerBlock[1];
-  Params.blockDimZ = ThreadsPerBlock[2];
-  Params.sharedMemBytes = KernelCommandHandle->Kernel->getLocalSize();
+  Params.gridDimX = static_cast<unsigned int>(BlocksPerGrid[0]);
+  Params.gridDimY = static_cast<unsigned int>(BlocksPerGrid[1]);
+  Params.gridDimZ = static_cast<unsigned int>(BlocksPerGrid[2]);
+  Params.blockDimX = static_cast<unsigned int>(ThreadsPerBlock[0]);
+  Params.blockDimY = static_cast<unsigned int>(ThreadsPerBlock[1]);
+  Params.blockDimZ = static_cast<unsigned int>(ThreadsPerBlock[2]);
+  Params.sharedMemBytes = LocalSize;
   Params.kernelParams =
       const_cast<void **>(KernelCommandHandle->Kernel->getArgIndices().data());
 

test/conformance/device_code/saxpy_usm_local_mem.cpp

@@ -15,15 +15,27 @@ int main() {
     uint32_t A = 42;
 
     sycl_queue.submit([&](sycl::handler &cgh) {
-        sycl::local_accessor<uint32_t, 1> local_mem(local_size, cgh);
+        sycl::local_accessor<uint32_t, 1> local_mem_A(local_size, cgh);
+        sycl::local_accessor<uint32_t, 1> local_mem_B(1, cgh);
+
         cgh.parallel_for<class saxpy_usm_local_mem>(
             sycl::nd_range<1>{{array_size}, {local_size}},
             [=](sycl::nd_item<1> itemId) {
                 auto i = itemId.get_global_linear_id();
                 auto local_id = itemId.get_local_linear_id();
-                local_mem[local_id] = i;
-                Z[i] = A * X[i] + Y[i] + local_mem[local_id] +
+
+                local_mem_A[local_id] = i;
+                if (i == 0) {
+                    local_mem_B[0] = 0xA;
+                }
+
+                Z[i] = A * X[i] + Y[i] + local_mem_A[local_id] +
                        itemId.get_local_range(0);
+
+                if (i == 0) {
+                    Z[i] += local_mem_B[0];
+                }
+
             });
     });
     return 0;