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 all 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
48 changes: 28 additions & 20 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,29 @@ 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.


std::unique_ptr<RTDeviceBinaryImage> DevImg;
if (IsDeviceImageCompressed) {
// Decompress the image.
CheckAndDecompressImage(Img.get());
DevImg = 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);
DevImg =
std::make_unique<DynRTDeviceBinaryImage>(std::move(Data), ImgSize);
}

// Register export symbols for bfloat16 device library image.
auto ESPropSet = getExportedSymbolPS(RawImg);
sycl_device_binary_property ESProp;
for (ESProp = ESPropSet->PropertiesBegin;
for (auto ESProp = ESPropSet->PropertiesBegin;
ESProp != ESPropSet->PropertiesEnd; ++ESProp) {
m_ExportedSymbolImages.insert(
{ESProp->Name, DynBfloat16DeviceLibImg.get()});
m_ExportedSymbolImages.insert({ESProp->Name, DevImg.get()});
}
m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion] =
std::move(DynBfloat16DeviceLibImg);
m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion] = std::move(DevImg);

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
62 changes: 1 addition & 61 deletions 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 @@ -8,70 +8,10 @@

// REQUIRES: linux
// 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

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

#include <sycl/detail/core.hpp>
#include <sycl/ext/oneapi/bfloat16.hpp>

using namespace sycl;

constexpr access::mode sycl_read = access::mode::read;
constexpr access::mode sycl_write = access::mode::write;

using BFP = sycl::ext::oneapi::bfloat16;

#ifdef BUILD_LIB
void foo(queue &deviceQueue) {
BFP bf16_v;
float fp32_v = 16.5f;
{
buffer<float, 1> fp32_buffer{&fp32_v, 1};
buffer<BFP, 1> bf16_buffer{&bf16_v, 1};
deviceQueue
.submit([&](handler &cgh) {
auto fp32_acc = fp32_buffer.template get_access<sycl_read>(cgh);
auto bf16_acc = bf16_buffer.template get_access<sycl_write>(cgh);
cgh.single_task([=]() { bf16_acc[0] = BFP{fp32_acc[0]}; });
})
.wait();
}
std::cout << "In foo: " << bf16_v << std::endl;
}
#endif

#ifdef BUILD_EXE
void foo(queue &deviceQueue);
#endif

int main() {
BFP bf16_array[3];
float fp32_array[3] = {7.0f, 8.5f, 0.5f};

sycl::queue deviceQueue;
{
buffer<float, 1> fp32_buffer{fp32_array, 3};
buffer<BFP, 1> bf16_buffer{bf16_array, 3};
deviceQueue
.submit([&](handler &cgh) {
auto fp32_acc = fp32_buffer.template get_access<sycl_read>(cgh);
auto bf16_acc = bf16_buffer.template get_access<sycl_write>(cgh);
cgh.single_task([=]() {
bf16_acc[0] = BFP{fp32_acc[0]};
bf16_acc[1] = BFP{fp32_acc[1]};
bf16_acc[2] = BFP{fp32_acc[2]};
});
})
.wait();
}
std::cout << bf16_array[0] << " " << bf16_array[1] << " " << bf16_array[2]
<< std::endl;
#ifdef BUILD_EXE
foo(deviceQueue);
#endif
return 0;
}
#include "bfloat16_conversion_test.hpp"
60 changes: 60 additions & 0 deletions sycl/test-e2e/DeviceLib/bfloat16_conversion_test.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,60 @@
#include <sycl/detail/core.hpp>
#include <sycl/ext/oneapi/bfloat16.hpp>

using namespace sycl;

constexpr access::mode sycl_read = access::mode::read;
constexpr access::mode sycl_write = access::mode::write;

using BFP = sycl::ext::oneapi::bfloat16;

#ifdef BUILD_LIB
void foo(queue &deviceQueue) {
BFP bf16_v;
float fp32_v = 16.5f;
{
buffer<float, 1> fp32_buffer{&fp32_v, 1};
buffer<BFP, 1> bf16_buffer{&bf16_v, 1};
deviceQueue
.submit([&](handler &cgh) {
auto fp32_acc = fp32_buffer.template get_access<sycl_read>(cgh);
auto bf16_acc = bf16_buffer.template get_access<sycl_write>(cgh);
cgh.single_task([=]() { bf16_acc[0] = BFP{fp32_acc[0]}; });
})
.wait();
}
std::cout << "In foo: " << bf16_v << std::endl;
}
#endif

#ifdef BUILD_EXE
void foo(queue &deviceQueue);
#endif

int main() {
BFP bf16_array[3];
float fp32_array[3] = {7.0f, 8.5f, 0.5f};

sycl::queue deviceQueue;
{
buffer<float, 1> fp32_buffer{fp32_array, 3};
buffer<BFP, 1> bf16_buffer{bf16_array, 3};
deviceQueue
.submit([&](handler &cgh) {
auto fp32_acc = fp32_buffer.template get_access<sycl_read>(cgh);
auto bf16_acc = bf16_buffer.template get_access<sycl_write>(cgh);
cgh.single_task([=]() {
bf16_acc[0] = BFP{fp32_acc[0]};
bf16_acc[1] = BFP{fp32_acc[1]};
bf16_acc[2] = BFP{fp32_acc[2]};
});
})
.wait();
}
std::cout << bf16_array[0] << " " << bf16_array[1] << " " << bf16_array[2]
<< std::endl;
#ifdef BUILD_EXE
foo(deviceQueue);
#endif
return 0;
}
19 changes: 19 additions & 0 deletions sycl/test-e2e/DeviceLib/bfloat16_conversion_test_compress.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,19 @@
//==-------------- bf1oat16 devicelib test for SYCL JIT --------------------==//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

// Check bfloat16 devicelib device image compression.

// REQUIRES: linux, zstd
// 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.

#include "bfloat16_conversion_test.hpp"