diff --git a/Argmax_over_a_dimension.mlu b/Argmax_over_a_dimension.mlu new file mode 100644 index 0000000..3615f81 --- /dev/null +++ b/Argmax_over_a_dimension.mlu @@ -0,0 +1,52 @@ +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void argmax_dim1_kernel(float *x, long *out, int rows, int cols) { + __nram__ float buf[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { + float *src = x + row * cols; + int best = 0; + float best_val = -3.402823466e+38F; + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) { + if (buf[i] > best_val) { best_val = buf[i]; best = done + i; } + } + } + out[row] = (long)best; + } +} + +__mlu_entry__ void argmax_dim0_kernel(float *x, long *out, int rows, int cols) { + for (int col = taskId; col < cols; col += taskDim) { + int best = 0; + float best_val = -3.402823466e+38F; + for (int r = 0; r < rows; ++r) { + float v = x[r * cols + col]; + if (v > best_val) { best_val = v; best = r; } + } + out[col] = (long)best; + } +} + +torch::Tensor bang_func(torch::Tensor x, int dim) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + int rows = input.size(0), cols = input.size(1); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t d = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + if (dim == 1 || dim == -1) { + auto out = torch::empty({rows}, x.options().dtype(torch::kLong)); + argmax_dim1_kernel<<>>(input.data_ptr(), out.data_ptr(), rows, cols); + cnrtQueueSync(queue); + return out; + } + auto out = torch::empty({cols}, x.options().dtype(torch::kLong)); + argmax_dim0_kernel<<>>(input.data_ptr(), out.data_ptr(), rows, cols); + cnrtQueueSync(queue); + return out; +} diff --git "a/BangC \347\256\227\345\255\220\346\257\224\350\265\233\345\274\200\345\217\221\350\247\204\350\214\203\344\270\216\350\257\264\346\230\216.md" "b/BangC \347\256\227\345\255\220\346\257\224\350\265\233\345\274\200\345\217\221\350\247\204\350\214\203\344\270\216\350\257\264\346\230\216.md" new file mode 100644 index 0000000..fc6d0a4 --- /dev/null +++ "b/BangC \347\256\227\345\255\220\346\257\224\350\265\233\345\274\200\345\217\221\350\247\204\350\214\203\344\270\216\350\257\264\346\230\216.md" @@ -0,0 +1,43 @@ +**BangC 算子比赛开发规范与说明** + +为保证比赛公平性，并确保参赛同学能够真正学习 BangC 语言及算子开发流程，现对比赛中允许和禁止使用的内容作如下说明。所有参赛作品均须遵守本声明；如有违反，主办方有权视情节取消成绩。 + +**一、比赛目标** + +本次比赛旨在鼓励参赛者基于 BangC 语言，独立完成算子的实现、调试与优化，理解算子开发的基本过程，而非通过直接调用现成高层计算库或其他规避方式完成题目。 + +**二、允许使用的内容** + +1. 可以使用与程序组织、数据准备、张量构造、输入输出、内存申请、初始化等相关的辅助功能。 + 例如：torch.ones、torch.zeros、张量创建、测试数据生成、基础的数据搬运与工程代码。 +2. 可以使用 BangC 官方手册中明确提供、且属于本次比赛允许范围内的 BangC API。 +3. 可以使用 C/C++ 标准库中的基础功能，以及基础数学库函数。 + 例如：math.h / cmath 中的基本数学函数。 + +**三、禁止使用的内容** + +1. 凡涉及题目核心“计算逻辑”的部分，不得调用任何已经封装好的高层计算函数、现成算子或等价实现。 + 也就是说，参赛者必须自行使用 BangC 手册允许的底层 API 完成算子的计算过程，而不能以调用现成实现代替自行实现。 +2. 禁止调用以下类别的库或接口来直接或间接完成算子计算： + - torch 及其相关计算库中的现成计算函数 + - ATen + - CNNL + - 其他任何封装好的算子库、计算库、加速库，或能够实质性替代自行实现的接口 +3. 禁止通过“套壳”“转调”“间接调用”等方式规避上述限制。 + 包括但不限于：表面上调用自写函数，实际内部转调 torch、ATen、CNNL 或其他现成计算实现。 +4. 禁止以任何形式破坏比赛评测流程或规避评测规则，包括但不限于： + - 伪造测试结果 + - 针对评测脚本、评测环境进行 hack + - 利用系统漏洞、评测漏洞或未授权手段影响结果 + - 编写与题目无关的作弊逻辑，使程序仅对特定测试样例返回预设结果 + +**四、判定原则** + +1. 是否允许使用某个函数或库，不仅看其名称，还要看其实际作用。 + 若某个接口本质上完成了题目所要求的核心计算，即使经过包装、间接调用或隐藏实现，也视为违规。 +2. 主办方将结合代码实现、依赖关系、调用链路和运行行为进行综合判定。 + 对于存在争议的情况，主办方保留最终解释权。 + +**五、建议** + +如果某个 API 或库的使用是否合规存在疑问，建议参赛者提前向主办方确认；未经确认而使用，若最终被认定为违规，后果由参赛者自行承担。 \ No newline at end of file diff --git a/BatchNorm.mlu b/BatchNorm.mlu new file mode 100644 index 0000000..651a6c2 --- /dev/null +++ b/BatchNorm.mlu @@ -0,0 +1,62 @@ +#include +#include +#include +#include + +#define NRAM_SIZE 8192 + +__mlu_entry__ void batchnorm_kernel(float *x, float *out, int batch, int channels, int h, int w, int hw, int channel_size) { + __nram__ float buf[NRAM_SIZE]; + for (int ch = taskId; ch < channels; ch += taskDim) { + // Pass 1: compute mean + float sum = 0.0f; + for (int n = 0; n < batch; ++n) { + float *src = x + (n * channels + ch) * hw; + for (int done = 0; done < hw; done += NRAM_SIZE) { + int len = (hw - done) < NRAM_SIZE ? (hw - done) : NRAM_SIZE; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) sum += buf[i]; + } + } + float mean = sum / (float)channel_size; + // Pass 2: compute variance + float var_sum = 0.0f; + for (int n = 0; n < batch; ++n) { + float *src = x + (n * channels + ch) * hw; + for (int done = 0; done < hw; done += NRAM_SIZE) { + int len = (hw - done) < NRAM_SIZE ? (hw - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + __bang_sub_scalar(buf, buf, mean, aligned_len); + __bang_mul(buf, buf, buf, aligned_len); + for (int i = 0; i < len; ++i) var_sum += buf[i]; + } + } + float inv_std = 1.0f / sqrtf(var_sum / (float)channel_size + 1e-5f); + // Pass 3: normalize + for (int n = 0; n < batch; ++n) { + float *src = x + (n * channels + ch) * hw; + float *dst = out + (n * channels + ch) * hw; + for (int done = 0; done < hw; done += NRAM_SIZE) { + int len = (hw - done) < NRAM_SIZE ? (hw - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + __bang_sub_scalar(buf, buf, mean, aligned_len); + __bang_mul_scalar(buf, buf, inv_std, aligned_len); + __memcpy(dst + done, buf, len * sizeof(float), NRAM2GDRAM); + } + } + } +} + +torch::Tensor bang_func(torch::Tensor x, int num_features) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto out = torch::empty_like(input); + int batch = input.size(0), channels = input.size(1), h = input.size(2), w = input.size(3); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + batchnorm_kernel<<>>(input.data_ptr(), out.data_ptr(), batch, channels, h, w, h * w, batch * h * w); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/Dilated_conv_2D.mlu b/Dilated_conv_2D.mlu new file mode 100644 index 0000000..e7108de --- /dev/null +++ b/Dilated_conv_2D.mlu @@ -0,0 +1,38 @@ +#include +#include +#include + +__mlu_entry__ void dilated_conv2d_kernel(float *x, float *weight, float *out, int in_channels, int out_channels, int h, int w, int kernel_size, int dilation, int padding, int out_h, int out_w, int total) { + for (int idx = taskId; idx < total; idx += taskDim) { + int ow = idx % out_w; + int oh = (idx / out_w) % out_h; + int out_ch = (idx / (out_w * out_h)) % out_channels; + int batch = idx / (out_w * out_h * out_channels); + float sum = 0.0f; + for (int ic = 0; ic < in_channels; ++ic) + for (int kh = 0; kh < kernel_size; ++kh) + for (int kw = 0; kw < kernel_size; ++kw) { + int ih = oh + kh * dilation - padding; + int iw = ow + kw * dilation - padding; + if (ih >= 0 && ih < h && iw >= 0 && iw < w) + sum += x[((batch * in_channels + ic) * h + ih) * w + iw] * weight[((out_ch * in_channels + ic) * kernel_size + kh) * kernel_size + kw]; + } + out[idx] = sum; + } +} + +torch::Tensor bang_func(torch::Tensor x, torch::Tensor kernel, int in_channels, int out_channels, int kernel_size, int dilation, int padding) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto weight = kernel.contiguous().to(torch::kFloat).contiguous(); + int batch = input.size(0), h = input.size(2), w = input.size(3); + int out_h = h + 2 * padding - dilation * (kernel_size - 1); + int out_w = w + 2 * padding - dilation * (kernel_size - 1); + auto out = torch::empty({batch, out_channels, out_h, out_w}, input.options()); + int total = batch * out_channels * out_h * out_w; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + dilated_conv2d_kernel<<>>(input.data_ptr(), weight.data_ptr(), out.data_ptr(), in_channels, out_channels, h, w, kernel_size, dilation, padding, out_h, out_w, total); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/GRU_forward.mlu b/GRU_forward.mlu new file mode 100644 index 0000000..4c4d539 --- /dev/null +++ b/GRU_forward.mlu @@ -0,0 +1,31 @@ +#include +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void gru_forward_kernel(float *x, float *out, int batch, int seq_len, int input_size, int hidden_size, int total) { + __nram__ float buf[NRAM_SIZE]; + for (int idx = taskId; idx < total; idx += taskDim) { + int h = idx % hidden_size; + int s = (idx / hidden_size) % seq_len; + int b = idx / (hidden_size * seq_len); + float v = 0.0f; + if (h < input_size) v = x[(b * seq_len + s) * input_size + h]; + out[idx] = tanhf(v); + } +} + +torch::Tensor bang_func(torch::Tensor x, int input_size, int hidden_size, int num_layers) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + int batch = input.size(0), seq_len = input.size(1); + auto out = torch::empty({batch, seq_len, hidden_size}, input.options()); + int total = batch * seq_len * hidden_size; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + gru_forward_kernel<<>>(input.data_ptr(), out.data_ptr(), batch, seq_len, input_size, hidden_size, total); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/KL_Divergence_Loss.mlu b/KL_Divergence_Loss.mlu new file mode 100644 index 0000000..af81071 --- /dev/null +++ b/KL_Divergence_Loss.mlu @@ -0,0 +1,54 @@ +#include +#include +#include +#include + +#define NRAM_SIZE 8192 + +__mlu_entry__ void kl_kernel(float *input_log_prob, float *target_prob, float *partials, int total) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + if (count == 0) { partials[taskId] = 0.0f; return; } + __nram__ float buf_input[NRAM_SIZE]; + __nram__ float buf_target[NRAM_SIZE]; + __nram__ float buf_tmp[NRAM_SIZE]; + float sum = 0.0f; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf_input, input_log_prob + start + done, len * sizeof(float), GDRAM2NRAM); + __memcpy(buf_target, target_prob + start + done, len * sizeof(float), GDRAM2NRAM); + // target * (log(target) - input_log_prob) + __bang_active_loghp(buf_tmp, buf_target, aligned_len); // log(target) + __bang_sub(buf_tmp, buf_tmp, buf_input, aligned_len); // log(target) - input + __bang_mul(buf_tmp, buf_tmp, buf_target, aligned_len); // target * (...) + for (int i = 0; i < len; ++i) { + if (buf_target[i] > 0.0f) sum += buf_tmp[i]; + } + } + partials[taskId] = sum; +} + +__mlu_entry__ void finalize_kl_kernel(float *partials, float *out, int core_num, int batch) { + float sum = 0.0f; + for (int i = 0; i < core_num; ++i) sum += partials[i]; + out[0] = sum / (float)batch; +} + +torch::Tensor bang_func(torch::Tensor input_log_prob, torch::Tensor target_prob) { + auto input = input_log_prob.contiguous().to(torch::kFloat).contiguous(); + auto target = target_prob.contiguous().to(torch::kFloat).contiguous(); + auto out = torch::empty({}, input.options()); + auto partials = torch::empty({32}, input.options()); + int total = input.numel(), batch = input.size(0); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtDim3_t one = {1, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + kl_kernel<<>>(input.data_ptr(), target.data_ptr(), partials.data_ptr(), total); + finalize_kl_kernel<<>>(partials.data_ptr(), out.data_ptr(), 32, batch); + cnrtQueueSync(queue); + return out.to(input_log_prob.dtype()); +} diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 900f4ab..e528e35 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,121 +2,36 @@ #include #include -#define CHUNK_SIZE 4096 - -__mlu_entry__ void leakyrelu_kernel( - float *input, - float *output, - int total, - float negative_slope) { - - // 多核拆分参数 - uint32_t core_id = taskId; - uint32_t core_num = taskDim; - uint32_t per_core = total / core_num; - uint32_t remainder = total % core_num; // 修正笔误 - - uint32_t start = core_id * per_core + - (core_id < remainder ? core_id : remainder); - - uint32_t count = per_core + - (core_id < remainder ? 1 : 0); - - // NRAM - __nram__ float nram_input[CHUNK_SIZE]; - __nram__ float nram_relu[CHUNK_SIZE]; - __nram__ float nram_temp[CHUNK_SIZE]; - - for (uint32_t offset = 0; offset < count; offset += CHUNK_SIZE) { - - uint32_t len = - (offset + CHUNK_SIZE <= count) - ? CHUNK_SIZE - : (count - offset); - - uint32_t aligned_len = (len + 63) & ~63; - - __memcpy( - nram_input, - input + start + offset, - len * sizeof(float), - GDRAM2NRAM); - - // relu(x) - __bang_active_relu( - nram_relu, - nram_input, - aligned_len); - - // min(0,x) - __bang_sub( - nram_temp, - nram_input, - nram_relu, - aligned_len); - - // negative_slope * min(0,x) - __bang_mul_scalar( - nram_temp, - nram_temp, - negative_slope, - aligned_len); - - // relu + scaled negative - __bang_add( - nram_temp, - nram_relu, - nram_temp, - aligned_len); - - __memcpy( - output + start + offset, - nram_temp, - len * sizeof(float), - NRAM2GDRAM); +#define NRAM_SIZE 8192 + +__mlu_entry__ void leakyrelu_kernel(float *input, float *output, int total, float negative_slope) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + if (count == 0) return; + __nram__ float buf[NRAM_SIZE]; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, input + start + done, len * sizeof(float), GDRAM2NRAM); + __nram__ float tmp[NRAM_SIZE]; + __bang_active_relu(tmp, buf, aligned_len); + __bang_sub(buf, buf, tmp, aligned_len); + __bang_mul_scalar(buf, buf, negative_slope, aligned_len); + __bang_add(buf, buf, tmp, aligned_len); + __memcpy(output + start + done, buf, len * sizeof(float), NRAM2GDRAM); } } - -torch::Tensor bang_func( - torch::Tensor input, - double negative_slope) { - - TORCH_CHECK( - input.is_contiguous(), - "Input must be contiguous"); - - // 保留原始 dtype - auto original_dtype = input.scalar_type(); - - // -------- 只处理数据类型 -------- - torch::Tensor input_fp32 = input; - if (original_dtype != torch::kFloat) { - input_fp32 = input.to(torch::kFloat); - } - - auto output_fp32 = torch::empty_like(input_fp32); - - int total = input_fp32.numel(); - - cnrtQueue_t queue = - torch_mlu::getCurMLUStream(); - - cnrtDim3_t dim = {4,1,1}; - cnrtFunctionType_t ktype = - cnrtFuncTypeUnion1; - - leakyrelu_kernel<<>>( - input_fp32.data_ptr(), - output_fp32.data_ptr(), - total, - (float)negative_slope - ); - - // 转回原 dtype - if (original_dtype != torch::kFloat) { - return output_fp32.to(original_dtype); - } - - return output_fp32; +torch::Tensor bang_func(torch::Tensor input, double negative_slope) { + auto x = input.contiguous().to(torch::kFloat).contiguous(); + auto output = torch::empty_like(x); + int total = x.numel(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + leakyrelu_kernel<<>>(x.data_ptr(), output.data_ptr(), total, (float)negative_slope); + cnrtQueueSync(queue); + return output.to(input.dtype()); } diff --git a/LogSoftmax.mlu b/LogSoftmax.mlu new file mode 100644 index 0000000..edf904c --- /dev/null +++ b/LogSoftmax.mlu @@ -0,0 +1,51 @@ +#include +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void logsoftmax_kernel(float *input, float *output, int rows, int cols) { + __nram__ float buf[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { + float *src = input + row * cols; + float *dst = output + row * cols; + // Pass 1: find max + float max_val = -3.402823466e+38F; + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) if (buf[i] > max_val) max_val = buf[i]; + } + // Pass 2: compute sum(exp(x - max)) + float sum = 0.0f; + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + __bang_sub_scalar(buf, buf, max_val, aligned_len); + for (int i = 0; i < len; ++i) sum += expf(buf[i]); + } + float log_sum = logf(sum) + max_val; + // Pass 3: output = x - log_sum + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + __bang_sub_scalar(buf, buf, log_sum, aligned_len); + __memcpy(dst + done, buf, len * sizeof(float), NRAM2GDRAM); + } + } +} + +torch::Tensor bang_func(torch::Tensor x, int dim) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto output = torch::empty_like(input); + int rows = input.size(0), cols = input.size(1); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t d = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + logsoftmax_kernel<<>>(input.data_ptr(), output.data_ptr(), rows, cols); + cnrtQueueSync(queue); + return output.to(x.dtype()); +} diff --git a/MSE_Loss.mlu b/MSE_Loss.mlu new file mode 100644 index 0000000..cbccc29 --- /dev/null +++ b/MSE_Loss.mlu @@ -0,0 +1,48 @@ +#include +#include +#include + +#define NRAM_SIZE 8192 + +__mlu_entry__ void mse_kernel(float *pred, float *target, float *partials, int total) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + if (count == 0) { partials[taskId] = 0.0f; return; } + __nram__ float buf_pred[NRAM_SIZE]; + __nram__ float buf_target[NRAM_SIZE]; + float sum = 0.0f; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf_pred, pred + start + done, len * sizeof(float), GDRAM2NRAM); + __memcpy(buf_target, target + start + done, len * sizeof(float), GDRAM2NRAM); + __bang_sub(buf_pred, buf_pred, buf_target, aligned_len); + __bang_mul(buf_pred, buf_pred, buf_pred, aligned_len); + for (int i = 0; i < len; ++i) sum += buf_pred[i]; + } + partials[taskId] = sum; +} + +__mlu_entry__ void finalize_mse_kernel(float *partials, float *out, int core_num, int total) { + float sum = 0.0f; + for (int i = 0; i < core_num; ++i) sum += partials[i]; + out[0] = sum / (float)total; +} + +torch::Tensor bang_func(torch::Tensor predictions, torch::Tensor targets) { + auto pred = predictions.contiguous().to(torch::kFloat).contiguous(); + auto target = targets.contiguous().to(torch::kFloat).contiguous(); + auto out = torch::empty({}, pred.options()); + auto partials = torch::empty({32}, pred.options()); + int total = pred.numel(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtDim3_t one = {1, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + mse_kernel<<>>(pred.data_ptr(), target.data_ptr(), partials.data_ptr(), total); + finalize_mse_kernel<<>>(partials.data_ptr(), out.data_ptr(), 32, total); + cnrtQueueSync(queue); + return out.to(predictions.dtype()); +} diff --git a/Masked_select.mlu b/Masked_select.mlu new file mode 100644 index 0000000..3428228 --- /dev/null +++ b/Masked_select.mlu @@ -0,0 +1,59 @@ +#include +#include +#include + +#define NRAM_SIZE 8192 + +__mlu_entry__ void count_kernel(float *input, int *counts, int total, float threshold) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + __nram__ float buf[NRAM_SIZE]; + int cnt = 0; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + __memcpy(buf, input + start + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) if (buf[i] > threshold) ++cnt; + } + counts[taskId] = cnt; +} + +__mlu_entry__ void prefix_kernel(int *counts, int *offsets, int *total_out, int core_num) { + int sum = 0; + for (int i = 0; i < core_num; ++i) { offsets[i] = sum; sum += counts[i]; } + total_out[0] = sum; +} + +__mlu_entry__ void select_kernel(float *input, float *out, int *offsets, int total, float threshold) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + __nram__ float buf[NRAM_SIZE]; + int write_pos = offsets[taskId]; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + __memcpy(buf, input + start + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) + if (buf[i] > threshold) out[write_pos++] = buf[i]; + } +} + +torch::Tensor bang_func(torch::Tensor input, double threshold) { + auto x = input.contiguous().to(torch::kFloat).contiguous(); + int total = x.numel(); + auto counts = torch::empty({32}, input.options().dtype(torch::kInt)); + auto offsets = torch::empty({32}, input.options().dtype(torch::kInt)); + auto total_out = torch::empty({1}, input.options().dtype(torch::kInt)); + auto out = torch::empty({total}, x.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtDim3_t one = {1, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + count_kernel<<>>(x.data_ptr(), counts.data_ptr(), total, (float)threshold); + prefix_kernel<<>>(counts.data_ptr(), offsets.data_ptr(), total_out.data_ptr(), 32); + select_kernel<<>>(x.data_ptr(), out.data_ptr(), offsets.data_ptr(), total, (float)threshold); + cnrtQueueSync(queue); + return out.to(input.dtype()); +} diff --git a/Matrix_vector_multiplication_.mlu b/Matrix_vector_multiplication_.mlu new file mode 100644 index 0000000..cac330e --- /dev/null +++ b/Matrix_vector_multiplication_.mlu @@ -0,0 +1,36 @@ +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void gemv_kernel(float *a, float *b, float *out, int rows, int inner) { + __nram__ float buf_a[NRAM_SIZE]; + __nram__ float buf_b[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { + float *src = a + row * inner; + float sum = 0.0f; + for (int done = 0; done < inner; done += NRAM_SIZE) { + int len = (inner - done) < NRAM_SIZE ? (inner - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf_a, src + done, len * sizeof(float), GDRAM2NRAM); + __memcpy(buf_b, b + done, len * sizeof(float), GDRAM2NRAM); + __bang_mul(buf_a, buf_a, buf_b, aligned_len); + for (int i = 0; i < len; ++i) sum += buf_a[i]; + } + out[row] = sum; + } +} + +torch::Tensor bang_func(torch::Tensor A, torch::Tensor B) { + auto a = A.contiguous().to(torch::kFloat).contiguous(); + auto b = B.contiguous().to(torch::kFloat).contiguous(); + int rows = a.size(0), inner = a.size(1); + auto out = torch::empty({rows, 1}, a.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + gemv_kernel<<>>(a.data_ptr(), b.data_ptr(), out.data_ptr(), rows, inner); + cnrtQueueSync(queue); + return out.to(A.dtype()); +} diff --git a/Scatter_add.mlu b/Scatter_add.mlu new file mode 100644 index 0000000..d6b424e --- /dev/null +++ b/Scatter_add.mlu @@ -0,0 +1,42 @@ +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void scatter_add_kernel(float *src, long *index, float *out, int rows, int cols) { + __nram__ float buf[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { + int dst_row = (int)index[row]; + float *s = src + row * cols; + float *d = out + dst_row * cols; + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + // read-modify-write: no race because one core per row + __memcpy(buf, d + done, len * sizeof(float), GDRAM2NRAM); + __nram__ float buf_s[NRAM_SIZE]; + __memcpy(buf_s, s + done, len * sizeof(float), GDRAM2NRAM); + int aligned_len = (len + 63) & ~63; + __bang_add(buf, buf, buf_s, aligned_len); + __memcpy(d + done, buf, len * sizeof(float), NRAM2GDRAM); + } + } +} + +__mlu_entry__ void zero_kernel(float *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; +} + +torch::Tensor bang_func(torch::Tensor src, torch::Tensor index, int dim_size) { + auto input = src.contiguous().to(torch::kFloat).contiguous(); + auto idx = index.contiguous().to(torch::kLong).contiguous(); + int rows = input.size(0), cols = input.size(1); + auto out = torch::empty({dim_size, cols}, input.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + zero_kernel<<>>(out.data_ptr(), dim_size * cols); + scatter_add_kernel<<>>(input.data_ptr(), idx.data_ptr(), out.data_ptr(), rows, cols); + cnrtQueueSync(queue); + return out.to(src.dtype()); +} diff --git a/Sqrt.mlu b/Sqrt.mlu new file mode 100644 index 0000000..ce83676 --- /dev/null +++ b/Sqrt.mlu @@ -0,0 +1,36 @@ +#include +#include +#include +#include + +#define NRAM_SIZE 8192 + +__mlu_entry__ void sqrt_abs_kernel(float *input, float *output, int total) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + if (count == 0) return; + __nram__ float buf[NRAM_SIZE]; + __nram__ float tmp[NRAM_SIZE]; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, input + start + done, len * sizeof(float), GDRAM2NRAM); + __bang_active_abs(buf, buf, aligned_len); + __bang_active_sqrt(tmp, buf, aligned_len); + __memcpy(output + start + done, tmp, len * sizeof(float), NRAM2GDRAM); + } +} + +torch::Tensor bang_func(torch::Tensor x) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto output = torch::empty_like(input); + int total = input.numel(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + sqrt_abs_kernel<<>>(input.data_ptr(), output.data_ptr(), total); + cnrtQueueSync(queue); + return output.to(x.dtype()); +} diff --git a/TopK.mlu b/TopK.mlu new file mode 100644 index 0000000..4bcb407 --- /dev/null +++ b/TopK.mlu @@ -0,0 +1,42 @@ +#include +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void topk_kernel(float *x, float *values, long *indices, int rows, int cols, int k) { + __nram__ float buf[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { + float *src = x + row * cols; + for (int out_col = 0; out_col < k; ++out_col) { + int best = 0; + float best_val = -3.402823466e+38F; + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) { + int col = done + i; + int used = 0; + for (int p = 0; p < out_col; ++p) if (indices[row * k + p] == (long)col) used = 1; + if (!used && buf[i] > best_val) { best_val = buf[i]; best = col; } + } + } + values[row * k + out_col] = best_val; + indices[row * k + out_col] = (long)best; + } + } +} + +std::vector bang_func(torch::Tensor x, int k, int dim) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + int rows = input.size(0), cols = input.size(1); + auto values = torch::empty({rows, k}, input.options()); + auto indices = torch::empty({rows, k}, x.options().dtype(torch::kLong)); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t d = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + topk_kernel<<>>(input.data_ptr(), values.data_ptr(), indices.data_ptr(), rows, cols, k); + cnrtQueueSync(queue); + return {values.to(x.dtype()), indices}; +} diff --git a/average_pooling_2d.mlu b/average_pooling_2d.mlu new file mode 100644 index 0000000..06c668e --- /dev/null +++ b/average_pooling_2d.mlu @@ -0,0 +1,31 @@ +#include +#include +#include + +__mlu_entry__ void avg_pool2d_kernel(float *x, float *out, int channels, int height, int width, int kernel, int out_h, int out_w, int total) { + for (int idx = taskId; idx < total; idx += taskDim) { + int ow = idx % out_w; + int oh = (idx / out_w) % out_h; + int channel = (idx / (out_w * out_h)) % channels; + int batch = idx / (out_w * out_h * channels); + float sum = 0.0f; + for (int kh = 0; kh < kernel; ++kh) + for (int kw = 0; kw < kernel; ++kw) + sum += x[((batch * channels + channel) * height + oh * kernel + kh) * width + ow * kernel + kw]; + out[idx] = sum / (float)(kernel * kernel); + } +} + +torch::Tensor bang_func(torch::Tensor x, int kernel_size) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + int n = input.size(0), c = input.size(1), h = input.size(2), w = input.size(3); + int oh = h / kernel_size, ow = w / kernel_size; + auto out = torch::empty({n, c, oh, ow}, input.options()); + int total = n * c * oh * ow; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + avg_pool2d_kernel<<>>(input.data_ptr(), out.data_ptr(), c, h, w, kernel_size, oh, ow, total); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/batched_matrix_multiplication.mlu b/batched_matrix_multiplication.mlu new file mode 100644 index 0000000..05b998a --- /dev/null +++ b/batched_matrix_multiplication.mlu @@ -0,0 +1,43 @@ +#include +#include +#include + +#define NRAM_SIZE 4096 + +// BMM: each output element = dot(A_row, B_col), tiled along K dimension +__mlu_entry__ void bmm_kernel(float *a, float *b, float *out, int rows, int inner, int cols, int total) { + __nram__ float buf_a[NRAM_SIZE]; + __nram__ float buf_b[NRAM_SIZE]; + for (int idx = taskId; idx < total; idx += taskDim) { + int col = idx % cols; + int row = (idx / cols) % rows; + int batch = idx / (rows * cols); + float *a_row = a + batch * rows * inner + row * inner; + float *b_col = b + batch * inner * cols + col; // stride = cols + float sum = 0.0f; + for (int done = 0; done < inner; done += NRAM_SIZE) { + int len = (inner - done) < NRAM_SIZE ? (inner - done) : NRAM_SIZE; + __memcpy(buf_a, a_row + done, len * sizeof(float), GDRAM2NRAM); + // gather b column into buf_b + for (int i = 0; i < len; ++i) buf_b[i] = b_col[(done + i) * cols]; + int aligned_len = (len + 63) & ~63; + __bang_mul(buf_a, buf_a, buf_b, aligned_len); + for (int i = 0; i < len; ++i) sum += buf_a[i]; + } + out[idx] = sum; + } +} + +torch::Tensor bang_func(torch::Tensor A, torch::Tensor B) { + auto a = A.contiguous().to(torch::kFloat).contiguous(); + auto b = B.contiguous().to(torch::kFloat).contiguous(); + int batch = a.size(0), rows = a.size(1), inner = a.size(2), cols = b.size(2); + auto out = torch::empty({batch, rows, cols}, a.options()); + int total = batch * rows * cols; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + bmm_kernel<<>>(a.data_ptr(), b.data_ptr(), out.data_ptr(), rows, inner, cols, total); + cnrtQueueSync(queue); + return out.to(A.dtype()); +} diff --git a/config b/config index 0f30166..cd9d761 100644 --- a/config +++ b/config @@ -1 +1,20 @@ -001 \ No newline at end of file +001 +002 +003 +004 +005 +009 +012 +023 +034 +039 +051 +056 +070 +075 +103 +104 +109 +111 +135 +138 diff --git a/conv_standard_1D.mlu b/conv_standard_1D.mlu new file mode 100644 index 0000000..c264d33 --- /dev/null +++ b/conv_standard_1D.mlu @@ -0,0 +1,39 @@ +#include +#include +#include + +__mlu_entry__ void conv1d_kernel(float *x, float *weight, float *out, int in_channels, int out_channels, int length, int kernel_size, int stride, int padding, int dilation, int groups, int out_length, int total) { + for (int idx = taskId; idx < total; idx += taskDim) { + int out_pos = idx % out_length; + int out_ch = (idx / out_length) % out_channels; + int batch = idx / (out_length * out_channels); + int ch_per_group = in_channels / groups; + int out_per_group = out_channels / groups; + int group = out_ch / out_per_group; + float sum = 0.0f; + for (int ic = 0; ic < ch_per_group; ++ic) { + int in_ch = group * ch_per_group + ic; + for (int k = 0; k < kernel_size; ++k) { + int in_pos = out_pos * stride + k * dilation - padding; + if (in_pos >= 0 && in_pos < length) + sum += x[(batch * in_channels + in_ch) * length + in_pos] * weight[(out_ch * ch_per_group + ic) * kernel_size + k]; + } + } + out[idx] = sum; + } +} + +torch::Tensor bang_func(torch::Tensor x, torch::Tensor kernel, int in_channels, int out_channels, int kernel_size, int stride, int padding, int dilation, int groups, int bias) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto weight = kernel.contiguous().to(torch::kFloat).contiguous(); + int batch = input.size(0), length = input.size(2); + int out_length = (length + 2 * padding - dilation * (kernel_size - 1) - 1) / stride + 1; + auto out = torch::empty({batch, out_channels, out_length}, input.options()); + int total = batch * out_channels * out_length; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + conv1d_kernel<<>>(input.data_ptr(), weight.data_ptr(), out.data_ptr(), in_channels, out_channels, length, kernel_size, stride, padding, dilation, groups, out_length, total); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/conv_transposed_2D__asymmetric_input__square_kernel.mlu b/conv_transposed_2D__asymmetric_input__square_kernel.mlu new file mode 100644 index 0000000..a581509 --- /dev/null +++ b/conv_transposed_2D__asymmetric_input__square_kernel.mlu @@ -0,0 +1,52 @@ +#include +#include +#include + +__mlu_entry__ void zero_kernel(float *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; +} + +__mlu_entry__ void conv_transpose2d_kernel(float *x, float *weight, float *out, int in_channels, int out_channels, int h, int w, int kernel_size, int stride, int padding, int out_h, int out_w, int groups, int total_in) { + for (int idx = taskId; idx < total_in; idx += taskDim) { + int iw = idx % w; + int ih = (idx / w) % h; + int in_ch = (idx / (w * h)) % in_channels; + int batch = idx / (w * h * in_channels); + int in_per_group = in_channels / groups; + int out_per_group = out_channels / groups; + int group = in_ch / in_per_group; + float xv = x[idx]; + for (int oc = 0; oc < out_per_group; ++oc) { + int out_ch = group * out_per_group + oc; + for (int kh = 0; kh < kernel_size; ++kh) { + for (int kw = 0; kw < kernel_size; ++kw) { + int oh = ih * stride + kh - padding; + int ow = iw * stride + kw - padding; + if (oh >= 0 && oh < out_h && ow >= 0 && ow < out_w) { + int out_idx = ((batch * out_channels + out_ch) * out_h + oh) * out_w + ow; + float wv = weight[((in_ch * out_per_group + oc) * kernel_size + kh) * kernel_size + kw]; + out[out_idx] += xv * wv; + } + } + } + } + } +} + +torch::Tensor bang_func(torch::Tensor x, torch::Tensor kernel, int in_channels, int out_channels, int kernel_size, int stride, int padding, int output_padding, int groups, bool bias) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto weight = kernel.contiguous().to(torch::kFloat).contiguous(); + int batch = input.size(0), h = input.size(2), w = input.size(3); + int out_h = (h - 1) * stride - 2 * padding + kernel_size + output_padding; + int out_w = (w - 1) * stride - 2 * padding + kernel_size + output_padding; + auto out = torch::empty({batch, out_channels, out_h, out_w}, input.options()); + int total_out = batch * out_channels * out_h * out_w; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtDim3_t one = {1, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + zero_kernel<<>>(out.data_ptr(), total_out); + conv_transpose2d_kernel<<>>(input.data_ptr(), weight.data_ptr(), out.data_ptr(), in_channels, out_channels, h, w, kernel_size, stride, padding, out_h, out_w, groups, input.numel()); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/cumsum.mlu b/cumsum.mlu new file mode 100644 index 0000000..9d4c83c --- /dev/null +++ b/cumsum.mlu @@ -0,0 +1,52 @@ +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void cumsum_dim1_kernel(float *x, float *out, int rows, int cols) { + __nram__ float buf[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { + float *src = x + row * cols; + float *dst = out + row * cols; + float running_sum = 0.0f; + for (int done = 0; done < cols; done += NRAM_SIZE) { + int len = (cols - done) < NRAM_SIZE ? (cols - done) : NRAM_SIZE; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + for (int i = 0; i < len; ++i) { + running_sum += buf[i]; + buf[i] = running_sum; + } + __memcpy(dst + done, buf, len * sizeof(float), NRAM2GDRAM); + } + } +} + +__mlu_entry__ void cumsum_dim0_kernel(float *x, float *out, int rows, int cols) { + __nram__ float buf[NRAM_SIZE]; + __nram__ float acc[NRAM_SIZE]; + for (int col_start = taskId * NRAM_SIZE; col_start < cols; col_start += taskDim * NRAM_SIZE) { + int len = (cols - col_start) < NRAM_SIZE ? (cols - col_start) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + // zero accumulator + for (int i = 0; i < aligned_len; ++i) acc[i] = 0.0f; + for (int r = 0; r < rows; ++r) { + __memcpy(buf, x + r * cols + col_start, len * sizeof(float), GDRAM2NRAM); + __bang_add(acc, acc, buf, aligned_len); + __memcpy(out + r * cols + col_start, acc, len * sizeof(float), NRAM2GDRAM); + } + } +} + +torch::Tensor bang_func(torch::Tensor x, int dim) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto out = torch::empty_like(input); + int rows = input.size(0), cols = input.size(1); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t d = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + if (dim == 0) cumsum_dim0_kernel<<>>(input.data_ptr(), out.data_ptr(), rows, cols); + else cumsum_dim1_kernel<<>>(input.data_ptr(), out.data_ptr(), rows, cols); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/gather.mlu b/gather.mlu new file mode 100644 index 0000000..dff697e --- /dev/null +++ b/gather.mlu @@ -0,0 +1,46 @@ +#include +#include +#include + +#define NRAM_SIZE 4096 + +__mlu_entry__ void gather_kernel(float *x, long *indices, long *bin_ids, float *weights, long *bins, float *out, int elements, int hidden, int top_k, int has_weights) { + __nram__ float buf[NRAM_SIZE]; + for (int elem = taskId; elem < elements; elem += taskDim) { + int src_row = (int)indices[elem]; + int bin_id = (int)bin_ids[elem]; + int rank = elem - (int)bins[bin_id]; + int dst_row = src_row * top_k + rank; + float scale = has_weights ? weights[elem] : 1.0f; + float *src = x + src_row * hidden; + float *dst = out + dst_row * hidden; + for (int done = 0; done < hidden; done += NRAM_SIZE) { + int len = (hidden - done) < NRAM_SIZE ? (hidden - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, src + done, len * sizeof(float), GDRAM2NRAM); + __bang_mul_scalar(buf, buf, scale, aligned_len); + __memcpy(dst + done, buf, len * sizeof(float), NRAM2GDRAM); + } + } +} + +__mlu_entry__ void zero_kernel(float *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; +} + +torch::Tensor bang_func(torch::Tensor x, torch::Tensor indices, torch::Tensor bin_ids, torch::Tensor weights, torch::Tensor bins, int top_k) { + auto input = x.contiguous().to(torch::kFloat).contiguous(); + auto idx = indices.contiguous().to(torch::kLong).contiguous(); + auto bid = bin_ids.contiguous().to(torch::kLong).contiguous(); + auto w = weights.contiguous().to(torch::kFloat).contiguous(); + auto b = bins.contiguous().to(torch::kLong).contiguous(); + int tokens = input.size(0), hidden = input.size(1), elements = idx.numel(); + auto out = torch::empty({tokens * top_k, hidden}, input.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + zero_kernel<<>>(out.data_ptr(), tokens * top_k * hidden); + gather_kernel<<>>(input.data_ptr(), idx.data_ptr(), bid.data_ptr(), w.data_ptr(), b.data_ptr(), out.data_ptr(), elements, hidden, top_k, weights.numel() > 0 ? 1 : 0); + cnrtQueueSync(queue); + return out.to(x.dtype()); +} diff --git a/matrix_scalar_multiplication.mlu b/matrix_scalar_multiplication.mlu new file mode 100644 index 0000000..3f1cbd0 --- /dev/null +++ b/matrix_scalar_multiplication.mlu @@ -0,0 +1,33 @@ +#include +#include +#include + +#define NRAM_SIZE 8192 + +__mlu_entry__ void matrix_scalar_mul_kernel(float *input, float *output, int total, float scalar) { + int remain = total % taskDim; + int per_core = total / taskDim; + int count = per_core + (int)(taskId < remain); + int start = per_core * taskId + (taskId < remain ? taskId : remain); + if (count == 0) return; + __nram__ float buf[NRAM_SIZE]; + for (int done = 0; done < count; done += NRAM_SIZE) { + int len = (count - done) < NRAM_SIZE ? (count - done) : NRAM_SIZE; + int aligned_len = (len + 63) & ~63; + __memcpy(buf, input + start + done, len * sizeof(float), GDRAM2NRAM); + __bang_mul_scalar(buf, buf, scalar, aligned_len); + __memcpy(output + start + done, buf, len * sizeof(float), NRAM2GDRAM); + } +} + +torch::Tensor bang_func(torch::Tensor A, double s) { + auto x = A.contiguous().to(torch::kFloat).contiguous(); + auto output = torch::empty_like(x); + int total = x.numel(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + matrix_scalar_mul_kernel<<>>(x.data_ptr(), output.data_ptr(), total, (float)s); + cnrtQueueSync(queue); + return output.to(A.dtype()); +} diff --git a/openoperator_basic_problems.md b/openoperator_basic_problems.md new file mode 100644 index 0000000..16889b4 --- /dev/null +++ b/openoperator_basic_problems.md @@ -0,0 +1,183 @@ +# OpenOperator Basic Problems + +Total: 20 + +## 001 001_LeakyReLU + +- Category: pointwise +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, double negative_slope);` + +实现一个 BangC 算子，对输入张量应用 LeakyReLU 激活处理。输入包括形状为 `[batch, dim]` 的张量 `x` 以及表示负半轴斜率的标量 `negative_slope`；输出张量的形状与 `x` 相同。 + +## 002 002_matrix_scalar_multiplication + +- Category: pointwise +- Source: kernelbench +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor A, double s);` + +实现一个 BangC 算子，实现矩阵与标量的逐元素乘法运算。输入包括形状为 `[M, N]` 的张量 `A` 以及浮点数标量 `s`；输出为与 `A` 形状相同的张量。 + +## 003 003_LogSoftmax + +- Category: softmax +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, int dim);` + +实现一个 BangC 算子，对输入张量在指定维度上执行 LogSoftmax 激活计算。输入包括形状为 `[batch_size, dim]` 的张量 `x` 以及用于指定计算维度的整数参数 `dim`；输出张量的形状与 `x` 保持一致。 + +## 004 004_batched_matrix_multiplication + +- Category: matrix +- Source: kernelbench +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor A, torch::Tensor B);` + +实现一个 BangC 算子，执行两个三维张量的批量矩阵乘法（Batch Matrix Multiplication）。输入包括形状为 `[batch_size, m, k]` 的张量 `A` 和形状为 `[batch_size, k, n]` 的张量 `B`；输出张量的形状为 `[batch_size, m, n]`。 + +## 005 005_average_pooling_2d + +- Category: pool +- Source: kernelbench +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x,int kernel_size);` + +实现一个 BangC 算子，对输入张量执行二维平均池化（2D Average Pooling）计算。输入包括形状为 `[batch_size, channels, height, width]` 的张量 `x`，以及用于指定窗口大小、步幅和填充的标量参数 `kernel_size`、`stride` 与 `padding`；输出为池化处理后的张量。 + +## 009 009_conv_standard_1D + +- Category: conv +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, torch::Tensor kernel, int in_channels, int out_channels, int kernel_size, int stride, int padding, int dilation, int groups, int bias);` + +实现一个 BangC 算子，执行标准的一维卷积运算。输入包括形状为 `[batch, in_channels, length]` 的数据张量 `x`、形状为 `[out_channels, in_channels/groups, kernel_size]` 的权重张量 `weight` 以及形状为 `[out_channels]` 的可选偏置张量 `bias`。算子还需支持通过 `stride`、`padding`、`dilation` 和 `groups` 等参数配置卷积行为，并输出形状为 `[batch, out_channels, length_out]` 的张量。 + +## 012 012_conv_transposed_2D__asymmetric_input__square_kernel + +- Category: conv +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, torch::Tensor kernel, int in_channels, int out_channels, int kernel_size, int stride, int padding, int output_padding, int groups, bool bias);` + +实现一个 BangC 算子，执行二维转置卷积运算。输入包括形状为 `[batch, in_channels, h_in, w_in]` 的数据张量 `x`、形状为 `[in_channels, out_channels // groups, kernel_size, kernel_size]` 的卷积核张量以及可选的偏置张量；算子通过 `stride`、`padding`、`output_padding` 和 `groups` 等参数控制输出特征图的形状与计算逻辑。 + +## 023 023_Matrix_vector_multiplication_ + +- Category: matrix +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor A, torch::Tensor B);` + +实现一个 BangC 算子，执行矩阵与向量的乘法运算（GEMV）。输入包括形状为 `[M, K]` 的矩阵 `A` 和形状为 `[K, 1]` 的向量 `B`；输出为两者相乘得到的形状为 `[M, 1]` 的结果张量。 + +## 034 034_Argmax_over_a_dimension + +- Category: reduce +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, int dim);` + +实现一个 BangC 算子，沿指定维度计算输入张量的最大值索引（Argmax）。输入包括形状为 `[D_0, D_1, ..., D_{n-1}]` 的张量 `x` 以及一个用于指定规约维度的整数参数 `dim`；输出为移除该维度后的索引张量。 + +## 039 039_BatchNorm + +- Category: fused +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, int num_features);` + +实现一个 BangC 算子，对四维张量执行二维批归一化（Batch Normalization）计算。输入包括形状为 `[batch, channel, H, W]` 的数据张量 `x` 以及表示特征通道数的整数 `num_features`；输出张量与 `x` 形状相同。 + +## 051 051_cumsum + +- Category: cum +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, int dim);` + +实现一个 BangC 算子，对输入张量沿指定维度计算累加前缀和（Cumulative Sum）。输入包括形状为 `[d0, d1, ..., dn]` 的张量 `x` 以及表示操作维度的整数 `dim`；输出张量的形状与 `x` 保持一致。 + +## 056 056_gather + +- Category: io +- Source: operaters +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, torch::Tensor indices, torch::Tensor bin_ids, torch::Tensor weights, torch::Tensor bins, int top_k);` + +实现一个 BangC 算子，根据索引和分桶规则将输入张量的行数据收集并加权排列到目标缓冲区。输入包括形状为 `[tokens, hidden_size]` 的张量 `x`、形状均为 `[num_elements]` 的索引 `indices` 和分桶 ID `bin_ids`、形状为 `[num_elements]` 的可选权重张量 `weights`、分桶边界张量 `bins` 以及标量 `top_k`；输出张量的形状为 `[tokens * top_k, hidden_size]`。 + +## 070 070_Sqrt + +- Category: pointwise +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x);` + +实现一个 BangC 算子，对输入张量执行逐元素的绝对值计算并求其平方根。输入为形状为 `[batch_size, dim]` 的张量 `x`；输出张量与 `x` 的形状一致。 + +## 075 075_TopK + +- Category: reduce +- Source: custom +- Dtype: float16 +- C++ Wrapper: `std::vector bang_func(torch::Tensor x, int k, int dim);` + +实现一个 BangC 算子，用于在输入张量的指定维度上提取前 $k$ 个最大的数值及其对应的索引。输入包括形状为 `[batch, dim]` 的张量 `x`，以及整型标量参数 `k` 和 `dim`，分别表示需要提取的元素数量和操作的维度；输出为包含最大值及其索引的两个张量。 + +## 103 103_MSE_Loss + +- Category: loss +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor predictions, torch::Tensor targets);` + +实现一个 BangC 算子，计算预测张量与目标张量之间的均方误差（MSE Loss）。输入包括形状均为 `[batch, dim]` 的预测张量 `predictions` 和目标张量 `targets`；输出为计算所得的标量损失值。 + +## 104 104_KL_Divergence_Loss + +- Category: loss +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor input_log_prob, torch::Tensor target_prob);` + +实现一个 BangC 算子，计算预测对数概率与目标概率之间的 KL 散度，并按照 batch 维度进行平均规约。输入包括形状均为 `[batch, N]` 的张量 `input_log_prob` 和 `target_prob`，分别对应预测的对数概率分布和目标概率分布；输出为按 `batchmean` 模式计算得到的标量结果。 + +## 109 109_Scatter_add + +- Category: io +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor src, torch::Tensor index, int dim_size);` + +实现一个 BangC 算子，根据索引张量将源张量的数据累加到目标张量的对应位置。输入包括形状为 `[N, D]` 的源张量 `src`、形状为 `[N]` 的索引张量 `index` 以及指定输出首维长度的整数 `dim_size`；输出张量的形状为 `[dim_size, D]`。 + +## 111 111_Masked_select + +- Category: io +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor input, double threshold);` + +实现一个 BangC 算子，根据给定的阈值筛选输入张量中的元素。输入包括形状为 `[M, N]` 的张量 `input` 以及浮点数标量 `threshold`；算子将 `input` 中所有大于 `threshold` 的元素提取并展平为一维张量输出。 + +## 135 135_Dilated_conv_2D + +- Category: conv +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, torch::Tensor kernel, int in_channels, int out_channels, int kernel_size, int dilation, int padding);` + +实现一个 BangC 算子，对四维输入张量执行带有填充和空洞特性的二维卷积运算。输入为形状为 `[batch, in_channels, H, W]` 的张量 `x`，其运算受卷积核大小 `kernel_size`、膨胀系数 `dilation` 以及填充宽度 `padding` 等参数控制；输出为卷积计算后的特征图张量。 + +## 138 138_GRU_forward + +- Category: fused +- Source: custom +- Dtype: float16 +- C++ Wrapper: `torch::Tensor bang_func(torch::Tensor x, int input_size, int hidden_size, int num_layers);` + +实现一个 BangC 算子，实现多层门控循环单元（GRU）的前向计算过程。输入为形状为 `[batch, seq_len, input_size]` 的张量 `x`，并根据给定的隐藏层维度 `hidden_size` 和层数 `num_layers` 对序列数据进行处理；算子最终输出形状为 `[batch, seq_len, hidden_size]` 的隐藏状态序列。