SWDEV-549247 - Fix stream checks for hipModuleLaunchCooperativeKernel… (#858)

SWDEV-549247 - Fix stream validation checks in hipModuleLaunchCooperativeKernelMultiDevice

In CooperativeKernelMultiDevice mode, the stream's device
should be validated against the grid's assigned device
rather than the current device.

[ROCm/clr commit: 2aa2212d93]

projects/clr/hipamd/src/hip_module.cpp

@@ -429,7 +429,11 @@ hipError_t ihipModuleLaunchKernel(hipFunction_t f, amd::LaunchParams& launch_par
                                   uint64_t prevGridSum = 0, uint64_t allGridSum = 0,
                                   uint32_t firstDevice = 0) {
   int deviceId = hip::Stream::DeviceId(hStream);
-  if (deviceId != ihipGetDevice()) {
+
+  // Ensure the stream's device matches the current device,
+  // or the grid's assigned device in CooperativeKernelMultiDevice mode
+  int targetDevice = (numGrids == 0) ? ihipGetDevice() : gridId;
+  if (deviceId != targetDevice) {
     return hipErrorInvalidResourceHandle;
   }
   HIP_RETURN_ONFAIL(PlatformState::instance().initStatManagedVarDevicePtr(deviceId));