From 913486590ec18c553b1867c58a81efac8bba5b6a Mon Sep 17 00:00:00 2001 From: Sicheng <15533151020@163.com> Date: Sat, 6 Jun 2026 11:10:29 +0800 Subject: [PATCH 1/4] 056 --- gather.mlu | 240 +++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 240 insertions(+) create mode 100644 gather.mlu diff --git a/gather.mlu b/gather.mlu new file mode 100644 index 0000000..8dfb741 --- /dev/null +++ b/gather.mlu @@ -0,0 +1,240 @@ +#include +#include +#include +#include +#include + +#define TILE_ELEMS 4096 +#define MAX_TASKS 64 + +__mlu_entry__ void gather_with_bins_float_kernel(const float *x, + const int64_t *indices, + const int64_t *bin_ids, + const float *weights, + const int64_t *bins, + float *out, + int64_t tokens, + int64_t hidden_size, + int64_t num_elements, + int64_t top_k, + int has_weights) { + __nram__ float tile[TILE_ELEMS]; + + for (int64_t row = taskId; row < num_elements; row += taskDim) { + int64_t flat_index = indices[row]; + int64_t bin_id = bin_ids[row]; + int64_t bin_start = bin_id > 0 ? bins[bin_id - 1] : 0; + int64_t dst_row = row - bin_start + bin_start; + + if (flat_index < 0 || flat_index >= num_elements || + dst_row < 0 || dst_row >= num_elements || top_k <= 0) { + for (int64_t col = 0; col < hidden_size; col += TILE_ELEMS) { + uint32_t len = static_cast(hidden_size - col); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + __bang_write_value(tile, len, 0.0f); + __memcpy(out + row * hidden_size + col, + tile, + len * sizeof(float), + NRAM2GDRAM); + } + continue; + } + + int64_t src_row = flat_index / top_k; + if (src_row < 0 || src_row >= tokens) { + for (int64_t col = 0; col < hidden_size; col += TILE_ELEMS) { + uint32_t len = static_cast(hidden_size - col); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + __bang_write_value(tile, len, 0.0f); + __memcpy(out + dst_row * hidden_size + col, + tile, + len * sizeof(float), + NRAM2GDRAM); + } + continue; + } + + float scale = has_weights ? weights[flat_index] : 1.0f; + for (int64_t col = 0; col < hidden_size; col += TILE_ELEMS) { + uint32_t len = static_cast(hidden_size - col); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + + __memcpy(tile, + x + src_row * hidden_size + col, + len * sizeof(float), + GDRAM2NRAM); + if (has_weights) { + __bang_mul_scalar(tile, tile, scale, len); + } + __memcpy(out + dst_row * hidden_size + col, + tile, + len * sizeof(float), + NRAM2GDRAM); + } + } +} + +__mlu_entry__ void gather_with_bins_half_kernel(const half *x, + const int64_t *indices, + const int64_t *bin_ids, + const float *weights, + const int64_t *bins, + half *out, + int64_t tokens, + int64_t hidden_size, + int64_t num_elements, + int64_t top_k, + int has_weights) { + __nram__ half tile[TILE_ELEMS]; + + for (int64_t row = taskId; row < num_elements; row += taskDim) { + int64_t flat_index = indices[row]; + int64_t bin_id = bin_ids[row]; + int64_t bin_start = bin_id > 0 ? bins[bin_id - 1] : 0; + int64_t dst_row = row - bin_start + bin_start; + + if (flat_index < 0 || flat_index >= num_elements || + dst_row < 0 || dst_row >= num_elements || top_k <= 0) { + for (int64_t col = 0; col < hidden_size; col += TILE_ELEMS) { + uint32_t len = static_cast(hidden_size - col); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + __bang_write_value(tile, len, (half)0.0f); + __memcpy(out + row * hidden_size + col, + tile, + len * sizeof(half), + NRAM2GDRAM); + } + continue; + } + + int64_t src_row = flat_index / top_k; + if (src_row < 0 || src_row >= tokens) { + for (int64_t col = 0; col < hidden_size; col += TILE_ELEMS) { + uint32_t len = static_cast(hidden_size - col); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + __bang_write_value(tile, len, (half)0.0f); + __memcpy(out + dst_row * hidden_size + col, + tile, + len * sizeof(half), + NRAM2GDRAM); + } + continue; + } + + half scale = has_weights ? (half)weights[flat_index] : (half)1.0f; + for (int64_t col = 0; col < hidden_size; col += TILE_ELEMS) { + uint32_t len = static_cast(hidden_size - col); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + + __memcpy(tile, + x + src_row * hidden_size + col, + len * sizeof(half), + GDRAM2NRAM); + if (has_weights) { + __bang_mul_scalar(tile, tile, scale, len); + } + __memcpy(out + dst_row * hidden_size + col, + tile, + len * sizeof(half), + NRAM2GDRAM); + } + } +} + +torch::Tensor bang_func(torch::Tensor x, + torch::Tensor indices, + torch::Tensor bin_ids, + c10::optional weights, + torch::Tensor bins, + int64_t top_k) { + TORCH_CHECK(x.dim() == 2, "x must be a 2D tensor"); + TORCH_CHECK(indices.dim() == 1, "indices must be a 1D tensor"); + TORCH_CHECK(bin_ids.dim() == 1, "bin_ids must be a 1D tensor"); + TORCH_CHECK(bins.dim() == 1, "bins must be a 1D tensor"); + TORCH_CHECK(top_k > 0, "top_k must be positive"); + TORCH_CHECK(x.scalar_type() == torch::kFloat32 || + x.scalar_type() == torch::kHalf, + "x must be float32 or float16"); + + auto x_c = x.contiguous(); + auto indices_c = indices.to(torch::kInt64).contiguous(); + auto bin_ids_c = bin_ids.to(torch::kInt64).contiguous(); + auto bins_c = bins.to(torch::kInt64).contiguous(); + + int64_t tokens = x_c.size(0); + int64_t hidden_size = x_c.size(1); + int64_t num_elements = indices_c.numel(); + TORCH_CHECK(tokens == 0 || top_k <= INT64_MAX / tokens, + "output size overflow"); + int64_t output_rows = tokens * top_k; + + TORCH_CHECK(num_elements == output_rows, + "indices length must equal tokens * top_k"); + TORCH_CHECK(bin_ids_c.numel() == num_elements, + "bin_ids length must match indices length"); + int has_weights = 0; + torch::Tensor weights_c; + if (weights.has_value() && weights.value().defined() && + weights.value().numel() > 0) { + TORCH_CHECK(weights.value().dim() == 1, "weights must be a 1D tensor"); + TORCH_CHECK(weights.value().numel() == num_elements, + "weights length must match indices length"); + weights_c = weights.value().to(torch::kFloat32).contiguous(); + has_weights = 1; + } else { + weights_c = torch::empty({1}, x_c.options().dtype(torch::kFloat32)); + } + + auto out = torch::empty({output_rows, hidden_size}, x_c.options()); + if (num_elements == 0 || hidden_size == 0) { + return out; + } + + uint32_t task_num = static_cast( + num_elements < MAX_TASKS ? num_elements : MAX_TASKS); + cnrtDim3_t dim = {task_num, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeBlock; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + + if (x_c.scalar_type() == torch::kFloat32) { + gather_with_bins_float_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(indices_c.data_ptr()), + reinterpret_cast(bin_ids_c.data_ptr()), + reinterpret_cast(weights_c.data_ptr()), + reinterpret_cast(bins_c.data_ptr()), + reinterpret_cast(out.data_ptr()), + tokens, + hidden_size, + num_elements, + top_k, + has_weights); + } else { + gather_with_bins_half_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(indices_c.data_ptr()), + reinterpret_cast(bin_ids_c.data_ptr()), + reinterpret_cast(weights_c.data_ptr()), + reinterpret_cast(bins_c.data_ptr()), + reinterpret_cast(out.data_ptr()), + tokens, + hidden_size, + num_elements, + top_k, + has_weights); + } + + return out; +} From bd52f5fad65dfed94de0a88a5ccb1d21c6f092e6 Mon Sep 17 00:00:00 2001 From: Sicheng <15533151020@163.com> Date: Sat, 6 Jun 2026 11:10:37 +0800 Subject: [PATCH 2/4] 070 --- Sqrt.mlu | 132 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 132 insertions(+) create mode 100644 Sqrt.mlu diff --git a/Sqrt.mlu b/Sqrt.mlu new file mode 100644 index 0000000..b7d7130 --- /dev/null +++ b/Sqrt.mlu @@ -0,0 +1,132 @@ +#include +#include +#include +#include +#include + +#define TILE_ELEMS 4096 +#define TASK_ELEMS (TILE_ELEMS * 8) + +__mlu_entry__ void sqrt_abs_float_kernel(const float *x, + float *out, + int64_t total) { + __nram__ float tile[TILE_ELEMS]; + + int64_t stride = static_cast(taskDim) * TASK_ELEMS; + if (stride <= 0) { + return; + } + + for (int64_t start = taskId * TASK_ELEMS; + start < total; + start += stride) { + int64_t count = total - start; + if (count > TASK_ELEMS) { + count = TASK_ELEMS; + } + + for (int64_t offset = 0; offset < count; offset += TILE_ELEMS) { + uint32_t len = static_cast(count - offset); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + + __memcpy(tile, + x + start + offset, + len * sizeof(float), + GDRAM2NRAM); + + for (uint32_t i = 0; i < len; ++i) { + float value = tile[i]; + tile[i] = sqrtf(value < 0.0f ? -value : value); + } + + __memcpy(out + start + offset, + tile, + len * sizeof(float), + NRAM2GDRAM); + } + } +} + +__mlu_entry__ void sqrt_abs_half_kernel(const half *x, + half *out, + int64_t total) { + __nram__ half half_tile[TILE_ELEMS]; + __nram__ float float_tile[TILE_ELEMS]; + + int64_t stride = static_cast(taskDim) * TASK_ELEMS; + if (stride <= 0) { + return; + } + + for (int64_t start = taskId * TASK_ELEMS; + start < total; + start += stride) { + int64_t count = total - start; + if (count > TASK_ELEMS) { + count = TASK_ELEMS; + } + + for (int64_t offset = 0; offset < count; offset += TILE_ELEMS) { + uint32_t len = static_cast(count - offset); + if (len > TILE_ELEMS) { + len = TILE_ELEMS; + } + + __memcpy(half_tile, + x + start + offset, + len * sizeof(half), + GDRAM2NRAM); + __bang_half2float(float_tile, half_tile, len); + + for (uint32_t i = 0; i < len; ++i) { + float value = float_tile[i]; + float_tile[i] = sqrtf(value < 0.0f ? -value : value); + } + + __bang_float2half_rn(half_tile, float_tile, len); + __memcpy(out + start + offset, + half_tile, + len * sizeof(half), + NRAM2GDRAM); + } + } +} + +torch::Tensor bang_func(torch::Tensor x) { + TORCH_CHECK(x.dim() == 2, "x must be a 2D tensor with shape [batch_size, dim]"); + TORCH_CHECK(x.scalar_type() == torch::kFloat32 || + x.scalar_type() == torch::kHalf, + "x must be float32 or float16"); + + auto x_c = x.contiguous(); + auto out = torch::empty_like(x_c); + int64_t total = x_c.numel(); + if (total == 0) { + return out; + } + + int64_t task_num_64 = (total + TASK_ELEMS - 1) / TASK_ELEMS; + if (task_num_64 > 64) { + task_num_64 = 64; + } + + cnrtDim3_t dim = {static_cast(task_num_64), 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeBlock; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + + if (x_c.scalar_type() == torch::kFloat32) { + sqrt_abs_float_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(out.data_ptr()), + total); + } else { + sqrt_abs_half_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(out.data_ptr()), + total); + } + + return out; +} From 0a67a84f3f715dd6b3dc5610fcdf57bc876fbee3 Mon Sep 17 00:00:00 2001 From: Sicheng <15533151020@163.com> Date: Sat, 6 Jun 2026 11:10:45 +0800 Subject: [PATCH 3/4] 075 --- TopK.mlu | 175 +++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 175 insertions(+) create mode 100644 TopK.mlu diff --git a/TopK.mlu b/TopK.mlu new file mode 100644 index 0000000..66b26b1 --- /dev/null +++ b/TopK.mlu @@ -0,0 +1,175 @@ +#include +#include +#include +#include +#include + +#define MAX_TASKS 64 +#define NEG_INF (-3.4028234663852886e+38F) + +__mlu_entry__ void topk_float_kernel(const float *x, + float *values, + int64_t *indices, + int64_t batch, + int64_t dim_size, + int64_t k, + int64_t op_dim) { + int64_t work_items = op_dim == 1 ? batch : dim_size; + + for (int64_t item = taskId; item < work_items; item += taskDim) { + for (int64_t rank = 0; rank < k; ++rank) { + float best_value = NEG_INF; + int64_t best_index = 0; + int found = 0; + int64_t reduce_size = op_dim == 1 ? dim_size : batch; + + for (int64_t i = 0; i < reduce_size; ++i) { + int selected = 0; + for (int64_t prev = 0; prev < rank; ++prev) { + int64_t prev_index = op_dim == 1 + ? indices[item * k + prev] + : indices[prev * dim_size + item]; + if (prev_index == i) { + selected = 1; + break; + } + } + if (selected) { + continue; + } + + float candidate = op_dim == 1 ? x[item * dim_size + i] + : x[i * dim_size + item]; + if (!found || candidate > best_value || + (candidate == best_value && i < best_index)) { + best_value = candidate; + best_index = i; + found = 1; + } + } + + if (op_dim == 1) { + values[item * k + rank] = best_value; + indices[item * k + rank] = best_index; + } else { + values[rank * dim_size + item] = best_value; + indices[rank * dim_size + item] = best_index; + } + } + } +} + +__mlu_entry__ void topk_half_kernel(const half *x, + half *values, + int64_t *indices, + int64_t batch, + int64_t dim_size, + int64_t k, + int64_t op_dim) { + int64_t work_items = op_dim == 1 ? batch : dim_size; + + for (int64_t item = taskId; item < work_items; item += taskDim) { + for (int64_t rank = 0; rank < k; ++rank) { + float best_value = NEG_INF; + half best_half = (half)0.0f; + int64_t best_index = 0; + int found = 0; + int64_t reduce_size = op_dim == 1 ? dim_size : batch; + + for (int64_t i = 0; i < reduce_size; ++i) { + int selected = 0; + for (int64_t prev = 0; prev < rank; ++prev) { + int64_t prev_index = op_dim == 1 + ? indices[item * k + prev] + : indices[prev * dim_size + item]; + if (prev_index == i) { + selected = 1; + break; + } + } + if (selected) { + continue; + } + + half candidate_half = op_dim == 1 ? x[item * dim_size + i] + : x[i * dim_size + item]; + float candidate = (float)candidate_half; + if (!found || candidate > best_value || + (candidate == best_value && i < best_index)) { + best_value = candidate; + best_half = candidate_half; + best_index = i; + found = 1; + } + } + + if (op_dim == 1) { + values[item * k + rank] = best_half; + indices[item * k + rank] = best_index; + } else { + values[rank * dim_size + item] = best_half; + indices[rank * dim_size + item] = best_index; + } + } + } +} + +std::vector bang_func(torch::Tensor x, int64_t k, int64_t dim) { + TORCH_CHECK(x.dim() == 2, "x must be a 2D tensor with shape [batch, dim]"); + TORCH_CHECK(x.scalar_type() == torch::kFloat32 || + x.scalar_type() == torch::kHalf, + "x must be float32 or float16"); + TORCH_CHECK(dim == 0 || dim == 1 || dim == -1 || dim == -2, + "dim must be 0, 1, -1, or -2"); + int64_t op_dim = dim < 0 ? dim + 2 : dim; + + int64_t batch = x.size(0); + int64_t dim_size = x.size(1); + int64_t reduce_size = op_dim == 1 ? dim_size : batch; + + TORCH_CHECK(k >= 0, "k must be non-negative"); + TORCH_CHECK(k <= reduce_size, "k must not exceed the selected dimension size"); + + auto x_c = x.contiguous(); + std::vector out_shape; + if (op_dim == 1) { + out_shape = {batch, k}; + } else { + out_shape = {k, dim_size}; + } + + auto values = torch::empty(out_shape, x_c.options()); + auto indices = torch::empty(out_shape, x_c.options().dtype(torch::kInt64)); + if (k == 0 || batch == 0 || dim_size == 0) { + return {values, indices}; + } + + int64_t work_items = op_dim == 1 ? batch : dim_size; + uint32_t task_num = static_cast( + work_items < MAX_TASKS ? work_items : MAX_TASKS); + cnrtDim3_t k_dim = {task_num, 1, 1}; + cnrtFunctionType_t ktype = cnrtFuncTypeBlock; + cnrtQueue_t queue = torch_mlu::getCurMLUStream(); + + if (x_c.scalar_type() == torch::kFloat32) { + topk_float_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(values.data_ptr()), + reinterpret_cast(indices.data_ptr()), + batch, + dim_size, + k, + op_dim); + } else { + topk_half_kernel<<>>( + reinterpret_cast(x_c.data_ptr()), + reinterpret_cast(values.data_ptr()), + reinterpret_cast(indices.data_ptr()), + batch, + dim_size, + k, + op_dim); + } + + return {values, indices}; +} From 96c12266e0f9eb97557baaedbfae94d238859467 Mon Sep 17 00:00:00 2001 From: Sicheng <15533151020@163.com> Date: Sat, 6 Jun 2026 11:10:53 +0800 Subject: [PATCH 4/4] config --- config | 3 +++ 1 file changed, 3 insertions(+) diff --git a/config b/config index 7cd8208..0d5f384 100644 --- a/config +++ b/config @@ -2,3 +2,6 @@ 109 135 111 +056 +070 +075