wehub-resource-sync
136 lines
4.3 KiB
C++
136 lines
4.3 KiB
C++
// Copyright (c) Microsoft Corporation.
|
|
// SPDX-License-Identifier: Apache-2.0
|
|
|
|
// DeepSpeed Team
|
|
|
|
/*
|
|
Functionality for swapping tensors to/from (NVMe) storage devices.
|
|
*/
|
|
|
|
#include "deepspeed_py_copy.h"
|
|
#include <omp.h>
|
|
|
|
#define ROUND_DOWN(size, step) ((size) & ~((step) - 1))
|
|
|
|
#if defined(__AVX512__) or defined(__AVX256__)
|
|
union AVX_Data {
|
|
#if defined(__AVX512__)
|
|
__m512 data;
|
|
#else
|
|
__m256 data;
|
|
#endif
|
|
};
|
|
#endif
|
|
|
|
static void helper_memcpy_1(float* dest, float* src, size_t param_size)
|
|
{
|
|
size_t rounded_size = 0;
|
|
|
|
#if defined(__AVX512__) or defined(__AVX256__)
|
|
|
|
rounded_size = ROUND_DOWN(param_size, SIMD_WIDTH);
|
|
|
|
for (size_t t = 0; t < rounded_size; t += TILE) {
|
|
size_t copy_size = TILE;
|
|
if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
|
|
size_t offset = copy_size + t;
|
|
#pragma omp parallel for
|
|
for (size_t i = t; i < offset; i += SIMD_WIDTH) {
|
|
AVX_Data src_4;
|
|
src_4.data = SIMD_LOAD(src + i);
|
|
|
|
SIMD_STORE(dest + i, src_4.data);
|
|
}
|
|
}
|
|
|
|
#endif
|
|
|
|
if (param_size > rounded_size) {
|
|
#pragma omp parallel for
|
|
for (size_t k = rounded_size; k < param_size; k++) { dest[k] = src[k]; }
|
|
}
|
|
}
|
|
|
|
static void helper_memcpy_4(float* dest, float* src, size_t param_size)
|
|
{
|
|
size_t rounded_size = 0;
|
|
|
|
#if defined(__AVX512__) or defined(__AVX256__)
|
|
|
|
rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2));
|
|
|
|
for (size_t t = 0; t < rounded_size; t += TILE) {
|
|
size_t copy_size = TILE;
|
|
if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
|
|
size_t offset = copy_size + t;
|
|
#pragma omp parallel for
|
|
for (size_t i = t; i < offset; i += (SIMD_WIDTH << 2)) {
|
|
AVX_Data src_4[4];
|
|
src_4[0].data = SIMD_LOAD(src + i);
|
|
src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH);
|
|
src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1));
|
|
src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3);
|
|
|
|
SIMD_STORE(dest + i, src_4[0].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data);
|
|
SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data);
|
|
}
|
|
}
|
|
#endif
|
|
if (param_size > rounded_size)
|
|
helper_memcpy_1((dest + rounded_size), (src + rounded_size), (param_size - rounded_size));
|
|
}
|
|
|
|
static void helper_mempcy_8(float* dest, float* src, size_t param_size)
|
|
{
|
|
size_t rounded_size = 0;
|
|
|
|
#if defined(__AVX512__) or defined(__AVX256__)
|
|
|
|
rounded_size = ROUND_DOWN(param_size, (SIMD_WIDTH << 2));
|
|
|
|
for (size_t t = 0; t < rounded_size; t += TILE) {
|
|
size_t copy_size = TILE;
|
|
if ((t + TILE) > rounded_size) copy_size = rounded_size - t;
|
|
size_t offset = copy_size + t;
|
|
#pragma omp parallel for
|
|
for (size_t i = t; i < offset; i += (SIMD_WIDTH << 3)) {
|
|
AVX_Data src_4[8];
|
|
src_4[0].data = SIMD_LOAD(src + i);
|
|
src_4[1].data = SIMD_LOAD(src + i + SIMD_WIDTH);
|
|
src_4[2].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 1));
|
|
src_4[3].data = SIMD_LOAD(src + i + SIMD_WIDTH * 3);
|
|
src_4[4].data = SIMD_LOAD(src + i + (SIMD_WIDTH << 2));
|
|
src_4[5].data = SIMD_LOAD(src + i + SIMD_WIDTH * 5);
|
|
src_4[6].data = SIMD_LOAD(src + i + SIMD_WIDTH * 6);
|
|
src_4[7].data = SIMD_LOAD(src + i + SIMD_WIDTH * 7);
|
|
|
|
SIMD_STORE(dest + i, src_4[0].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH, src_4[1].data);
|
|
SIMD_STORE(dest + i + (SIMD_WIDTH << 1), src_4[2].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH * 3, src_4[3].data);
|
|
SIMD_STORE(dest + i + (SIMD_WIDTH << 2), src_4[4].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH * 5, src_4[5].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH * 6, src_4[6].data);
|
|
SIMD_STORE(dest + i + SIMD_WIDTH * 7, src_4[7].data);
|
|
}
|
|
}
|
|
#endif
|
|
if (param_size > rounded_size)
|
|
helper_memcpy_4((dest + rounded_size), (src + rounded_size), (param_size - rounded_size));
|
|
}
|
|
|
|
int deepspeed_py_memcpy(torch::Tensor& dest, const torch::Tensor& src)
|
|
{
|
|
auto dest_c = dest.contiguous();
|
|
auto src_c = src.contiguous();
|
|
|
|
float* dest_ptr = (float*)dest_c.data_ptr();
|
|
float* src_ptr = (float*)src_c.data_ptr();
|
|
|
|
helper_mempcy_8(dest_ptr, src_ptr, dest_c.size(0));
|
|
|
|
return 0;
|
|
}
|