From c2b2dcaab87b5bd84153f39e5274c8802be65254 Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 17:20:35 +0800 Subject: [PATCH 01/18] 001 --- LeakyReLU.mlu | 104 ++++++++++++++++++++++---------------------------- 1 file changed, 46 insertions(+), 58 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 900f4ab..f0641fe 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,7 +2,7 @@ #include #include -#define CHUNK_SIZE 4096 +#define CHUNK_SIZE 16384 __mlu_entry__ void leakyrelu_kernel( float *input, @@ -10,70 +10,58 @@ __mlu_entry__ void leakyrelu_kernel( 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 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); + // 双缓冲流水线: ping-pong buffers + __nram__ float nram_buf[CHUNK_SIZE * 2]; + __nram__ float nram_tmp[CHUNK_SIZE]; + + float *ping = nram_buf; + float *pong = nram_buf + CHUNK_SIZE; + + uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; + + if (num_chunks == 0) return; + + // Load first chunk + uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; + __memcpy(ping, input + start, len0 * sizeof(float), GDRAM2NRAM); + + for (uint32_t i = 0; i < num_chunks; i++) { + uint32_t cur_offset = i * CHUNK_SIZE; + uint32_t cur_len = ((cur_offset + CHUNK_SIZE) <= count) + ? CHUNK_SIZE : (count - cur_offset); + uint32_t cur_aligned = (cur_len + 63) & ~63; + + float *cur_buf = (i % 2 == 0) ? ping : pong; + float *nxt_buf = (i % 2 == 0) ? pong : ping; + + // Prefetch next chunk (overlap with compute) + if (i + 1 < num_chunks) { + uint32_t nxt_offset = (i + 1) * CHUNK_SIZE; + uint32_t nxt_len = ((nxt_offset + CHUNK_SIZE) <= count) + ? CHUNK_SIZE : (count - nxt_offset); + __memcpy_async(nxt_buf, input + start + nxt_offset, + nxt_len * sizeof(float), GDRAM2NRAM); + } + + // Compute LeakyReLU: relu(x) + negative_slope * min(0,x) + __bang_active_relu(nram_tmp, cur_buf, cur_aligned); + __bang_sub(cur_buf, cur_buf, nram_tmp, cur_aligned); + __bang_mul_scalar(cur_buf, cur_buf, negative_slope, cur_aligned); + __bang_add(cur_buf, nram_tmp, cur_buf, cur_aligned); + + // Store result + __memcpy(output + start + cur_offset, cur_buf, + cur_len * sizeof(float), NRAM2GDRAM); } } @@ -102,9 +90,9 @@ torch::Tensor bang_func( cnrtQueue_t queue = torch_mlu::getCurMLUStream(); - cnrtDim3_t dim = {4,1,1}; + cnrtDim3_t dim = {16,1,1}; cnrtFunctionType_t ktype = - cnrtFuncTypeUnion1; + cnrtFuncTypeUnion4; leakyrelu_kernel<<>>( input_fp32.data_ptr(), From 73887f6ad94da94a2941d36cf67f56e59c07c57e Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:10:24 +0800 Subject: [PATCH 02/18] 001-V2 --- LeakyReLU.mlu | 165 +++++++++++++++++++++----------------------------- 1 file changed, 69 insertions(+), 96 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index f0641fe..4fc540c 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -1,110 +1,83 @@ -#include +#include "bang.h" +#include "cnrt.h" #include -#include - -#define CHUNK_SIZE 16384 - -__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); - - // 双缓冲流水线: ping-pong buffers - __nram__ float nram_buf[CHUNK_SIZE * 2]; - __nram__ float nram_tmp[CHUNK_SIZE]; - - float *ping = nram_buf; - float *pong = nram_buf + CHUNK_SIZE; - - uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; - - if (num_chunks == 0) return; - - // Load first chunk - uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; - __memcpy(ping, input + start, len0 * sizeof(float), GDRAM2NRAM); - - for (uint32_t i = 0; i < num_chunks; i++) { - uint32_t cur_offset = i * CHUNK_SIZE; - uint32_t cur_len = ((cur_offset + CHUNK_SIZE) <= count) - ? CHUNK_SIZE : (count - cur_offset); - uint32_t cur_aligned = (cur_len + 63) & ~63; - - float *cur_buf = (i % 2 == 0) ? ping : pong; - float *nxt_buf = (i % 2 == 0) ? pong : ping; - - // Prefetch next chunk (overlap with compute) - if (i + 1 < num_chunks) { - uint32_t nxt_offset = (i + 1) * CHUNK_SIZE; - uint32_t nxt_len = ((nxt_offset + CHUNK_SIZE) <= count) - ? CHUNK_SIZE : (count - nxt_offset); - __memcpy_async(nxt_buf, input + start + nxt_offset, - nxt_len * sizeof(float), GDRAM2NRAM); - } - - // Compute LeakyReLU: relu(x) + negative_slope * min(0,x) - __bang_active_relu(nram_tmp, cur_buf, cur_aligned); - __bang_sub(cur_buf, cur_buf, nram_tmp, cur_aligned); - __bang_mul_scalar(cur_buf, cur_buf, negative_slope, cur_aligned); - __bang_add(cur_buf, nram_tmp, cur_buf, cur_aligned); - - // Store result - __memcpy(output + start + cur_offset, cur_buf, - cur_len * sizeof(float), NRAM2GDRAM); - } -} +const int NRAM_SIZE = 1024 * 128; +const int NRAM_FLOAT_CAP = NRAM_SIZE / sizeof(float); -torch::Tensor bang_func( - torch::Tensor input, - double negative_slope) { +__mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elements, float negative_slope) { + int per_task = num_elements / taskDim; + int remainder = num_elements % taskDim; + int start = taskId * per_task + (taskId < remainder ? taskId : remainder); + int len = per_task + (taskId < remainder ? 1 : 0); + if (len <= 0) return; - TORCH_CHECK( - input.is_contiguous(), - "Input must be contiguous"); + // 3 buffers: src, scaled, mask + int max_deal = NRAM_FLOAT_CAP / 3; + max_deal = max_deal / 64 * 64; - // 保留原始 dtype - auto original_dtype = input.scalar_type(); + __nram__ float nram_src[NRAM_FLOAT_CAP / 3]; + __nram__ float nram_scaled[NRAM_FLOAT_CAP / 3]; + __nram__ float nram_mask[NRAM_FLOAT_CAP / 3]; - // -------- 只处理数据类型 -------- - torch::Tensor input_fp32 = input; - if (original_dtype != torch::kFloat) { - input_fp32 = input.to(torch::kFloat); - } + int offset = start; + int remain = len; - auto output_fp32 = torch::empty_like(input_fp32); + while (remain > 0) { + int deal_num = remain > max_deal ? max_deal : remain; + int align_num = (deal_num + 63) / 64 * 64; - int total = input_fp32.numel(); + __memcpy(nram_src, input + offset, deal_num * sizeof(float), GDRAM2NRAM); - cnrtQueue_t queue = - torch_mlu::getCurMLUStream(); + // scaled = x * negative_slope + __bang_mul_scalar(nram_scaled, nram_src, negative_slope, align_num); - cnrtDim3_t dim = {16,1,1}; - cnrtFunctionType_t ktype = - cnrtFuncTypeUnion4; + // mask = (x >= 0) ? 1.0 : 0.0 + __bang_ge_scalar(nram_mask, nram_src, 0.0f, align_num); - leakyrelu_kernel<<>>( - input_fp32.data_ptr(), - output_fp32.data_ptr(), - total, - (float)negative_slope - ); + // result = mask * x + (1 - mask) * scaled + // nram_src = mask * x + __bang_mul(nram_src, nram_src, nram_mask, align_num); - // 转回原 dtype - if (original_dtype != torch::kFloat) { - return output_fp32.to(original_dtype); - } + // nram_mask = 1 - mask (invert) + __bang_mul_scalar(nram_mask, nram_mask, -1.0f, align_num); + __bang_add_scalar(nram_mask, nram_mask, 1.0f, align_num); + + // nram_scaled = (1-mask) * scaled + __bang_mul(nram_scaled, nram_scaled, nram_mask, align_num); + + // result = nram_src + nram_scaled + __bang_add(nram_src, nram_src, nram_scaled, align_num); + + __memcpy(output + offset, nram_src, deal_num * sizeof(float), NRAM2GDRAM); - return output_fp32; + offset += deal_num; + remain -= deal_num; + } } + +torch::Tensor bang_func(torch::Tensor input, float negative_slope) { + auto output = torch::empty_like(input); + int num_elements = input.numel(); + + cnrtDim3_t dim; + cnrtFunctionType_t func_type = CNRT_FUNC_TYPE_UNION1; + dim.x = 4; + dim.y = 1; + dim.z = 1; + + cnrtQueue_t queue; + cnrtQueueCreate(&queue); + + leaky_relu_kernel<<>>( + (float *)input.data_ptr(), + (float *)output.data_ptr(), + num_elements, + negative_slope + ); + + cnrtQueueSync(queue); + cnrtQueueDestroy(queue); + + return output; +} \ No newline at end of file From 7f445c531f9a329a019a8216f182aa7deb23d412 Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:13:22 +0800 Subject: [PATCH 03/18] 001-V3 --- LeakyReLU.mlu | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 4fc540c..bf88df4 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -1,6 +1,7 @@ #include "bang.h" #include "cnrt.h" #include +#include const int NRAM_SIZE = 1024 * 128; const int NRAM_FLOAT_CAP = NRAM_SIZE / sizeof(float); @@ -66,8 +67,7 @@ torch::Tensor bang_func(torch::Tensor input, float negative_slope) { dim.y = 1; dim.z = 1; - cnrtQueue_t queue; - cnrtQueueCreate(&queue); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); leaky_relu_kernel<<>>( (float *)input.data_ptr(), @@ -77,7 +77,6 @@ torch::Tensor bang_func(torch::Tensor input, float negative_slope) { ); cnrtQueueSync(queue); - cnrtQueueDestroy(queue); return output; } \ No newline at end of file From 9df0de5faa80c59dd06f47450d430233de705552 Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:19:05 +0800 Subject: [PATCH 04/18] 001-V4 --- LeakyReLU.mlu | 30 +++++++++++++++++------------- 1 file changed, 17 insertions(+), 13 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index bf88df4..b6cc1f3 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -3,8 +3,8 @@ #include #include -const int NRAM_SIZE = 1024 * 128; -const int NRAM_FLOAT_CAP = NRAM_SIZE / sizeof(float); +#define NRAM_SIZE (1024 * 128) +#define MAX_DEAL (NRAM_SIZE / sizeof(float) / 3 / 64 * 64) __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elements, float negative_slope) { int per_task = num_elements / taskDim; @@ -13,19 +13,15 @@ __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elemen int len = per_task + (taskId < remainder ? 1 : 0); if (len <= 0) return; - // 3 buffers: src, scaled, mask - int max_deal = NRAM_FLOAT_CAP / 3; - max_deal = max_deal / 64 * 64; - - __nram__ float nram_src[NRAM_FLOAT_CAP / 3]; - __nram__ float nram_scaled[NRAM_FLOAT_CAP / 3]; - __nram__ float nram_mask[NRAM_FLOAT_CAP / 3]; + __nram__ float nram_src[MAX_DEAL]; + __nram__ float nram_scaled[MAX_DEAL]; + __nram__ float nram_mask[MAX_DEAL]; int offset = start; int remain = len; while (remain > 0) { - int deal_num = remain > max_deal ? max_deal : remain; + int deal_num = remain > (int)MAX_DEAL ? (int)MAX_DEAL : remain; int align_num = (deal_num + 63) / 64 * 64; __memcpy(nram_src, input + offset, deal_num * sizeof(float), GDRAM2NRAM); @@ -34,13 +30,21 @@ __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elemen __bang_mul_scalar(nram_scaled, nram_src, negative_slope, align_num); // mask = (x >= 0) ? 1.0 : 0.0 - __bang_ge_scalar(nram_mask, nram_src, 0.0f, align_num); + __bang_active_relu(nram_mask, nram_src, align_num); + __bang_gt_scalar(nram_mask, nram_mask, 0.0f, align_num); + // For x > 0: relu(x) > 0 -> mask=1; for x <= 0: relu(x)=0 -> mask=0 + // But x==0 should give mask=1. Use different approach: + + // mask: x >= 0 -> 1, x < 0 -> 0 + // Use: mask = (x + |x|) / (2*x) won't work for x=0 + // Better: use bang_ge_scalar if available, otherwise: + __bang_write_zero(nram_mask, align_num); + __bang_ge_scalar(nram_mask, nram_src, (float)0, align_num); - // result = mask * x + (1 - mask) * scaled // nram_src = mask * x __bang_mul(nram_src, nram_src, nram_mask, align_num); - // nram_mask = 1 - mask (invert) + // nram_mask = 1 - mask __bang_mul_scalar(nram_mask, nram_mask, -1.0f, align_num); __bang_add_scalar(nram_mask, nram_mask, 1.0f, align_num); From feddc37034c9797732acac4692b128afdbd68a3a Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:32:17 +0800 Subject: [PATCH 05/18] 001-V5 --- LeakyReLU.mlu | 47 +++++++++++++++++------------------------------ 1 file changed, 17 insertions(+), 30 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index b6cc1f3..40782b2 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -4,7 +4,8 @@ #include #define NRAM_SIZE (1024 * 128) -#define MAX_DEAL (NRAM_SIZE / sizeof(float) / 3 / 64 * 64) +#define NRAM_ELEM (NRAM_SIZE / 4) +#define BUF_SIZE (NRAM_ELEM / 4) __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elements, float negative_slope) { int per_task = num_elements / taskDim; @@ -13,48 +14,34 @@ __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elemen int len = per_task + (taskId < remainder ? 1 : 0); if (len <= 0) return; - __nram__ float nram_src[MAX_DEAL]; - __nram__ float nram_scaled[MAX_DEAL]; - __nram__ float nram_mask[MAX_DEAL]; + __nram__ float nram_src[BUF_SIZE]; + __nram__ float nram_dst[BUF_SIZE]; + __nram__ float nram_tmp[BUF_SIZE]; + __nram__ float nram_mask[BUF_SIZE]; + int max_deal = BUF_SIZE / 64 * 64; int offset = start; int remain = len; while (remain > 0) { - int deal_num = remain > (int)MAX_DEAL ? (int)MAX_DEAL : remain; + int deal_num = remain > max_deal ? max_deal : remain; int align_num = (deal_num + 63) / 64 * 64; __memcpy(nram_src, input + offset, deal_num * sizeof(float), GDRAM2NRAM); - // scaled = x * negative_slope - __bang_mul_scalar(nram_scaled, nram_src, negative_slope, align_num); + // dst = relu(x), gets max(0, x) + __bang_active_relu(nram_dst, nram_src, align_num); - // mask = (x >= 0) ? 1.0 : 0.0 - __bang_active_relu(nram_mask, nram_src, align_num); - __bang_gt_scalar(nram_mask, nram_mask, 0.0f, align_num); - // For x > 0: relu(x) > 0 -> mask=1; for x <= 0: relu(x)=0 -> mask=0 - // But x==0 should give mask=1. Use different approach: + // tmp = x - relu(x), gets min(0, x) i.e. negative part + __bang_sub(nram_tmp, nram_src, nram_dst, align_num); - // mask: x >= 0 -> 1, x < 0 -> 0 - // Use: mask = (x + |x|) / (2*x) won't work for x=0 - // Better: use bang_ge_scalar if available, otherwise: - __bang_write_zero(nram_mask, align_num); - __bang_ge_scalar(nram_mask, nram_src, (float)0, align_num); + // tmp = negative_part * negative_slope + __bang_mul_scalar(nram_tmp, nram_tmp, negative_slope, align_num); - // nram_src = mask * x - __bang_mul(nram_src, nram_src, nram_mask, align_num); + // result = relu(x) + negative_part * slope + __bang_add(nram_dst, nram_dst, nram_tmp, align_num); - // nram_mask = 1 - mask - __bang_mul_scalar(nram_mask, nram_mask, -1.0f, align_num); - __bang_add_scalar(nram_mask, nram_mask, 1.0f, align_num); - - // nram_scaled = (1-mask) * scaled - __bang_mul(nram_scaled, nram_scaled, nram_mask, align_num); - - // result = nram_src + nram_scaled - __bang_add(nram_src, nram_src, nram_scaled, align_num); - - __memcpy(output + offset, nram_src, deal_num * sizeof(float), NRAM2GDRAM); + __memcpy(output + offset, nram_dst, deal_num * sizeof(float), NRAM2GDRAM); offset += deal_num; remain -= deal_num; From 8f936aacff68cf7101d99ceb46504dc7a776e17a Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:34:58 +0800 Subject: [PATCH 06/18] 1 --- LeakyReLU.mlu | 19 ++++++++----------- 1 file changed, 8 insertions(+), 11 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 40782b2..4a3b18c 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -1,11 +1,7 @@ #include "bang.h" #include "cnrt.h" -#include -#include -#define NRAM_SIZE (1024 * 128) -#define NRAM_ELEM (NRAM_SIZE / 4) -#define BUF_SIZE (NRAM_ELEM / 4) +#define BUF_SIZE 8192 __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elements, float negative_slope) { int per_task = num_elements / taskDim; @@ -17,22 +13,20 @@ __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elemen __nram__ float nram_src[BUF_SIZE]; __nram__ float nram_dst[BUF_SIZE]; __nram__ float nram_tmp[BUF_SIZE]; - __nram__ float nram_mask[BUF_SIZE]; - int max_deal = BUF_SIZE / 64 * 64; int offset = start; int remain = len; while (remain > 0) { - int deal_num = remain > max_deal ? max_deal : remain; + int deal_num = remain > BUF_SIZE ? BUF_SIZE : remain; int align_num = (deal_num + 63) / 64 * 64; __memcpy(nram_src, input + offset, deal_num * sizeof(float), GDRAM2NRAM); - // dst = relu(x), gets max(0, x) + // dst = relu(x) = max(0, x) __bang_active_relu(nram_dst, nram_src, align_num); - // tmp = x - relu(x), gets min(0, x) i.e. negative part + // tmp = x - relu(x) = min(0, x), the negative part __bang_sub(nram_tmp, nram_src, nram_dst, align_num); // tmp = negative_part * negative_slope @@ -48,6 +42,9 @@ __mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elemen } } +#include +#include "torch_mlu/csrc/aten/utils/cnnl_util.h" + torch::Tensor bang_func(torch::Tensor input, float negative_slope) { auto output = torch::empty_like(input); int num_elements = input.numel(); @@ -58,7 +55,7 @@ torch::Tensor bang_func(torch::Tensor input, float negative_slope) { dim.y = 1; dim.z = 1; - cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + auto queue = torch_mlu::getQueueFromPool(); leaky_relu_kernel<<>>( (float *)input.data_ptr(), From 56a52cf72a9345a3b9c08ba95ec9bd6977d8df6d Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:44:22 +0800 Subject: [PATCH 07/18] 2 --- LeakyReLU.mlu | 68 +++++++++++++++++++-------------------------------- 1 file changed, 25 insertions(+), 43 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 4a3b18c..28aaa2c 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -1,70 +1,52 @@ #include "bang.h" #include "cnrt.h" +#include -#define BUF_SIZE 8192 +#define NRAM_MAX_SIZE 1024 * 128 +#define ELEM_COUNT (NRAM_MAX_SIZE / sizeof(float)) -__mlu_entry__ void leaky_relu_kernel(float *input, float *output, int num_elements, float negative_slope) { - int per_task = num_elements / taskDim; - int remainder = num_elements % taskDim; - int start = taskId * per_task + (taskId < remainder ? taskId : remainder); - int len = per_task + (taskId < remainder ? 1 : 0); - if (len <= 0) return; - - __nram__ float nram_src[BUF_SIZE]; - __nram__ float nram_dst[BUF_SIZE]; - __nram__ float nram_tmp[BUF_SIZE]; - - int offset = start; - int remain = len; - - while (remain > 0) { - int deal_num = remain > BUF_SIZE ? BUF_SIZE : remain; - int align_num = (deal_num + 63) / 64 * 64; - - __memcpy(nram_src, input + offset, deal_num * sizeof(float), GDRAM2NRAM); +__mlu_global__ void leaky_relu_kernel(float* input, float* output, int num_elements, float negative_slope) { + __nram__ float nram_input[ELEM_COUNT]; + __nram__ float nram_temp[ELEM_COUNT]; - // dst = relu(x) = max(0, x) - __bang_active_relu(nram_dst, nram_src, align_num); - - // tmp = x - relu(x) = min(0, x), the negative part - __bang_sub(nram_tmp, nram_src, nram_dst, align_num); + int task_size = num_elements / taskDim; + int remainder = num_elements % taskDim; + int start = taskId * task_size + (taskId < remainder ? taskId : remainder); + int len = task_size + (taskId < remainder ? 1 : 0); - // tmp = negative_part * negative_slope - __bang_mul_scalar(nram_tmp, nram_tmp, negative_slope, align_num); + for (int offset = 0; offset < len; offset += ELEM_COUNT) { + int cur_len = (offset + ELEM_COUNT <= len) ? ELEM_COUNT : (len - offset); + int align_len = (cur_len + 63) / 64 * 64; - // result = relu(x) + negative_part * slope - __bang_add(nram_dst, nram_dst, nram_tmp, align_num); + __memcpy(nram_input, input + start + offset, cur_len * sizeof(float), GDRAM2NRAM); - __memcpy(output + offset, nram_dst, deal_num * sizeof(float), NRAM2GDRAM); + // nram_temp = nram_input * negative_slope + __bang_mul_const(nram_temp, nram_input, negative_slope, align_len); + // output = (input >= 0) ? input : input * negative_slope + __bang_maxequal(nram_temp, nram_temp, nram_input, align_len); - offset += deal_num; - remain -= deal_num; + __memcpy(output + start + offset, nram_temp, cur_len * sizeof(float), NRAM2GDRAM); } } -#include -#include "torch_mlu/csrc/aten/utils/cnnl_util.h" - torch::Tensor bang_func(torch::Tensor input, float negative_slope) { auto output = torch::empty_like(input); int num_elements = input.numel(); + cnrtQueue_t queue; + cnrtQueueCreate(&queue); + cnrtDim3_t dim; - cnrtFunctionType_t func_type = CNRT_FUNC_TYPE_UNION1; dim.x = 4; dim.y = 1; dim.z = 1; - - auto queue = torch_mlu::getQueueFromPool(); + cnrtFunctionType_t func_type = CNRT_FUNC_TYPE_BLOCK; leaky_relu_kernel<<>>( - (float *)input.data_ptr(), - (float *)output.data_ptr(), - num_elements, - negative_slope - ); + input.data_ptr(), output.data_ptr(), num_elements, negative_slope); cnrtQueueSync(queue); + cnrtQueueDestroy(queue); return output; } \ No newline at end of file From 273c4b514181ff32e3f0c0a099b6b3add558ad4c Mon Sep 17 00:00:00 2001 From: duhang01 Date: Wed, 10 Jun 2026 19:46:46 +0800 Subject: [PATCH 08/18] 3 --- LeakyReLU.mlu | 90 ++++++++++++++++++++++----------------------------- 1 file changed, 38 insertions(+), 52 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 28aaa2c..5d39d58 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -1,52 +1,38 @@ -#include "bang.h" -#include "cnrt.h" -#include - -#define NRAM_MAX_SIZE 1024 * 128 -#define ELEM_COUNT (NRAM_MAX_SIZE / sizeof(float)) - -__mlu_global__ void leaky_relu_kernel(float* input, float* output, int num_elements, float negative_slope) { - __nram__ float nram_input[ELEM_COUNT]; - __nram__ float nram_temp[ELEM_COUNT]; - - int task_size = num_elements / taskDim; - int remainder = num_elements % taskDim; - int start = taskId * task_size + (taskId < remainder ? taskId : remainder); - int len = task_size + (taskId < remainder ? 1 : 0); - - for (int offset = 0; offset < len; offset += ELEM_COUNT) { - int cur_len = (offset + ELEM_COUNT <= len) ? ELEM_COUNT : (len - offset); - int align_len = (cur_len + 63) / 64 * 64; - - __memcpy(nram_input, input + start + offset, cur_len * sizeof(float), GDRAM2NRAM); - - // nram_temp = nram_input * negative_slope - __bang_mul_const(nram_temp, nram_input, negative_slope, align_len); - // output = (input >= 0) ? input : input * negative_slope - __bang_maxequal(nram_temp, nram_temp, nram_input, align_len); - - __memcpy(output + start + offset, nram_temp, cur_len * sizeof(float), NRAM2GDRAM); - } -} - -torch::Tensor bang_func(torch::Tensor input, float negative_slope) { - auto output = torch::empty_like(input); - int num_elements = input.numel(); - - cnrtQueue_t queue; - cnrtQueueCreate(&queue); - - cnrtDim3_t dim; - dim.x = 4; - dim.y = 1; - dim.z = 1; - cnrtFunctionType_t func_type = CNRT_FUNC_TYPE_BLOCK; - - leaky_relu_kernel<<>>( - input.data_ptr(), output.data_ptr(), num_elements, negative_slope); - - cnrtQueueSync(queue); - cnrtQueueDestroy(queue); - - return output; -} \ No newline at end of file +import torch +import torch.nn as nn + +class Model(nn.Module): + """ + Simple model that performs a LeakyReLU activation. + """ + def __init__(self, negative_slope: float = 0.01): + """ + Initializes the LeakyReLU module. + + Args: + negative_slope (float, optional): The negative slope of the activation function. Defaults to 0.01. + """ + super(Model, self).__init__() + self.negative_slope = negative_slope + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Applies LeakyReLU activation to the input tensor. + + Args: + x (torch.Tensor): Input tensor of any shape. + + Returns: + torch.Tensor: Output tensor with LeakyReLU applied, same shape as input. + """ + return torch.nn.functional.leaky_relu(x, negative_slope=self.negative_slope) + +batch_size = 16 +dim = 16384 + +def get_inputs(): + x = torch.randn(batch_size, dim) + return [x] + +def get_init_inputs(): + return [0.01] # No special initialization inputs needed \ No newline at end of file From 6f19dd24d786b97a74281aa35933d8b87471a81c Mon Sep 17 00:00:00 2001 From: "Unicorn." <119285775+dh-Unicorn@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:49:01 +0800 Subject: [PATCH 09/18] 5 --- LeakyReLU.mlu | 160 ++++++++++++++++++++++++++++++++++++++------------ 1 file changed, 122 insertions(+), 38 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 5d39d58..900f4ab 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -1,38 +1,122 @@ -import torch -import torch.nn as nn - -class Model(nn.Module): - """ - Simple model that performs a LeakyReLU activation. - """ - def __init__(self, negative_slope: float = 0.01): - """ - Initializes the LeakyReLU module. - - Args: - negative_slope (float, optional): The negative slope of the activation function. Defaults to 0.01. - """ - super(Model, self).__init__() - self.negative_slope = negative_slope - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Applies LeakyReLU activation to the input tensor. - - Args: - x (torch.Tensor): Input tensor of any shape. - - Returns: - torch.Tensor: Output tensor with LeakyReLU applied, same shape as input. - """ - return torch.nn.functional.leaky_relu(x, negative_slope=self.negative_slope) - -batch_size = 16 -dim = 16384 - -def get_inputs(): - x = torch.randn(batch_size, dim) - return [x] - -def get_init_inputs(): - return [0.01] # No special initialization inputs needed \ No newline at end of file +#include +#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); + } +} + + +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; +} From dbd7868eaed61b412adb1132029f1c652abbf9d3 Mon Sep 17 00:00:00 2001 From: "Unicorn." <119285775+dh-Unicorn@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:51:38 +0800 Subject: [PATCH 10/18] qew --- LeakyReLU.mlu | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 900f4ab..f3ecad3 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,7 +2,7 @@ #include #include -#define CHUNK_SIZE 4096 +#define CHUNK_SIZE 6144 __mlu_entry__ void leakyrelu_kernel( float *input, @@ -24,8 +24,7 @@ __mlu_entry__ void leakyrelu_kernel( // NRAM __nram__ float nram_input[CHUNK_SIZE]; - __nram__ float nram_relu[CHUNK_SIZE]; - __nram__ float nram_temp[CHUNK_SIZE]; + __nram__ float nram_output[CHUNK_SIZE]; for (uint32_t offset = 0; offset < count; offset += CHUNK_SIZE) { @@ -44,34 +43,34 @@ __mlu_entry__ void leakyrelu_kernel( // relu(x) __bang_active_relu( - nram_relu, + nram_output, nram_input, aligned_len); - // min(0,x) + // x - relu(x) = min(0, x) __bang_sub( - nram_temp, nram_input, - nram_relu, + nram_input, + nram_output, aligned_len); // negative_slope * min(0,x) __bang_mul_scalar( - nram_temp, - nram_temp, + nram_input, + nram_input, negative_slope, aligned_len); - // relu + scaled negative + // relu(x) + negative_slope * min(0,x) __bang_add( - nram_temp, - nram_relu, - nram_temp, + nram_output, + nram_output, + nram_input, aligned_len); __memcpy( output + start + offset, - nram_temp, + nram_output, len * sizeof(float), NRAM2GDRAM); } From 06935750ef7e6a7996f51da97e320e568694b5a4 Mon Sep 17 00:00:00 2001 From: "Unicorn." <119285775+dh-Unicorn@users.noreply.github.com> Date: Wed, 10 Jun 2026 19:56:40 +0800 Subject: [PATCH 11/18] asdasfd --- LeakyReLU.mlu | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index f3ecad3..31d1bfa 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,7 +2,7 @@ #include #include -#define CHUNK_SIZE 6144 +#define CHUNK_SIZE 4096 __mlu_entry__ void leakyrelu_kernel( float *input, @@ -41,13 +41,14 @@ __mlu_entry__ void leakyrelu_kernel( len * sizeof(float), GDRAM2NRAM); - // relu(x) + // leaky_relu: max(0,x) + negative_slope * min(0,x) + // 用 ge 掩码替代 relu+sub，减少一次向量运算 __bang_active_relu( nram_output, nram_input, aligned_len); - // x - relu(x) = min(0, x) + // nram_input = x - relu(x) = min(0, x) __bang_sub( nram_input, nram_input, @@ -101,9 +102,9 @@ torch::Tensor bang_func( cnrtQueue_t queue = torch_mlu::getCurMLUStream(); - cnrtDim3_t dim = {4,1,1}; + cnrtDim3_t dim = {16,1,1}; cnrtFunctionType_t ktype = - cnrtFuncTypeUnion1; + cnrtFuncTypeUnion4; leakyrelu_kernel<<>>( input_fp32.data_ptr(), From bf7ac403f7cdec9bb61dddd4adf8be3632eb10e5 Mon Sep 17 00:00:00 2001 From: "Unicorn." <119285775+dh-Unicorn@users.noreply.github.com> Date: Wed, 10 Jun 2026 20:03:17 +0800 Subject: [PATCH 12/18] asdqwdqwd --- LeakyReLU.mlu | 117 +++++++++++++++++++++++--------------------------- 1 file changed, 53 insertions(+), 64 deletions(-) diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 31d1bfa..2787971 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,7 +2,8 @@ #include #include -#define CHUNK_SIZE 4096 +// 最大化 NRAM 利用：双缓冲，每个 buffer 32KB / 4B = 8192 floats +#define CHUNK_SIZE 8192 __mlu_entry__ void leakyrelu_kernel( float *input, @@ -10,71 +11,67 @@ __mlu_entry__ void leakyrelu_kernel( 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 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_output[CHUNK_SIZE]; + // 双缓冲：计算与 DMA 重叠 + __nram__ float buf_in[2][CHUNK_SIZE]; + __nram__ float buf_out[CHUNK_SIZE]; + __nram__ float buf_tmp[CHUNK_SIZE]; - for (uint32_t offset = 0; offset < count; offset += CHUNK_SIZE) { + uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; + int cur = 0; - uint32_t len = - (offset + CHUNK_SIZE <= count) - ? CHUNK_SIZE - : (count - offset); + // prefetch 第一块 + if (num_chunks > 0) { + uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; + __memcpy_async(buf_in[0], input + start, + len0 * sizeof(float), GDRAM2NRAM); + } + for (uint32_t i = 0; i < num_chunks; i++) { + uint32_t offset = i * 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); - - // leaky_relu: max(0,x) + negative_slope * min(0,x) - // 用 ge 掩码替代 relu+sub，减少一次向量运算 - __bang_active_relu( - nram_output, - nram_input, - aligned_len); - - // nram_input = x - relu(x) = min(0, x) - __bang_sub( - nram_input, - nram_input, - nram_output, - aligned_len); - - // negative_slope * min(0,x) - __bang_mul_scalar( - nram_input, - nram_input, - negative_slope, - aligned_len); - - // relu(x) + negative_slope * min(0,x) - __bang_add( - nram_output, - nram_output, - nram_input, - aligned_len); - - __memcpy( - output + start + offset, - nram_output, - len * sizeof(float), - NRAM2GDRAM); + // 等待当前块 DMA 完成 + __sync_move(); + + // prefetch 下一块（与计算重叠） + if (i + 1 < num_chunks) { + uint32_t next_offset = (i + 1) * CHUNK_SIZE; + uint32_t next_len = (next_offset + CHUNK_SIZE <= count) + ? CHUNK_SIZE : (count - next_offset); + __memcpy_async(buf_in[1 - cur], + input + start + next_offset, + next_len * sizeof(float), GDRAM2NRAM); + } + + // 计算 LeakyReLU：融合运算 + // relu(x) + __bang_active_relu(buf_out, buf_in[cur], aligned_len); + // min(0,x) = x - relu(x) + __bang_sub(buf_tmp, buf_in[cur], buf_out, aligned_len); + // negative_slope * min(0,x) + relu(x) => 用 mul_scalar + add 融合 + __bang_mul_scalar(buf_tmp, buf_tmp, negative_slope, aligned_len); + __bang_add(buf_out, buf_out, buf_tmp, aligned_len); + + // 写回（异步） + __memcpy_async(output + start + offset, buf_out, + len * sizeof(float), NRAM2GDRAM); + + cur = 1 - cur; } + // 等待最后一次写回完成 + __sync_move(); } @@ -82,29 +79,23 @@ torch::Tensor bang_func( torch::Tensor input, double negative_slope) { - TORCH_CHECK( - input.is_contiguous(), - "Input must be contiguous"); + 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(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); - cnrtDim3_t dim = {16,1,1}; - cnrtFunctionType_t ktype = - cnrtFuncTypeUnion4; + // 激进并行：使用 Union8（32核）最大化多核利用 + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; leakyrelu_kernel<<>>( input_fp32.data_ptr(), @@ -113,10 +104,8 @@ torch::Tensor bang_func( (float)negative_slope ); - // 转回原 dtype if (original_dtype != torch::kFloat) { return output_fp32.to(original_dtype); } - return output_fp32; } From 0f50e63037f36a302984bb475c28886e8ef5737a Mon Sep 17 00:00:00 2001 From: duhang01 Date: Thu, 11 Jun 2026 10:51:54 +0800 Subject: [PATCH 13/18] sadasdasd --- LeakyReLU.mlu | 4 ++ config | 2 +- matrix_scalar_multiplication.mlu | 92 ++++++++++++++++++++++++++++++++ 3 files changed, 97 insertions(+), 1 deletion(-) create mode 100644 matrix_scalar_multiplication.mlu diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index 2787971..c4c55ed 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -108,4 +108,8 @@ torch::Tensor bang_func( return output_fp32.to(original_dtype); } return output_fp32; +<<<<<<< HEAD } +======= +} +>>>>>>> 12714b2 (sadasdasd) diff --git a/config b/config index 0f30166..3b653d0 100644 --- a/config +++ b/config @@ -1 +1 @@ -001 \ No newline at end of file +002 \ No newline at end of file diff --git a/matrix_scalar_multiplication.mlu b/matrix_scalar_multiplication.mlu new file mode 100644 index 0000000..914aa09 --- /dev/null +++ b/matrix_scalar_multiplication.mlu @@ -0,0 +1,92 @@ +#include +#include +#include + +#define CHUNK_SIZE 8192 + +__mlu_entry__ void matrix_scalar_mul_kernel( + float *input, + float *output, + int total, + float scalar) { + + 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__ float buf_in[2][CHUNK_SIZE]; + __nram__ float buf_out[CHUNK_SIZE]; + + uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; + int cur = 0; + + if (num_chunks > 0) { + uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; + __memcpy_async(buf_in[0], input + start, + len0 * sizeof(float), GDRAM2NRAM); + } + + for (uint32_t i = 0; i < num_chunks; i++) { + uint32_t offset = i * CHUNK_SIZE; + uint32_t len = (offset + CHUNK_SIZE <= count) + ? CHUNK_SIZE : (count - offset); + uint32_t aligned_len = (len + 63) & ~63; + + __sync_move(); + + if (i + 1 < num_chunks) { + uint32_t next_offset = (i + 1) * CHUNK_SIZE; + uint32_t next_len = (next_offset + CHUNK_SIZE <= count) + ? CHUNK_SIZE : (count - next_offset); + __memcpy_async(buf_in[1 - cur], + input + start + next_offset, + next_len * sizeof(float), GDRAM2NRAM); + } + + __bang_mul_scalar(buf_out, buf_in[cur], scalar, aligned_len); + + __memcpy_async(output + start + offset, buf_out, + len * sizeof(float), NRAM2GDRAM); + + cur = 1 - cur; + } + __sync_move(); +} + + +torch::Tensor bang_func(torch::Tensor A, double s) { + TORCH_CHECK(A.is_contiguous(), "Input must be contiguous"); + + auto original_dtype = A.scalar_type(); + + torch::Tensor A_fp32 = A; + if (original_dtype != torch::kFloat) { + A_fp32 = A.to(torch::kFloat); + } + + auto output = torch::empty_like(A_fp32); + int total = A_fp32.numel(); + + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + + matrix_scalar_mul_kernel<<>>( + A_fp32.data_ptr(), + output.data_ptr(), + total, + (float)s + ); + + if (original_dtype != torch::kFloat) { + return output.to(original_dtype); + } + return output; +} From 28cc2d900a8692ef201cf21c5e5154e6f7ce8a4e Mon Sep 17 00:00:00 2001 From: "Unicorn." <119285775+dh-Unicorn@users.noreply.github.com> Date: Thu, 11 Jun 2026 14:15:28 +0800 Subject: [PATCH 14/18] wdfasdasdasdadv --- matrix_scalar_multiplication.mlu | 16 +++++++++------- 1 file changed, 9 insertions(+), 7 deletions(-) diff --git a/matrix_scalar_multiplication.mlu b/matrix_scalar_multiplication.mlu index 914aa09..369f25c 100644 --- a/matrix_scalar_multiplication.mlu +++ b/matrix_scalar_multiplication.mlu @@ -20,23 +20,23 @@ __mlu_entry__ void matrix_scalar_mul_kernel( uint32_t count = per_core + (core_id < remainder ? 1 : 0); + if (count == 0) return; + __nram__ float buf_in[2][CHUNK_SIZE]; __nram__ float buf_out[CHUNK_SIZE]; uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; int cur = 0; - if (num_chunks > 0) { - uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; - __memcpy_async(buf_in[0], input + start, - len0 * sizeof(float), GDRAM2NRAM); - } + uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; + __memcpy_async(buf_in[0], input + start, + len0 * sizeof(float), GDRAM2NRAM); for (uint32_t i = 0; i < num_chunks; i++) { uint32_t offset = i * CHUNK_SIZE; uint32_t len = (offset + CHUNK_SIZE <= count) ? CHUNK_SIZE : (count - offset); - uint32_t aligned_len = (len + 63) & ~63; + uint32_t aligned_len = (len + 127) & ~127; __sync_move(); @@ -67,7 +67,7 @@ torch::Tensor bang_func(torch::Tensor A, double s) { torch::Tensor A_fp32 = A; if (original_dtype != torch::kFloat) { - A_fp32 = A.to(torch::kFloat); + A_fp32 = A.to(torch::kFloat).contiguous(); } auto output = torch::empty_like(A_fp32); @@ -85,6 +85,8 @@ torch::Tensor bang_func(torch::Tensor A, double s) { (float)s ); + cnrtQueueSync(queue); + if (original_dtype != torch::kFloat) { return output.to(original_dtype); } From 8f267c266010135a42c0b750a10ce0fd02bfdf2b Mon Sep 17 00:00:00 2001 From: "Unicorn." <119285775+dh-Unicorn@users.noreply.github.com> Date: Thu, 11 Jun 2026 14:20:59 +0800 Subject: [PATCH 15/18] Refactor matrix scalar multiplication kernel --- matrix_scalar_multiplication.mlu | 70 +++++++------------------------- 1 file changed, 15 insertions(+), 55 deletions(-) diff --git a/matrix_scalar_multiplication.mlu b/matrix_scalar_multiplication.mlu index 369f25c..dd875ac 100644 --- a/matrix_scalar_multiplication.mlu +++ b/matrix_scalar_multiplication.mlu @@ -2,7 +2,7 @@ #include #include -#define CHUNK_SIZE 8192 +#define NRAM_SIZE 8192 __mlu_entry__ void matrix_scalar_mul_kernel( float *input, @@ -10,65 +10,29 @@ __mlu_entry__ void matrix_scalar_mul_kernel( int total, float scalar) { - 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); + 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_in[2][CHUNK_SIZE]; - __nram__ float buf_out[CHUNK_SIZE]; - - uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; - int cur = 0; - - uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; - __memcpy_async(buf_in[0], input + start, - len0 * sizeof(float), GDRAM2NRAM); - - for (uint32_t i = 0; i < num_chunks; i++) { - uint32_t offset = i * CHUNK_SIZE; - uint32_t len = (offset + CHUNK_SIZE <= count) - ? CHUNK_SIZE : (count - offset); - uint32_t aligned_len = (len + 127) & ~127; - - __sync_move(); - - if (i + 1 < num_chunks) { - uint32_t next_offset = (i + 1) * CHUNK_SIZE; - uint32_t next_len = (next_offset + CHUNK_SIZE <= count) - ? CHUNK_SIZE : (count - next_offset); - __memcpy_async(buf_in[1 - cur], - input + start + next_offset, - next_len * sizeof(float), GDRAM2NRAM); - } + __nram__ float buf[NRAM_SIZE]; - __bang_mul_scalar(buf_out, buf_in[cur], scalar, aligned_len); + 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_async(output + start + offset, buf_out, - len * sizeof(float), NRAM2GDRAM); - - cur = 1 - cur; + __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); } - __sync_move(); } torch::Tensor bang_func(torch::Tensor A, double s) { - TORCH_CHECK(A.is_contiguous(), "Input must be contiguous"); - - auto original_dtype = A.scalar_type(); - - torch::Tensor A_fp32 = A; - if (original_dtype != torch::kFloat) { - A_fp32 = A.to(torch::kFloat).contiguous(); - } + auto A_contig = A.contiguous(); + auto A_fp32 = A_contig.to(torch::kFloat).contiguous(); auto output = torch::empty_like(A_fp32); int total = A_fp32.numel(); @@ -86,9 +50,5 @@ torch::Tensor bang_func(torch::Tensor A, double s) { ); cnrtQueueSync(queue); - - if (original_dtype != torch::kFloat) { - return output.to(original_dtype); - } - return output; + return output.to(A.dtype()); } From 15fce570182b427c7cb50219ae18c11fc8fd81ab Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=AD=E7=AB=A3=E8=B1=AA?= Date: Thu, 11 Jun 2026 19:09:21 +0800 Subject: [PATCH 16/18] Add basic BangC operator implementations --- Argmax_over_a_dimension.mlu | 49 +++++ ...03\344\270\216\350\257\264\346\230\216.md" | 43 ++++ BatchNorm.mlu | 39 ++++ Dilated_conv_2D.mlu | 38 ++++ GRU_forward.mlu | 25 +++ KL_Divergence_Loss.mlu | 34 ++++ LeakyReLU.mlu | 113 +---------- LogSoftmax.mlu | 29 +++ MSE_Loss.mlu | 32 +++ Masked_select.mlu | 41 ++++ Matrix_vector_multiplication_.mlu | 22 +++ Scatter_add.mlu | 28 +++ Sqrt.mlu | 21 ++ TopK.mlu | 37 ++++ average_pooling_2d.mlu | 32 +++ batched_matrix_multiplication.mlu | 25 +++ config | 21 +- conv_standard_1D.mlu | 38 ++++ ...ed_2D__asymmetric_input__square_kernel.mlu | 49 +++++ cumsum.mlu | 35 ++++ gather.mlu | 35 ++++ matrix_scalar_multiplication.mlu | 53 +---- openoperator_basic_problems.md | 183 ++++++++++++++++++ 23 files changed, 874 insertions(+), 148 deletions(-) create mode 100644 Argmax_over_a_dimension.mlu create mode 100644 "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" create mode 100644 BatchNorm.mlu create mode 100644 Dilated_conv_2D.mlu create mode 100644 GRU_forward.mlu create mode 100644 KL_Divergence_Loss.mlu create mode 100644 LogSoftmax.mlu create mode 100644 MSE_Loss.mlu create mode 100644 Masked_select.mlu create mode 100644 Matrix_vector_multiplication_.mlu create mode 100644 Scatter_add.mlu create mode 100644 Sqrt.mlu create mode 100644 TopK.mlu create mode 100644 average_pooling_2d.mlu create mode 100644 batched_matrix_multiplication.mlu create mode 100644 conv_standard_1D.mlu create mode 100644 conv_transposed_2D__asymmetric_input__square_kernel.mlu create mode 100644 cumsum.mlu create mode 100644 gather.mlu create mode 100644 openoperator_basic_problems.md diff --git a/Argmax_over_a_dimension.mlu b/Argmax_over_a_dimension.mlu new file mode 100644 index 0000000..d6f5174 --- /dev/null +++ b/Argmax_over_a_dimension.mlu @@ -0,0 +1,49 @@ +#include +#include +#include + +__mlu_entry__ void argmax_dim1_kernel(half *x, long *out, int rows, int cols) { + for (int row = taskId; row < rows; row += taskDim) { + int best = 0; + float best_value = (float)x[row * cols]; + for (int col = 1; col < cols; ++col) { + float value = (float)x[row * cols + col]; + if (value > best_value) { + best_value = value; + best = col; + } + } + out[row] = (long)best; + } +} + +__mlu_entry__ void argmax_dim0_kernel(half *x, long *out, int rows, int cols) { + for (int col = taskId; col < cols; col += taskDim) { + int best = 0; + float best_value = (float)x[col]; + for (int row = 1; row < rows; ++row) { + float value = (float)x[row * cols + col]; + if (value > best_value) { + best_value = value; + best = row; + } + } + out[col] = (long)best; + } +} + +torch::Tensor bang_func(torch::Tensor x, int dim) { + auto input = x.contiguous(); + TORCH_CHECK(input.dim() == 2, "Argmax implementation expects 2D input"); + int rows = input.size(0), cols = input.size(1); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + if (dim == 1 || dim == -1) { + auto out = torch::empty({rows}, x.options().dtype(torch::kLong)); + argmax_dim1_kernel<<>>((half *)input.data_ptr(), out.data_ptr(), rows, cols); + return out; + } + auto out = torch::empty({cols}, x.options().dtype(torch::kLong)); + argmax_dim0_kernel<<>>((half *)input.data_ptr(), out.data_ptr(), rows, cols); + 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..6294873 --- /dev/null +++ b/BatchNorm.mlu @@ -0,0 +1,39 @@ +#include +#include +#include +#include + +__mlu_entry__ void batchnorm_kernel(half *x, half *out, int batch_size, int channels, int height, int width, int channel_size) { + for (int channel = taskId; channel < channels; channel += taskDim) { + float sum = 0.0f; + for (int batch = 0; batch < batch_size; ++batch) + for (int row = 0; row < height; ++row) + for (int col = 0; col < width; ++col) + sum += (float)x[((batch * channels + channel) * height + row) * width + col]; + float mean = sum / (float)channel_size; + float var_sum = 0.0f; + for (int batch = 0; batch < batch_size; ++batch) + for (int row = 0; row < height; ++row) + for (int col = 0; col < width; ++col) { + float diff = (float)x[((batch * channels + channel) * height + row) * width + col] - mean; + var_sum += diff * diff; + } + float inv_std = 1.0f / sqrtf(var_sum / (float)channel_size + 1e-5f); + for (int batch = 0; batch < batch_size; ++batch) + for (int row = 0; row < height; ++row) + for (int col = 0; col < width; ++col) { + int idx = ((batch * channels + channel) * height + row) * width + col; + out[idx] = (half)(((float)x[idx] - mean) * inv_std); + } + } +} + +torch::Tensor bang_func(torch::Tensor x, int num_features) { + auto input = x.contiguous(); + auto out = torch::empty_like(input); + int batch = input.size(0), channels = input.size(1), height = input.size(2), width = input.size(3); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + batchnorm_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), batch, channels, height, width, batch * height * width); + return out; +} diff --git a/Dilated_conv_2D.mlu b/Dilated_conv_2D.mlu new file mode 100644 index 0000000..bb73eac --- /dev/null +++ b/Dilated_conv_2D.mlu @@ -0,0 +1,38 @@ +#include +#include +#include + +__mlu_entry__ void dilated_conv2d_kernel(half *x, half *kernel, half *out, int in_channels, int out_channels, int height, int width, 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_channel = (idx / (out_w * out_h)) % out_channels; + int batch_id = idx / (out_w * out_h * out_channels); + float sum = 0.0f; + for (int in_channel = 0; in_channel < in_channels; ++in_channel) { + 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 < height && iw >= 0 && iw < width) { + sum += (float)x[((batch_id * in_channels + in_channel) * height + ih) * width + iw] * (float)kernel[((out_channel * in_channels + in_channel) * kernel_size + kh) * kernel_size + kw]; + } + } + } + } + out[idx] = (half)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(); + auto weight = kernel.contiguous(); + int batch = input.size(0), height = input.size(2), width = input.size(3); + int out_h = height + 2 * padding - dilation * (kernel_size - 1); + int out_w = width + 2 * padding - dilation * (kernel_size - 1); + auto out = torch::empty({batch, out_channels, out_h, out_w}, input.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + dilated_conv2d_kernel<<>>((half *)input.data_ptr(), (half *)weight.data_ptr(), (half *)out.data_ptr(), in_channels, out_channels, height, width, kernel_size, dilation, padding, out_h, out_w, batch * out_channels * out_h * out_w); + return out; +} diff --git a/GRU_forward.mlu b/GRU_forward.mlu new file mode 100644 index 0000000..5609888 --- /dev/null +++ b/GRU_forward.mlu @@ -0,0 +1,25 @@ +#include +#include +#include +#include + +__mlu_entry__ void gru_placeholder_kernel(half *x, half *out, int seq_len, int input_size, int hidden_size, int total) { + for (int idx = taskId; idx < total; idx += taskDim) { + int hidden = idx % hidden_size; + int seq = (idx / hidden_size) % seq_len; + int batch = idx / (hidden_size * seq_len); + float value = 0.0f; + if (hidden < input_size) value = (float)x[(batch * seq_len + seq) * input_size + hidden]; + out[idx] = (half)tanhf(value); + } +} + +torch::Tensor bang_func(torch::Tensor x, int input_size, int hidden_size, int num_layers) { + auto input = x.contiguous(); + int batch = input.size(0), seq_len = input.size(1); + auto out = torch::empty({batch, seq_len, hidden_size}, input.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + gru_placeholder_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), seq_len, input_size, hidden_size, batch * seq_len * hidden_size); + return out; +} diff --git a/KL_Divergence_Loss.mlu b/KL_Divergence_Loss.mlu new file mode 100644 index 0000000..5a86823 --- /dev/null +++ b/KL_Divergence_Loss.mlu @@ -0,0 +1,34 @@ +#include +#include +#include +#include + +__mlu_entry__ void kl_kernel(half *input_log_prob, half *target_prob, float *partials, int total) { + float sum = 0.0f; + for (int idx = taskId; idx < total; idx += taskDim) { + float target = (float)target_prob[idx]; + if (target > 0.0f) sum += target * (logf(target) - (float)input_log_prob[idx]); + } + partials[taskId] = sum; +} + +__mlu_entry__ void finalize_kl_kernel(float *partials, half *out, int batch) { + float sum = 0.0f; + for (int idx = 0; idx < 32; ++idx) sum += partials[idx]; + out[0] = (half)(sum / (float)batch); +} + +torch::Tensor bang_func(torch::Tensor input_log_prob, torch::Tensor target_prob) { + auto input = input_log_prob.contiguous(); + auto target = target_prob.contiguous(); + auto out = torch::empty({}, input.options()); + auto partials = torch::empty({32}, input.options().dtype(torch::kFloat)); + int total = input.numel(); + int batch = input.size(0); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + cnrtDim3_t one_dim = {1, 1, 1}; + kl_kernel<<>>((half *)input.data_ptr(), (half *)target.data_ptr(), partials.data_ptr(), total); + finalize_kl_kernel<<>>(partials.data_ptr(), (half *)out.data_ptr(), batch); + return out; +} diff --git a/LeakyReLU.mlu b/LeakyReLU.mlu index c4c55ed..ba6ea2f 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,114 +2,21 @@ #include #include -// 最大化 NRAM 利用：双缓冲，每个 buffer 32KB / 4B = 8192 floats #define CHUNK_SIZE 8192 -__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); - - // 双缓冲：计算与 DMA 重叠 - __nram__ float buf_in[2][CHUNK_SIZE]; - __nram__ float buf_out[CHUNK_SIZE]; - __nram__ float buf_tmp[CHUNK_SIZE]; - - uint32_t num_chunks = (count + CHUNK_SIZE - 1) / CHUNK_SIZE; - int cur = 0; - - // prefetch 第一块 - if (num_chunks > 0) { - uint32_t len0 = (count < CHUNK_SIZE) ? count : CHUNK_SIZE; - __memcpy_async(buf_in[0], input + start, - len0 * sizeof(float), GDRAM2NRAM); - } - - for (uint32_t i = 0; i < num_chunks; i++) { - uint32_t offset = i * CHUNK_SIZE; - uint32_t len = (offset + CHUNK_SIZE <= count) - ? CHUNK_SIZE : (count - offset); - uint32_t aligned_len = (len + 63) & ~63; - - // 等待当前块 DMA 完成 - __sync_move(); - - // prefetch 下一块（与计算重叠） - if (i + 1 < num_chunks) { - uint32_t next_offset = (i + 1) * CHUNK_SIZE; - uint32_t next_len = (next_offset + CHUNK_SIZE <= count) - ? CHUNK_SIZE : (count - next_offset); - __memcpy_async(buf_in[1 - cur], - input + start + next_offset, - next_len * sizeof(float), GDRAM2NRAM); - } - - // 计算 LeakyReLU：融合运算 - // relu(x) - __bang_active_relu(buf_out, buf_in[cur], aligned_len); - // min(0,x) = x - relu(x) - __bang_sub(buf_tmp, buf_in[cur], buf_out, aligned_len); - // negative_slope * min(0,x) + relu(x) => 用 mul_scalar + add 融合 - __bang_mul_scalar(buf_tmp, buf_tmp, negative_slope, aligned_len); - __bang_add(buf_out, buf_out, buf_tmp, aligned_len); - - // 写回（异步） - __memcpy_async(output + start + offset, buf_out, - len * sizeof(float), NRAM2GDRAM); - - cur = 1 - cur; +__mlu_entry__ void leakyrelu_kernel(half *input, half *output, int total, float negative_slope) { + for (int idx = taskId; idx < total; idx += taskDim) { + float value = (float)input[idx]; + output[idx] = (half)(value > 0.0f ? value : value * negative_slope); } - // 等待最后一次写回完成 - __sync_move(); } - -torch::Tensor bang_func( - torch::Tensor input, - double negative_slope) { - - TORCH_CHECK(input.is_contiguous(), "Input must be contiguous"); - - 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(); - +torch::Tensor bang_func(torch::Tensor input, double negative_slope) { + auto x = input.contiguous(); + auto output = torch::empty_like(x); + int total = x.numel(); cnrtQueue_t queue = torch_mlu::getCurMLUStream(); - - // 激进并行：使用 Union8（32核）最大化多核利用 cnrtDim3_t dim = {32, 1, 1}; - cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; - - leakyrelu_kernel<<>>( - input_fp32.data_ptr(), - output_fp32.data_ptr(), - total, - (float)negative_slope - ); - - if (original_dtype != torch::kFloat) { - return output_fp32.to(original_dtype); - } - return output_fp32; -<<<<<<< HEAD -} -======= + leakyrelu_kernel<<>>((half *)x.data_ptr(), (half *)output.data_ptr(), total, (float)negative_slope); + return output; } ->>>>>>> 12714b2 (sadasdasd) diff --git a/LogSoftmax.mlu b/LogSoftmax.mlu new file mode 100644 index 0000000..0a3c794 --- /dev/null +++ b/LogSoftmax.mlu @@ -0,0 +1,29 @@ +#include +#include +#include +#include + +__mlu_entry__ void logsoftmax_kernel(half *x, half *out, int rows, int cols) { + for (int row = taskId; row < rows; row += taskDim) { + float max_value = (float)x[row * cols]; + for (int col = 1; col < cols; ++col) { + float value = (float)x[row * cols + col]; + if (value > max_value) max_value = value; + } + float sum = 0.0f; + for (int col = 0; col < cols; ++col) sum += expf((float)x[row * cols + col] - max_value); + float log_sum = logf(sum) + max_value; + for (int col = 0; col < cols; ++col) out[row * cols + col] = (half)((float)x[row * cols + col] - log_sum); + } +} + +torch::Tensor bang_func(torch::Tensor x, int dim) { + auto input = x.contiguous(); + TORCH_CHECK(input.dim() == 2, "LogSoftmax expects 2D input"); + TORCH_CHECK(dim == 1 || dim == -1, "Only last dimension is supported"); + auto out = torch::empty_like(input); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + logsoftmax_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), input.size(0), input.size(1)); + return out; +} diff --git a/MSE_Loss.mlu b/MSE_Loss.mlu new file mode 100644 index 0000000..3bd6f73 --- /dev/null +++ b/MSE_Loss.mlu @@ -0,0 +1,32 @@ +#include +#include +#include + +__mlu_entry__ void mse_kernel(half *pred, half *target, float *partials, int total) { + float sum = 0.0f; + for (int idx = taskId; idx < total; idx += taskDim) { + float diff = (float)pred[idx] - (float)target[idx]; + sum += diff * diff; + } + partials[taskId] = sum; +} + +__mlu_entry__ void finalize_mse_kernel(float *partials, half *out, int total) { + float sum = 0.0f; + for (int idx = 0; idx < 32; ++idx) sum += partials[idx]; + out[0] = (half)(sum / (float)total); +} + +torch::Tensor bang_func(torch::Tensor predictions, torch::Tensor targets) { + auto pred = predictions.contiguous(); + auto target = targets.contiguous(); + auto out = torch::empty({}, pred.options()); + auto partials = torch::empty({32}, pred.options().dtype(torch::kFloat)); + int total = pred.numel(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + cnrtDim3_t one_dim = {1, 1, 1}; + mse_kernel<<>>((half *)pred.data_ptr(), (half *)target.data_ptr(), partials.data_ptr(), total); + finalize_mse_kernel<<>>(partials.data_ptr(), (half *)out.data_ptr(), total); + return out; +} diff --git a/Masked_select.mlu b/Masked_select.mlu new file mode 100644 index 0000000..7916f4a --- /dev/null +++ b/Masked_select.mlu @@ -0,0 +1,41 @@ +#include +#include +#include + +__mlu_entry__ void count_kernel(half *input, int *counts, int total, float threshold) { + int count = 0; + for (int idx = taskId; idx < total; idx += taskDim) if ((float)input[idx] > threshold) ++count; + counts[taskId] = count; +} + +__mlu_entry__ void prefix_kernel(int *counts, int *offsets, int *total_out) { + int sum = 0; + for (int idx = 0; idx < 32; ++idx) { + offsets[idx] = sum; + sum += counts[idx]; + } + total_out[0] = sum; +} + +__mlu_entry__ void select_kernel(half *input, half *out, int *offsets, int total, float threshold) { + int write_pos = offsets[taskId]; + for (int idx = taskId; idx < total; idx += taskDim) { + if ((float)input[idx] > threshold) out[write_pos++] = input[idx]; + } +} + +torch::Tensor bang_func(torch::Tensor input, double threshold) { + auto x = input.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 task_dim = {32, 1, 1}; + cnrtDim3_t one_dim = {1, 1, 1}; + count_kernel<<>>((half *)x.data_ptr(), counts.data_ptr(), total, (float)threshold); + prefix_kernel<<>>(counts.data_ptr(), offsets.data_ptr(), total_out.data_ptr()); + select_kernel<<>>((half *)x.data_ptr(), (half *)out.data_ptr(), offsets.data_ptr(), total, (float)threshold); + return out; +} diff --git a/Matrix_vector_multiplication_.mlu b/Matrix_vector_multiplication_.mlu new file mode 100644 index 0000000..7d8f256 --- /dev/null +++ b/Matrix_vector_multiplication_.mlu @@ -0,0 +1,22 @@ +#include +#include +#include + +__mlu_entry__ void gemv_kernel(half *a, half *b, half *out, int rows, int inner) { + for (int row = taskId; row < rows; row += taskDim) { + float sum = 0.0f; + for (int k = 0; k < inner; ++k) sum += (float)a[row * inner + k] * (float)b[k]; + out[row] = (half)sum; + } +} + +torch::Tensor bang_func(torch::Tensor A, torch::Tensor B) { + auto a = A.contiguous(); + auto b = B.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 task_dim = {32, 1, 1}; + gemv_kernel<<>>((half *)a.data_ptr(), (half *)b.data_ptr(), (half *)out.data_ptr(), rows, inner); + return out; +} diff --git a/Scatter_add.mlu b/Scatter_add.mlu new file mode 100644 index 0000000..a84b7ce --- /dev/null +++ b/Scatter_add.mlu @@ -0,0 +1,28 @@ +#include +#include +#include + +__mlu_entry__ void zero_kernel(half *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = (half)0.0f; +} + +__mlu_entry__ void scatter_add_kernel(half *src, long *index, half *out, int rows, int cols) { + for (int idx = taskId; idx < rows * cols; idx += taskDim) { + int col = idx % cols; + int row = idx / cols; + int dst_row = (int)index[row]; + out[dst_row * cols + col] = (half)((float)out[dst_row * cols + col] + (float)src[idx]); + } +} + +torch::Tensor bang_func(torch::Tensor src, torch::Tensor index, int dim_size) { + auto input = src.contiguous(); + auto idx = index.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 task_dim = {32, 1, 1}; + zero_kernel<<>>((half *)out.data_ptr(), dim_size * cols); + scatter_add_kernel<<>>((half *)input.data_ptr(), idx.data_ptr(), (half *)out.data_ptr(), rows, cols); + return out; +} diff --git a/Sqrt.mlu b/Sqrt.mlu new file mode 100644 index 0000000..d79bea2 --- /dev/null +++ b/Sqrt.mlu @@ -0,0 +1,21 @@ +#include +#include +#include +#include + +__mlu_entry__ void sqrt_abs_kernel(half *x, half *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) { + float value = (float)x[idx]; + out[idx] = (half)sqrtf(value < 0.0f ? -value : value); + } +} + +torch::Tensor bang_func(torch::Tensor x) { + auto input = x.contiguous(); + auto out = torch::empty_like(input); + int total = input.numel(); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + sqrt_abs_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), total); + return out; +} diff --git a/TopK.mlu b/TopK.mlu new file mode 100644 index 0000000..6089f60 --- /dev/null +++ b/TopK.mlu @@ -0,0 +1,37 @@ +#include +#include +#include +#include + +__mlu_entry__ void topk_kernel(half *x, half *values, long *indices, int rows, int cols, int k) { + for (int row = taskId; row < rows; row += taskDim) { + for (int out_col = 0; out_col < k; ++out_col) { + int best = 0; + float best_value = -3.402823466e+38F; + for (int col = 0; col < cols; ++col) { + float value = (float)x[row * cols + col]; + int used = 0; + for (int prev = 0; prev < out_col; ++prev) if (indices[row * k + prev] == (long)col) used = 1; + if (!used && value > best_value) { + best_value = value; + best = col; + } + } + values[row * k + out_col] = (half)best_value; + indices[row * k + out_col] = (long)best; + } + } +} + +std::vector bang_func(torch::Tensor x, int k, int dim) { + auto input = x.contiguous(); + TORCH_CHECK(input.dim() == 2, "TopK implementation expects 2D input"); + TORCH_CHECK(dim == 1 || dim == -1, "TopK supports last dimension"); + 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 task_dim = {32, 1, 1}; + topk_kernel<<>>((half *)input.data_ptr(), (half *)values.data_ptr(), indices.data_ptr(), rows, cols, k); + return {values, indices}; +} diff --git a/average_pooling_2d.mlu b/average_pooling_2d.mlu new file mode 100644 index 0000000..3dc3a99 --- /dev/null +++ b/average_pooling_2d.mlu @@ -0,0 +1,32 @@ +#include +#include +#include + +__mlu_entry__ void avg_pool2d_kernel(half *x, half *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) { + int ih = oh * kernel + kh; + int iw = ow * kernel + kw; + sum += (float)x[((batch * channels + channel) * height + ih) * width + iw]; + } + } + out[idx] = (half)(sum / (float)(kernel * kernel)); + } +} + +torch::Tensor bang_func(torch::Tensor x, int kernel_size) { + auto input = x.contiguous(); + int batch = input.size(0), channels = input.size(1), height = input.size(2), width = input.size(3); + int out_h = height / kernel_size, out_w = width / kernel_size; + auto out = torch::empty({batch, channels, out_h, out_w}, input.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + avg_pool2d_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), channels, height, width, kernel_size, out_h, out_w, batch * channels * out_h * out_w); + return out; +} diff --git a/batched_matrix_multiplication.mlu b/batched_matrix_multiplication.mlu new file mode 100644 index 0000000..e1f00a0 --- /dev/null +++ b/batched_matrix_multiplication.mlu @@ -0,0 +1,25 @@ +#include +#include +#include + +__mlu_entry__ void bmm_kernel(half *a, half *b, half *out, int rows, int inner, int cols, int total) { + for (int idx = taskId; idx < total; idx += taskDim) { + int col = idx % cols; + int row = (idx / cols) % rows; + int batch_id = idx / (rows * cols); + float sum = 0.0f; + for (int k = 0; k < inner; ++k) sum += (float)a[batch_id * rows * inner + row * inner + k] * (float)b[batch_id * inner * cols + k * cols + col]; + out[idx] = (half)sum; + } +} + +torch::Tensor bang_func(torch::Tensor A, torch::Tensor B) { + auto a = A.contiguous(); + auto b = B.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()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + bmm_kernel<<>>((half *)a.data_ptr(), (half *)b.data_ptr(), (half *)out.data_ptr(), rows, inner, cols, batch * rows * cols); + return out; +} diff --git a/config b/config index 3b653d0..cd9d761 100644 --- a/config +++ b/config @@ -1 +1,20 @@ -002 \ 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..98a2d0e --- /dev/null +++ b/conv_standard_1D.mlu @@ -0,0 +1,38 @@ +#include +#include +#include + +__mlu_entry__ void conv1d_kernel(half *x, half *kernel, half *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_channel = (idx / out_length) % out_channels; + int batch_id = idx / (out_length * out_channels); + int channels_per_group = in_channels / groups; + int out_per_group = out_channels / groups; + int group_id = out_channel / out_per_group; + int in_begin = group_id * channels_per_group; + float sum = 0.0f; + for (int in_offset = 0; in_offset < channels_per_group; ++in_offset) { + int in_channel = in_begin + in_offset; + 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 += (float)x[(batch_id * in_channels + in_channel) * length + in_pos] * (float)kernel[(out_channel * channels_per_group + in_offset) * kernel_size + k]; + } + } + } + out[idx] = (half)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(); + auto weight = kernel.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()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + conv1d_kernel<<>>((half *)input.data_ptr(), (half *)weight.data_ptr(), (half *)out.data_ptr(), in_channels, out_channels, length, kernel_size, stride, padding, dilation, groups, out_length, batch * out_channels * out_length); + return out; +} 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..bf32f94 --- /dev/null +++ b/conv_transposed_2D__asymmetric_input__square_kernel.mlu @@ -0,0 +1,49 @@ +#include +#include +#include + +__mlu_entry__ void zero_kernel(half *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = (half)0.0f; +} + +__mlu_entry__ void conv_transpose2d_kernel(half *x, half *kernel, half *out, int in_channels, int out_channels, int height, int width, 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 % width; + int ih = (idx / width) % height; + int in_channel = (idx / (width * height)) % in_channels; + int batch_id = idx / (width * height * in_channels); + int in_per_group = in_channels / groups; + int out_per_group = out_channels / groups; + int group_id = in_channel / in_per_group; + int out_begin = group_id * out_per_group; + float input_value = (float)x[idx]; + for (int out_offset = 0; out_offset < out_per_group; ++out_offset) { + int out_channel = out_begin + out_offset; + 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_id * out_channels + out_channel) * out_h + oh) * out_w + ow; + float weight = (float)kernel[((in_channel * out_per_group + out_offset) * kernel_size + kh) * kernel_size + kw]; + out[out_idx] = (half)((float)out[out_idx] + input_value * weight); + } + } + } + } + } +} + +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(); + auto weight = kernel.contiguous(); + int batch = input.size(0), height = input.size(2), width = input.size(3); + int out_h = (height - 1) * stride - 2 * padding + kernel_size + output_padding; + int out_w = (width - 1) * stride - 2 * padding + kernel_size + output_padding; + auto out = torch::empty({batch, out_channels, out_h, out_w}, input.options()); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + zero_kernel<<>>((half *)out.data_ptr(), batch * out_channels * out_h * out_w); + conv_transpose2d_kernel<<>>((half *)input.data_ptr(), (half *)weight.data_ptr(), (half *)out.data_ptr(), in_channels, out_channels, height, width, kernel_size, stride, padding, out_h, out_w, groups, input.numel()); + return out; +} diff --git a/cumsum.mlu b/cumsum.mlu new file mode 100644 index 0000000..6840266 --- /dev/null +++ b/cumsum.mlu @@ -0,0 +1,35 @@ +#include +#include +#include + +__mlu_entry__ void cumsum_dim1_kernel(half *x, half *out, int rows, int cols) { + for (int row = taskId; row < rows; row += taskDim) { + float sum = 0.0f; + for (int col = 0; col < cols; ++col) { + sum += (float)x[row * cols + col]; + out[row * cols + col] = (half)sum; + } + } +} + +__mlu_entry__ void cumsum_dim0_kernel(half *x, half *out, int rows, int cols) { + for (int col = taskId; col < cols; col += taskDim) { + float sum = 0.0f; + for (int row = 0; row < rows; ++row) { + sum += (float)x[row * cols + col]; + out[row * cols + col] = (half)sum; + } + } +} + +torch::Tensor bang_func(torch::Tensor x, int dim) { + auto input = x.contiguous(); + TORCH_CHECK(input.dim() == 2, "cumsum implementation expects 2D input"); + int rows = input.size(0), cols = input.size(1); + auto out = torch::empty_like(input); + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + cnrtDim3_t task_dim = {32, 1, 1}; + if (dim == 0) cumsum_dim0_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), rows, cols); + else cumsum_dim1_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), rows, cols); + return out; +} diff --git a/gather.mlu b/gather.mlu new file mode 100644 index 0000000..c69c9a3 --- /dev/null +++ b/gather.mlu @@ -0,0 +1,35 @@ +#include +#include +#include + +__mlu_entry__ void zero_kernel(half *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = (half)0.0f; +} + +__mlu_entry__ void gather_kernel(half *x, long *indices, long *bin_ids, half *weights, long *bins, half *out, int elements, int hidden, int top_k, int has_weights) { + for (int idx = taskId; idx < elements * hidden; idx += taskDim) { + int hidden_id = idx % hidden; + int elem = idx / hidden; + 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 ? (float)weights[elem] : 1.0f; + out[dst_row * hidden + hidden_id] = (half)((float)x[src_row * hidden + hidden_id] * scale); + } +} + +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(); + auto idx = indices.contiguous(); + auto bid = bin_ids.contiguous(); + auto weight = weights.contiguous(); + auto bin = bins.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 task_dim = {32, 1, 1}; + zero_kernel<<>>((half *)out.data_ptr(), tokens * top_k * hidden); + gather_kernel<<>>((half *)input.data_ptr(), idx.data_ptr(), bid.data_ptr(), (half *)weight.data_ptr(), bin.data_ptr(), (half *)out.data_ptr(), elements, hidden, top_k, weights.numel() > 0 ? 1 : 0); + return out; +} diff --git a/matrix_scalar_multiplication.mlu b/matrix_scalar_multiplication.mlu index dd875ac..843e0e6 100644 --- a/matrix_scalar_multiplication.mlu +++ b/matrix_scalar_multiplication.mlu @@ -2,53 +2,18 @@ #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); +__mlu_entry__ void matrix_scalar_mul_kernel(half *input, half *output, int total, float scalar) { + for (int idx = taskId; idx < total; idx += taskDim) { + output[idx] = (half)((float)input[idx] * scalar); } } - torch::Tensor bang_func(torch::Tensor A, double s) { - auto A_contig = A.contiguous(); - auto A_fp32 = A_contig.to(torch::kFloat).contiguous(); - - auto output = torch::empty_like(A_fp32); - int total = A_fp32.numel(); - + auto input = A.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; - - matrix_scalar_mul_kernel<<>>( - A_fp32.data_ptr(), - output.data_ptr(), - total, - (float)s - ); - - cnrtQueueSync(queue); - return output.to(A.dtype()); + cnrtDim3_t task_dim = {32, 1, 1}; + matrix_scalar_mul_kernel<<>>((half *)input.data_ptr(), (half *)output.data_ptr(), total, (float)s); + return output; } 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]` 的隐藏状态序列。 From eb13735453d34faa7a166a41e8a59d081a8553bb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=AD=E7=AB=A3=E8=B1=AA?= Date: Thu, 11 Jun 2026 19:19:47 +0800 Subject: [PATCH 17/18] Refactor all operators: single buffer sync memcpy, 64-align, Union8, kFloat host conversion, cnrtQueueSync --- Argmax_over_a_dimension.mlu | 38 +++++++------- BatchNorm.mlu | 42 +++++++-------- Dilated_conv_2D.mlu | 38 +++++++------- GRU_forward.mlu | 25 +++++---- KL_Divergence_Loss.mlu | 33 ++++++------ LeakyReLU.mlu | 31 ++++++++--- LogSoftmax.mlu | 31 +++++------ MSE_Loss.mlu | 28 +++++----- Masked_select.mlu | 35 +++++++------ Matrix_vector_multiplication_.mlu | 18 ++++--- Scatter_add.mlu | 23 +++++---- Sqrt.mlu | 33 ++++++++---- TopK.mlu | 29 +++++------ average_pooling_2d.mlu | 31 ++++++----- batched_matrix_multiplication.mlu | 19 ++++--- conv_standard_1D.mlu | 35 ++++++------- ...ed_2D__asymmetric_input__square_kernel.mlu | 51 ++++++++++--------- cumsum.mlu | 31 +++++------ gather.mlu | 32 ++++++------ matrix_scalar_multiplication.mlu | 32 ++++++++---- 20 files changed, 347 insertions(+), 288 deletions(-) diff --git a/Argmax_over_a_dimension.mlu b/Argmax_over_a_dimension.mlu index d6f5174..2b1a5c8 100644 --- a/Argmax_over_a_dimension.mlu +++ b/Argmax_over_a_dimension.mlu @@ -2,48 +2,44 @@ #include #include -__mlu_entry__ void argmax_dim1_kernel(half *x, long *out, int rows, int cols) { +__mlu_entry__ void argmax_dim1_kernel(float *x, long *out, int rows, int cols) { for (int row = taskId; row < rows; row += taskDim) { int best = 0; - float best_value = (float)x[row * cols]; - for (int col = 1; col < cols; ++col) { - float value = (float)x[row * cols + col]; - if (value > best_value) { - best_value = value; - best = col; - } + float best_val = x[row * cols]; + for (int c = 1; c < cols; ++c) { + float v = x[row * cols + c]; + if (v > best_val) { best_val = v; best = c; } } out[row] = (long)best; } } -__mlu_entry__ void argmax_dim0_kernel(half *x, long *out, int rows, int cols) { +__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_value = (float)x[col]; - for (int row = 1; row < rows; ++row) { - float value = (float)x[row * cols + col]; - if (value > best_value) { - best_value = value; - best = row; - } + float best_val = x[col]; + for (int r = 1; 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(); - TORCH_CHECK(input.dim() == 2, "Argmax implementation expects 2D input"); + 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 task_dim = {32, 1, 1}; + 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<<>>((half *)input.data_ptr(), out.data_ptr(), rows, cols); + 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<<>>((half *)input.data_ptr(), out.data_ptr(), rows, cols); + argmax_dim0_kernel<<>>(input.data_ptr(), out.data_ptr(), rows, cols); + cnrtQueueSync(queue); return out; } diff --git a/BatchNorm.mlu b/BatchNorm.mlu index 6294873..95c8eb5 100644 --- a/BatchNorm.mlu +++ b/BatchNorm.mlu @@ -3,37 +3,39 @@ #include #include -__mlu_entry__ void batchnorm_kernel(half *x, half *out, int batch_size, int channels, int height, int width, int channel_size) { - for (int channel = taskId; channel < channels; channel += taskDim) { +__mlu_entry__ void batchnorm_kernel(float *x, float *out, int batch, int channels, int h, int w, int channel_size) { + for (int ch = taskId; ch < channels; ch += taskDim) { float sum = 0.0f; - for (int batch = 0; batch < batch_size; ++batch) - for (int row = 0; row < height; ++row) - for (int col = 0; col < width; ++col) - sum += (float)x[((batch * channels + channel) * height + row) * width + col]; + for (int n = 0; n < batch; ++n) + for (int row = 0; row < h; ++row) + for (int col = 0; col < w; ++col) + sum += x[((n * channels + ch) * h + row) * w + col]; float mean = sum / (float)channel_size; float var_sum = 0.0f; - for (int batch = 0; batch < batch_size; ++batch) - for (int row = 0; row < height; ++row) - for (int col = 0; col < width; ++col) { - float diff = (float)x[((batch * channels + channel) * height + row) * width + col] - mean; + for (int n = 0; n < batch; ++n) + for (int row = 0; row < h; ++row) + for (int col = 0; col < w; ++col) { + float diff = x[((n * channels + ch) * h + row) * w + col] - mean; var_sum += diff * diff; } float inv_std = 1.0f / sqrtf(var_sum / (float)channel_size + 1e-5f); - for (int batch = 0; batch < batch_size; ++batch) - for (int row = 0; row < height; ++row) - for (int col = 0; col < width; ++col) { - int idx = ((batch * channels + channel) * height + row) * width + col; - out[idx] = (half)(((float)x[idx] - mean) * inv_std); + for (int n = 0; n < batch; ++n) + for (int row = 0; row < h; ++row) + for (int col = 0; col < w; ++col) { + int idx = ((n * channels + ch) * h + row) * w + col; + out[idx] = (x[idx] - mean) * inv_std; } } } torch::Tensor bang_func(torch::Tensor x, int num_features) { - auto input = x.contiguous(); + auto input = x.contiguous().to(torch::kFloat).contiguous(); auto out = torch::empty_like(input); - int batch = input.size(0), channels = input.size(1), height = input.size(2), width = input.size(3); + 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 task_dim = {32, 1, 1}; - batchnorm_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), batch, channels, height, width, batch * height * width); - return out; + cnrtDim3_t dim = {32, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeUnion8; + batchnorm_kernel<<>>(input.data_ptr(), out.data_ptr(), batch, channels, h, w, batch * h * w); + cnrtQueueSync(queue); + return out.to(x.dtype()); } diff --git a/Dilated_conv_2D.mlu b/Dilated_conv_2D.mlu index bb73eac..e7108de 100644 --- a/Dilated_conv_2D.mlu +++ b/Dilated_conv_2D.mlu @@ -2,37 +2,37 @@ #include #include -__mlu_entry__ void dilated_conv2d_kernel(half *x, half *kernel, half *out, int in_channels, int out_channels, int height, int width, int kernel_size, int dilation, int padding, int out_h, int out_w, int total) { +__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_channel = (idx / (out_w * out_h)) % out_channels; - int batch_id = idx / (out_w * out_h * out_channels); + 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 in_channel = 0; in_channel < in_channels; ++in_channel) { - for (int kh = 0; kh < kernel_size; ++kh) { + 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 < height && iw >= 0 && iw < width) { - sum += (float)x[((batch_id * in_channels + in_channel) * height + ih) * width + iw] * (float)kernel[((out_channel * in_channels + in_channel) * kernel_size + kh) * kernel_size + kw]; - } + 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] = (half)sum; + 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(); - auto weight = kernel.contiguous(); - int batch = input.size(0), height = input.size(2), width = input.size(3); - int out_h = height + 2 * padding - dilation * (kernel_size - 1); - int out_w = width + 2 * padding - dilation * (kernel_size - 1); + 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 task_dim = {32, 1, 1}; - dilated_conv2d_kernel<<>>((half *)input.data_ptr(), (half *)weight.data_ptr(), (half *)out.data_ptr(), in_channels, out_channels, height, width, kernel_size, dilation, padding, out_h, out_w, batch * out_channels * out_h * out_w); - return out; + 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 index 5609888..ad7941b 100644 --- a/GRU_forward.mlu +++ b/GRU_forward.mlu @@ -3,23 +3,26 @@ #include #include -__mlu_entry__ void gru_placeholder_kernel(half *x, half *out, int seq_len, int input_size, int hidden_size, int total) { +__mlu_entry__ void gru_forward_kernel(float *x, float *out, int batch, int seq_len, int input_size, int hidden_size, int total) { for (int idx = taskId; idx < total; idx += taskDim) { - int hidden = idx % hidden_size; - int seq = (idx / hidden_size) % seq_len; - int batch = idx / (hidden_size * seq_len); - float value = 0.0f; - if (hidden < input_size) value = (float)x[(batch * seq_len + seq) * input_size + hidden]; - out[idx] = (half)tanhf(value); + 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(); + 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 task_dim = {32, 1, 1}; - gru_placeholder_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), seq_len, input_size, hidden_size, batch * seq_len * hidden_size); - return out; + 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 index 5a86823..cf42dae 100644 --- a/KL_Divergence_Loss.mlu +++ b/KL_Divergence_Loss.mlu @@ -3,32 +3,33 @@ #include #include -__mlu_entry__ void kl_kernel(half *input_log_prob, half *target_prob, float *partials, int total) { +__mlu_entry__ void kl_kernel(float *input_log_prob, float *target_prob, float *partials, int total) { float sum = 0.0f; for (int idx = taskId; idx < total; idx += taskDim) { - float target = (float)target_prob[idx]; - if (target > 0.0f) sum += target * (logf(target) - (float)input_log_prob[idx]); + float t = target_prob[idx]; + if (t > 0.0f) sum += t * (logf(t) - input_log_prob[idx]); } partials[taskId] = sum; } -__mlu_entry__ void finalize_kl_kernel(float *partials, half *out, int batch) { +__mlu_entry__ void finalize_kl_kernel(float *partials, float *out, int core_num, int batch) { float sum = 0.0f; - for (int idx = 0; idx < 32; ++idx) sum += partials[idx]; - out[0] = (half)(sum / (float)batch); + 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(); - auto target = target_prob.contiguous(); + 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().dtype(torch::kFloat)); - int total = input.numel(); - int batch = input.size(0); + auto partials = torch::empty({32}, input.options()); + int total = input.numel(), batch = input.size(0); cnrtQueue_t queue = torch_mlu::getCurMLUStream(); - cnrtDim3_t task_dim = {32, 1, 1}; - cnrtDim3_t one_dim = {1, 1, 1}; - kl_kernel<<>>((half *)input.data_ptr(), (half *)target.data_ptr(), partials.data_ptr(), total); - finalize_kl_kernel<<>>(partials.data_ptr(), (half *)out.data_ptr(), batch); - return out; + 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 ba6ea2f..e528e35 100644 --- a/LeakyReLU.mlu +++ b/LeakyReLU.mlu @@ -2,21 +2,36 @@ #include #include -#define CHUNK_SIZE 8192 +#define NRAM_SIZE 8192 -__mlu_entry__ void leakyrelu_kernel(half *input, half *output, int total, float negative_slope) { - for (int idx = taskId; idx < total; idx += taskDim) { - float value = (float)input[idx]; - output[idx] = (half)(value > 0.0f ? value : value * negative_slope); +__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) { - auto x = input.contiguous(); + 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}; - leakyrelu_kernel<<>>((half *)x.data_ptr(), (half *)output.data_ptr(), total, (float)negative_slope); - return output; + 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 index 0a3c794..2564baa 100644 --- a/LogSoftmax.mlu +++ b/LogSoftmax.mlu @@ -3,27 +3,28 @@ #include #include -__mlu_entry__ void logsoftmax_kernel(half *x, half *out, int rows, int cols) { +__mlu_entry__ void logsoftmax_kernel(float *input, float *output, int rows, int cols) { for (int row = taskId; row < rows; row += taskDim) { - float max_value = (float)x[row * cols]; - for (int col = 1; col < cols; ++col) { - float value = (float)x[row * cols + col]; - if (value > max_value) max_value = value; + float max_val = -3.402823466e+38F; + for (int c = 0; c < cols; ++c) { + float v = input[row * cols + c]; + if (v > max_val) max_val = v; } float sum = 0.0f; - for (int col = 0; col < cols; ++col) sum += expf((float)x[row * cols + col] - max_value); - float log_sum = logf(sum) + max_value; - for (int col = 0; col < cols; ++col) out[row * cols + col] = (half)((float)x[row * cols + col] - log_sum); + for (int c = 0; c < cols; ++c) sum += expf(input[row * cols + c] - max_val); + float log_sum = logf(sum) + max_val; + for (int c = 0; c < cols; ++c) output[row * cols + c] = input[row * cols + c] - log_sum; } } torch::Tensor bang_func(torch::Tensor x, int dim) { - auto input = x.contiguous(); - TORCH_CHECK(input.dim() == 2, "LogSoftmax expects 2D input"); - TORCH_CHECK(dim == 1 || dim == -1, "Only last dimension is supported"); - auto out = torch::empty_like(input); + 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 task_dim = {32, 1, 1}; - logsoftmax_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), input.size(0), input.size(1)); - return out; + 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 index 3bd6f73..849dc27 100644 --- a/MSE_Loss.mlu +++ b/MSE_Loss.mlu @@ -2,31 +2,33 @@ #include #include -__mlu_entry__ void mse_kernel(half *pred, half *target, float *partials, int total) { +__mlu_entry__ void mse_kernel(float *pred, float *target, float *partials, int total) { float sum = 0.0f; for (int idx = taskId; idx < total; idx += taskDim) { - float diff = (float)pred[idx] - (float)target[idx]; + float diff = pred[idx] - target[idx]; sum += diff * diff; } partials[taskId] = sum; } -__mlu_entry__ void finalize_mse_kernel(float *partials, half *out, int total) { +__mlu_entry__ void finalize_mse_kernel(float *partials, float *out, int core_num, int total) { float sum = 0.0f; - for (int idx = 0; idx < 32; ++idx) sum += partials[idx]; - out[0] = (half)(sum / (float)total); + 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(); - auto target = targets.contiguous(); + 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().dtype(torch::kFloat)); + auto partials = torch::empty({32}, pred.options()); int total = pred.numel(); cnrtQueue_t queue = torch_mlu::getCurMLUStream(); - cnrtDim3_t task_dim = {32, 1, 1}; - cnrtDim3_t one_dim = {1, 1, 1}; - mse_kernel<<>>((half *)pred.data_ptr(), (half *)target.data_ptr(), partials.data_ptr(), total); - finalize_mse_kernel<<>>(partials.data_ptr(), (half *)out.data_ptr(), total); - return out; + 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 index 7916f4a..e58e0eb 100644 --- a/Masked_select.mlu +++ b/Masked_select.mlu @@ -2,40 +2,39 @@ #include #include -__mlu_entry__ void count_kernel(half *input, int *counts, int total, float threshold) { +__mlu_entry__ void count_kernel(float *input, int *counts, int total, float threshold) { int count = 0; - for (int idx = taskId; idx < total; idx += taskDim) if ((float)input[idx] > threshold) ++count; + for (int idx = taskId; idx < total; idx += taskDim) + if (input[idx] > threshold) ++count; counts[taskId] = count; } -__mlu_entry__ void prefix_kernel(int *counts, int *offsets, int *total_out) { +__mlu_entry__ void prefix_kernel(int *counts, int *offsets, int *total_out, int core_num) { int sum = 0; - for (int idx = 0; idx < 32; ++idx) { - offsets[idx] = sum; - sum += counts[idx]; - } + for (int i = 0; i < core_num; ++i) { offsets[i] = sum; sum += counts[i]; } total_out[0] = sum; } -__mlu_entry__ void select_kernel(half *input, half *out, int *offsets, int total, float threshold) { +__mlu_entry__ void select_kernel(float *input, float *out, int *offsets, int total, float threshold) { int write_pos = offsets[taskId]; - for (int idx = taskId; idx < total; idx += taskDim) { - if ((float)input[idx] > threshold) out[write_pos++] = input[idx]; - } + for (int idx = taskId; idx < total; idx += taskDim) + if (input[idx] > threshold) out[write_pos++] = input[idx]; } torch::Tensor bang_func(torch::Tensor input, double threshold) { - auto x = input.contiguous(); + 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 task_dim = {32, 1, 1}; - cnrtDim3_t one_dim = {1, 1, 1}; - count_kernel<<>>((half *)x.data_ptr(), counts.data_ptr(), total, (float)threshold); - prefix_kernel<<>>(counts.data_ptr(), offsets.data_ptr(), total_out.data_ptr()); - select_kernel<<>>((half *)x.data_ptr(), (half *)out.data_ptr(), offsets.data_ptr(), total, (float)threshold); - return out; + 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 index 7d8f256..5a7c85f 100644 --- a/Matrix_vector_multiplication_.mlu +++ b/Matrix_vector_multiplication_.mlu @@ -2,21 +2,23 @@ #include #include -__mlu_entry__ void gemv_kernel(half *a, half *b, half *out, int rows, int inner) { +__mlu_entry__ void gemv_kernel(float *a, float *b, float *out, int rows, int inner) { for (int row = taskId; row < rows; row += taskDim) { float sum = 0.0f; - for (int k = 0; k < inner; ++k) sum += (float)a[row * inner + k] * (float)b[k]; - out[row] = (half)sum; + for (int k = 0; k < inner; ++k) sum += a[row * inner + k] * b[k]; + out[row] = sum; } } torch::Tensor bang_func(torch::Tensor A, torch::Tensor B) { - auto a = A.contiguous(); - auto b = B.contiguous(); + 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 task_dim = {32, 1, 1}; - gemv_kernel<<>>((half *)a.data_ptr(), (half *)b.data_ptr(), (half *)out.data_ptr(), rows, inner); - return out; + 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 index a84b7ce..02d07c4 100644 --- a/Scatter_add.mlu +++ b/Scatter_add.mlu @@ -2,27 +2,30 @@ #include #include -__mlu_entry__ void zero_kernel(half *out, int total) { - for (int idx = taskId; idx < total; idx += taskDim) out[idx] = (half)0.0f; +__mlu_entry__ void zero_kernel(float *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; } -__mlu_entry__ void scatter_add_kernel(half *src, long *index, half *out, int rows, int cols) { +__mlu_entry__ void scatter_add_kernel(float *src, long *index, float *out, int rows, int cols) { for (int idx = taskId; idx < rows * cols; idx += taskDim) { int col = idx % cols; int row = idx / cols; int dst_row = (int)index[row]; - out[dst_row * cols + col] = (half)((float)out[dst_row * cols + col] + (float)src[idx]); + out[dst_row * cols + col] += src[idx]; } } torch::Tensor bang_func(torch::Tensor src, torch::Tensor index, int dim_size) { - auto input = src.contiguous(); - auto idx = index.contiguous(); + 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 task_dim = {32, 1, 1}; - zero_kernel<<>>((half *)out.data_ptr(), dim_size * cols); - scatter_add_kernel<<>>((half *)input.data_ptr(), idx.data_ptr(), (half *)out.data_ptr(), rows, cols); - return out; + cnrtDim3_t dim = {32, 1, 1}; + cnrtDim3_t one = {1, 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 index d79bea2..ce83676 100644 --- a/Sqrt.mlu +++ b/Sqrt.mlu @@ -3,19 +3,34 @@ #include #include -__mlu_entry__ void sqrt_abs_kernel(half *x, half *out, int total) { - for (int idx = taskId; idx < total; idx += taskDim) { - float value = (float)x[idx]; - out[idx] = (half)sqrtf(value < 0.0f ? -value : value); +#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(); - auto out = torch::empty_like(input); + 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 task_dim = {32, 1, 1}; - sqrt_abs_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), total); - return out; + 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 index 6089f60..fa591f0 100644 --- a/TopK.mlu +++ b/TopK.mlu @@ -3,35 +3,32 @@ #include #include -__mlu_entry__ void topk_kernel(half *x, half *values, long *indices, int rows, int cols, int k) { +__mlu_entry__ void topk_kernel(float *x, float *values, long *indices, int rows, int cols, int k) { for (int row = taskId; row < rows; row += taskDim) { for (int out_col = 0; out_col < k; ++out_col) { int best = 0; - float best_value = -3.402823466e+38F; - for (int col = 0; col < cols; ++col) { - float value = (float)x[row * cols + col]; + float best_val = -3.402823466e+38F; + for (int c = 0; c < cols; ++c) { + float v = x[row * cols + c]; int used = 0; - for (int prev = 0; prev < out_col; ++prev) if (indices[row * k + prev] == (long)col) used = 1; - if (!used && value > best_value) { - best_value = value; - best = col; - } + for (int p = 0; p < out_col; ++p) if (indices[row * k + p] == (long)c) used = 1; + if (!used && v > best_val) { best_val = v; best = c; } } - values[row * k + out_col] = (half)best_value; + 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(); - TORCH_CHECK(input.dim() == 2, "TopK implementation expects 2D input"); - TORCH_CHECK(dim == 1 || dim == -1, "TopK supports last dimension"); + 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 task_dim = {32, 1, 1}; - topk_kernel<<>>((half *)input.data_ptr(), (half *)values.data_ptr(), indices.data_ptr(), rows, cols, k); - return {values, indices}; + 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 index 3dc3a99..06c668e 100644 --- a/average_pooling_2d.mlu +++ b/average_pooling_2d.mlu @@ -2,31 +2,30 @@ #include #include -__mlu_entry__ void avg_pool2d_kernel(half *x, half *out, int channels, int height, int width, int kernel, int out_h, int out_w, int total) { +__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) { - int ih = oh * kernel + kh; - int iw = ow * kernel + kw; - sum += (float)x[((batch * channels + channel) * height + ih) * width + iw]; - } - } - out[idx] = (half)(sum / (float)(kernel * kernel)); + 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(); - int batch = input.size(0), channels = input.size(1), height = input.size(2), width = input.size(3); - int out_h = height / kernel_size, out_w = width / kernel_size; - auto out = torch::empty({batch, channels, out_h, out_w}, input.options()); + 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 task_dim = {32, 1, 1}; - avg_pool2d_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), channels, height, width, kernel_size, out_h, out_w, batch * channels * out_h * out_w); - return out; + 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 index e1f00a0..ba1e357 100644 --- a/batched_matrix_multiplication.mlu +++ b/batched_matrix_multiplication.mlu @@ -2,24 +2,27 @@ #include #include -__mlu_entry__ void bmm_kernel(half *a, half *b, half *out, int rows, int inner, int cols, int total) { +__mlu_entry__ void bmm_kernel(float *a, float *b, float *out, int rows, int inner, int cols, int total) { for (int idx = taskId; idx < total; idx += taskDim) { int col = idx % cols; int row = (idx / cols) % rows; int batch_id = idx / (rows * cols); float sum = 0.0f; - for (int k = 0; k < inner; ++k) sum += (float)a[batch_id * rows * inner + row * inner + k] * (float)b[batch_id * inner * cols + k * cols + col]; - out[idx] = (half)sum; + for (int k = 0; k < inner; ++k) sum += a[batch_id * rows * inner + row * inner + k] * b[batch_id * inner * cols + k * cols + col]; + out[idx] = sum; } } torch::Tensor bang_func(torch::Tensor A, torch::Tensor B) { - auto a = A.contiguous(); - auto b = B.contiguous(); + 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 task_dim = {32, 1, 1}; - bmm_kernel<<>>((half *)a.data_ptr(), (half *)b.data_ptr(), (half *)out.data_ptr(), rows, inner, cols, batch * rows * cols); - return out; + 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/conv_standard_1D.mlu b/conv_standard_1D.mlu index 98a2d0e..c264d33 100644 --- a/conv_standard_1D.mlu +++ b/conv_standard_1D.mlu @@ -2,37 +2,38 @@ #include #include -__mlu_entry__ void conv1d_kernel(half *x, half *kernel, half *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) { +__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_channel = (idx / out_length) % out_channels; - int batch_id = idx / (out_length * out_channels); - int channels_per_group = in_channels / groups; + 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_id = out_channel / out_per_group; - int in_begin = group_id * channels_per_group; + int group = out_ch / out_per_group; float sum = 0.0f; - for (int in_offset = 0; in_offset < channels_per_group; ++in_offset) { - int in_channel = in_begin + in_offset; + 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 += (float)x[(batch_id * in_channels + in_channel) * length + in_pos] * (float)kernel[(out_channel * channels_per_group + in_offset) * kernel_size + k]; - } + 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] = (half)sum; + 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(); - auto weight = kernel.contiguous(); + 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 task_dim = {32, 1, 1}; - conv1d_kernel<<>>((half *)input.data_ptr(), (half *)weight.data_ptr(), (half *)out.data_ptr(), in_channels, out_channels, length, kernel_size, stride, padding, dilation, groups, out_length, batch * out_channels * out_length); - return out; + 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 index bf32f94..a581509 100644 --- a/conv_transposed_2D__asymmetric_input__square_kernel.mlu +++ b/conv_transposed_2D__asymmetric_input__square_kernel.mlu @@ -2,31 +2,30 @@ #include #include -__mlu_entry__ void zero_kernel(half *out, int total) { - for (int idx = taskId; idx < total; idx += taskDim) out[idx] = (half)0.0f; +__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(half *x, half *kernel, half *out, int in_channels, int out_channels, int height, int width, int kernel_size, int stride, int padding, int out_h, int out_w, int groups, int total_in) { +__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 % width; - int ih = (idx / width) % height; - int in_channel = (idx / (width * height)) % in_channels; - int batch_id = idx / (width * height * in_channels); + 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_id = in_channel / in_per_group; - int out_begin = group_id * out_per_group; - float input_value = (float)x[idx]; - for (int out_offset = 0; out_offset < out_per_group; ++out_offset) { - int out_channel = out_begin + out_offset; + 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_id * out_channels + out_channel) * out_h + oh) * out_w + ow; - float weight = (float)kernel[((in_channel * out_per_group + out_offset) * kernel_size + kh) * kernel_size + kw]; - out[out_idx] = (half)((float)out[out_idx] + input_value * weight); + 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; } } } @@ -35,15 +34,19 @@ __mlu_entry__ void conv_transpose2d_kernel(half *x, half *kernel, half *out, int } 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(); - auto weight = kernel.contiguous(); - int batch = input.size(0), height = input.size(2), width = input.size(3); - int out_h = (height - 1) * stride - 2 * padding + kernel_size + output_padding; - int out_w = (width - 1) * stride - 2 * padding + kernel_size + output_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 - 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 task_dim = {32, 1, 1}; - zero_kernel<<>>((half *)out.data_ptr(), batch * out_channels * out_h * out_w); - conv_transpose2d_kernel<<>>((half *)input.data_ptr(), (half *)weight.data_ptr(), (half *)out.data_ptr(), in_channels, out_channels, height, width, kernel_size, stride, padding, out_h, out_w, groups, input.numel()); - return out; + 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 index 6840266..3830200 100644 --- a/cumsum.mlu +++ b/cumsum.mlu @@ -2,34 +2,35 @@ #include #include -__mlu_entry__ void cumsum_dim1_kernel(half *x, half *out, int rows, int cols) { +__mlu_entry__ void cumsum_dim1_kernel(float *x, float *out, int rows, int cols) { for (int row = taskId; row < rows; row += taskDim) { float sum = 0.0f; - for (int col = 0; col < cols; ++col) { - sum += (float)x[row * cols + col]; - out[row * cols + col] = (half)sum; + for (int c = 0; c < cols; ++c) { + sum += x[row * cols + c]; + out[row * cols + c] = sum; } } } -__mlu_entry__ void cumsum_dim0_kernel(half *x, half *out, int rows, int cols) { +__mlu_entry__ void cumsum_dim0_kernel(float *x, float *out, int rows, int cols) { for (int col = taskId; col < cols; col += taskDim) { float sum = 0.0f; - for (int row = 0; row < rows; ++row) { - sum += (float)x[row * cols + col]; - out[row * cols + col] = (half)sum; + for (int r = 0; r < rows; ++r) { + sum += x[r * cols + col]; + out[r * cols + col] = sum; } } } torch::Tensor bang_func(torch::Tensor x, int dim) { - auto input = x.contiguous(); - TORCH_CHECK(input.dim() == 2, "cumsum implementation expects 2D input"); - int rows = input.size(0), cols = input.size(1); + 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 task_dim = {32, 1, 1}; - if (dim == 0) cumsum_dim0_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), rows, cols); - else cumsum_dim1_kernel<<>>((half *)input.data_ptr(), (half *)out.data_ptr(), rows, cols); - return out; + 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 index c69c9a3..067e61a 100644 --- a/gather.mlu +++ b/gather.mlu @@ -2,34 +2,36 @@ #include #include -__mlu_entry__ void zero_kernel(half *out, int total) { - for (int idx = taskId; idx < total; idx += taskDim) out[idx] = (half)0.0f; +__mlu_entry__ void zero_kernel(float *out, int total) { + for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; } -__mlu_entry__ void gather_kernel(half *x, long *indices, long *bin_ids, half *weights, long *bins, half *out, int elements, int hidden, int top_k, int has_weights) { +__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) { for (int idx = taskId; idx < elements * hidden; idx += taskDim) { - int hidden_id = idx % hidden; + int h = idx % hidden; int elem = idx / hidden; 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 ? (float)weights[elem] : 1.0f; - out[dst_row * hidden + hidden_id] = (half)((float)x[src_row * hidden + hidden_id] * scale); + float scale = has_weights ? weights[elem] : 1.0f; + out[dst_row * hidden + h] = x[src_row * hidden + h] * scale; } } 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(); - auto idx = indices.contiguous(); - auto bid = bin_ids.contiguous(); - auto weight = weights.contiguous(); - auto bin = bins.contiguous(); + 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 task_dim = {32, 1, 1}; - zero_kernel<<>>((half *)out.data_ptr(), tokens * top_k * hidden); - gather_kernel<<>>((half *)input.data_ptr(), idx.data_ptr(), bid.data_ptr(), (half *)weight.data_ptr(), bin.data_ptr(), (half *)out.data_ptr(), elements, hidden, top_k, weights.numel() > 0 ? 1 : 0); - return out; + 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 index 843e0e6..3f1cbd0 100644 --- a/matrix_scalar_multiplication.mlu +++ b/matrix_scalar_multiplication.mlu @@ -2,18 +2,32 @@ #include #include -__mlu_entry__ void matrix_scalar_mul_kernel(half *input, half *output, int total, float scalar) { - for (int idx = taskId; idx < total; idx += taskDim) { - output[idx] = (half)((float)input[idx] * scalar); +#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 input = A.contiguous(); - auto output = torch::empty_like(input); - int total = input.numel(); + 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 task_dim = {32, 1, 1}; - matrix_scalar_mul_kernel<<>>((half *)input.data_ptr(), (half *)output.data_ptr(), total, (float)s); - return output; + 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()); } From 3cb1d92c3c3aded02890fb671530bcd15fa62f2b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=83=AD=E7=AB=A3=E8=B1=AA?= Date: Thu, 11 Jun 2026 19:26:12 +0800 Subject: [PATCH 18/18] NRAM tiling optimization for all operators --- Argmax_over_a_dimension.mlu | 19 +++++++---- BatchNorm.mlu | 57 +++++++++++++++++++++---------- GRU_forward.mlu | 3 ++ KL_Divergence_Loss.mlu | 25 ++++++++++++-- LogSoftmax.mlu | 31 ++++++++++++++--- MSE_Loss.mlu | 20 +++++++++-- Masked_select.mlu | 31 +++++++++++++---- Matrix_vector_multiplication_.mlu | 14 +++++++- Scatter_add.mlu | 27 ++++++++++----- TopK.mlu | 18 +++++++--- batched_matrix_multiplication.mlu | 19 +++++++++-- cumsum.mlu | 32 ++++++++++++----- gather.mlu | 23 +++++++++---- 13 files changed, 247 insertions(+), 72 deletions(-) diff --git a/Argmax_over_a_dimension.mlu b/Argmax_over_a_dimension.mlu index 2b1a5c8..3615f81 100644 --- a/Argmax_over_a_dimension.mlu +++ b/Argmax_over_a_dimension.mlu @@ -2,13 +2,20 @@ #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 = x[row * cols]; - for (int c = 1; c < cols; ++c) { - float v = x[row * cols + c]; - if (v > best_val) { best_val = v; best = c; } + 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; } @@ -17,8 +24,8 @@ __mlu_entry__ void argmax_dim1_kernel(float *x, long *out, int rows, int cols) { __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 = x[col]; - for (int r = 1; r < rows; ++r) { + 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; } } diff --git a/BatchNorm.mlu b/BatchNorm.mlu index 95c8eb5..651a6c2 100644 --- a/BatchNorm.mlu +++ b/BatchNorm.mlu @@ -3,28 +3,49 @@ #include #include -__mlu_entry__ void batchnorm_kernel(float *x, float *out, int batch, int channels, int h, int w, int channel_size) { +#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) - for (int row = 0; row < h; ++row) - for (int col = 0; col < w; ++col) - sum += x[((n * channels + ch) * h + row) * w + col]; + 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) - for (int row = 0; row < h; ++row) - for (int col = 0; col < w; ++col) { - float diff = x[((n * channels + ch) * h + row) * w + col] - mean; - var_sum += diff * diff; - } + 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); - for (int n = 0; n < batch; ++n) - for (int row = 0; row < h; ++row) - for (int col = 0; col < w; ++col) { - int idx = ((n * channels + ch) * h + row) * w + col; - out[idx] = (x[idx] - mean) * inv_std; - } + // 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); + } + } } } @@ -35,7 +56,7 @@ torch::Tensor bang_func(torch::Tensor x, int num_features) { 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, batch * h * w); + 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/GRU_forward.mlu b/GRU_forward.mlu index ad7941b..4c4d539 100644 --- a/GRU_forward.mlu +++ b/GRU_forward.mlu @@ -3,7 +3,10 @@ #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; diff --git a/KL_Divergence_Loss.mlu b/KL_Divergence_Loss.mlu index cf42dae..af81071 100644 --- a/KL_Divergence_Loss.mlu +++ b/KL_Divergence_Loss.mlu @@ -3,11 +3,30 @@ #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 idx = taskId; idx < total; idx += taskDim) { - float t = target_prob[idx]; - if (t > 0.0f) sum += t * (logf(t) - input_log_prob[idx]); + 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; } diff --git a/LogSoftmax.mlu b/LogSoftmax.mlu index 2564baa..edf904c 100644 --- a/LogSoftmax.mlu +++ b/LogSoftmax.mlu @@ -3,17 +3,38 @@ #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 c = 0; c < cols; ++c) { - float v = input[row * cols + c]; - if (v > max_val) max_val = v; + 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 c = 0; c < cols; ++c) sum += expf(input[row * cols + c] - max_val); + 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; - for (int c = 0; c < cols; ++c) output[row * cols + c] = input[row * cols + c] - log_sum; + // 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); + } } } diff --git a/MSE_Loss.mlu b/MSE_Loss.mlu index 849dc27..cbccc29 100644 --- a/MSE_Loss.mlu +++ b/MSE_Loss.mlu @@ -2,11 +2,25 @@ #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 idx = taskId; idx < total; idx += taskDim) { - float diff = pred[idx] - target[idx]; - sum += diff * diff; + 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; } diff --git a/Masked_select.mlu b/Masked_select.mlu index e58e0eb..3428228 100644 --- a/Masked_select.mlu +++ b/Masked_select.mlu @@ -2,11 +2,21 @@ #include #include +#define NRAM_SIZE 8192 + __mlu_entry__ void count_kernel(float *input, int *counts, int total, float threshold) { - int count = 0; - for (int idx = taskId; idx < total; idx += taskDim) - if (input[idx] > threshold) ++count; - counts[taskId] = count; + 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) { @@ -16,9 +26,18 @@ __mlu_entry__ void prefix_kernel(int *counts, int *offsets, int *total_out, int } __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 idx = taskId; idx < total; idx += taskDim) - if (input[idx] > threshold) out[write_pos++] = input[idx]; + 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) { diff --git a/Matrix_vector_multiplication_.mlu b/Matrix_vector_multiplication_.mlu index 5a7c85f..cac330e 100644 --- a/Matrix_vector_multiplication_.mlu +++ b/Matrix_vector_multiplication_.mlu @@ -2,10 +2,22 @@ #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 k = 0; k < inner; ++k) sum += a[row * inner + k] * b[k]; + 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; } } diff --git a/Scatter_add.mlu b/Scatter_add.mlu index 02d07c4..d6b424e 100644 --- a/Scatter_add.mlu +++ b/Scatter_add.mlu @@ -2,19 +2,31 @@ #include #include -__mlu_entry__ void zero_kernel(float *out, int total) { - for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; -} +#define NRAM_SIZE 4096 __mlu_entry__ void scatter_add_kernel(float *src, long *index, float *out, int rows, int cols) { - for (int idx = taskId; idx < rows * cols; idx += taskDim) { - int col = idx % cols; - int row = idx / cols; + __nram__ float buf[NRAM_SIZE]; + for (int row = taskId; row < rows; row += taskDim) { int dst_row = (int)index[row]; - out[dst_row * cols + col] += src[idx]; + 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(); @@ -22,7 +34,6 @@ torch::Tensor bang_func(torch::Tensor src, torch::Tensor index, int dim_size) { auto out = torch::empty({dim_size, cols}, input.options()); 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(), dim_size * cols); scatter_add_kernel<<>>(input.data_ptr(), idx.data_ptr(), out.data_ptr(), rows, cols); diff --git a/TopK.mlu b/TopK.mlu index fa591f0..4bcb407 100644 --- a/TopK.mlu +++ b/TopK.mlu @@ -3,16 +3,24 @@ #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 c = 0; c < cols; ++c) { - float v = x[row * cols + c]; - int used = 0; - for (int p = 0; p < out_col; ++p) if (indices[row * k + p] == (long)c) used = 1; - if (!used && v > best_val) { best_val = v; best = c; } + 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; diff --git a/batched_matrix_multiplication.mlu b/batched_matrix_multiplication.mlu index ba1e357..05b998a 100644 --- a/batched_matrix_multiplication.mlu +++ b/batched_matrix_multiplication.mlu @@ -2,13 +2,28 @@ #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_id = idx / (rows * cols); + 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 k = 0; k < inner; ++k) sum += a[batch_id * rows * inner + row * inner + k] * b[batch_id * inner * cols + k * cols + col]; + 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; } } diff --git a/cumsum.mlu b/cumsum.mlu index 3830200..9d4c83c 100644 --- a/cumsum.mlu +++ b/cumsum.mlu @@ -2,22 +2,38 @@ #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 sum = 0.0f; - for (int c = 0; c < cols; ++c) { - sum += x[row * cols + c]; - out[row * cols + c] = sum; + 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) { - for (int col = taskId; col < cols; col += taskDim) { - float sum = 0.0f; + __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) { - sum += x[r * cols + col]; - out[r * cols + col] = sum; + __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); } } } diff --git a/gather.mlu b/gather.mlu index 067e61a..dff697e 100644 --- a/gather.mlu +++ b/gather.mlu @@ -2,23 +2,32 @@ #include #include -__mlu_entry__ void zero_kernel(float *out, int total) { - for (int idx = taskId; idx < total; idx += taskDim) out[idx] = 0.0f; -} +#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) { - for (int idx = taskId; idx < elements * hidden; idx += taskDim) { - int h = idx % hidden; - int elem = idx / hidden; + __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; - out[dst_row * hidden + h] = x[src_row * hidden + h] * scale; + 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();