add ops
This commit is contained in:
Chranos
214
torch_mlu_ops-v1.3.2/csrc/kernels/generate_mask.mlu
Normal file
214
torch_mlu_ops-v1.3.2/csrc/kernels/generate_mask.mlu
Normal file
@@ -0,0 +1,214 @@
|
||||
#include <cstddef>
|
||||
#include <iostream>
|
||||
#include "cn_api.h"
|
||||
#include "cnnl.h"
|
||||
#include "cnrt.h"
|
||||
#include "generate_mask.mluh"
|
||||
// clang-format off
|
||||
#include <mlu.h>
|
||||
// clang-format on
|
||||
|
||||
namespace tmo {
|
||||
namespace kernels {
|
||||
template <typename T>
|
||||
__mlu_func__ void write_value(void *dst, unsigned int elem_count, T value) {
|
||||
__bang_write_value(dst, elem_count, value);
|
||||
}
|
||||
|
||||
template <>
|
||||
__mlu_func__ void write_value(void *dst, unsigned int elem_count, bfloat16_t value) {
|
||||
#if __BANG_ARCH__ >= 500
|
||||
__bang_write_value(dst, elem_count, value);
|
||||
#endif
|
||||
}
|
||||
|
||||
// [once_len, once_len]
|
||||
__nram__ int8_t nram_small[(__MLU_NRAM_SIZE__ * 1 / 4 * 1024)];
|
||||
// [1 + once_len, 2 * once_len]
|
||||
__nram__ int8_t nram_large[(__MLU_NRAM_SIZE__ * 2 / 4 * 1024 + 1024)];
|
||||
// [once_len * 2 + 1]
|
||||
__nram__ int8_t nram_tiny[2048];
|
||||
template <typename T>
|
||||
class GenerateMask {
|
||||
constexpr static int once_len = sizeof(T) == 4 ? 160 : 256;
|
||||
// [once_len, once_len]
|
||||
T *nram_upper = (T *)(nram_small);
|
||||
// [1 + once_len, 2 * once_len]
|
||||
T *nram_buf = (T *)(nram_large);
|
||||
// [once_len, once_len], reuse upper part of nram_buf
|
||||
T *nram_filled = nram_buf;
|
||||
// [once_len, once_len], reuse lower part of nram_buf
|
||||
T *nram_zeros = nram_buf + once_len * once_len;
|
||||
// [once_len]
|
||||
T *nram_ones_zeros = (T *)nram_tiny;
|
||||
|
||||
__mlu_func__ void initBuffers(T fill_value = -10000) {
|
||||
/* nram_buf:
|
||||
|---once_len---||---once_len---|
|
||||
0, 1, 1, 1, ..., 1, 0, 0, 0, ...
|
||||
0, 0, 1, 1, ..., 1, 1, 0, 0, ...
|
||||
0, 0, 0, 1, ..., 1, 1, 1, 0, ...
|
||||
... */
|
||||
nram_buf[0] = 0;
|
||||
constexpr static int copy_size = (once_len * 2 + 1) * sizeof(T);
|
||||
__memcpy(nram_buf + 1, nram_ones_zeros, copy_size, NRAM2NRAM, copy_size, 0, once_len - 1);
|
||||
__memcpy(nram_upper, nram_buf, once_len * sizeof(T), NRAM2NRAM, once_len * sizeof(T),
|
||||
once_len * 2 * sizeof(T), once_len - 1);
|
||||
// nram_buf is nolonger needed
|
||||
write_value(nram_filled, once_len * once_len, (T)fill_value);
|
||||
write_value(nram_zeros, once_len * once_len, (T)0);
|
||||
}
|
||||
|
||||
__mlu_func__ void dealOneBatch(T *output, // [max_seq_len, max_seq_len]
|
||||
int max_seq_len,
|
||||
int seq_len) {
|
||||
/*
|
||||
| once_len |
|
||||
+----------+-----------------------------------+
|
||||
| | | |
|
||||
| upper | fill_value | |
|
||||
| | | |
|
||||
+----------+----------+ | |
|
||||
| | | | |
|
||||
| | upper | | fill |
|
||||
| | | | value |
|
||||
| +----------+----------+ | |
|
||||
| | | | |
|
||||
| | upper | | |
|
||||
| 0 | | | |
|
||||
| +----------+---+ |
|
||||
| | u | |
|
||||
|--------------------------------+---+ |
|
||||
| |
|
||||
| fill_value |
|
||||
| |
|
||||
+----------------------------------------------+
|
||||
*/
|
||||
int tile_count = seq_len / once_len;
|
||||
int tile_remain = seq_len % once_len;
|
||||
int boarder_len = max_seq_len - seq_len;
|
||||
int row = 0;
|
||||
for (; row < tile_count * once_len; row += once_len) {
|
||||
// fill left with zeros
|
||||
// assume that max_seq_len <= once_len^2
|
||||
if (row > 0) {
|
||||
__memcpy_async(output + (size_t)row * max_seq_len, nram_zeros, row * sizeof(T), NRAM2GDRAM,
|
||||
max_seq_len * sizeof(T), 0, once_len - 1);
|
||||
}
|
||||
// fill middle with upper
|
||||
__memcpy_async(output + (size_t)row * max_seq_len + row, nram_upper, once_len * sizeof(T),
|
||||
NRAM2GDRAM, max_seq_len * sizeof(T), once_len * sizeof(T), once_len - 1);
|
||||
// fill right with fill_value
|
||||
if (row + once_len < max_seq_len) {
|
||||
__memcpy_async(output + (size_t)row * max_seq_len + row + once_len, nram_filled,
|
||||
(max_seq_len - row - once_len) * sizeof(T), NRAM2GDRAM,
|
||||
max_seq_len * sizeof(T), 0, once_len - 1);
|
||||
}
|
||||
}
|
||||
|
||||
if (tile_remain) {
|
||||
// fill left with zeros
|
||||
if (row > 0) {
|
||||
__memcpy_async(output + (size_t)row * max_seq_len, nram_zeros, row * sizeof(T), NRAM2GDRAM,
|
||||
max_seq_len * sizeof(T), 0, tile_remain - 1);
|
||||
}
|
||||
// fill middle with upper
|
||||
__memcpy_async(output + (size_t)row * max_seq_len + row, nram_upper, tile_remain * sizeof(T),
|
||||
NRAM2GDRAM, max_seq_len * sizeof(T), once_len * sizeof(T), tile_remain - 1);
|
||||
// fill right with fill_value
|
||||
if (row + tile_remain < max_seq_len) {
|
||||
__memcpy_async(output + (size_t)row * max_seq_len + row + tile_remain, nram_filled,
|
||||
(max_seq_len - row - tile_remain) * sizeof(T), NRAM2GDRAM,
|
||||
max_seq_len * sizeof(T), 0, tile_remain - 1);
|
||||
}
|
||||
}
|
||||
|
||||
if (boarder_len) {
|
||||
// fill right boarder with fill_value
|
||||
__memcpy_async(output + seq_len, nram_filled, boarder_len * sizeof(T), NRAM2GDRAM,
|
||||
max_seq_len * sizeof(T), 0, (max_seq_len - boarder_len) - 1);
|
||||
// fill bottom boarder with fill_value
|
||||
__memcpy_async(output + (size_t)seq_len * max_seq_len, nram_filled, max_seq_len * sizeof(T),
|
||||
NRAM2GDRAM, max_seq_len * sizeof(T), 0, boarder_len - 1);
|
||||
}
|
||||
__sync_io();
|
||||
}
|
||||
|
||||
public:
|
||||
__mlu_func__ void execute(T *output_ddr, // [total_batch, max_seq_len, max_seq_len]
|
||||
int *batch_seq_len,
|
||||
int total_batch,
|
||||
int max_seq_len,
|
||||
T fill_value = -10000) {
|
||||
int batch_each = total_batch / taskDimY;
|
||||
int batch_remain = total_batch % taskDimY;
|
||||
int batch_start = taskIdY * batch_each + (taskIdY < batch_remain ? taskIdY : batch_remain);
|
||||
int batch_count = batch_each + (taskIdY < batch_remain ? 1 : 0);
|
||||
write_value(nram_ones_zeros, once_len, (T)fill_value);
|
||||
write_value(nram_ones_zeros + once_len, once_len + 1, (T)0);
|
||||
initBuffers();
|
||||
|
||||
for (int n = batch_start; n < batch_start + batch_count; n++) {
|
||||
T *output = output_ddr + (size_t)n * max_seq_len * max_seq_len;
|
||||
int seq_len = batch_seq_len[n];
|
||||
dealOneBatch(output, max_seq_len, seq_len);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__mlu_global__ void MLUUnion1GenerateMask(T *output_ddr, // [total_batch, max_seq_len, max_seq_len]
|
||||
int *batch_seq_len,
|
||||
int total_batch,
|
||||
int max_seq_len,
|
||||
T fill_value = -10000) {
|
||||
if (coreId != 0) {
|
||||
return; // we only use 1 core in a cluster
|
||||
}
|
||||
GenerateMask<T>().execute(output_ddr, batch_seq_len, total_batch, max_seq_len, fill_value);
|
||||
}
|
||||
} // namespace kernels
|
||||
|
||||
KernelStatus invokeGenerateMask(cnnlHandle_t handle,
|
||||
void *output_ddr,
|
||||
int *batch_seq_len,
|
||||
int total_batch,
|
||||
int max_seq_len,
|
||||
cnnlDataType_t data_type,
|
||||
float fill_value) {
|
||||
cnrtQueue_t queue;
|
||||
cnnlGetQueue(handle, &queue);
|
||||
CNdev dev;
|
||||
cnnlGetDevice(handle, &dev);
|
||||
int cluster_num;
|
||||
CNRT_CHECK(cnrtDeviceGetAttribute(&cluster_num, cnrtAttrClusterCount, dev));
|
||||
cnrtDim3_t dim;
|
||||
dim.x = 4;
|
||||
dim.y = cluster_num;
|
||||
dim.z = 1;
|
||||
if (data_type == CNNL_DTYPE_FLOAT) {
|
||||
kernels::MLUUnion1GenerateMask<float><<<dim, cnrtFuncTypeUnion1, queue>>>(
|
||||
static_cast<float *>(output_ddr), batch_seq_len, total_batch, max_seq_len,
|
||||
static_cast<float>(fill_value));
|
||||
} else if (data_type == CNNL_DTYPE_HALF) {
|
||||
kernels::MLUUnion1GenerateMask<half><<<dim, cnrtFuncTypeUnion1, queue>>>(
|
||||
static_cast<half *>(output_ddr), batch_seq_len, total_batch, max_seq_len,
|
||||
static_cast<half>(fill_value));
|
||||
} else if (data_type == CNNL_DTYPE_BFLOAT16) {
|
||||
if (!isBf16Supported()) {
|
||||
std::cerr << "[invokeGenerateMask]: MLU300 devices do not support bfloat16." << std::endl;
|
||||
return KernelStatus::KERNEL_STATUS_FAILED;
|
||||
}
|
||||
kernels::MLUUnion1GenerateMask<bfloat16_t><<<dim, cnrtFuncTypeUnion1, queue>>>(
|
||||
static_cast<bfloat16_t *>(output_ddr), batch_seq_len, total_batch, max_seq_len,
|
||||
static_cast<bfloat16_t>(fill_value));
|
||||
} else {
|
||||
std::cerr << "[invokeGenerateMask]: invokeGenerateMask: data_type is not supported"
|
||||
<< std::endl;
|
||||
return KernelStatus::KERNEL_STATUS_FAILED;
|
||||
}
|
||||
|
||||
return KernelStatus::KERNEL_STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
} // namespace tmo
|
||||
Reference in New Issue
Block a user