diff --git a/source/adapters/hip/device.hpp b/source/adapters/hip/device.hpp index 181c5a7bdb..3474b313b6 100644 --- a/source/adapters/hip/device.hpp +++ b/source/adapters/hip/device.hpp @@ -26,6 +26,8 @@ struct ur_device_handle_t_ { ur_platform_handle_t Platform; hipCtx_t HIPContext; uint32_t DeviceIndex; + int DeviceMaxLocalMem; + int ManagedMemSupport; public: ur_device_handle_t_(native_type HipDevice, hipCtx_t Context, @@ -35,6 +37,13 @@ struct ur_device_handle_t_ { ~ur_device_handle_t_() noexcept(false) { UR_CHECK_ERROR(hipDevicePrimaryCtxRelease(HIPDevice)); + + UR_CHECK_ERROR(hipDeviceGetAttribute( + &DeviceMaxLocalMem, hipDeviceAttributeMaxSharedMemoryPerBlock, + HIPDevice)); + UR_CHECK_ERROR(hipDeviceGetAttribute( + &ManagedMemSupport, hipDeviceAttributeManagedMemory, + HIPDevice)); } native_type get() const noexcept { return HIPDevice; }; @@ -48,6 +57,10 @@ struct ur_device_handle_t_ { // Returns the index of the device relative to the other devices in the same // platform uint32_t getIndex() const noexcept { return DeviceIndex; }; + + int getDeviceMaxLocalMem() const noexcept { return DeviceMaxLocalMem; }; + + bool getManagedMemSupport() const noexcept { return ManagedMemSupport; }; }; int getAttribute(ur_device_handle_t Device, hipDeviceAttribute_t Attribute); diff --git a/source/adapters/hip/enqueue.cpp b/source/adapters/hip/enqueue.cpp index c24287749e..cd8b60e7a2 100644 --- a/source/adapters/hip/enqueue.cpp +++ b/source/adapters/hip/enqueue.cpp @@ -423,11 +423,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch( : (LocalMemSzPtrPI ? LocalMemSzPtrPI : nullptr); if (LocalMemSzPtr) { - int DeviceMaxLocalMem = 0; - UR_CHECK_ERROR(hipDeviceGetAttribute( - &DeviceMaxLocalMem, hipDeviceAttributeMaxSharedMemoryPerBlock, - Dev->get())); - + int DeviceMaxLocalMem = Dev->getDeviceMaxLocalMem(); static const int EnvVal = std::atoi(LocalMemSzPtr); if (EnvVal <= 0 || EnvVal > DeviceMaxLocalMem) { setErrorMessage(LocalMemSzPtrUR ? "Invalid value specified for "