Skip to content

[SYCL] Add support for compressed BF16 device lib images in runtime #18108

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Apr 22, 2025
Merged

[SYCL] Add support for compressed BF16 device lib images in runtime #18108

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
da6c4d6
Add reproducer
uditagarwal97 Apr 21, 2025
86ca5cc
Check for compressed BF16 conversion functions
uditagarwal97 Apr 21, 2025
a7b72ee
Address code reviews
uditagarwal97 Apr 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 34 additions & 22 deletions sycl/source/detail/program_manager/program_manager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1912,10 +1912,13 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
if (EntriesB == EntriesE && shouldSkipEmptyImage(RawImg))
return;

std::unique_ptr<RTDeviceBinaryImage> Img;
bool IsBfloat16DeviceLib = false;
uint32_t Bfloat16DeviceLibVersion = 0;
if (isDeviceImageCompressed(RawImg))
const bool IsBfloat16DeviceLib =
isBfloat16DeviceLibImage(RawImg, &Bfloat16DeviceLibVersion);
const bool IsDeviceImageCompressed = isDeviceImageCompressed(RawImg);

std::unique_ptr<RTDeviceBinaryImage> Img;
if (IsDeviceImageCompressed) {
#ifndef SYCL_RT_ZSTD_NOT_AVAIABLE
Img = std::make_unique<CompressedRTDeviceBinaryImage>(RawImg);
#else
Expand All @@ -1924,11 +1927,8 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
"SYCL RT was built without ZSTD support."
"Aborting. ");
#endif
else {
IsBfloat16DeviceLib =
isBfloat16DeviceLibImage(RawImg, &Bfloat16DeviceLibVersion);
if (!IsBfloat16DeviceLib)
Img = std::make_unique<RTDeviceBinaryImage>(RawImg);
} else if (!IsBfloat16DeviceLib) {
Img = std::make_unique<RTDeviceBinaryImage>(RawImg);
}

// If an output image is requested, set it to the newly allocated image.
Expand Down Expand Up @@ -1966,21 +1966,33 @@ void ProgramManager::addImage(sycl_device_binary RawImg,
"Invalid Bfloat16 Device Library Index.");
if (m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion].get())
return;
size_t ImgSize =
static_cast<size_t>(RawImg->BinaryEnd - RawImg->BinaryStart);
std::unique_ptr<char[]> Data(new char[ImgSize]);
std::memcpy(Data.get(), RawImg->BinaryStart, ImgSize);
auto DynBfloat16DeviceLibImg =
std::make_unique<DynRTDeviceBinaryImage>(std::move(Data), ImgSize);
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi @jinge90 ,
I'm wondering why do we need to use dynamic RT device image and not a static one (RTDeviceBinaryImage)?

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, @uditagarwal97
The reason is RTDeviceBinaryImage does not own the sycl_device_binary_struct, which itself does not own the raw binary data (e.g. SPIRV) at sycl_device_binary_struct::BinaryStart. The previous discussion is in: #17461 (comment)

Suppose following scenario, if main program dlopen 2 sycl dynamic library lib1 and lib2, both of them depend on bfloat16 devicelib. SYCL RT calls addImages for lib1 and keeps the required bfloat16 devicelib RTDeviceBinaryImages in program manager's m_Bfloat16DeviceLibImages for further dynamic symbol resolve. Then, SYCL RT loads lib2 and resolves dependent bfloat16 symbols with m_Bfloat16DeviceLibImages. The RTDeviceBinaryImage stored in m_Bfloat16DeviceLibImages is from lib1 and SYCL RT may dlclose lib1 when lib2 device code is still alive. To avoid this, we need to make sure bfloat16 devicelib image is alive during the lifetime of the sycl programs, so we use dynamic device binary image.

I know little about 'CompressedDeviceImage', I think we need to do same thing if the compressed device image may be removed by SYCL RT.

Adding @jopperm for some suggestion.

