Fix format and type conversion errors

Author: xu-zhengzhong

python/infinicore/__init__.py

@@ -64,14 +64,14 @@
 from infinicore.ops.axpy import axpy
 from infinicore.ops.baddbmm import baddbmm
 from infinicore.ops.bilinear import bilinear
-from infinicore.ops.blas_amax import blas_amax
-from infinicore.ops.blas_amin import blas_amin
-from infinicore.ops.blas_copy import blas_copy
-from infinicore.ops.blas_dot import blas_dot
 from infinicore.ops.binary_cross_entropy_with_logits import (
     binary_cross_entropy_with_logits,
 )
 from infinicore.ops.bitwise_right_shift import bitwise_right_shift
+from infinicore.ops.blas_amax import blas_amax
+from infinicore.ops.blas_amin import blas_amin
+from infinicore.ops.blas_copy import blas_copy
+from infinicore.ops.blas_dot import blas_dot
 from infinicore.ops.block_diag import block_diag
 from infinicore.ops.broadcast_to import broadcast_to
 from infinicore.ops.cat import cat
@@ -121,9 +121,9 @@
 from infinicore.ops.scal import scal
 from infinicore.ops.scatter import scatter
 from infinicore.ops.sinh import sinh
-from infinicore.ops.swap import swap
 from infinicore.ops.squeeze import squeeze
 from infinicore.ops.sum import sum
+from infinicore.ops.swap import swap
 from infinicore.ops.take import take
 from infinicore.ops.tan import tan
 from infinicore.ops.topk import topk

src/infiniop/ops/asum/bang/asum_bang.mlu

@@ -33,8 +33,8 @@ infiniStatus_t calculateAsum(
     Tdata *result,
     cnrtQueue_t queue) {
 
-    const size_t n = info.n;
-    const ptrdiff_t incx = info.incx;
+    const int n = utils::cast<int>(info.n);
+    const int incx = utils::cast<int>(info.incx);
 
     cnrtDim3_t k_dim;
     cnrtFunctionType_t k_type;

src/infiniop/ops/asum/bang/asum_bang_kernel.mlu

@@ -1,90 +1,140 @@
 #include "../../../devices/bang/common_bang.h"
 #include "asum_bang.h"
 
+#include <type_traits>
+
 __nram__ char nram_buffer[NRAM_MAX_SIZE];
 
+template <typename Tdata>
+__mlu_device__ void asumToCompute(float *dst, const Tdata *src, int size) {
+    if constexpr (std::is_same_v<Tdata, half>) {
+        __bang_half2float(dst, src, size);
+    } else if constexpr (std::is_same_v<Tdata, bfloat16_t>) {
+        __bang_bfloat162float(dst, src, size);
+    } else {
+        __memcpy(dst, src, size * sizeof(float), NRAM2NRAM);
+    }
+}
+
+template <typename Tdata>
+__mlu_device__ float asumToCompute(Tdata value) {
+    if constexpr (std::is_same_v<Tdata, half>) {
+        return __half2float(value);
+    } else if constexpr (std::is_same_v<Tdata, bfloat16_t>) {
+        return __bfloat162float(value);
+    } else {
+        return static_cast<float>(value);
+    }
+}
+
+template <typename Tdata>
+__mlu_device__ void asumStoreResult(Tdata *result, Tdata *nram_result, float *nram_compute, float value) {
+    nram_compute[0] = value;
+    if constexpr (std::is_same_v<Tdata, half>) {
+        __bang_float2half(nram_result, nram_compute, 1);
+        result[0] = nram_result[0];
+    } else if constexpr (std::is_same_v<Tdata, bfloat16_t>) {
+        __bang_float2bfloat16(nram_result, nram_compute, 1);
+        result[0] = nram_result[0];
+    } else {
+        result[0] = nram_compute[0];
+    }
+}
+
 template <typename Tdata>
 __mlu_global__ void asumKernelContiguous(
-    size_t n,
+    int n,
     const Tdata *x,
     Tdata *result) {
 
-    __mlu_shared__ Tdata shared_partial_sum[4];
+    __mlu_shared__ float shared_partial_sum[4];
 
-    Tdata *nram_x = (Tdata *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1));
+    char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1));
 
