diff --git a/sycl/source/detail/program_manager/program_manager.cpp b/sycl/source/detail/program_manager/program_manager.cpp index 3f1ae09a7966f..6375415c16adf 100644 --- a/sycl/source/detail/program_manager/program_manager.cpp +++ b/sycl/source/detail/program_manager/program_manager.cpp @@ -509,7 +509,7 @@ std::pair ProgramManager::getOrCreateURProgram( const std::vector &AllImages, const context &Context, const std::vector &Devices, const std::string &CompileAndLinkOptions, SerializedObj SpecConsts) { - ur_program_handle_t NativePrg; // TODO: Or native? + ur_program_handle_t NativePrg; // Get binaries for each device (1:1 correpsondence with input Devices). auto Binaries = PersistentDeviceCodeCache::getItemFromDisc( @@ -768,7 +768,8 @@ setSpecializationConstants(const std::shared_ptr &InputImpl, } } -static inline void CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) { +static inline void +CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) { #ifndef SYCL_RT_ZSTD_NOT_AVAIABLE if (auto CompImg = dynamic_cast(Img)) if (CompImg->IsCompressed()) @@ -913,6 +914,11 @@ ur_program_handle_t ProgramManager::getBuiltURProgram( { std::lock_guard Lock(MNativeProgramsMutex); + // NativePrograms map does not intend to keep reference to program handle, + // so keys in the map can be invalid (reference count went to zero and the + // underlying program disposed of). Protecting from incorrect values by + // removal of map entries with same handle (obviously invalid entries). + std::ignore = NativePrograms.erase(BuiltProgram.get()); for (const RTDeviceBinaryImage *Img : ImgWithDeps) { NativePrograms.insert({BuiltProgram.get(), Img}); } @@ -2738,6 +2744,11 @@ ProgramManager::link(const DevImgPlainWithDeps &ImgWithDeps, { std::lock_guard Lock(MNativeProgramsMutex); + // NativePrograms map does not intend to keep reference to program handle, + // so keys in the map can be invalid (reference count went to zero and the + // underlying program disposed of). Protecting from incorrect values by + // removal of map entries with same handle (obviously invalid entries). + std::ignore = NativePrograms.erase(LinkedProg); for (const device_image_plain &Img : ImgWithDeps) { NativePrograms.insert( {LinkedProg, getSyclObjImpl(Img)->get_bin_image_ref()}); diff --git a/sycl/source/detail/program_manager/program_manager.hpp b/sycl/source/detail/program_manager/program_manager.hpp index 717eebbc99cd7..61a3240c1ddd4 100644 --- a/sycl/source/detail/program_manager/program_manager.hpp +++ b/sycl/source/detail/program_manager/program_manager.hpp @@ -46,6 +46,9 @@ extern "C" __SYCL_EXPORT void __sycl_unregister_lib(sycl_device_binaries desc); // +++ } +// For testing purposes +class ProgramManagerTest; + namespace sycl { inline namespace _V1 { class context; @@ -494,6 +497,8 @@ class ProgramManager { using MaterializedEntries = std::map, ur_kernel_handle_t>; std::unordered_map m_MaterializedKernels; + + friend class ::ProgramManagerTest; }; } // namespace detail } // namespace _V1 diff --git a/sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp b/sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp index 02c007c7d27e7..74e5dc5fecf42 100644 --- a/sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp +++ b/sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp @@ -6,14 +6,17 @@ // //===----------------------------------------------------------------------===// +#include #include #include +#include #include #include #include #include #include +#include #include #include @@ -98,10 +101,10 @@ static sycl::unittest::MockDeviceImageArray<1> EAMImgArray{&EAMImg}; static sycl::unittest::MockDeviceImageArray<1> EAM2ImgArray{&EAM2Img}; static sycl::unittest::MockDeviceImageArray<1> EAM3ImgArray{&EAM3Img}; -// ur_program_handle_t address is used as a key for ProgramManager::NativePrograms -// storage. redefinedProgramLinkCommon makes ur_program_handle_t address equal to 0x1. -// Make sure that size of Bin is different for device images used in these tests -// and greater than 1. +// ur_program_handle_t address is used as a key for +// ProgramManager::NativePrograms storage. redefinedProgramLinkCommon makes +// ur_program_handle_t address equal to 0x1. Make sure that size of Bin is +// different for device images used in these tests and greater than 1. inline ur_result_t redefinedProgramCreateEAM(void *pParams) { auto params = *static_cast(pParams); static size_t UrProgramAddr = 2; @@ -109,17 +112,6 @@ inline ur_result_t redefinedProgramCreateEAM(void *pParams) { return UR_RESULT_SUCCESS; } -mock::dummy_handle_t_ FixedHandle; -inline ur_result_t setFixedProgramPtr(void *pParams) { - auto params = *static_cast(pParams); - **params.pphProgram = reinterpret_cast(&FixedHandle); - return UR_RESULT_SUCCESS; -} -inline ur_result_t releaseFixedProgramPtr(void *pParams) { - // Do nothing - return UR_RESULT_SUCCESS; -} - class MockHandler : public sycl::handler { public: @@ -218,6 +210,53 @@ TEST(EliminatedArgMask, KernelBundleWith2Kernels) { EXPECT_EQ(*EliminatedArgMask, ExpElimArgMask); } +std::vector> UsedProgramHandles; +std::vector> ProgramHandlesToReuse; +inline ur_result_t setFixedProgramPtr(void *pParams) { + auto params = *static_cast(pParams); + if (!ProgramHandlesToReuse.empty()) { + auto it = ProgramHandlesToReuse.begin() + 1; + std::move(ProgramHandlesToReuse.begin(), it, + std::back_inserter(UsedProgramHandles)); + ProgramHandlesToReuse.erase(ProgramHandlesToReuse.begin(), it); + } else + UsedProgramHandles.push_back( + std::make_unique(sizeof(unsigned))); + **params.pphProgram = + reinterpret_cast(UsedProgramHandles.back().get()); + return UR_RESULT_SUCCESS; +} +inline ur_result_t releaseFixedProgramPtr(void *pParams) { + auto params = *static_cast(pParams); + { + auto it = std::find_if( + UsedProgramHandles.begin(), UsedProgramHandles.end(), + [¶ms](const std::unique_ptr &item) { + return reinterpret_cast(item.get()) == + *params.phProgram; + }); + if (it == UsedProgramHandles.end()) + return UR_RESULT_SUCCESS; + std::move(it, it + 1, std::back_inserter(ProgramHandlesToReuse)); + UsedProgramHandles.erase(it, it + 1); + } + return UR_RESULT_SUCCESS; +} + +inline ur_result_t customProgramRetain(void *pParams) { + // do nothing + return UR_RESULT_SUCCESS; +} + +class ProgramManagerTest { +public: + static std::unordered_multimap & + getNativePrograms() { + return sycl::detail::ProgramManager::getInstance().NativePrograms; + } +}; + // It's possible for the same handle to be reused for multiple distinct programs // This can happen if a program is released (freeing underlying memory) and then // a new program happens to get given that same memory for its handle. @@ -227,6 +266,7 @@ TEST(EliminatedArgMask, KernelBundleWith2Kernels) { TEST(EliminatedArgMask, ReuseOfHandleValues) { sycl::detail::ProgramManager &PM = sycl::detail::ProgramManager::getInstance(); + auto &NativePrograms = ProgramManagerTest::getNativePrograms(); ur_program_handle_t ProgBefore = nullptr; ur_program_handle_t ProgAfter = nullptr; @@ -238,6 +278,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) { &setFixedProgramPtr); mock::getCallbacks().set_replace_callback("urProgramRelease", &releaseFixedProgramPtr); + mock::getCallbacks().set_replace_callback("urProgramRetain", + &customProgramRetain); const sycl::device Dev = Plt.get_devices()[0]; sycl::queue Queue{Dev}; @@ -247,8 +289,12 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) { auto Mask = PM.getEliminatedKernelArgMask(ProgBefore, Name); EXPECT_NE(Mask, nullptr); EXPECT_EQ(Mask->at(0), 1); + EXPECT_EQ(UsedProgramHandles.size(), 1u); + EXPECT_EQ(NativePrograms.count(ProgBefore), 1u); } + EXPECT_EQ(UsedProgramHandles.size(), 0u); + { auto Name = sycl::detail::KernelInfo::getName(); sycl::unittest::UrMock<> Mock; @@ -257,6 +303,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) { &setFixedProgramPtr); mock::getCallbacks().set_replace_callback("urProgramRelease", &releaseFixedProgramPtr); + mock::getCallbacks().set_replace_callback("urProgramRetain", + &customProgramRetain); const sycl::device Dev = Plt.get_devices()[0]; sycl::queue Queue{Dev}; @@ -266,6 +314,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) { auto Mask = PM.getEliminatedKernelArgMask(ProgAfter, Name); EXPECT_NE(Mask, nullptr); EXPECT_EQ(Mask->at(0), 0); + EXPECT_EQ(UsedProgramHandles.size(), 1u); + EXPECT_EQ(NativePrograms.count(ProgBefore), 1u); } // Verify that the test is behaving correctly and that the pointer is being