Thanks very much.

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I see, makes sense then. CompressedDeviceImage is same as DynRTDeviceImage, i.e., it also allocates data on heap but, the former can also handle compressed images.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hi, @uditagarwal97
Just double confirm, do you mean the CompressedDeviceImage holds the compressed device image which is located on head and it be alive during the whole lifetime of current sycl program? If so, the patch is OK to be merged.

Thanks very much.

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

After a quick look, my understanding is that CompressedRTDeviceBinaryImage owns a copy of the original image's sycl_device_binary struct, and it owns the decompressed bytes, so that would be fine 👍

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Yes, just like DynRTDeviceImage, CompressedRTDeviceBinaryImage also owns the device image bytes.

auto ESPropSet = getExportedSymbolPS(RawImg);
sycl_device_binary_property ESProp;
for (ESProp = ESPropSet->PropertiesBegin;
ESProp != ESPropSet->PropertiesEnd; ++ESProp) {
m_ExportedSymbolImages.insert(
{ESProp->Name, DynBfloat16DeviceLibImg.get()});

if (IsDeviceImageCompressed) {
// Decompress the image.
auto ESPropSet = getExportedSymbolPS(RawImg);
for (auto ESProp = ESPropSet->PropertiesBegin;
ESProp != ESPropSet->PropertiesEnd; ++ESProp) {
m_ExportedSymbolImages.insert({ESProp->Name, Img.get()});
}
CheckAndDecompressImage(Img.get());
m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion] = std::move(Img);
} else {
size_t ImgSize =
static_cast<size_t>(RawImg->BinaryEnd - RawImg->BinaryStart);
std::unique_ptr<char[]> Data(new char[ImgSize]);
std::memcpy(Data.get(), RawImg->BinaryStart, ImgSize);
auto DynBfloat16DeviceLibImg =
std::make_unique<DynRTDeviceBinaryImage>(std::move(Data), ImgSize);
auto ESPropSet = getExportedSymbolPS(RawImg);
sycl_device_binary_property ESProp;
for (ESProp = ESPropSet->PropertiesBegin;
ESProp != ESPropSet->PropertiesEnd; ++ESProp) {
m_ExportedSymbolImages.insert(
{ESProp->Name, DynBfloat16DeviceLibImg.get()});
}
m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion] =
std::move(DynBfloat16DeviceLibImg);
}
m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion] =
std::move(DynBfloat16DeviceLibImg);
return;
}
}
Expand Down
2 changes: 1 addition & 1 deletion sycl/source/detail/program_manager/program_manager.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -541,7 +541,7 @@ class ProgramManager {
// version and 2nd is for native version. These bfloat16 device library
// images are provided by compiler long time ago, we expect no further
// update, so keeping 1 copy should be OK.
std::array<DynRTDeviceBinaryImageUPtr, 2> m_Bfloat16DeviceLibImages;
std::array<RTDeviceBinaryImageUPtr, 2> m_Bfloat16DeviceLibImages;

friend class ::ProgramManagerTest;
};
Expand Down
7 changes: 6 additions & 1 deletion sycl/test-e2e/DeviceLib/bfloat16_conversion_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,12 +6,17 @@
//
//===----------------------------------------------------------------------===//

// REQUIRES: linux
// REQUIRES: linux, zstd
// RUN: %{build} -DBUILD_LIB -fPIC -shared -o %T/lib%basename_t.so

// RUN: %{build} -DBUILD_EXE -L%T -o %t1.out -l%basename_t -Wl,-rpath=%T
// RUN: %{run} %t1.out

// Check with device image compression.
// RUN: %{build} --offload-compress -DBUILD_LIB -fPIC -shared -o %T/lib%basename_t_compress.so
// RUN: %{build} --offload-compress -DBUILD_EXE -L%T -o %t1.out -l%basename_t_compress -Wl,-rpath=%T
// RUN: %{run} %t1.out

// UNSUPPORTED: target-nvidia || target-amd
// UNSUPPORTED-INTENDED: bfloat16 device library is not used on AMD and Nvidia.

Expand Down