From e598fd1ba0819f28b30e989b84a1477261949ec8 Mon Sep 17 00:00:00 2001
From: Zhiyuan Chen <136175529+Jassicia@users.noreply.github.com>
Date: Mon, 8 Jun 2026 19:30:10 +0800
Subject: [PATCH] Add files via upload

---
 Sqrt.mlu | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/Sqrt.mlu b/Sqrt.mlu
index 1e11f98..e7a7921 100644
--- a/Sqrt.mlu
+++ b/Sqrt.mlu
@@ -11,6 +11,7 @@ __mlu_entry__ void sqrt_kernel(half *input, half *output, int total) {
     uint32_t stride = task_num * BLOCK_SIZE;
 
     __nram__ half buffer[BLOCK_SIZE];
+    __nram__ float float_buffer[BLOCK_SIZE];
 
     for (uint32_t offset = start; offset < (uint32_t)total; offset += stride) {
         uint32_t remain = (uint32_t)total - offset;
@@ -19,7 +20,9 @@ __mlu_entry__ void sqrt_kernel(half *input, half *output, int total) {
 
         __memcpy(buffer, input + offset, len * sizeof(half), GDRAM2NRAM);
         __bang_abs(buffer, buffer, aligned_len);
-        __bang_sqrt(buffer, buffer, aligned_len);
+        __bang_half2float(float_buffer, buffer, aligned_len);
+        __bang_sqrt(float_buffer, float_buffer, aligned_len);
+        __bang_float2half(buffer, float_buffer, aligned_len);
         __memcpy(output + offset, buffer, len * sizeof(half), NRAM2GDRAM);
     }
 }