From 399851ce8eadaee0ca4020dfd1829b2087c83feb Mon Sep 17 00:00:00 2001
From: Zhiyuan Chen <136175529+Jassicia@users.noreply.github.com>
Date: Mon, 8 Jun 2026 18:36:45 +0800
Subject: [PATCH] Add files via upload

---
 Sqrt.mlu | 44 ++++++++++++++++++++++++++++++++++++++++++++
 config   |  3 +--
 2 files changed, 45 insertions(+), 2 deletions(-)
 create mode 100644 Sqrt.mlu

diff --git a/Sqrt.mlu b/Sqrt.mlu
new file mode 100644
index 0000000..1e11f98
--- /dev/null
+++ b/Sqrt.mlu
@@ -0,0 +1,44 @@
+#include <bang.h>
+#include <torch/extension.h>
+#include <cnrt.h>
+
+#define BLOCK_SIZE 256
+
+__mlu_entry__ void sqrt_kernel(half *input, half *output, int total) {
+    uint32_t task_id = taskId;
+    uint32_t task_num = taskDim;
+    uint32_t start = task_id * BLOCK_SIZE;
+    uint32_t stride = task_num * BLOCK_SIZE;
+
+    __nram__ half buffer[BLOCK_SIZE];
+
+    for (uint32_t offset = start; offset < (uint32_t)total; offset += stride) {
+        uint32_t remain = (uint32_t)total - offset;
+        uint32_t len = remain > BLOCK_SIZE ? BLOCK_SIZE : remain;
+        uint32_t aligned_len = (len + 63) & ~63;
+
+        __memcpy(buffer, input + offset, len * sizeof(half), GDRAM2NRAM);
+        __bang_abs(buffer, buffer, aligned_len);
+        __bang_sqrt(buffer, buffer, aligned_len);
+        __memcpy(output + offset, buffer, len * sizeof(half), NRAM2GDRAM);
+    }
+}
+
+torch::Tensor bang_func(torch::Tensor input) {
+    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+    TORCH_CHECK(input.scalar_type() == torch::kFloat16, "input must be float16");
+
+    auto output = torch::empty_like(input);
+    int total = input.numel();
+
+    cnrtQueue_t queue = torch_mlu::getCurMLUStream();
+    cnrtDim3_t dim = {4, 1, 1};
+    cnrtFunctionType_t ktype = cnrtFuncTypeUnion1;
+
+    sqrt_kernel<<<dim, ktype, queue>>>(
+        reinterpret_cast<half *>(input.data_ptr<at::Half>()),
+        reinterpret_cast<half *>(output.data_ptr<at::Half>()),
+        total);
+
+    return output;
+}
diff --git a/config b/config
index e14f913..49a5bb2 100644
--- a/config
+++ b/config
@@ -1,2 +1 @@
-012
-104
\ No newline at end of file
+070