-    size_t nram_usable = NRAM_MAX_SIZE - ((char *)nram_x - nram_buffer);
-    size_t max_chunk_elements = nram_usable / sizeof(Tdata);
+    size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer);
+    size_t max_chunk_elements = nram_usable / (sizeof(Tdata) + sizeof(float));
 
-    int align_elements = ALIGN_SIZE / sizeof(Tdata);
+    size_t align_elements = ALIGN_SIZE / sizeof(Tdata);
     if (align_elements == 0) {
         align_elements = 1;
     }
-    max_chunk_elements = (max_chunk_elements / align_elements) * align_elements;
+    int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements);
+
+    Tdata *nram_x = (Tdata *)nram_aligned;
+    float *nram_compute = (float *)(nram_x + chunk_size);
 
     int elements_per_core = n / taskDim;
     int remain = n % taskDim;
     int core_elements = elements_per_core + (taskId < remain ? 1 : 0);
     int core_offset = taskId < remain ? taskId * core_elements : taskId * elements_per_core + remain;
 
-    int chunks = core_elements / max_chunk_elements;
-    int chunk_rem = core_elements % max_chunk_elements;
+    int chunks = core_elements / chunk_size;
+    int chunk_rem = core_elements % chunk_size;
 
-    Tdata partial_sum = static_cast<Tdata>(0);
+    float partial_sum = 0.0f;
 
     for (int c = 0; c < chunks; c++) {
-        size_t current_offset = core_offset + c * max_chunk_elements;
-        __memcpy(nram_x, x + current_offset, max_chunk_elements * sizeof(Tdata), GDRAM2NRAM);
+        int current_offset = core_offset + c * chunk_size;
 
-        __bang_abs(nram_x, nram_x, max_chunk_elements);
+        __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM);
 
-        partial_sum += __bang_sum(nram_x, max_chunk_elements);
+        asumToCompute(nram_compute, nram_x, chunk_size);
+        __bang_abs(nram_compute, nram_compute, chunk_size);
+
+        partial_sum += __bang_sum(nram_compute, chunk_size);
     }
 
     if (chunk_rem > 0) {
-        size_t current_offset = core_offset + chunks * max_chunk_elements;
+        int current_offset = core_offset + chunks * chunk_size;
 
         __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM);
 
-        __bang_abs(nram_x, nram_x, chunk_rem);
+        asumToCompute(nram_compute, nram_x, chunk_rem);
+        __bang_abs(nram_compute, nram_compute, chunk_rem);
 
-        partial_sum += __bang_sum(nram_x, chunk_rem);
+        partial_sum += __bang_sum(nram_compute, chunk_rem);
     }
 
     shared_partial_sum[coreId] = partial_sum;
 
     __sync_cluster();
 
     if (coreId == 0) {
-        Tdata cluster_sum = static_cast<Tdata>(0);
+        float cluster_sum = 0.0f;
 
         for (int i = 0; i < coreDim; i++) {
             cluster_sum += shared_partial_sum[i];
         }
 
-        result[0] = cluster_sum;
+        asumStoreResult(result, nram_x, nram_compute, cluster_sum);
     }
 }
 
 template <typename Tdata>
 __mlu_global__ void asumKernelStrided(
-    size_t n,
+    int n,
     const Tdata *x,
-    size_t incx,
+    int incx,
     Tdata *result) {
 
-    __mlu_shared__ Tdata shared_partial_sum[4];
+    __mlu_shared__ float shared_partial_sum[4];
+
+    char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1));
+
+    float *nram_compute = (float *)nram_aligned;
+    Tdata *nram_result = (Tdata *)(nram_compute + 1);
 
     int elements_per_core = n / taskDim;
     int remain = n % taskDim;
     int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0);
     int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain;
 
-    Tdata partial_sum = static_cast<Tdata>(0);
+    float partial_sum = 0.0f;
 
     for (int i = start_idx; i < start_idx + actual_tasks; ++i) {
-        size_t offset = i * incx;
-        Tdata abs_val = x[offset] > static_cast<Tdata>(0) ? x[offset] : -x[offset];
+        int offset = i * incx;
+        float x_val = asumToCompute(x[offset]);
+        float abs_val = x_val > 0.0f ? x_val : -x_val;
 
         partial_sum += abs_val;
     }
