wehub-resource-sync
93 lines
3.4 KiB
C++
93 lines
3.4 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include <sycl/sycl.hpp>
|
|
|
|
inline void has_capability_or_fail(const sycl::device& dev,
|
|
const std::initializer_list<sycl::aspect>& props)
|
|
{
|
|
for (const auto& it : props) {
|
|
if (dev.has(it)) continue;
|
|
switch (it) {
|
|
case sycl::aspect::fp64:
|
|
throw std::runtime_error("'double' is not supported in '" +
|
|
dev.get_info<sycl::info::device::name>() + "' device");
|
|
break;
|
|
case sycl::aspect::fp16:
|
|
throw std::runtime_error("'half' is not supported in '" +
|
|
dev.get_info<sycl::info::device::name>() + "' device");
|
|
break;
|
|
default:
|
|
#define __SYCL_ASPECT(ASPECT, ID) \
|
|
case sycl::aspect::ASPECT: return #ASPECT;
|
|
#define __SYCL_ASPECT_DEPRECATED(ASPECT, ID, MESSAGE) __SYCL_ASPECT(ASPECT, ID)
|
|
#define __SYCL_ASPECT_DEPRECATED_ALIAS(ASPECT, ID, MESSAGE)
|
|
auto getAspectNameStr = [](sycl::aspect AspectNum) -> std::string {
|
|
switch (AspectNum) {
|
|
#include <sycl/info/aspects.def>
|
|
#include <sycl/info/aspects_deprecated.def>
|
|
default: return "unknown aspect";
|
|
}
|
|
};
|
|
#undef __SYCL_ASPECT_DEPRECATED_ALIAS
|
|
#undef __SYCL_ASPECT_DEPRECATED
|
|
#undef __SYCL_ASPECT
|
|
throw std::runtime_error("'" + getAspectNameStr(it) + "' is not supported in '" +
|
|
dev.get_info<sycl::info::device::name>() + "' device");
|
|
}
|
|
break;
|
|
}
|
|
}
|
|
|
|
void param_update_kernel(const float* input, sycl::half* output, int size)
|
|
{
|
|
auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>();
|
|
int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2);
|
|
|
|
if (id < size) { output[id] = (sycl::half)input[id]; }
|
|
}
|
|
|
|
void launch_param_update(const float* input, sycl::half* output, int size, sycl::queue* stream)
|
|
{
|
|
int threads = 1024;
|
|
|
|
sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1);
|
|
sycl::range<3> block_dim(1, 1, threads);
|
|
|
|
{
|
|
has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16});
|
|
stream->parallel_for(
|
|
sycl::nd_range<3>(grid_dim * block_dim, block_dim),
|
|
[=](sycl::nd_item<3> item_ct1) { param_update_kernel(input, output, size); });
|
|
}
|
|
}
|
|
|
|
void param_update_kernel_half(const float* input, sycl::half* output, int size)
|
|
{
|
|
auto item_ct1 = sycl::ext::oneapi::experimental::this_nd_item<3>();
|
|
int id = item_ct1.get_group(2) * item_ct1.get_local_range(2) + item_ct1.get_local_id(2);
|
|
sycl::half2* output_cast = reinterpret_cast<sycl::half2*>(output);
|
|
if (id < size) {
|
|
float input_f = input[id];
|
|
sycl::half2* input_h = reinterpret_cast<sycl::half2*>(&input_f);
|
|
output_cast[id] = *input_h;
|
|
}
|
|
}
|
|
|
|
void launch_param_update_half(const float* input, sycl::half* output, int size, sycl::queue* stream)
|
|
{
|
|
int threads = 1024;
|
|
size /= 2;
|
|
sycl::range<3> grid_dim(1, 1, (size - 1) / threads + 1);
|
|
sycl::range<3> block_dim(1, 1, threads);
|
|
|
|
{
|
|
has_capability_or_fail(stream->get_device(), {sycl::aspect::fp16});
|
|
stream->parallel_for(
|
|
sycl::nd_range<3>(grid_dim * block_dim, block_dim),
|
|
[=](sycl::nd_item<3> item_ct1) { param_update_kernel_half(input, output, size); });
|
|
}
|
|
}
|