wehub-resource-sync
98 lines
3.1 KiB
C++
98 lines
3.1 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
#include <c10/xpu/XPUStream.h>
|
|
#include <torch/extension.h>
|
|
#include <iostream>
|
|
#include <sycl/sycl.hpp>
|
|
|
|
using namespace sycl;
|
|
|
|
void packbitskernel(const float* input, uint8_t* output, const int input_size, id<1> item_ct1)
|
|
{
|
|
// get the sign bit of each float and pack them into byte
|
|
int i = item_ct1;
|
|
for (int j = 0; j < 8; ++j) {
|
|
int k = i * 8 + j;
|
|
int bit = k < input_size && (!sycl::signbit(input[k]));
|
|
output[i] |= bit << (7 - j);
|
|
}
|
|
}
|
|
|
|
void unpackbitskernel(const uint8_t* input, float* output, id<1> item_ct1)
|
|
{
|
|
// use the bit value to set float, bit 0 -> float -1, bit 1 -> float 1
|
|
int i = item_ct1;
|
|
output[i] = (float((input[i / 8] >> (7 - i % 8)) & 1) - 0.5) * 2;
|
|
}
|
|
|
|
sycl::queue get_current_queue(at::Device device)
|
|
{
|
|
c10::xpu::XPUStream stream = c10::xpu::getCurrentXPUStream(device.index());
|
|
return stream.queue();
|
|
}
|
|
|
|
/*
|
|
pack float tensor into uint8 tensor. Every eight float elements get packed into one uint8
|
|
if float x >= 0, will be packed as a '1' bit, or will be packed as '0'
|
|
Arguments:
|
|
tensor: A bool tensor that get packed.
|
|
input_size: numel of input tensor
|
|
rank: device id in order to get corresponding stream
|
|
*/
|
|
at::Tensor packbits(at::Tensor tensor, int input_size, int rank)
|
|
{
|
|
at::Device device = "xpu:" + std::to_string(rank);
|
|
sycl::queue q = get_current_queue(device);
|
|
|
|
int packed_size = (input_size + 7) / 8;
|
|
auto unit8_options = at::TensorOptions().dtype(at::kByte).device(at::kXPU);
|
|
at::Tensor packed = torch::zeros({packed_size}, unit8_options);
|
|
|
|
float* input = (float*)tensor.data_ptr();
|
|
uint8_t* output = (uint8_t*)packed.data_ptr();
|
|
|
|
auto event = q.submit([&](sycl::handler& cgh) {
|
|
cgh.parallel_for<>(range(packed_size), [=](id<1> item_ct1) {
|
|
packbitskernel(input, output, input_size, item_ct1);
|
|
});
|
|
});
|
|
|
|
return packed;
|
|
}
|
|
|
|
/*
|
|
unpack uint8 tensor into float tensor. Every uint8 element get unpacked into eight float
|
|
a '1' bit will be converted to a float(1), a '0' bit will be converted to a float(-1).
|
|
Arguments:
|
|
tensor: A uint8 tensor that get unpacked.
|
|
input_size: numel of input tensor
|
|
rank: device id in order to get corresponding stream
|
|
*/
|
|
at::Tensor unpackbits(at::Tensor tensor, int input_size, int rank)
|
|
{
|
|
at::Device device = "xpu:" + std::to_string(rank);
|
|
sycl::queue q = get_current_queue(device);
|
|
|
|
auto float_options = at::TensorOptions().dtype(at::kFloat).device(at::kXPU);
|
|
at::Tensor unpacked = torch::empty({input_size * 8}, float_options);
|
|
|
|
uint8_t* input = (uint8_t*)tensor.data_ptr();
|
|
float* output = (float*)unpacked.data_ptr();
|
|
|
|
auto event = q.submit([&](sycl::handler& cgh) {
|
|
cgh.parallel_for<>(range(input_size * 8),
|
|
[=](id<1> item_ct1) { unpackbitskernel(input, output, item_ct1); });
|
|
});
|
|
|
|
return unpacked;
|
|
}
|
|
|
|
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m)
|
|
{
|
|
m.def("packbits", &packbits, "DeepSpeed XPU packbits (C++)");
|
|
m.def("unpackbits", &unpackbits, "DeepSpeed XPU unpackbits (C++)");
|
|
}
|