@@ -94,12 +144,12 @@ __mlu_global__ void asumKernelStrided(
     __sync_cluster();
 
     if (coreId == 0) {
-        Tdata cluster_sum = static_cast<Tdata>(0);
+        float cluster_sum = 0.0f;
 
         for (int i = 0; i < coreDim; i++) {
             cluster_sum += shared_partial_sum[i];
         }
 
-        result[0] = cluster_sum;
+        asumStoreResult(result, nram_result, nram_compute, cluster_sum);
     }
-}
+}

src/infiniop/ops/asum/cpu/asum_cpu.cc

@@ -31,22 +31,24 @@ infiniStatus_t calculateAsum(
     const Tdata *x,
     Tdata *result) {
 
-    const ptrdiff_t n = info.n;
+    const size_t n = info.n;
     const ptrdiff_t incx = info.incx;
 
     if constexpr (std::is_same<Tdata, fp16_t>::value || std::is_same<Tdata, bf16_t>::value) {
         float total_sum = 0.0;
 
-        for (ptrdiff_t i = 0; i < n; ++i) {
-            total_sum += std::abs(utils::cast<float>(x[i * incx]));
+        for (size_t i = 0; i < n; ++i) {
+            const ptrdiff_t idx = utils::cast<ptrdiff_t>(i) * incx;
+            total_sum += std::abs(utils::cast<float>(x[idx]));
         }
 
         result[0] = utils::cast<Tdata>(total_sum);
     } else {
         Tdata total_sum = 0.0;
 
-        for (ptrdiff_t i = 0; i < n; ++i) {
-            total_sum += std::abs(x[i * incx]);
+        for (size_t i = 0; i < n; ++i) {
+            const ptrdiff_t idx = utils::cast<ptrdiff_t>(i) * incx;
+            total_sum += std::abs(x[idx]);
         }
 
         result[0] = total_sum;

src/infiniop/ops/asum/metax/asum_metax.cc

@@ -42,8 +42,8 @@ infiniStatus_t Descriptor::calculate(
     (void)workspace;
     (void)workspace_size;
 
-    const size_t n = _info.n;
-    const ptrdiff_t incx = _info.incx;
+    const int n = utils::cast<int>(_info.n);
+    const int incx = utils::cast<int>(_info.incx);
     const infiniDtype_t data_type = _info.data_type;
 
     CHECK_STATUS(_opaque->internal->useMcblas(

src/infiniop/ops/axpy/bang/axpy_bang.mlu

@@ -36,9 +36,9 @@ infiniStatus_t calculateAxpy(
     Tdata *y,
     cnrtQueue_t queue) {
 
-    const size_t size = info.n;
-    const ptrdiff_t incx = info.incx;
-    const ptrdiff_t incy = info.incy;
+    const int n = utils::cast<int>(info.n);
+    const int incx = utils::cast<int>(info.incx);
+    const int incy = utils::cast<int>(info.incy);
 
     cnrtDim3_t k_dim;
     cnrtFunctionType_t k_type;
@@ -50,13 +50,13 @@ infiniStatus_t calculateAxpy(
 
     if (incx == 1 && incy == 1) {
         axpyKernelContiguous<Tdata><<<k_dim, k_type, queue>>>(
-            size,
+            n,
             alpha,
             x,
             y);
     } else {
         axpyKernelStrided<Tdata><<<k_dim, k_type, queue>>>(
-            size,
+            n,
             alpha,
             x,
             incx,

src/infiniop/ops/axpy/bang/axpy_bang_kernel.mlu

@@ -5,24 +5,24 @@ __nram__ char nram_buffer[NRAM_MAX_SIZE];
 
 template <typename Tdata>
 __mlu_global__ void axpyKernelContiguous(
-    size_t n,
+    int n,
     const Tdata *alpha,
     const Tdata *x,
     Tdata *y) {
 
-    Tdata *nram_align = (Tdata *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1));
+    char *nram_aligned = (char *)(((size_t)nram_buffer + ALIGN_SIZE - 1) & ~(ALIGN_SIZE - 1));
 
-    size_t nram_usable = NRAM_MAX_SIZE - ((char *)nram_align - nram_buffer);
+    size_t nram_usable = NRAM_MAX_SIZE - (nram_aligned - nram_buffer);
     size_t max_chunk_elements = nram_usable / (2 * sizeof(Tdata));
 
-    int align_elements = ALIGN_SIZE / sizeof(Tdata);
+    size_t align_elements = ALIGN_SIZE / sizeof(Tdata);
     if (align_elements == 0) {
         align_elements = 1;
     }
-    max_chunk_elements = (max_chunk_elements / align_elements) * align_elements;
+    int chunk_size = (int)((max_chunk_elements / align_elements) * align_elements);
 
-    Tdata *nram_x = nram_align;
-    Tdata *nram_y = nram_align + max_chunk_elements;
+    Tdata *nram_x = (Tdata *)nram_aligned;
+    Tdata *nram_y = nram_x + chunk_size;
 
     int elements_per_core = n / taskDim;
     int remain = n % taskDim;
@@ -33,22 +33,23 @@ __mlu_global__ void axpyKernelContiguous(
         return;
     }
 
-    int chunks = core_elements / max_chunk_elements;
-    int chunk_rem = core_elements % max_chunk_elements;
+    int chunks = core_elements / chunk_size;
+    int chunk_rem = core_elements % chunk_size;
 
     for (int c = 0; c < chunks; c++) {
-        size_t current_offset = core_offset + c * max_chunk_elements;
-        __memcpy(nram_x, x + current_offset, max_chunk_elements * sizeof(Tdata), GDRAM2NRAM);
-        __memcpy(nram_y, y + current_offset, max_chunk_elements * sizeof(Tdata), GDRAM2NRAM);
+        int current_offset = core_offset + c * chunk_size;
 
-        __bang_mul_scalar(nram_x, nram_x, alpha[0], max_chunk_elements);
-        __bang_add(nram_y, nram_y, nram_x, max_chunk_elements);
+        __memcpy(nram_x, x + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM);
+        __memcpy(nram_y, y + current_offset, chunk_size * sizeof(Tdata), GDRAM2NRAM);
 
-        __memcpy(y + current_offset, nram_y, max_chunk_elements * sizeof(Tdata), NRAM2GDRAM);
+        __bang_mul_scalar(nram_x, nram_x, alpha[0], chunk_size);
+        __bang_add(nram_y, nram_y, nram_x, chunk_size);
+
+        __memcpy(y + current_offset, nram_y, chunk_size * sizeof(Tdata), NRAM2GDRAM);
     }
 
     if (chunk_rem > 0) {
-        size_t current_offset = core_offset + chunks * max_chunk_elements;
+        int current_offset = core_offset + chunks * chunk_size;
         int align_rem = ((chunk_rem + align_elements - 1) / align_elements) * align_elements;
 
         __memcpy(nram_x, x + current_offset, chunk_rem * sizeof(Tdata), GDRAM2NRAM);
@@ -63,22 +64,22 @@ __mlu_global__ void axpyKernelContiguous(
 
 template <typename Tdata>
 __mlu_global__ void axpyKernelStrided(
-    size_t n,
+    int n,
     const Tdata *alpha,
     const Tdata *x,
-    size_t incx,
+    int incx,
     Tdata *y,
-    size_t incy) {
+    int incy) {
 
     int elements_per_core = n / taskDim;
     int remain = n % taskDim;
     int actual_tasks = elements_per_core + (taskId < remain ? 1 : 0);
     int start_idx = taskId < remain ? taskId * actual_tasks : taskId * elements_per_core + remain;
 
     for (int i = start_idx; i < start_idx + actual_tasks; ++i) {
-        size_t idx_x = i * incx;
-        size_t idx_y = i * incy;
+        int idx_x = i * incx;
+        int idx_y = i * incy;
 
         y[idx_y] += alpha[0] * x[idx_x];
     }
-}
+}