mirror of
				https://github.com/ggml-org/llama.cpp.git
				synced 2025-10-28 08:31:25 +00:00 
			
		
		
		
	[CANN] get_rows and dup optimization (#12671)
* [CANN]get_rows and dup optimization. Co-authored-by: hipudding <huafengchun@gmail.com> Signed-off-by: noemotiovon <noemotiovon@gmail.com> * [CANN]GET_ROWS and CPY/DUP optimization Co-authored-by: hipudding <huafengchun@gmail.com> Signed-off-by: noemotiovon <noemotiovon@gmail.com> * [CANN]code style adjustment Signed-off-by: noemotiovon <noemotiovon@gmail.com> * [CANN]code style adjustment Signed-off-by: noemotiovon <noemotiovon@gmail.com> * [CANN]code style adjustment Signed-off-by: noemotiovon <noemotiovon@gmail.com> * [CANN]code style adjustment Signed-off-by: noemotiovon <noemotiovon@gmail.com> --------- Signed-off-by: noemotiovon <noemotiovon@gmail.com> Co-authored-by: noemotiovon <noemotiovon@gmail.com> Co-authored-by: hipudding <huafengchun@gmail.com>
This commit is contained in:
		| @@ -51,13 +51,11 @@ if (CANN_INSTALL_DIR) | |||||||
|         ${CANN_INSTALL_DIR}/acllib/include |         ${CANN_INSTALL_DIR}/acllib/include | ||||||
|     ) |     ) | ||||||
|  |  | ||||||
|     add_subdirectory(kernels) |  | ||||||
|     list(APPEND CANN_LIBRARIES |     list(APPEND CANN_LIBRARIES | ||||||
|         ascendcl |         ascendcl | ||||||
|         nnopbase |         nnopbase | ||||||
|         opapi |         opapi | ||||||
|         acl_op_compiler |         acl_op_compiler | ||||||
|         ascendc_kernels |  | ||||||
|     ) |     ) | ||||||
|  |  | ||||||
|     file(GLOB GGML_SOURCES_CANN "*.cpp") |     file(GLOB GGML_SOURCES_CANN "*.cpp") | ||||||
|   | |||||||
| @@ -30,6 +30,7 @@ | |||||||
| #include <aclnnop/aclnn_copy.h> | #include <aclnnop/aclnn_copy.h> | ||||||
| #include <aclnnop/aclnn_cos.h> | #include <aclnnop/aclnn_cos.h> | ||||||
| #include <aclnnop/aclnn_div.h> | #include <aclnnop/aclnn_div.h> | ||||||
|  | #include <aclnnop/aclnn_embedding.h> | ||||||
| #include <aclnnop/aclnn_exp.h> | #include <aclnnop/aclnn_exp.h> | ||||||
| #include <aclnnop/aclnn_fill_scalar.h> | #include <aclnnop/aclnn_fill_scalar.h> | ||||||
| #include <aclnnop/aclnn_group_norm.h> | #include <aclnnop/aclnn_group_norm.h> | ||||||
| @@ -58,7 +59,6 @@ | |||||||
| #include <vector> | #include <vector> | ||||||
|  |  | ||||||
| #include "ggml-impl.h" | #include "ggml-impl.h" | ||||||
| #include "kernels/ascendc_kernels.h" |  | ||||||
|  |  | ||||||
| #define GGML_COMMON_DECL_C | #define GGML_COMMON_DECL_C | ||||||
|  |  | ||||||
| @@ -99,6 +99,35 @@ static void aclnn_repeat(ggml_backend_cann_context& ctx, aclTensor* acl_src, | |||||||
|     ACL_CHECK(aclDestroyIntArray(repeats)); |     ACL_CHECK(aclDestroyIntArray(repeats)); | ||||||
| } | } | ||||||
|  |  | ||||||
|  | /** | ||||||
|  |  * @brief Casts the elements of a tensor to a specified data type using the CANN backend. | ||||||
|  |  * | ||||||
|  |  * @details This function performs a type conversion on the elements of the input tensor `acl_src` | ||||||
|  |  *          and stores the results in the destination tensor `acl_dst`. The conversion type is | ||||||
|  |  *          determined based on the `dst` tensor's data type. | ||||||
|  |  * | ||||||
|  |  * @param ctx The context for the CANN backend operations. | ||||||
|  |  * @param acl_src The source tensor whose elements will be cast. | ||||||
|  |  * @param acl_dst The destination tensor that will store the casted elements. | ||||||
|  |  * @param dst The ggml tensor specifying the target data type. | ||||||
|  |  */ | ||||||
|  | static void aclnn_cast(ggml_backend_cann_context& ctx, aclTensor* acl_src, | ||||||
|  |                     aclTensor* acl_dst, ggml_tensor* dst) { | ||||||
|  |     uint64_t workspaceSize = 0; | ||||||
|  |     aclOpExecutor* executor; | ||||||
|  |     void* workspaceAddr = nullptr; | ||||||
|  |     ACL_CHECK(aclnnCastGetWorkspaceSize(acl_src, | ||||||
|  |                                         ggml_cann_type_mapping(dst->type), | ||||||
|  |                                         acl_dst, &workspaceSize, &executor)); | ||||||
|  |  | ||||||
|  |     if (workspaceSize > 0) { | ||||||
|  |         ggml_cann_pool_alloc workspace_allocator(ctx.pool(), workspaceSize); | ||||||
|  |         workspaceAddr = workspace_allocator.get(); | ||||||
|  |     } | ||||||
|  |  | ||||||
|  |     ACL_CHECK(aclnnCast(workspaceAddr, workspaceSize, executor, ctx.stream())); | ||||||
|  | } | ||||||
|  |  | ||||||
| void ggml_cann_repeat(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | void ggml_cann_repeat(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | ||||||
|     ggml_tensor* src = dst->src[0]; |     ggml_tensor* src = dst->src[0]; | ||||||
|     GGML_ASSERT(ggml_can_repeat(src, dst)); |     GGML_ASSERT(ggml_can_repeat(src, dst)); | ||||||
| @@ -889,173 +918,76 @@ static void cann_copy(ggml_backend_cann_context& ctx, aclTensor* acl_src, | |||||||
| } | } | ||||||
|  |  | ||||||
| void ggml_cann_dup(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | void ggml_cann_dup(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | ||||||
|     ggml_tensor* src = dst->src[0]; |     ggml_tensor* src0 = dst->src[0]; | ||||||
|  |  | ||||||
|     aclTensor* acl_src = ggml_cann_create_tensor(src); |     aclTensor* acl_src = ggml_cann_create_tensor(src0); | ||||||
|     aclTensor* acl_dst = ggml_cann_create_tensor(dst); |     aclTensor* acl_dst = ggml_cann_create_tensor(dst); | ||||||
|  |     if (ggml_are_same_shape(src0, dst)) { | ||||||
|     ggml_cann_pool_alloc src_extra_allocator(ctx.pool(), sizeof(ggml_tensor)); |         if (dst->type == src0->type) { | ||||||
|     ggml_cann_pool_alloc dst_extra_allocator(ctx.pool(), sizeof(ggml_tensor)); |  | ||||||
|     src->extra = src_extra_allocator.get(); |  | ||||||
|     dst->extra = dst_extra_allocator.get(); |  | ||||||
|     ACL_CHECK(aclrtMemcpyAsync(src->extra, sizeof(ggml_tensor), src, |  | ||||||
|                                sizeof(ggml_tensor), ACL_MEMCPY_HOST_TO_DEVICE, |  | ||||||
|                                ctx.stream())); |  | ||||||
|     ACL_CHECK(aclrtMemcpyAsync(dst->extra, sizeof(ggml_tensor), dst, |  | ||||||
|                                sizeof(ggml_tensor), ACL_MEMCPY_HOST_TO_DEVICE, |  | ||||||
|                                ctx.stream())); |  | ||||||
|  |  | ||||||
|     if ((dst->type == GGML_TYPE_F16 || dst->type == GGML_TYPE_F32) && |  | ||||||
|         ggml_are_same_shape(src, dst)) { |  | ||||||
|         cann_copy(ctx, acl_src, acl_dst); |  | ||||||
|         ACL_CHECK(aclDestroyTensor(acl_src)); |  | ||||||
|         ACL_CHECK(aclDestroyTensor(acl_dst)); |  | ||||||
|         return; |  | ||||||
|     } |  | ||||||
|     // TODO: simplify |  | ||||||
|     if (src->type == GGML_TYPE_F16) { |  | ||||||
|         if (dst->type == GGML_TYPE_Q8_0) { |  | ||||||
|             aclrtlaunch_ascendc_quantize_f16_q8_0( |  | ||||||
|                 24, ctx.stream(), src->data, dst->data, |  | ||||||
|                 ((ggml_tensor*)src->extra)->ne, ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                 ((ggml_tensor*)dst->extra)->ne); |  | ||||||
|             return; |  | ||||||
|         } |  | ||||||
|         if (dst->type == GGML_TYPE_Q4_0) { |  | ||||||
|             aclrtlaunch_ascendc_quantize_f16_to_q4_0( |  | ||||||
|                 24, ctx.stream(), src->data, dst->data, |  | ||||||
|                 ((ggml_tensor*)src->extra)->ne, ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                 ((ggml_tensor*)dst->extra)->ne); |  | ||||||
|             return; |  | ||||||
|         } |  | ||||||
|         if (dst->type == GGML_TYPE_F16) { |  | ||||||
|             if (ggml_are_same_shape(src, dst)) { |  | ||||||
|                 cann_copy(ctx, acl_src, acl_dst); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_src)); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_dst)); |  | ||||||
|                 return; |  | ||||||
|             } |  | ||||||
|             if (ggml_is_contiguous(dst)) { |  | ||||||
|                 const size_t src_type_size = ggml_type_size(src->type); |  | ||||||
|                 if (src->nb[0] == src_type_size) { |  | ||||||
|                     // src0 is contigous on first dimension, copy by rows |  | ||||||
|                     int64_t rows_num = ggml_nrows(src); |  | ||||||
|  |  | ||||||
|                     aclrtlaunch_ascendc_dup_by_rows_fp16( |  | ||||||
|                         rows_num, ctx.stream(), src->data, dst->data, |  | ||||||
|                         ((ggml_tensor*)src->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->nb); |  | ||||||
|                     return; |  | ||||||
|                 } |  | ||||||
|                 GGML_ABORT("fatal error"); |  | ||||||
|             } |  | ||||||
|             GGML_ABORT("fatal error"); |  | ||||||
|         } |  | ||||||
|         if (dst->type == GGML_TYPE_F32) { |  | ||||||
|             if (ggml_are_same_shape(src, dst)) { |  | ||||||
|                 cann_copy(ctx, acl_src, acl_dst); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_src)); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_dst)); |  | ||||||
|                 return; |  | ||||||
|             } |  | ||||||
|             if (ggml_is_contiguous(dst)) { |  | ||||||
|                 const size_t src_type_size = ggml_type_size(src->type); |  | ||||||
|                 if (src->nb[0] == src_type_size) { |  | ||||||
|                     // src0 is contigous on first dimension, copy by rows |  | ||||||
|                     int64_t rows_num = ggml_nrows(src); |  | ||||||
|                     aclrtlaunch_ascendc_dup_by_rows_fp16_to_fp32( |  | ||||||
|                         rows_num, ctx.stream(), src->data, dst->data, |  | ||||||
|                         ((ggml_tensor*)src->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->nb); |  | ||||||
|                     return; |  | ||||||
|                 } |  | ||||||
|                 GGML_ABORT("fatal error"); |  | ||||||
|             } |  | ||||||
|             GGML_ABORT("fatal error"); |  | ||||||
|         } |  | ||||||
|         // TODO |  | ||||||
|         GGML_ABORT("fatal error"); |  | ||||||
|     } else if (src->type == GGML_TYPE_F32) { |  | ||||||
|         // TODO: if (src0->type == dst->type && ne00 == ne0 && nb00 == type_size |  | ||||||
|         //          && nb0 == type_size) |  | ||||||
|         if (dst->type == GGML_TYPE_Q8_0) { |  | ||||||
|             aclrtlaunch_ascendc_quantize_f32_q8_0( |  | ||||||
|                 24, ctx.stream(), src->data, dst->data, |  | ||||||
|                 ((ggml_tensor*)src->extra)->ne, ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                 ((ggml_tensor*)dst->extra)->ne); |  | ||||||
|             return; |  | ||||||
|         } |  | ||||||
|         if (dst->type == GGML_TYPE_Q4_0) { |  | ||||||
|             aclrtlaunch_ascendc_quantize_f32_to_q4_0( |  | ||||||
|                 24, ctx.stream(), src->data, dst->data, |  | ||||||
|                 ((ggml_tensor*)src->extra)->ne, ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                 ((ggml_tensor*)dst->extra)->ne); |  | ||||||
|             return; |  | ||||||
|         } |  | ||||||
|         if (dst->type == GGML_TYPE_F32) { |  | ||||||
|             if (ggml_are_same_shape(src, dst)) { |  | ||||||
|                 cann_copy(ctx, acl_src, acl_dst); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_src)); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_dst)); |  | ||||||
|                 return; |  | ||||||
|             } |  | ||||||
|             if (ggml_is_contiguous(dst)) { |  | ||||||
|                 const size_t src_type_size = ggml_type_size(src->type); |  | ||||||
|                 if (src->nb[0] == src_type_size) { |  | ||||||
|                     // src0 is contigous on first dimension, copy by rows |  | ||||||
|                     int64_t rows_num = ggml_nrows(src); |  | ||||||
|                     aclrtlaunch_ascendc_dup_by_rows_fp32( |  | ||||||
|                         rows_num, ctx.stream(), src->data, dst->data, |  | ||||||
|                         ((ggml_tensor*)src->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->nb); |  | ||||||
|                     return; |  | ||||||
|                 } |  | ||||||
|                 GGML_ABORT("fatal error"); |  | ||||||
|             } else { |  | ||||||
|                 // TODO: dst not contiguous |  | ||||||
|                 GGML_ABORT("fatal error"); |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|         if (dst->type == GGML_TYPE_F16) { |  | ||||||
|             if (ggml_are_same_shape(src, dst)) { |  | ||||||
|                 cann_copy(ctx, acl_src, acl_dst); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_src)); |  | ||||||
|                 ACL_CHECK(aclDestroyTensor(acl_dst)); |  | ||||||
|                 return; |  | ||||||
|             } |  | ||||||
|             if (ggml_is_contiguous(dst)) { |  | ||||||
|                 const size_t src_type_size = ggml_type_size(src->type); |  | ||||||
|                 if (src->nb[0] == src_type_size) { |  | ||||||
|                     // src0 is contigous on first dimension, copy by rows |  | ||||||
|                     int64_t rows_num = ggml_nrows(src); |  | ||||||
|                     aclrtlaunch_ascendc_dup_by_rows_fp32_to_fp16( |  | ||||||
|                         rows_num, ctx.stream(), src->data, dst->data, |  | ||||||
|                         ((ggml_tensor*)src->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)src->extra)->nb, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->ne, |  | ||||||
|                         ((ggml_tensor*)dst->extra)->nb); |  | ||||||
|                     return; |  | ||||||
|                 } |  | ||||||
|                 GGML_ABORT("fatal error"); |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|         // TODO |  | ||||||
|         GGML_ABORT("fatal error"); |  | ||||||
|     } else { |  | ||||||
|         if (ggml_are_same_shape(src, dst)) { |  | ||||||
|             cann_copy(ctx, acl_src, acl_dst); |             cann_copy(ctx, acl_src, acl_dst); | ||||||
|             ACL_CHECK(aclDestroyTensor(acl_src)); |         } else { | ||||||
|             ACL_CHECK(aclDestroyTensor(acl_dst)); |             aclnn_cast(ctx, acl_src, acl_dst, dst); | ||||||
|             return; |         } | ||||||
|  |     } else { | ||||||
|  |         if (ggml_is_contiguous(src0) && ggml_is_contiguous(dst)) { | ||||||
|  |             if (dst->type == src0->type) { | ||||||
|  |                 size_t cpy_size = ggml_nbytes(dst); | ||||||
|  |                 ACL_CHECK(aclrtMemcpyAsync( | ||||||
|  |                     dst->data, cpy_size, src0->data, cpy_size, | ||||||
|  |                     ACL_MEMCPY_DEVICE_TO_DEVICE, ctx.stream())); | ||||||
|  |                 return; | ||||||
|  |             } else { | ||||||
|  |                 ggml_cann_pool_alloc src_buffer_allocator( | ||||||
|  |                     ctx.pool(), | ||||||
|  |                     ggml_nelements(dst) * ggml_type_size(dst->type)); | ||||||
|  |                 void* src_trans_buffer = src_buffer_allocator.get(); | ||||||
|  |                 size_t src_trans_nb[GGML_MAX_DIMS]; | ||||||
|  |                 src_trans_nb[0] = ggml_type_size(dst->type); | ||||||
|  |                 for (int i = 1; i < GGML_MAX_DIMS; i++) { | ||||||
|  |                     src_trans_nb[i] = src_trans_nb[i - 1] * src0->ne[i - 1]; | ||||||
|  |                 } | ||||||
|  |                 aclTensor* src_trans_tensor = ggml_cann_create_tensor( | ||||||
|  |                     src_trans_buffer, ggml_cann_type_mapping(dst->type), | ||||||
|  |                     ggml_type_size(dst->type), src0->ne, src_trans_nb, | ||||||
|  |                     GGML_MAX_DIMS); | ||||||
|  |  | ||||||
|  |                 aclnn_cast(ctx, acl_src, src_trans_tensor, dst); | ||||||
|  |                 size_t cpy_size = ggml_nbytes(dst); | ||||||
|  |                 ACL_CHECK(aclrtMemcpyAsync( | ||||||
|  |                     dst->data, cpy_size, src_trans_buffer, cpy_size, | ||||||
|  |                     ACL_MEMCPY_DEVICE_TO_DEVICE, ctx.stream())); | ||||||
|  |                 ACL_CHECK(aclDestroyTensor(src_trans_tensor)); | ||||||
|  |                 return; | ||||||
|  |             } | ||||||
|  |         } else if (ggml_is_contiguous(dst)) { | ||||||
|  |             ggml_cann_pool_alloc src_buffer_allocator( | ||||||
|  |                 ctx.pool(), ggml_nelements(dst) * ggml_type_size(dst->type)); | ||||||
|  |             void* src_trans_buffer = src_buffer_allocator.get(); | ||||||
|  |             size_t src_trans_nb[GGML_MAX_DIMS]; | ||||||
|  |             src_trans_nb[0] = ggml_type_size(dst->type); | ||||||
|  |             for (int i = 1; i < GGML_MAX_DIMS; i++) { | ||||||
|  |                 src_trans_nb[i] = src_trans_nb[i - 1] * src0->ne[i - 1]; | ||||||
|  |             } | ||||||
|  |             aclTensor* src_trans_tensor = ggml_cann_create_tensor( | ||||||
|  |                 src_trans_buffer, ggml_cann_type_mapping(dst->type), | ||||||
|  |                 ggml_type_size(dst->type), src0->ne, src_trans_nb, | ||||||
|  |                 GGML_MAX_DIMS); | ||||||
|  |  | ||||||
|  |             aclnn_cast(ctx, acl_src, src_trans_tensor, dst); | ||||||
|  |  | ||||||
|  |             size_t cpy_size = ggml_nbytes(dst); | ||||||
|  |             ACL_CHECK(aclrtMemcpyAsync(dst->data, cpy_size, src_trans_buffer, | ||||||
|  |                                        cpy_size, ACL_MEMCPY_DEVICE_TO_DEVICE, | ||||||
|  |                                        ctx.stream())); | ||||||
|  |             ACL_CHECK(aclDestroyTensor(src_trans_tensor)); | ||||||
|  |             return; | ||||||
|  |         } else { | ||||||
|  |             GGML_ABORT("Unsupport dst is not tontiguous."); | ||||||
|         } |         } | ||||||
|         GGML_ABORT("fatal error"); |  | ||||||
|     } |     } | ||||||
|  |  | ||||||
|  |     ACL_CHECK(aclDestroyTensor(acl_src)); | ||||||
|  |     ACL_CHECK(aclDestroyTensor(acl_dst)); | ||||||
| } | } | ||||||
|  |  | ||||||
| #ifdef __cplusplus | #ifdef __cplusplus | ||||||
| @@ -2378,85 +2310,168 @@ void ggml_cann_softmax(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | |||||||
|     ACL_CHECK(aclDestroyTensor(tmp_mask_tensor)); |     ACL_CHECK(aclDestroyTensor(tmp_mask_tensor)); | ||||||
| } | } | ||||||
|  |  | ||||||
| void ggml_cann_get_rows(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | /** | ||||||
|     ggml_tensor* src0 = dst->src[0]; |  * @brief Performs embedding operation on a 4D tensor using the CANN backend. | ||||||
|     ggml_tensor* src1 = dst->src[1]; |  * | ||||||
|  |  * This function extracts slices from the source tensor (`src_buffer`), | ||||||
|  |  * index tensor (`index`), and destination tensor (`dst`), and performs an | ||||||
|  |  * embedding operation on them. The embedding operation is applied by iterating | ||||||
|  |  * over the last two dimensions of the source tensor, creating the necessary | ||||||
|  |  * tensors for the source, index, and output, and executing the embedding operation. | ||||||
|  |  * | ||||||
|  |  * @param ctx The context for CANN backend operations. | ||||||
|  |  * @param src_buffer The source buffer holding the data for the source tensor. | ||||||
|  |  * @param src_ne The dimensions of the source tensor. | ||||||
|  |  * @param src_nb The strides (byte offsets) of the source tensor. | ||||||
|  |  * @param index The index tensor used in the embedding operation. | ||||||
|  |  * @param dst The destination tensor where the result will be stored. | ||||||
|  |  */ | ||||||
|  | static void aclnn_embedding_4d(ggml_backend_cann_context& ctx, void* src_buffer, | ||||||
|  |                             int64_t* src_ne, size_t* src_nb, ggml_tensor* index, | ||||||
|  |                             ggml_tensor* dst) { | ||||||
|  |     for (int64_t i = 0; i < src_ne[3]; i++) { | ||||||
|  |         for (int64_t j = 0; j < src_ne[2]; j++) { | ||||||
|  |             // src | ||||||
|  |             int64_t acl_src_ne[2] = {src_ne[0], src_ne[1]}; | ||||||
|  |             size_t acl_src_nb[2] = {src_nb[0], src_nb[1]}; | ||||||
|  |             aclTensor* acl_src_tensor = ggml_cann_create_tensor( | ||||||
|  |                 (char*)src_buffer + i * src_nb[3] + j * src_nb[2], | ||||||
|  |                 ggml_cann_type_mapping(dst->type), ggml_element_size(dst), | ||||||
|  |                 acl_src_ne, acl_src_nb, 2); | ||||||
|  |  | ||||||
|     ggml_cann_pool_alloc src0_extra_allocator(ctx.pool(), sizeof(ggml_tensor)); |             // index | ||||||
|     ggml_cann_pool_alloc src1_extra_allocator(ctx.pool(), sizeof(ggml_tensor)); |             int64_t acl_index_ne[1] = {index->ne[0]}; | ||||||
|     ggml_cann_pool_alloc dst_extra_allocator(ctx.pool(), sizeof(ggml_tensor)); |             size_t acl_index_nb[1] = {index->nb[0]}; | ||||||
|     src0->extra = src0_extra_allocator.get(); |             aclTensor* acl_index = ggml_cann_create_tensor( | ||||||
|     src1->extra = src1_extra_allocator.get(); |                 (char*)index->data + i * index->nb[2] + j * index->nb[1], | ||||||
|     dst->extra = dst_extra_allocator.get(); |                 ggml_cann_type_mapping(index->type), ggml_element_size(index), | ||||||
|     ACL_CHECK(aclrtMemcpyAsync(src0->extra, sizeof(ggml_tensor), src0, |                 acl_index_ne, acl_index_nb, 1); | ||||||
|                                sizeof(ggml_tensor), ACL_MEMCPY_HOST_TO_DEVICE, |  | ||||||
|                                ctx.stream())); |             // out | ||||||
|     ACL_CHECK(aclrtMemcpyAsync(src1->extra, sizeof(ggml_tensor), src1, |             int64_t acl_out_ne[2] = {dst->ne[0], dst->ne[1]}; | ||||||
|                                sizeof(ggml_tensor), ACL_MEMCPY_HOST_TO_DEVICE, |             size_t acl_out_nb[2] = {dst->nb[0], dst->nb[1]}; | ||||||
|                                ctx.stream())); |             aclTensor* acl_out = ggml_cann_create_tensor( | ||||||
|     ACL_CHECK(aclrtMemcpyAsync(dst->extra, sizeof(ggml_tensor), dst, |                 (char*)dst->data + i * dst->nb[3] + j * dst->nb[2], | ||||||
|                                sizeof(ggml_tensor), ACL_MEMCPY_HOST_TO_DEVICE, |                 ggml_cann_type_mapping(dst->type), ggml_element_size(dst), | ||||||
|                                ctx.stream())); |                 acl_out_ne, acl_out_nb, 2); | ||||||
|  |  | ||||||
|  |             uint64_t workspaceSize = 0; | ||||||
|  |             aclOpExecutor* executor; | ||||||
|  |             void* workspaceAddr = nullptr; | ||||||
|  |  | ||||||
|  |             ACL_CHECK(aclnnEmbeddingGetWorkspaceSize( | ||||||
|  |                 acl_src_tensor, acl_index, acl_out, &workspaceSize, &executor)); | ||||||
|  |  | ||||||
|  |             if (workspaceSize > 0) { | ||||||
|  |                 ggml_cann_pool_alloc workspace_allocator(ctx.pool(), | ||||||
|  |                                                          workspaceSize); | ||||||
|  |                 workspaceAddr = workspace_allocator.get(); | ||||||
|  |             } | ||||||
|  |  | ||||||
|  |             ACL_CHECK(aclnnEmbedding(workspaceAddr, workspaceSize, executor, | ||||||
|  |                                      ctx.stream())); | ||||||
|  |  | ||||||
|  |             ACL_CHECK(aclDestroyTensor(acl_src_tensor)); | ||||||
|  |             ACL_CHECK(aclDestroyTensor(acl_index)); | ||||||
|  |             ACL_CHECK(aclDestroyTensor(acl_out)); | ||||||
|  |         } | ||||||
|  |     } | ||||||
|  | } | ||||||
|  |  | ||||||
|  | void ggml_cann_get_rows(ggml_backend_cann_context& ctx, ggml_tensor* dst) { | ||||||
|  |     ggml_tensor* src0 = dst->src[0];  // src | ||||||
|  |     ggml_tensor* src1 = dst->src[1];  // index | ||||||
|  |  | ||||||
|     switch (src0->type) { |     switch (src0->type) { | ||||||
|         case GGML_TYPE_F32: { |         case GGML_TYPE_F32: { | ||||||
| #ifdef ASCEND_310P |             aclnn_embedding_4d(ctx, src0->data, src0->ne, src0->nb, src1, | ||||||
|             // Special operation for get_row_f32 kernel of 310P: clear the |                                    dst); | ||||||
|             // content of dest data buffer when row is not aligned to 32 bytes |  | ||||||
|             if ((src0->ne[0] % 8) != 0) { |  | ||||||
|                 size_t dst_len = src1->ne[0] * src1->ne[1] * src1->ne[2] * |  | ||||||
|                                  src0->ne[0] * ggml_type_size(GGML_TYPE_F32); |  | ||||||
|                 ACL_CHECK(aclrtMemset((char*)dst->data, dst_len, 0, dst_len)); |  | ||||||
|             } |  | ||||||
| #endif |  | ||||||
|             aclrtlaunch_ascendc_get_row_f32( |  | ||||||
|                 24, ctx.stream(), src0->data, src1->data, dst->data, |  | ||||||
|                 ((ggml_tensor*)src0->extra)->ne, |  | ||||||
|                 ((ggml_tensor*)src0->extra)->nb, |  | ||||||
|                 ((ggml_tensor*)src1->extra)->ne, |  | ||||||
|                 ((ggml_tensor*)src1->extra)->nb, ((ggml_tensor*)dst->extra)->ne, |  | ||||||
|                 ((ggml_tensor*)dst->extra)->nb); |  | ||||||
|             break; |             break; | ||||||
|         } |         } | ||||||
|         case GGML_TYPE_F16: { |         case GGML_TYPE_F16: { | ||||||
| #ifdef ASCEND_310P |             aclTensor* acl_src0 = ggml_cann_create_tensor(src0); | ||||||
|             // Special operation for get_row_f16 kernel of 310P: clear the |             ggml_cann_pool_alloc src_buffer_allocator( | ||||||
|             // content of dest data buffer when row is not aligned to 32 bytes |                 ctx.pool(), ggml_nelements(src0) * sizeof(float_t)); | ||||||
|             if ((src0->ne[0] % 16) != 0) { |             void* src_trans_buffer = src_buffer_allocator.get(); | ||||||
|                 size_t dst_len = |             size_t src_trans_nb[GGML_MAX_DIMS]; | ||||||
|                     src1->ne[0] * src1->ne[1] * src1->ne[2] * src0->ne[0] * |             src_trans_nb[0] = sizeof(float_t); | ||||||
|                     ggml_type_size( |             for (int i = 1; i < GGML_MAX_DIMS; i++) { | ||||||
|                         GGML_TYPE_F32);  // out is also f32, even input is f16 |                 src_trans_nb[i] = src_trans_nb[i - 1] * src0->ne[i - 1]; | ||||||
|                 ACL_CHECK(aclrtMemset((char*)dst->data, dst_len, 0, dst_len)); |  | ||||||
|             } |             } | ||||||
| #endif |             aclTensor* src_trans_tensor = ggml_cann_create_tensor( | ||||||
|             aclrtlaunch_ascendc_get_row_f16( |                 src_trans_buffer, ACL_FLOAT, ggml_type_size(dst->type), | ||||||
|                 24, ctx.stream(), src0->data, src1->data, dst->data, |                 src0->ne, src_trans_nb, GGML_MAX_DIMS); | ||||||
|                 ((ggml_tensor*)src0->extra)->ne, |             aclnn_cast(ctx, acl_src0, src_trans_tensor, dst); | ||||||
|                 ((ggml_tensor*)src0->extra)->nb, |             aclnn_embedding_4d(ctx, src_trans_buffer, src0->ne, | ||||||
|                 ((ggml_tensor*)src1->extra)->ne, |                                    src_trans_nb, src1, dst); | ||||||
|                 ((ggml_tensor*)src1->extra)->nb, ((ggml_tensor*)dst->extra)->ne, |             ACL_CHECK(aclDestroyTensor(acl_src0)); | ||||||
|                 ((ggml_tensor*)dst->extra)->nb); |             ACL_CHECK(aclDestroyTensor(src_trans_tensor)); | ||||||
|             break; |             break; | ||||||
|         } |         } | ||||||
|         case GGML_TYPE_Q4_0: |         case GGML_TYPE_Q8_0: { | ||||||
|             aclrtlaunch_ascendc_get_row_q4_0( |             // add 1 dim for bcast mul. | ||||||
|                 24, ctx.stream(), src0->data, src1->data, dst->data, |             size_t weight_nb[GGML_MAX_DIMS + 1], scale_nb[GGML_MAX_DIMS + 1], | ||||||
|                 ((ggml_tensor*)src0->extra)->ne, |                 dequant_nb[GGML_MAX_DIMS + 1]; | ||||||
|                 ((ggml_tensor*)src1->extra)->ne, |             int64_t weight_ne[GGML_MAX_DIMS + 1], scale_ne[GGML_MAX_DIMS + 1], | ||||||
|                 ((ggml_tensor*)src1->extra)->nb, ((ggml_tensor*)dst->extra)->ne, |                 *dequant_ne; | ||||||
|                 ((ggml_tensor*)dst->extra)->nb); |             int64_t scale_offset = 0; | ||||||
|             break; |  | ||||||
|         case GGML_TYPE_Q8_0: |             // [3,4,5,64] -> [3,4,5,2,32] | ||||||
|             aclrtlaunch_ascendc_get_row_q8_0( |             weight_ne[0] = QK8_0; | ||||||
|                 24, ctx.stream(), src0->data, src1->data, dst->data, |             weight_ne[1] = src0->ne[0] / QK8_0; | ||||||
|                 ((ggml_tensor*)src0->extra)->ne, |             weight_nb[0] = sizeof(int8_t); | ||||||
|                 ((ggml_tensor*)src1->extra)->ne, |             weight_nb[1] = weight_nb[0] * weight_ne[0]; | ||||||
|                 ((ggml_tensor*)src1->extra)->nb, ((ggml_tensor*)dst->extra)->ne, |             for (int i = 2; i < GGML_MAX_DIMS + 1; i++) { | ||||||
|                 ((ggml_tensor*)dst->extra)->nb); |                 weight_ne[i] = src0->ne[i - 1]; | ||||||
|  |                 weight_nb[i] = weight_nb[i - 1] * weight_ne[i - 1]; | ||||||
|  |             } | ||||||
|  |  | ||||||
|  |             // [3,4,5,64] -> [3,4,5,2,1] | ||||||
|  |             scale_ne[0] = 1; | ||||||
|  |             scale_ne[1] = src0->ne[0] / QK8_0; | ||||||
|  |             scale_nb[0] = sizeof(uint16_t); | ||||||
|  |             scale_nb[1] = scale_nb[0] * scale_ne[0]; | ||||||
|  |             for (int i = 2; i < GGML_MAX_DIMS + 1; i++) { | ||||||
|  |                 scale_ne[i] = src0->ne[i - 1]; | ||||||
|  |                 scale_nb[i] = scale_nb[i - 1] * scale_ne[i - 1]; | ||||||
|  |             } | ||||||
|  |  | ||||||
|  |             // [3,4,5,64] -> [3,4,5,2,32] | ||||||
|  |             dequant_ne = weight_ne; | ||||||
|  |             dequant_nb[0] = sizeof(float_t); | ||||||
|  |             for (int i = 1; i < GGML_MAX_DIMS + 1; i++) { | ||||||
|  |                 dequant_nb[i] = dequant_nb[i - 1] * dequant_ne[i - 1]; | ||||||
|  |             } | ||||||
|  |  | ||||||
|  |             scale_offset = ggml_nelements(src0) * sizeof(int8_t); | ||||||
|  |             ggml_cann_pool_alloc dequant_buffer_allocator( | ||||||
|  |                 ctx.pool(), ggml_nelements(src0) * sizeof(float_t)); | ||||||
|  |  | ||||||
|  |             aclTensor* acl_weight_tensor = ggml_cann_create_tensor( | ||||||
|  |                 src0->data, ACL_INT8, sizeof(int8_t), weight_ne, weight_nb, | ||||||
|  |                 GGML_MAX_DIMS + 1); | ||||||
|  |             aclTensor* acl_scale_tensor = ggml_cann_create_tensor( | ||||||
|  |                 src0->data, ACL_FLOAT16, sizeof(float16_t), scale_ne, scale_nb, | ||||||
|  |                 GGML_MAX_DIMS + 1, ACL_FORMAT_ND, scale_offset); | ||||||
|  |             aclTensor* dequant_tensor = ggml_cann_create_tensor( | ||||||
|  |                 dequant_buffer_allocator.get(), ACL_FLOAT, sizeof(float_t), | ||||||
|  |                 dequant_ne, dequant_nb, GGML_MAX_DIMS + 1); | ||||||
|  |  | ||||||
|  |             aclnn_mul(ctx, acl_weight_tensor, acl_scale_tensor, dequant_tensor); | ||||||
|  |             dequant_nb[0] = sizeof(float_t); | ||||||
|  |             dequant_ne = src0->ne; | ||||||
|  |             for (int i = 1; i < GGML_MAX_DIMS; i++) { | ||||||
|  |                 dequant_nb[i] = dequant_nb[i - 1] * src0->ne[i - 1]; | ||||||
|  |             } | ||||||
|  |  | ||||||
|  |             aclnn_embedding_4d(ctx, dequant_buffer_allocator.get(), | ||||||
|  |                                    dequant_ne, dequant_nb, src1, dst); | ||||||
|  |  | ||||||
|  |             ACL_CHECK(aclDestroyTensor(dequant_tensor)); | ||||||
|             break; |             break; | ||||||
|  |         } | ||||||
|         default: |         default: | ||||||
|             GGML_ABORT("fatal error"); |             GGML_ABORT("Unsupported tensor type for GGML_OP_GET_ROWS"); | ||||||
|             break; |             break; | ||||||
|     } |     } | ||||||
| } | } | ||||||
| @@ -2797,8 +2812,8 @@ static void ggml_cann_mul_mat_quant(ggml_backend_cann_context& ctx, | |||||||
|  |  | ||||||
|             ACL_CHECK(aclnnWeightQuantBatchMatmulV2GetWorkspaceSize( |             ACL_CHECK(aclnnWeightQuantBatchMatmulV2GetWorkspaceSize( | ||||||
|                 acl_input_tensor, acl_weight_tensor, acl_scale_tensor, nullptr, |                 acl_input_tensor, acl_weight_tensor, acl_scale_tensor, nullptr, | ||||||
|                 nullptr, nullptr, nullptr, antiquantGroupSize, acl_output_tensor, |                 nullptr, nullptr, nullptr, antiquantGroupSize, | ||||||
|                 &workspaceSize, &executor)); |                 acl_output_tensor, &workspaceSize, &executor)); | ||||||
|             if (workspaceAddr == nullptr) { |             if (workspaceAddr == nullptr) { | ||||||
|                 workspaceAddr = workspace_allocator.alloc(workspaceSize); |                 workspaceAddr = workspace_allocator.alloc(workspaceSize); | ||||||
|             } |             } | ||||||
|   | |||||||
| @@ -1704,7 +1704,6 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, | |||||||
|             switch (op->src[0]->type) { |             switch (op->src[0]->type) { | ||||||
|                 case GGML_TYPE_F32: |                 case GGML_TYPE_F32: | ||||||
|                 case GGML_TYPE_F16: |                 case GGML_TYPE_F16: | ||||||
|                 case GGML_TYPE_Q4_0: |  | ||||||
|                 case GGML_TYPE_Q8_0: |                 case GGML_TYPE_Q8_0: | ||||||
|                     return true; |                     return true; | ||||||
|                 default: |                 default: | ||||||
| @@ -1712,16 +1711,21 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, | |||||||
|             } |             } | ||||||
|         } break; |         } break; | ||||||
|         case GGML_OP_CPY: { |         case GGML_OP_CPY: { | ||||||
|             switch (op->type) { |             ggml_tensor *src = op->src[0]; | ||||||
|                 case GGML_TYPE_F32: |             if ((op->type != GGML_TYPE_F32 && op->type != GGML_TYPE_F16) || | ||||||
|                 case GGML_TYPE_F16: |                   (src->type != GGML_TYPE_F32 && | ||||||
|                 case GGML_TYPE_Q8_0: |                     src->type != GGML_TYPE_F16)) { | ||||||
|                 case GGML_TYPE_Q4_0: |                 // only support F32 and F16. | ||||||
|                     return true; |                 return false; | ||||||
|                 default: |  | ||||||
|                     return false; |  | ||||||
|             } |             } | ||||||
|         } |  | ||||||
|  |             if (!ggml_are_same_shape(op, src) && !ggml_is_contiguous(op)) { | ||||||
|  |                 // unsupport dst is not contiguous. | ||||||
|  |                 return false; | ||||||
|  |             } | ||||||
|  |  | ||||||
|  |             return true; | ||||||
|  |         } break; | ||||||
|         case GGML_OP_CONT: { |         case GGML_OP_CONT: { | ||||||
|             // TODO: support GGML_TYPE_BF16 |             // TODO: support GGML_TYPE_BF16 | ||||||
|             switch (op->src[0]->type) { |             switch (op->src[0]->type) { | ||||||
| @@ -1762,9 +1766,9 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, | |||||||
|             } |             } | ||||||
|             return true; |             return true; | ||||||
|         } |         } | ||||||
|  |         case GGML_OP_DUP: | ||||||
|         case GGML_OP_IM2COL: |         case GGML_OP_IM2COL: | ||||||
|         case GGML_OP_CONCAT: |         case GGML_OP_CONCAT: | ||||||
|         case GGML_OP_DUP: |  | ||||||
|         case GGML_OP_REPEAT: |         case GGML_OP_REPEAT: | ||||||
|         case GGML_OP_NONE: |         case GGML_OP_NONE: | ||||||
|         case GGML_OP_RESHAPE: |         case GGML_OP_RESHAPE: | ||||||
|   | |||||||
| @@ -1,30 +0,0 @@ | |||||||
| file(GLOB SRC_FILES |  | ||||||
|     get_row_f32.cpp |  | ||||||
|     get_row_f16.cpp |  | ||||||
|     get_row_q4_0.cpp |  | ||||||
|     get_row_q8_0.cpp |  | ||||||
|     quantize_f32_q8_0.cpp |  | ||||||
|     quantize_f16_q8_0.cpp |  | ||||||
|     quantize_float_to_q4_0.cpp |  | ||||||
|     dup.cpp |  | ||||||
| ) |  | ||||||
|  |  | ||||||
| set(ASCEND_CANN_PACKAGE_PATH ${CANN_INSTALL_DIR}) |  | ||||||
| set(RUN_MODE "npu" CACHE STRING "run mode: npu/sim") |  | ||||||
|  |  | ||||||
| if(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/compiler/tikcpp/ascendc_kernel_cmake) |  | ||||||
|     set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/compiler/tikcpp/ascendc_kernel_cmake) |  | ||||||
| elseif(EXISTS ${ASCEND_CANN_PACKAGE_PATH}/ascendc_devkit/tikcpp/samples/cmake) |  | ||||||
|     set(ASCENDC_CMAKE_DIR ${ASCEND_CANN_PACKAGE_PATH}/ascendc_devkit/tikcpp/samples/cmake) |  | ||||||
| else() |  | ||||||
|     message(FATAL_ERROR "ascendc_kernel_cmake does not exist, please check whether the compiler package is installed.") |  | ||||||
| endif() |  | ||||||
| include(${ASCENDC_CMAKE_DIR}/ascendc.cmake) |  | ||||||
|  |  | ||||||
| ascendc_library(ascendc_kernels STATIC |  | ||||||
|     ${SRC_FILES} |  | ||||||
| ) |  | ||||||
|  |  | ||||||
| message(STATUS "CANN: compile ascend kernels witch SOC_TYPE:${SOC_TYPE}, SOC_VERSION:${SOC_VERSION}, compile macro:-D${SOC_TYPE_COMPILE_OPTION}.") |  | ||||||
| ascendc_compile_definitions(ascendc_kernels PRIVATE "-D${SOC_TYPE_COMPILE_OPTION}") |  | ||||||
| # ascendc_compile_definitions(ascendc_kernels PRIVATE -DASCENDC_DUMP) |  | ||||||
| @@ -1,19 +0,0 @@ | |||||||
| #ifndef ASCENDC_KERNELS_H |  | ||||||
| #define ASCENDC_KERNELS_H |  | ||||||
|  |  | ||||||
| #include "aclrtlaunch_ascendc_get_row_f32.h" |  | ||||||
| #include "aclrtlaunch_ascendc_get_row_f16.h" |  | ||||||
| #include "aclrtlaunch_ascendc_get_row_q8_0.h" |  | ||||||
| #include "aclrtlaunch_ascendc_get_row_q4_0.h" |  | ||||||
|  |  | ||||||
| #include "aclrtlaunch_ascendc_quantize_f32_q8_0.h" |  | ||||||
| #include "aclrtlaunch_ascendc_quantize_f16_q8_0.h" |  | ||||||
| #include "aclrtlaunch_ascendc_quantize_f16_to_q4_0.h" |  | ||||||
| #include "aclrtlaunch_ascendc_quantize_f32_to_q4_0.h" |  | ||||||
|  |  | ||||||
| #include "aclrtlaunch_ascendc_dup_by_rows_fp16.h" |  | ||||||
| #include "aclrtlaunch_ascendc_dup_by_rows_fp32.h" |  | ||||||
| #include "aclrtlaunch_ascendc_dup_by_rows_fp32_to_fp16.h" |  | ||||||
| #include "aclrtlaunch_ascendc_dup_by_rows_fp16_to_fp32.h" |  | ||||||
|  |  | ||||||
| #endif  // ASCENDC_KERNELS_H |  | ||||||
| @@ -1,234 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| using namespace AscendC; |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
| const int64_t SUPPORTED_MAX_DIM = 65535;  // currently the limit of max block dim supportted by dup kernel is 65535template <typename SRC_T, typename DST_T> |  | ||||||
|  |  | ||||||
| template <typename SRC_T, typename DST_T> |  | ||||||
| class DupByRows { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline DupByRows() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR src, GM_ADDR dst, int64_t *input_ne_ub, |  | ||||||
|                                 size_t *input_nb_ub) { |  | ||||||
|         /* Dup by rows when src is contigous on first dimension and dst is |  | ||||||
|         contiguous, each kernel process one row. |  | ||||||
|         */ |  | ||||||
|  |  | ||||||
|         // Input has four dims. |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         int64_t op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         // param |  | ||||||
|         num_rows = input_ne_ub[1] * input_ne_ub[2] * input_ne_ub[3]; |  | ||||||
|         num_elem = input_ne_ub[0]; |  | ||||||
|  |  | ||||||
|         // index for (ne[1], ne[2], ne[3]): (idx_ne1, idx_ne2, idx_ne3) |  | ||||||
|         idx_ne3 = op_block_idx / (input_ne_ub[1] * input_ne_ub[2]); |  | ||||||
|         idx_ne2 = (op_block_idx - idx_ne3 * (input_ne_ub[1] * input_ne_ub[2])) |  | ||||||
|                   / (input_ne_ub[1]); |  | ||||||
|         idx_ne1 = op_block_idx - idx_ne3 * (input_ne_ub[1] * input_ne_ub[2]) |  | ||||||
|                 - idx_ne2 * input_ne_ub[1]; |  | ||||||
|  |  | ||||||
|         // src may not contiguous in dim [1,2,3], so stride decited by ne&nb |  | ||||||
|         src_stride = input_nb_ub[3] * idx_ne3 + input_nb_ub[2] * idx_ne2 |  | ||||||
|                      + input_nb_ub[1] * idx_ne1; |  | ||||||
|  |  | ||||||
|         // dst is contiguous |  | ||||||
|         dst_stride = op_block_idx * (input_ne_ub[0] * sizeof(DST_T)); |  | ||||||
|  |  | ||||||
|         src_gm.SetGlobalBuffer(reinterpret_cast<__gm__ SRC_T *>(src + |  | ||||||
|                                                                 src_stride)); |  | ||||||
|         dst_gm.SetGlobalBuffer(reinterpret_cast<__gm__ DST_T *>(dst + |  | ||||||
|                                                                 dst_stride)); |  | ||||||
|  |  | ||||||
|         pipe.InitBuffer(src_queue, BUFFER_NUM, (sizeof(SRC_T) * num_elem + |  | ||||||
|                                                 32 - 1) / 32 * 32); |  | ||||||
|         pipe.InitBuffer(dst_queue, BUFFER_NUM, (sizeof(DST_T) * num_elem + |  | ||||||
|                                                 32 - 1) / 32 * 32); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in() { |  | ||||||
|         LocalTensor<SRC_T> src_local = src_queue.AllocTensor<SRC_T>(); |  | ||||||
|         const size_t elem_per_block = 32 / sizeof(SRC_T); |  | ||||||
|         size_t tail = num_elem % elem_per_block; |  | ||||||
|         size_t cpy_elements_len = tail > 0 ? num_elem + 1 : num_elem; |  | ||||||
|         DataCopy(src_local, src_gm, cpy_elements_len); |  | ||||||
|         src_queue.EnQue(src_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out() { |  | ||||||
|         LocalTensor<DST_T> dst_local = dst_queue.DeQue<DST_T>(); |  | ||||||
| #ifdef ASCEND_310P |  | ||||||
|         const size_t elem_per_block = 32 / sizeof(DST_T); |  | ||||||
|         size_t tail = num_elem % elem_per_block; |  | ||||||
|         size_t len = num_elem & ~(elem_per_block - 1); |  | ||||||
|         if (len > 0) { |  | ||||||
|             DataCopy(dst_gm, dst_local, len); |  | ||||||
|         } |  | ||||||
|         if(tail != 0) { |  | ||||||
|             for (size_t i = tail; i < elem_per_block; i++) { |  | ||||||
|                 dst_local[len + i].SetValue(0, 0); |  | ||||||
|             } |  | ||||||
|             SetAtomicAdd<float>(); |  | ||||||
|             DataCopy(dst_gm[len], dst_local[len], elem_per_block); |  | ||||||
|             SetAtomicNone(); |  | ||||||
|         } |  | ||||||
| #else |  | ||||||
|         DataCopyExtParams dataCopyParams; |  | ||||||
|         dataCopyParams.blockCount = 1; |  | ||||||
|         dataCopyParams.blockLen = num_elem * sizeof(DST_T); |  | ||||||
|         DataCopyPad(dst_gm, dst_local, dataCopyParams); |  | ||||||
| #endif |  | ||||||
|         dst_queue.FreeTensor(dst_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void dup() { |  | ||||||
|         // main process, copy one row data from src to dst. |  | ||||||
|         copy_in(); |  | ||||||
|  |  | ||||||
|         LocalTensor<SRC_T> src_local = src_queue.DeQue<SRC_T>(); |  | ||||||
|         LocalTensor<DST_T> dst_local = dst_queue.AllocTensor<DST_T>(); |  | ||||||
|  |  | ||||||
|         int32_t BLOCK_NUM = 32 / sizeof(DST_T); |  | ||||||
|         DataCopy(dst_local, src_local, (num_elem + BLOCK_NUM - 1) |  | ||||||
|                                         / BLOCK_NUM * BLOCK_NUM); |  | ||||||
|         dst_queue.EnQue<DST_T>(dst_local); |  | ||||||
|  |  | ||||||
|         src_queue.FreeTensor(src_local); |  | ||||||
|         copy_out(); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void dup_with_cast() { |  | ||||||
|         // main process, copy one row data from src to dst. |  | ||||||
|         // cast dtype from src to dst. |  | ||||||
|         copy_in(); |  | ||||||
|  |  | ||||||
|         LocalTensor<SRC_T> src_local = src_queue.DeQue<SRC_T>(); |  | ||||||
|         LocalTensor<DST_T> dst_local = dst_queue.AllocTensor<DST_T>(); |  | ||||||
|  |  | ||||||
|         Cast(dst_local, src_local, RoundMode::CAST_NONE, num_elem); |  | ||||||
|         dst_queue.EnQue<DST_T>(dst_local); |  | ||||||
|  |  | ||||||
|         src_queue.FreeTensor(src_local); |  | ||||||
|         copy_out(); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<SRC_T> src_gm; |  | ||||||
|     GlobalTensor<DST_T> dst_gm; |  | ||||||
|  |  | ||||||
|     int64_t num_rows; |  | ||||||
|     int64_t num_elem; |  | ||||||
|     int64_t idx_ne3; |  | ||||||
|     int64_t idx_ne2; |  | ||||||
|     int64_t idx_ne1; |  | ||||||
|     int64_t src_stride; |  | ||||||
|     int64_t dst_stride; |  | ||||||
|  |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> src_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> dst_queue; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_dup_by_rows_fp16( |  | ||||||
|                                                         GM_ADDR src_gm, |  | ||||||
|                                                         GM_ADDR dst_gm, |  | ||||||
|                                                         GM_ADDR input_ne_gm, |  | ||||||
|                                                         GM_ADDR input_nb_gm, |  | ||||||
|                                                         GM_ADDR output_ne_gm, |  | ||||||
|                                                         GM_ADDR output_nb_gm) { |  | ||||||
|  |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     DupByRows<half, half> op; |  | ||||||
|     op.init(src_gm, dst_gm, input_ne_ub, input_nb_ub); |  | ||||||
|     op.dup(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_dup_by_rows_fp32( |  | ||||||
|                                                         GM_ADDR src_gm, |  | ||||||
|                                                         GM_ADDR dst_gm, |  | ||||||
|                                                         GM_ADDR input_ne_gm, |  | ||||||
|                                                         GM_ADDR input_nb_gm, |  | ||||||
|                                                         GM_ADDR output_ne_gm, |  | ||||||
|                                                         GM_ADDR output_nb_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     DupByRows<float, float> op; |  | ||||||
|     op.init(src_gm, dst_gm, input_ne_ub, input_nb_ub); |  | ||||||
|     op.dup(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_dup_by_rows_fp32_to_fp16( |  | ||||||
|                                                         GM_ADDR src_gm, |  | ||||||
|                                                         GM_ADDR dst_gm, |  | ||||||
|                                                         GM_ADDR input_ne_gm, |  | ||||||
|                                                         GM_ADDR input_nb_gm, |  | ||||||
|                                                         GM_ADDR output_ne_gm, |  | ||||||
|                                                         GM_ADDR output_nb_gm) { |  | ||||||
|  |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     DupByRows<float, half> op; |  | ||||||
|     op.init(src_gm, dst_gm, input_ne_ub, input_nb_ub); |  | ||||||
|     op.dup_with_cast(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_dup_by_rows_fp16_to_fp32( |  | ||||||
|                                                         GM_ADDR src_gm, |  | ||||||
|                                                         GM_ADDR dst_gm, |  | ||||||
|                                                         GM_ADDR input_ne_gm, |  | ||||||
|                                                         GM_ADDR input_nb_gm, |  | ||||||
|                                                         GM_ADDR output_ne_gm, |  | ||||||
|                                                         GM_ADDR output_nb_gm) { |  | ||||||
|  |  | ||||||
|     // copy params from gm to ub. |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     DupByRows<half, float> op; |  | ||||||
|     op.init(src_gm, dst_gm, input_ne_ub, input_nb_ub); |  | ||||||
|     op.dup_with_cast(); |  | ||||||
| } |  | ||||||
| @@ -1,197 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| // optimize me. Use template to avoid copy code. |  | ||||||
| using namespace AscendC; |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
|  |  | ||||||
| class GET_ROW_F16 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline GET_ROW_F16() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR indices, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, size_t *input_nb_ub, |  | ||||||
|                                 int64_t *indices_ne_ub, size_t *indices_nb_ub, |  | ||||||
|                                 int64_t *output_ne_ub, size_t *output_nb_ub) { |  | ||||||
|         // TODO, use template for F16/f32 |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             input_stride[i] = input_nb_ub[i] / input_nb_ub[0]; |  | ||||||
|  |  | ||||||
|             indices_ne[i] = indices_ne_ub[i]; |  | ||||||
|             indices_stride[i] = indices_nb_ub[i] / indices_nb_ub[0]; |  | ||||||
|  |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|             output_stride[i] = output_nb_ub[i] / output_nb_ub[0]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // Indices has two dims. n_elements = all rows should get. |  | ||||||
|         // dr, all rows should this thread get. |  | ||||||
|         uint64_t n_elements = |  | ||||||
|             indices_ne[0] * indices_ne[1] * indices_ne[2] * indices_ne[3]; |  | ||||||
|         dr = n_elements / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = n_elements % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ half *)input); |  | ||||||
|         indices_gm.SetGlobalBuffer((__gm__ int32_t *)indices); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ float *)output); |  | ||||||
|  |  | ||||||
|         uint64_t input_local_buffer_size = ((input_ne[0] * sizeof(half) + 31) |  | ||||||
|                                              & ~31); |  | ||||||
|         uint64_t output_local_buffer_size = ((input_ne[0] * sizeof(float) + 31) |  | ||||||
|                                               & ~31); |  | ||||||
|  |  | ||||||
|         local_buffer_elems = input_local_buffer_size / sizeof(half); |  | ||||||
|  |  | ||||||
|         // TODO, consider long row that can't put in UB. |  | ||||||
|         // All data should asign to 32. It's ok because all data is align to 32. |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, input_local_buffer_size); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, output_local_buffer_size); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset, size_t len) { |  | ||||||
|         size_t origin_len = len; |  | ||||||
|         LocalTensor<half> input_local = input_queue.AllocTensor<half>(); |  | ||||||
|         const size_t elem_per_block = 32 / sizeof(half); |  | ||||||
|         size_t tail = len % elem_per_block; |  | ||||||
|         len = len & ~(elem_per_block - 1); |  | ||||||
|         if(tail != 0) { |  | ||||||
|             len += elem_per_block; |  | ||||||
|         } |  | ||||||
|         DataCopy(input_local, input_gm[offset], len); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset, size_t len) { |  | ||||||
|         LocalTensor<float> output_local = output_queue.DeQue<float>(); |  | ||||||
|         const size_t elem_per_block = 32 / sizeof(float); |  | ||||||
|         size_t tail = len % elem_per_block; |  | ||||||
|         len = len & ~(elem_per_block - 1); |  | ||||||
|         if (len > 0) { |  | ||||||
|             DataCopy(output_gm[offset], output_local, len); |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         if(tail != 0) { |  | ||||||
| #ifdef ASCEND_310P |  | ||||||
|             for (size_t i = tail; i < elem_per_block; i++) { |  | ||||||
|                 output_local[len + i].SetValue(0, 0); |  | ||||||
|             } |  | ||||||
|             SetAtomicAdd<float>(); |  | ||||||
|             DataCopy(output_gm[offset + len], output_local[len], elem_per_block); |  | ||||||
|             SetAtomicNone(); |  | ||||||
| #else |  | ||||||
|             DataCopyExtParams dataCopyParams; |  | ||||||
|             dataCopyParams.blockCount = 1; |  | ||||||
|             dataCopyParams.blockLen = tail * sizeof(float); |  | ||||||
|             DataCopyPad(output_gm[offset + len], output_local[len], |  | ||||||
|                         dataCopyParams); |  | ||||||
| #endif |  | ||||||
|         } |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate_row(int64_t idx) { |  | ||||||
|         const int64_t indices_ne2_idx = idx / (indices_ne[0] * indices_ne[1]); |  | ||||||
|         const int64_t indices_ne1_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1]) / |  | ||||||
|             indices_ne[0]; |  | ||||||
|         const int64_t indices_ne0_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1] - |  | ||||||
|              indices_ne1_idx * indices_ne[0]); |  | ||||||
|  |  | ||||||
|         const int64_t indices_offset = indices_ne0_idx * indices_stride[0] + |  | ||||||
|                                        indices_ne1_idx * indices_stride[1] + |  | ||||||
|                                        indices_ne2_idx * indices_stride[2]; |  | ||||||
|         const int32_t selected_row_idx = indices_gm.GetValue(indices_offset); |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = selected_row_idx * input_stride[1] + |  | ||||||
|                                      indices_ne1_idx * input_stride[2] + |  | ||||||
|                                      indices_ne2_idx * input_stride[3]; |  | ||||||
|  |  | ||||||
|         const int64_t output_offset = indices_ne0_idx * output_stride[1] + |  | ||||||
|                                       indices_ne1_idx * output_stride[2] + |  | ||||||
|                                       indices_ne2_idx * output_stride[3]; |  | ||||||
|  |  | ||||||
|         copy_in(input_offset, input_ne[0]); |  | ||||||
|         LocalTensor<half> input_local = input_queue.DeQue<half>(); |  | ||||||
|         LocalTensor<float> output_local = output_queue.AllocTensor<float>(); |  | ||||||
|  |  | ||||||
|         Cast(output_local, input_local, RoundMode::CAST_NONE, |  | ||||||
|              local_buffer_elems); |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|         copy_out(output_offset, input_ne[0]); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             calculate_row(i); |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t indices_ne[4]; |  | ||||||
|     size_t indices_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     size_t local_buffer_elems; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<half> input_gm; |  | ||||||
|     GlobalTensor<int32_t> indices_gm; |  | ||||||
|     GlobalTensor<float> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     int64_t op_block_idx; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_get_row_f16( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR indices_gm, GM_ADDR output_gm, |  | ||||||
|     GM_ADDR input_ne_gm, GM_ADDR input_nb_gm, GM_ADDR indices_ne_gm, |  | ||||||
|     GM_ADDR indices_nb_gm, GM_ADDR output_ne_gm, GM_ADDR output_nb_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t indices_ne_ub[4]; |  | ||||||
|     size_t indices_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(indices_ne_gm, indices_ne_ub, 32); |  | ||||||
|     copy_to_ub(indices_nb_gm, indices_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     GET_ROW_F16 op; |  | ||||||
|     op.init(input_gm, indices_gm, output_gm, input_ne_ub, input_nb_ub, |  | ||||||
|             indices_ne_ub, indices_nb_ub, output_ne_ub, output_nb_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
| @@ -1,190 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| // optimize me. Use template to avoid copy code. |  | ||||||
| using namespace AscendC; |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
|  |  | ||||||
| class GET_ROW_F32 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline GET_ROW_F32() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR indices, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, size_t *input_nb_ub, |  | ||||||
|                                 int64_t *indices_ne_ub, size_t *indices_nb_ub, |  | ||||||
|                                 int64_t *output_ne_ub, size_t *output_nb_ub) { |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             input_stride[i] = input_nb_ub[i] / input_nb_ub[0]; |  | ||||||
|  |  | ||||||
|             indices_ne[i] = indices_ne_ub[i]; |  | ||||||
|             indices_stride[i] = indices_nb_ub[i] / indices_nb_ub[0]; |  | ||||||
|  |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|             output_stride[i] = output_nb_ub[i] / output_nb_ub[0]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // Indices has two dims. n_elements = all rows should get. |  | ||||||
|         // dr, all rows should this thread get. |  | ||||||
|         uint64_t n_elements = |  | ||||||
|             indices_ne[0] * indices_ne[1] * indices_ne[2] * indices_ne[3]; |  | ||||||
|         dr = n_elements / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = n_elements % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ float *)input); |  | ||||||
|         indices_gm.SetGlobalBuffer((__gm__ int32_t *)indices); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ float *)output); |  | ||||||
|  |  | ||||||
|         uint64_t local_buffer_size = ((input_ne[0] * sizeof(float) + 31) & ~31); |  | ||||||
|         local_buffer_elems = local_buffer_size / sizeof(float); |  | ||||||
|  |  | ||||||
|         // TODO, consider long row that can't put in UB. |  | ||||||
|         // All data should asign to 32. It's ok because all data is align to 32. |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, local_buffer_size); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, local_buffer_size); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset, size_t len) { |  | ||||||
|         LocalTensor<float> input_local = input_queue.AllocTensor<float>(); |  | ||||||
|         const size_t elem_per_block = 32 / sizeof(float); |  | ||||||
|         size_t tail = len % elem_per_block; |  | ||||||
|         len = len & ~(elem_per_block - 1); |  | ||||||
|         if(tail != 0) { |  | ||||||
|             len += elem_per_block; |  | ||||||
|         } |  | ||||||
|         DataCopy(input_local, input_gm[offset], len); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset, size_t len) { |  | ||||||
|         LocalTensor<float> output_local = output_queue.DeQue<float>(); |  | ||||||
|         const size_t elem_per_block = 32 / sizeof(float); |  | ||||||
|         size_t tail = len % elem_per_block; |  | ||||||
|         len = len & ~(elem_per_block - 1); |  | ||||||
|         if (len > 0) { |  | ||||||
|             DataCopy(output_gm[offset], output_local, len); |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         if(tail != 0) { |  | ||||||
| #ifdef ASCEND_310P |  | ||||||
|             for (size_t i = tail; i < elem_per_block; i++) { |  | ||||||
|                 output_local[len + i].SetValue(0, 0); |  | ||||||
|             } |  | ||||||
|             SetAtomicAdd<float>(); |  | ||||||
|             DataCopy(output_gm[offset + len], output_local[len], elem_per_block); |  | ||||||
|             SetAtomicNone(); |  | ||||||
| #else |  | ||||||
|             DataCopyExtParams dataCopyParams; |  | ||||||
|             dataCopyParams.blockCount = 1; |  | ||||||
|             dataCopyParams.blockLen = tail * sizeof(float); |  | ||||||
|             DataCopyPad(output_gm[offset + len], output_local[len], |  | ||||||
|                         dataCopyParams); |  | ||||||
| #endif |  | ||||||
|         } |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate_row(int64_t idx) { |  | ||||||
|         const int64_t indices_ne2_idx = idx / (indices_ne[0] * indices_ne[1]); |  | ||||||
|         const int64_t indices_ne1_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1]) / |  | ||||||
|             indices_ne[0]; |  | ||||||
|         const int64_t indices_ne0_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1] - |  | ||||||
|              indices_ne1_idx * indices_ne[0]); |  | ||||||
|  |  | ||||||
|         const int64_t indices_offset = indices_ne0_idx * indices_stride[0] + |  | ||||||
|                                        indices_ne1_idx * indices_stride[1] + |  | ||||||
|                                        indices_ne2_idx * indices_stride[2]; |  | ||||||
|         const int32_t selected_row_idx = indices_gm.GetValue(indices_offset); |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = selected_row_idx * input_stride[1] + |  | ||||||
|                                      indices_ne1_idx * input_stride[2] + |  | ||||||
|                                      indices_ne2_idx * input_stride[3]; |  | ||||||
|  |  | ||||||
|         const int64_t output_offset = indices_ne0_idx * output_stride[1] + |  | ||||||
|                                       indices_ne1_idx * output_stride[2] + |  | ||||||
|                                       indices_ne2_idx * output_stride[3]; |  | ||||||
|  |  | ||||||
|         copy_in(input_offset, input_ne[0]); |  | ||||||
|         LocalTensor<float> input_local = input_queue.DeQue<float>(); |  | ||||||
|         LocalTensor<float> output_local = output_queue.AllocTensor<float>(); |  | ||||||
|  |  | ||||||
|         DataCopy(output_local, input_local, local_buffer_elems); |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|         copy_out(output_offset, input_ne[0]); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             calculate_row(i); |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t indices_ne[4]; |  | ||||||
|     size_t indices_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     size_t local_buffer_elems; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<float> input_gm; |  | ||||||
|     GlobalTensor<int32_t> indices_gm; |  | ||||||
|     GlobalTensor<float> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     int64_t op_block_idx; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_get_row_f32( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR indices_gm, GM_ADDR output_gm, |  | ||||||
|     GM_ADDR input_ne_gm, GM_ADDR input_nb_gm, GM_ADDR indices_ne_gm, |  | ||||||
|     GM_ADDR indices_nb_gm, GM_ADDR output_ne_gm, GM_ADDR output_nb_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t indices_ne_ub[4]; |  | ||||||
|     size_t indices_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(indices_ne_gm, indices_ne_ub, 32); |  | ||||||
|     copy_to_ub(indices_nb_gm, indices_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     GET_ROW_F32 op; |  | ||||||
|     op.init(input_gm, indices_gm, output_gm, input_ne_ub, input_nb_ub, |  | ||||||
|             indices_ne_ub, indices_nb_ub, output_ne_ub, output_nb_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
| @@ -1,204 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| // optimize me. Use template to avoid copy code. |  | ||||||
| using namespace AscendC; |  | ||||||
| #ifdef ASCEND_310P // 310P not support 4bit get row |  | ||||||
|     extern "C" __global__ __aicore__ void ascendc_get_row_q4_0( |  | ||||||
|         GM_ADDR input_gm, GM_ADDR indices_gm, GM_ADDR output_gm, |  | ||||||
|         GM_ADDR input_ne_gm, GM_ADDR indices_ne_gm, GM_ADDR indices_nb_gm, |  | ||||||
|         GM_ADDR output_ne_gm, GM_ADDR output_nb_gm) { |  | ||||||
|         // let following test cases can continue run, here just print error information. Of Cource the test case that call this operator is failed. |  | ||||||
|         printf("Ascend310P not support 4bit get row.\n"); |  | ||||||
|     } |  | ||||||
| #else |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
|  |  | ||||||
| #define QK4_0 32 |  | ||||||
|  |  | ||||||
| class GET_ROW_Q4_0 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline GET_ROW_Q4_0() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR indices, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, int64_t *indices_ne_ub, |  | ||||||
|                                 size_t *indices_nb_ub, int64_t *output_ne_ub, |  | ||||||
|                                 size_t *output_nb_ub) { |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         int64_t op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             indices_ne[i] = indices_ne_ub[i]; |  | ||||||
|             indices_stride[i] = indices_nb_ub[i] / indices_nb_ub[0]; |  | ||||||
|             scale_ne[i] = input_ne_ub[i]; |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|             output_stride[i] = output_nb_ub[i] / output_nb_ub[0]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // one scale for a group. |  | ||||||
|         scale_ne[0] /= QK4_0; |  | ||||||
|  |  | ||||||
|         input_stride[0] = 1; |  | ||||||
|         scale_stride[0] = 1; |  | ||||||
|         output_stride[0] = 1; |  | ||||||
|         for (int i = 1; i < 4; i++) { |  | ||||||
|             input_stride[i] = input_stride[i - 1] * input_ne[i - 1]; |  | ||||||
|             scale_stride[i] = scale_stride[i - 1] * scale_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         group_size_in_row = input_ne[0] / QK4_0; |  | ||||||
|         int64_t scale_offset = input_ne[0] * input_ne[1] * input_ne[2] * |  | ||||||
|                                input_ne[3] / 2; |  | ||||||
|  |  | ||||||
|         // Indices has two dims. n_elements = all rows should get. |  | ||||||
|         // dr, all rows should this thread get. |  | ||||||
|         uint64_t n_elements = |  | ||||||
|             indices_ne[0] * indices_ne[1] * indices_ne[2] * indices_ne[3]; |  | ||||||
|         dr = n_elements / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = n_elements % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ int4b_t *)input); |  | ||||||
|         scale_gm.SetGlobalBuffer((__gm__ half *)(input + scale_offset)); |  | ||||||
|         indices_gm.SetGlobalBuffer((__gm__ int32_t *)indices); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ float *)output); |  | ||||||
|  |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, QK4_0 * sizeof(int4b_t)); |  | ||||||
|         pipe.InitBuffer(cast_queue, BUFFER_NUM, QK4_0 * sizeof(half)); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, QK4_0 * sizeof(float)); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset) { |  | ||||||
|         LocalTensor<int4b_t> input_local = input_queue.AllocTensor<int4b_t>(); |  | ||||||
|         // 32 * sizeof(int4b_t) = 16, which is not aligned to 32, why no error? |  | ||||||
|         DataCopy(input_local, input_gm[offset], QK4_0); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset) { |  | ||||||
|         LocalTensor<float> output_local = output_queue.DeQue<float>(); |  | ||||||
|         DataCopy(output_gm[offset], output_local, QK4_0); |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate_group(int64_t idx, int64_t group) { |  | ||||||
|         const int64_t indices_ne2_idx = idx / (indices_ne[0] * indices_ne[1]); |  | ||||||
|         const int64_t indices_ne1_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1]) / |  | ||||||
|             indices_ne[0]; |  | ||||||
|         const int64_t indices_ne0_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1] - |  | ||||||
|              indices_ne1_idx * indices_ne[0]); |  | ||||||
|  |  | ||||||
|         const int64_t indices_offset = indices_ne0_idx * indices_stride[0] + |  | ||||||
|                                        indices_ne1_idx * indices_stride[1] + |  | ||||||
|                                        indices_ne2_idx * indices_stride[2]; |  | ||||||
|         const int32_t selected_row_idx = indices_gm.GetValue(indices_offset); |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = selected_row_idx * input_stride[1] + |  | ||||||
|                                      indices_ne1_idx * input_stride[2] + |  | ||||||
|                                      indices_ne2_idx * input_stride[3] + |  | ||||||
|                                      group * QK4_0; |  | ||||||
|         const int64_t scale_offset = selected_row_idx * scale_stride[1] + |  | ||||||
|                                      indices_ne1_idx * scale_stride[2] + |  | ||||||
|                                      indices_ne2_idx * scale_stride[3] + group; |  | ||||||
|         const int64_t output_offset = indices_ne0_idx * output_stride[1] + |  | ||||||
|                                       indices_ne1_idx * output_stride[2] + |  | ||||||
|                                       indices_ne2_idx * output_stride[3] + |  | ||||||
|                                       group * QK4_0; |  | ||||||
|  |  | ||||||
|         copy_in(input_offset); |  | ||||||
|         LocalTensor<int4b_t> input_local = input_queue.DeQue<int4b_t>(); |  | ||||||
|         LocalTensor<half> cast_local = cast_queue.AllocTensor<half>(); |  | ||||||
|         LocalTensor<float> output_local = output_queue.AllocTensor<float>(); |  | ||||||
|  |  | ||||||
|         // TODO: cast more data to speed up. |  | ||||||
|         Cast(cast_local, input_local, RoundMode::CAST_NONE, QK4_0); |  | ||||||
|         Cast(output_local, cast_local, RoundMode::CAST_NONE, QK4_0); |  | ||||||
|  |  | ||||||
|         // Only mul need compile by group. |  | ||||||
|         half scale = scale_gm.GetValue(scale_offset); |  | ||||||
|  |  | ||||||
|         Muls(output_local, output_local, (float)scale, QK4_0); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|         cast_queue.FreeTensor(cast_local); |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|  |  | ||||||
|         copy_out(output_offset); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             for (int64_t j = 0; j < group_size_in_row; j++) { |  | ||||||
|                 calculate_group(i, j); |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t scale_ne[4]; |  | ||||||
|     size_t scale_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t indices_ne[4]; |  | ||||||
|     size_t indices_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     int64_t group_size_in_row; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<int4b_t> input_gm; |  | ||||||
|     GlobalTensor<half> scale_gm; |  | ||||||
|     GlobalTensor<int32_t> indices_gm; |  | ||||||
|     GlobalTensor<float> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> cast_queue; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_get_row_q4_0( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR indices_gm, GM_ADDR output_gm, |  | ||||||
|     GM_ADDR input_ne_gm, GM_ADDR indices_ne_gm, GM_ADDR indices_nb_gm, |  | ||||||
|     GM_ADDR output_ne_gm, GM_ADDR output_nb_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     int64_t indices_ne_ub[4]; |  | ||||||
|     size_t indices_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(indices_ne_gm, indices_ne_ub, 32); |  | ||||||
|     copy_to_ub(indices_nb_gm, indices_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     GET_ROW_Q4_0 op; |  | ||||||
|     op.init(input_gm, indices_gm, output_gm, input_ne_ub, indices_ne_ub, |  | ||||||
|             indices_nb_ub, output_ne_ub, output_nb_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| #endif // #ifdef ASCEND_310P |  | ||||||
| @@ -1,191 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| // optimize me. Use template to avoid copy code. |  | ||||||
| using namespace AscendC; |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
|  |  | ||||||
| #define QK8_0 32 |  | ||||||
|  |  | ||||||
| class GET_ROW_Q8_0 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline GET_ROW_Q8_0() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR indices, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, int64_t *indices_ne_ub, |  | ||||||
|                                 size_t *indices_nb_ub, int64_t *output_ne_ub, |  | ||||||
|                                 size_t *output_nb_ub) { |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         int64_t op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             indices_ne[i] = indices_ne_ub[i]; |  | ||||||
|             indices_stride[i] = indices_nb_ub[i] / indices_nb_ub[0]; |  | ||||||
|             scale_ne[i] = input_ne_ub[i]; |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|             output_stride[i] = output_nb_ub[i] / output_nb_ub[0]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // one scale for a group. |  | ||||||
|         scale_ne[0] /= QK8_0; |  | ||||||
|  |  | ||||||
|         input_stride[0] = 1; |  | ||||||
|         scale_stride[0] = 1; |  | ||||||
|         output_stride[0] = 1; |  | ||||||
|         for (int i = 1; i < 4; i++) { |  | ||||||
|             input_stride[i] = input_stride[i - 1] * input_ne[i - 1]; |  | ||||||
|             scale_stride[i] = scale_stride[i - 1] * scale_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         group_size_in_row = input_ne[0] / QK8_0; |  | ||||||
|         int64_t scale_offset = input_ne[0] * input_ne[1] * input_ne[2] * |  | ||||||
|                                input_ne[3] * sizeof(int8_t); |  | ||||||
|  |  | ||||||
|         // Indices has two dims. n_elements = all rows should get. |  | ||||||
|         // dr, all rows should this thread get. |  | ||||||
|         uint64_t n_elements = |  | ||||||
|             indices_ne[0] * indices_ne[1] * indices_ne[2] * indices_ne[3]; |  | ||||||
|         dr = n_elements / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = n_elements % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ int8_t *)input); |  | ||||||
|         scale_gm.SetGlobalBuffer((__gm__ half *)(input + scale_offset)); |  | ||||||
|         indices_gm.SetGlobalBuffer((__gm__ int32_t *)indices); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ float *)output); |  | ||||||
|  |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, QK8_0 * sizeof(int8_t)); |  | ||||||
|         pipe.InitBuffer(cast_queue, BUFFER_NUM, QK8_0 * sizeof(half)); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, QK8_0 * sizeof(float)); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset) { |  | ||||||
|         LocalTensor<int8_t> input_local = input_queue.AllocTensor<int8_t>(); |  | ||||||
|         DataCopy(input_local, input_gm[offset], QK8_0); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset) { |  | ||||||
|         LocalTensor<float> output_local = output_queue.DeQue<float>(); |  | ||||||
|         DataCopy(output_gm[offset], output_local, QK8_0); |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate_group(int64_t idx, int64_t group) { |  | ||||||
|         const int64_t indices_ne2_idx = idx / (indices_ne[0] * indices_ne[1]); |  | ||||||
|         const int64_t indices_ne1_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1]) / |  | ||||||
|             indices_ne[0]; |  | ||||||
|         const int64_t indices_ne0_idx = |  | ||||||
|             (idx - indices_ne2_idx * indices_ne[0] * indices_ne[1] - |  | ||||||
|              indices_ne1_idx * indices_ne[0]); |  | ||||||
|  |  | ||||||
|         const int64_t indices_offset = indices_ne0_idx * indices_stride[0] + |  | ||||||
|                                        indices_ne1_idx * indices_stride[1] + |  | ||||||
|                                        indices_ne2_idx * indices_stride[2]; |  | ||||||
|         const int32_t selected_row_idx = indices_gm.GetValue(indices_offset); |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = selected_row_idx * input_stride[1] + |  | ||||||
|                                      indices_ne1_idx * input_stride[2] + |  | ||||||
|                                      indices_ne2_idx * input_stride[3] + |  | ||||||
|                                      group * QK8_0; |  | ||||||
|         const int64_t scale_offset = selected_row_idx * scale_stride[1] + |  | ||||||
|                                      indices_ne1_idx * scale_stride[2] + |  | ||||||
|                                      indices_ne2_idx * scale_stride[3] + group; |  | ||||||
|         const int64_t output_offset = indices_ne0_idx * output_stride[1] + |  | ||||||
|                                       indices_ne1_idx * output_stride[2] + |  | ||||||
|                                       indices_ne2_idx * output_stride[3] + |  | ||||||
|                                       group * QK8_0; |  | ||||||
|  |  | ||||||
|         copy_in(input_offset); |  | ||||||
|         LocalTensor<int8_t> input_local = input_queue.DeQue<int8_t>(); |  | ||||||
|         LocalTensor<half> cast_local = cast_queue.AllocTensor<half>(); |  | ||||||
|         LocalTensor<float> output_local = output_queue.AllocTensor<float>(); |  | ||||||
|  |  | ||||||
|         // TODO: cast more data to speed up. |  | ||||||
|         Cast(cast_local, input_local, RoundMode::CAST_NONE, QK8_0); |  | ||||||
|         Cast(output_local, cast_local, RoundMode::CAST_NONE, QK8_0); |  | ||||||
|  |  | ||||||
|         // Only mul need compile by group. |  | ||||||
|         half scale = scale_gm.GetValue(scale_offset); |  | ||||||
|         Muls(output_local, output_local, (float)scale, QK8_0); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|         cast_queue.FreeTensor(cast_local); |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|  |  | ||||||
|         copy_out(output_offset); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             for (int64_t j = 0; j < group_size_in_row; j++) { |  | ||||||
|                 calculate_group(i, j); |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t scale_ne[4]; |  | ||||||
|     size_t scale_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t indices_ne[4]; |  | ||||||
|     size_t indices_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     int64_t group_size_in_row; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<int8_t> input_gm; |  | ||||||
|     GlobalTensor<half> scale_gm; |  | ||||||
|     GlobalTensor<int32_t> indices_gm; |  | ||||||
|     GlobalTensor<float> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> cast_queue; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_get_row_q8_0( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR indices_gm, GM_ADDR output_gm, |  | ||||||
|     GM_ADDR input_ne_gm, GM_ADDR indices_ne_gm, GM_ADDR indices_nb_gm, |  | ||||||
|     GM_ADDR output_ne_gm, GM_ADDR output_nb_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     int64_t indices_ne_ub[4]; |  | ||||||
|     size_t indices_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|     size_t output_nb_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(indices_ne_gm, indices_ne_ub, 32); |  | ||||||
|     copy_to_ub(indices_nb_gm, indices_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|     copy_to_ub(output_nb_gm, output_nb_ub, 32); |  | ||||||
|  |  | ||||||
|     GET_ROW_Q8_0 op; |  | ||||||
|     op.init(input_gm, indices_gm, output_gm, input_ne_ub, indices_ne_ub, |  | ||||||
|             indices_nb_ub, output_ne_ub, output_nb_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
| @@ -1,218 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| using namespace AscendC; |  | ||||||
| #ifdef ASCEND_310P |  | ||||||
|     extern "C" __global__ __aicore__ void ascendc_quantize_f16_q8_0( |  | ||||||
|         GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|         GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|         // let following test cases can continue run, here just print error information. Of Cource the test case that call this operator is failed. |  | ||||||
|         printf("Ascend310P not support f16->8bit quantization.\n"); |  | ||||||
|     } |  | ||||||
| #else |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
| #define QK8_0 32 |  | ||||||
|  |  | ||||||
| class QUANTIZE_F16_Q8_0 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline QUANTIZE_F16_Q8_0() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, size_t *input_nb_ub, |  | ||||||
|                                 int64_t *output_ne_ub) { |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         int64_t op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             input_stride[i] = input_nb_ub[i] / input_nb_ub[0]; |  | ||||||
|  |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         output_stride[0] = 1; |  | ||||||
|         for (int i = 1; i < 4; i++) { |  | ||||||
|             output_stride[i] = output_stride[i - 1] * output_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         scale_ne = input_ne; |  | ||||||
|         scale_stride[0] = 1; |  | ||||||
|         scale_stride[1] = input_ne[0] / QK8_0; |  | ||||||
|         for (int i = 2; i < 4; i++) { |  | ||||||
|             scale_stride[i] = scale_stride[i - 1] * scale_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // split input tensor by rows. |  | ||||||
|         uint64_t nr = input_ne[1] * input_ne[2] * input_ne[3]; |  | ||||||
|         dr = nr / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = nr % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         group_size_in_row = scale_stride[1]; |  | ||||||
|         int64_t output_size = output_ne[0] * output_ne[1] * output_ne[2] * |  | ||||||
|                               output_ne[3] * sizeof(uint8_t); |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ half *)input); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ int8_t *)output); |  | ||||||
|         scale_gm.SetGlobalBuffer((__gm__ half *)(output + output_size + ir * |  | ||||||
|                                                  group_size_in_row * |  | ||||||
|                                                  sizeof(half))); |  | ||||||
|  |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, QK8_0 * sizeof(half)); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, QK8_0 * sizeof(int8_t)); |  | ||||||
|         pipe.InitBuffer(work_queue, 1, 32); |  | ||||||
|         pipe.InitBuffer(max_queue, 1, 32); |  | ||||||
|         pipe.InitBuffer(abs_queue, 1, QK8_0 * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(scale_queue, 1, 32); |  | ||||||
|         pipe.InitBuffer(cast_queue ,1 ,QK8_0 * sizeof(float)); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset) { |  | ||||||
|         LocalTensor<half> input_local = input_queue.AllocTensor<half>(); |  | ||||||
|         DataCopy(input_local, input_gm[offset], QK8_0); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset) { |  | ||||||
|         LocalTensor<int8_t> output_local = output_queue.DeQue<int8_t>(); |  | ||||||
|         DataCopy(output_gm[offset], output_local, QK8_0); |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline half calculate_group(int64_t row, int64_t group) { |  | ||||||
|         const int64_t i3 = row / (input_ne[1] * input_ne[2]); |  | ||||||
|         const int64_t i2 = (row - i3 * input_ne[1] * input_ne[2]) / input_ne[1]; |  | ||||||
|         const int64_t i1 = |  | ||||||
|             row - i3 * input_ne[1] * input_ne[2] - i2 * input_ne[1]; |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = i1 * input_stride[1] + |  | ||||||
|                                      i2 * input_stride[2] + |  | ||||||
|                                      i3 * input_stride[3] + QK8_0 * group; |  | ||||||
|  |  | ||||||
|         const int64_t output_offset = i1 * output_stride[1] + |  | ||||||
|                                       i2 * output_stride[2] + |  | ||||||
|                                       i3 * output_stride[3] + QK8_0 * group; |  | ||||||
|  |  | ||||||
|         copy_in(input_offset); |  | ||||||
|         LocalTensor<half> input_local = input_queue.DeQue<half>(); |  | ||||||
|         LocalTensor<int8_t> output_local = output_queue.AllocTensor<int8_t>(); |  | ||||||
|         LocalTensor<float> work_local = work_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> abs_local = abs_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> max_local = max_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> cast_local = cast_queue.AllocTensor<float>(); |  | ||||||
|  |  | ||||||
|         Cast(cast_local, input_local, RoundMode::CAST_NONE, QK8_0); |  | ||||||
|         Abs(abs_local, cast_local, QK8_0); |  | ||||||
|         ReduceMax(max_local, abs_local, work_local, QK8_0); |  | ||||||
|  |  | ||||||
|         pipe_barrier(PIPE_ALL); |  | ||||||
|         float d = max_local.GetValue(0); |  | ||||||
|         d = d / ((1 << 7) - 1); |  | ||||||
|         if (d != 0) { |  | ||||||
|             Muls(cast_local, cast_local, 1.0f / d, QK8_0); |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         Cast(cast_local, cast_local, RoundMode::CAST_ROUND, QK8_0); |  | ||||||
|         Cast(input_local, cast_local, RoundMode::CAST_ROUND, QK8_0); |  | ||||||
|         Cast(output_local, input_local, RoundMode::CAST_ROUND, QK8_0); |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|         copy_out(output_offset); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|         work_queue.FreeTensor(work_local); |  | ||||||
|         abs_queue.FreeTensor(abs_local); |  | ||||||
|         max_queue.FreeTensor(max_local); |  | ||||||
|         cast_queue.FreeTensor(cast_local); |  | ||||||
|         return (half)d; |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         LocalTensor<half> scale_local = scale_queue.AllocTensor<half>(); |  | ||||||
|         uint32_t scale_local_offset = 0; |  | ||||||
|         uint32_t scale_global_offset = 0; |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             for (int64_t j = 0; j < group_size_in_row; j++) { |  | ||||||
|                 half scale = calculate_group(i, j); |  | ||||||
|                 scale_local.SetValue(scale_local_offset++, scale); |  | ||||||
|                 if (scale_local_offset == 16) { |  | ||||||
|                     scale_local_offset = 0; |  | ||||||
|                     // TODO: OPTIMIZE ME |  | ||||||
|                     pipe_barrier(PIPE_ALL); |  | ||||||
|                     DataCopy(scale_gm[scale_global_offset], scale_local, 16); |  | ||||||
|                     pipe_barrier(PIPE_ALL); |  | ||||||
|                     scale_global_offset += 16; |  | ||||||
|                 } |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         if (scale_local_offset != 0) { |  | ||||||
|             pipe_barrier(PIPE_ALL); |  | ||||||
|             DataCopyExtParams dataCopyParams; |  | ||||||
|             dataCopyParams.blockCount = 1; |  | ||||||
|             dataCopyParams.blockLen = scale_local_offset * sizeof(half); |  | ||||||
|             DataCopyPad(scale_gm[scale_global_offset], scale_local, |  | ||||||
|                         dataCopyParams); |  | ||||||
|             pipe_barrier(PIPE_ALL); |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t *scale_ne; |  | ||||||
|     size_t scale_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t group_size_in_row; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<half> input_gm; |  | ||||||
|     GlobalTensor<half> scale_gm; |  | ||||||
|     GlobalTensor<int8_t> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     TQue<QuePosition::VECIN, 1> work_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, 1> max_queue; |  | ||||||
|     TQue<QuePosition::VECIN, 1> abs_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, 1> scale_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, 1> cast_queue; |  | ||||||
|  |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_quantize_f16_q8_0( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|     GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|  |  | ||||||
|     QUANTIZE_F16_Q8_0 op; |  | ||||||
|     op.init(input_gm, output_gm, input_ne_ub, input_nb_ub, output_ne_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| #endif // #ifdef ASCEND_310P |  | ||||||
| @@ -1,216 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| using namespace AscendC; |  | ||||||
| #ifdef ASCEND_310P // 310P not support f32->8bit quantization |  | ||||||
|     extern "C" __global__ __aicore__ void ascendc_quantize_f32_q8_0( |  | ||||||
|         GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|         GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|         // let following test cases can continue run, here just print error information. Of Cource the test case that call this operator is failed. |  | ||||||
|         printf("Ascend310P not support f32->8bit quantization.\n"); |  | ||||||
|     } |  | ||||||
| #else |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
| #define QK8_0 32 |  | ||||||
|  |  | ||||||
| class QUANTIZE_F32_Q8_0 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline QUANTIZE_F32_Q8_0() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, size_t *input_nb_ub, |  | ||||||
|                                 int64_t *output_ne_ub) { |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         int64_t op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             input_stride[i] = input_nb_ub[i] / input_nb_ub[0]; |  | ||||||
|  |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         output_stride[0] = 1; |  | ||||||
|         for (int i = 1; i < 4; i++) { |  | ||||||
|             output_stride[i] = output_stride[i - 1] * output_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         scale_ne = input_ne; |  | ||||||
|         scale_stride[0] = 1; |  | ||||||
|         scale_stride[1] = input_ne[0] / QK8_0; |  | ||||||
|         for (int i = 2; i < 4; i++) { |  | ||||||
|             scale_stride[i] = scale_stride[i - 1] * scale_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // split input tensor by rows. |  | ||||||
|         uint64_t nr = input_ne[1] * input_ne[2] * input_ne[3]; |  | ||||||
|         dr = nr / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = nr % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         group_size_in_row = scale_stride[1]; |  | ||||||
|         int64_t output_size = output_ne[0] * output_ne[1] * output_ne[2] * |  | ||||||
|                               output_ne[3] * sizeof(uint8_t); |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ float *)input); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ int8_t *)output); |  | ||||||
|         scale_gm.SetGlobalBuffer((__gm__ half *)(output + output_size + |  | ||||||
|                                                  ir * group_size_in_row * |  | ||||||
|                                                  sizeof(half))); |  | ||||||
|  |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, QK8_0 * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, QK8_0 * sizeof(int8_t)); |  | ||||||
|         pipe.InitBuffer(work_queue, 1, 32); |  | ||||||
|         pipe.InitBuffer(max_queue, 1, 32); |  | ||||||
|         pipe.InitBuffer(abs_queue, 1, QK8_0 * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(cast_queue, 1, QK8_0 * sizeof(half)); |  | ||||||
|         pipe.InitBuffer(scale_queue, 1, 32); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset) { |  | ||||||
|         LocalTensor<float> input_local = input_queue.AllocTensor<float>(); |  | ||||||
|         DataCopy(input_local, input_gm[offset], QK8_0); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset) { |  | ||||||
|         LocalTensor<int8_t> output_local = output_queue.DeQue<int8_t>(); |  | ||||||
|         DataCopy(output_gm[offset], output_local, QK8_0); |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline half calculate_group(int64_t row, int64_t group) { |  | ||||||
|         const int64_t i3 = row / (input_ne[1] * input_ne[2]); |  | ||||||
|         const int64_t i2 = (row - i3 * input_ne[1] * input_ne[2]) / input_ne[1]; |  | ||||||
|         const int64_t i1 = |  | ||||||
|             row - i3 * input_ne[1] * input_ne[2] - i2 * input_ne[1]; |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = i1 * input_stride[1] + |  | ||||||
|                                      i2 * input_stride[2] + |  | ||||||
|                                      i3 * input_stride[3] + QK8_0 * group; |  | ||||||
|  |  | ||||||
|         const int64_t output_offset = i1 * output_stride[1] + |  | ||||||
|                                       i2 * output_stride[2] + |  | ||||||
|                                       i3 * output_stride[3] + QK8_0 * group; |  | ||||||
|  |  | ||||||
|         copy_in(input_offset); |  | ||||||
|         LocalTensor<float> input_local = input_queue.DeQue<float>(); |  | ||||||
|         LocalTensor<int8_t> output_local = output_queue.AllocTensor<int8_t>(); |  | ||||||
|         LocalTensor<float> work_local = work_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> abs_local = abs_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> max_local = max_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<half> cast_local = cast_queue.AllocTensor<half>(); |  | ||||||
|  |  | ||||||
|         Abs(abs_local, input_local, QK8_0); |  | ||||||
|         ReduceMax(max_local, abs_local, work_local, QK8_0); |  | ||||||
|         pipe_barrier(PIPE_ALL); |  | ||||||
|         float d = max_local.GetValue(0); |  | ||||||
|         d = d / ((1 << 7) - 1); |  | ||||||
|         if (d != 0) { |  | ||||||
|             Muls(input_local, input_local, 1.0f / d, QK8_0); |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         Cast(input_local, input_local, RoundMode::CAST_ROUND, QK8_0); |  | ||||||
|         Cast(cast_local, input_local, RoundMode::CAST_ROUND, QK8_0); |  | ||||||
|         Cast(output_local, cast_local, RoundMode::CAST_ROUND, QK8_0); |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|         copy_out(output_offset); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|         work_queue.FreeTensor(work_local); |  | ||||||
|         abs_queue.FreeTensor(abs_local); |  | ||||||
|         max_queue.FreeTensor(max_local); |  | ||||||
|         cast_queue.FreeTensor(cast_local); |  | ||||||
|  |  | ||||||
|         return (half)d; |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         LocalTensor<half> scale_local = scale_queue.AllocTensor<half>(); |  | ||||||
|         uint32_t scale_local_offset = 0; |  | ||||||
|         uint32_t scale_global_offset = 0; |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             for (int64_t j = 0; j < group_size_in_row; j++) { |  | ||||||
|                 half scale = calculate_group(i, j); |  | ||||||
|                 scale_local.SetValue(scale_local_offset++, scale); |  | ||||||
|                 if (scale_local_offset == 16) { |  | ||||||
|                     scale_local_offset = 0; |  | ||||||
|                     // TODO: OPTIMIZE ME |  | ||||||
|                     pipe_barrier(PIPE_ALL); |  | ||||||
|                     DataCopy(scale_gm[scale_global_offset], scale_local, 16); |  | ||||||
|                     pipe_barrier(PIPE_ALL); |  | ||||||
|                     scale_global_offset += 16; |  | ||||||
|                 } |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         if (scale_local_offset != 0) { |  | ||||||
|             pipe_barrier(PIPE_ALL); |  | ||||||
|             DataCopyExtParams dataCopyParams; |  | ||||||
|             dataCopyParams.blockCount = 1; |  | ||||||
|             dataCopyParams.blockLen = scale_local_offset * sizeof(half); |  | ||||||
|             DataCopyPad(scale_gm[scale_global_offset], scale_local, |  | ||||||
|                         dataCopyParams); |  | ||||||
|             pipe_barrier(PIPE_ALL); |  | ||||||
|         } |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t *scale_ne; |  | ||||||
|     size_t scale_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t group_size_in_row; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<float> input_gm; |  | ||||||
|     GlobalTensor<half> scale_gm; |  | ||||||
|     GlobalTensor<int8_t> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     TQue<QuePosition::VECIN, 1> work_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, 1> max_queue; |  | ||||||
|     TQue<QuePosition::VECIN, 1> abs_queue; |  | ||||||
|     TQue<QuePosition::VECIN, 1> cast_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, 1> scale_queue; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_quantize_f32_q8_0( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|     GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|  |  | ||||||
|     QUANTIZE_F32_Q8_0 op; |  | ||||||
|     op.init(input_gm, output_gm, input_ne_ub, input_nb_ub, output_ne_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| #endif // #ifdef ASCEND_310P |  | ||||||
| @@ -1,295 +0,0 @@ | |||||||
| #include "kernel_operator.h" |  | ||||||
|  |  | ||||||
| using namespace AscendC; |  | ||||||
| #ifdef ASCEND_310P // 310P not support float->4bit quantization |  | ||||||
|     extern "C" __global__ __aicore__ void ascendc_quantize_f32_to_q4_0( |  | ||||||
|         GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|         GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|         // let following test cases can continue run, here just print error information. Of Cource the test case that call this operator is failed. |  | ||||||
|         printf("Ascend310P not support f32->4bit quantization.\n"); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     extern "C" __global__ __aicore__ void ascendc_quantize_f16_to_q4_0( |  | ||||||
|         GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|         GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|         // let following test cases can continue run, here just print error information. Of Cource the test case that call this operator is failed. |  | ||||||
|         printf("Ascend310P not support f16->4bit quantization.\n"); |  | ||||||
|     } |  | ||||||
| #else |  | ||||||
|  |  | ||||||
| #define BUFFER_NUM 2 |  | ||||||
| #define Group_Size 32 |  | ||||||
|  |  | ||||||
| template <typename SRC_T> |  | ||||||
| class QUANTIZE_FLOAT_TO_Q4_0 { |  | ||||||
|    public: |  | ||||||
|     __aicore__ inline QUANTIZE_FLOAT_TO_Q4_0() {} |  | ||||||
|     __aicore__ inline void init(GM_ADDR input, GM_ADDR output, |  | ||||||
|                                 int64_t *input_ne_ub, size_t *input_nb_ub, |  | ||||||
|                                 int64_t *output_ne_ub) { |  | ||||||
|         // TODO: fix test_case CPY(type_src=f16,type_dst=q4_0,ne=[256,4,4,4], |  | ||||||
|         //                         permute=[0,0,0,0]): |  | ||||||
|         // [CPY] NMSE = 0.000008343 > 0.000001000 FAIL |  | ||||||
|         int64_t op_block_num = GetBlockNum(); |  | ||||||
|         int64_t op_block_idx = GetBlockIdx(); |  | ||||||
|  |  | ||||||
|         // input stride of data elements |  | ||||||
|         for (int i = 0; i < 4; i++) { |  | ||||||
|             input_ne[i] = input_ne_ub[i]; |  | ||||||
|             input_stride[i] = input_nb_ub[i] / input_nb_ub[0]; |  | ||||||
|             output_ne[i] = output_ne_ub[i]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // output stride of data elements |  | ||||||
|         output_stride[0] = 1; |  | ||||||
|         for (int i = 1; i < 4; i++) { |  | ||||||
|             output_stride[i] = output_stride[i - 1] * output_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // scale saved one by one after data:. [group1_scale, group2_scale, ...] |  | ||||||
|         scale_ne = input_ne; |  | ||||||
|         scale_stride[0] = 1; |  | ||||||
|         scale_stride[1] = input_ne[0] / Group_Size; |  | ||||||
|         for (int i = 2; i < 4; i++) { |  | ||||||
|             scale_stride[i] = scale_stride[i - 1] * scale_ne[i - 1]; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // split input tensor by rows. |  | ||||||
|         uint64_t nr = input_ne[1] * input_ne[2] * input_ne[3]; |  | ||||||
|         dr = nr / op_block_num; |  | ||||||
|  |  | ||||||
|         uint64_t tails = nr % op_block_num; |  | ||||||
|         if (op_block_idx < tails) { |  | ||||||
|             dr += 1; |  | ||||||
|             ir = dr * op_block_idx; |  | ||||||
|         } else { |  | ||||||
|             ir = dr * op_block_idx + tails; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         group_size_in_row = scale_stride[1]; |  | ||||||
|         int64_t scale_offset = output_ne[0] * output_ne[1] * output_ne[2] * |  | ||||||
|                               output_ne[3] * sizeof(uint8_t) / 2; |  | ||||||
|  |  | ||||||
|         input_gm.SetGlobalBuffer((__gm__ SRC_T *)input); |  | ||||||
|         output_gm.SetGlobalBuffer((__gm__ int8_t *)output); |  | ||||||
|         scale_gm.SetGlobalBuffer((__gm__ half *)(output + scale_offset + ir * |  | ||||||
|                                                  group_size_in_row * |  | ||||||
|                                                  sizeof(half))); |  | ||||||
|  |  | ||||||
|         pipe.InitBuffer(input_queue, BUFFER_NUM, Group_Size * sizeof(SRC_T)); |  | ||||||
|         pipe.InitBuffer(output_queue, BUFFER_NUM, |  | ||||||
|                             Group_Size * sizeof(int8_t) / 2); |  | ||||||
|         pipe.InitBuffer(cast_queue , 1, Group_Size * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(work_queue, 1, Group_Size * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(max_queue, 1, Group_Size * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(min_queue, 1, Group_Size * sizeof(float)); |  | ||||||
|         pipe.InitBuffer(scale_queue, 1, Group_Size / 2 * sizeof(half)); |  | ||||||
|         pipe.InitBuffer(int8_queue, 1, Group_Size * sizeof(int8_t)); |  | ||||||
|         pipe.InitBuffer(half_queue, 1, Group_Size * sizeof(half)); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_in(uint32_t offset) { |  | ||||||
|         LocalTensor<SRC_T> input_local = input_queue.AllocTensor<SRC_T>(); |  | ||||||
|         DataCopy(input_local, input_gm[offset], Group_Size); |  | ||||||
|         input_queue.EnQue(input_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void copy_out(uint32_t offset) { |  | ||||||
|         // reinterpretcast Group_Size(32) * int4b_t to Group_Size / 2 * int8_t, |  | ||||||
|         // and using DataCopyPad to avoid 32 bits align. |  | ||||||
|         LocalTensor<int4b_t> output_local = output_queue.DeQue<int4b_t>(); |  | ||||||
|         LocalTensor<int8_t> output_int8_local = |  | ||||||
|                                     output_local.ReinterpretCast<int8_t>(); |  | ||||||
|  |  | ||||||
|         DataCopyExtParams dataCopyParams; |  | ||||||
|         dataCopyParams.blockCount = 1; |  | ||||||
|         dataCopyParams.blockLen = Group_Size / 2  * sizeof(int8_t); |  | ||||||
|         DataCopyPad(output_gm[offset], output_int8_local, dataCopyParams); |  | ||||||
|  |  | ||||||
|         output_queue.FreeTensor(output_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void input_to_cast(LocalTensor<float> cast_local, |  | ||||||
|                                          LocalTensor<float> input_local) { |  | ||||||
|         DataCopy(cast_local, input_local, Group_Size); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void input_to_cast(LocalTensor<float> cast_local, |  | ||||||
|                                          LocalTensor<half> input_local) { |  | ||||||
|         Cast(cast_local, input_local, RoundMode::CAST_NONE, Group_Size); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline half calculate_group(int64_t row, int64_t group) { |  | ||||||
|         const int64_t i3 = row / (input_ne[1] * input_ne[2]); |  | ||||||
|         const int64_t i2 = (row - i3 * input_ne[1] * input_ne[2]) / input_ne[1]; |  | ||||||
|         const int64_t i1 = |  | ||||||
|             row - i3 * input_ne[1] * input_ne[2] - i2 * input_ne[1]; |  | ||||||
|  |  | ||||||
|         const int64_t input_offset = i1 * input_stride[1] + |  | ||||||
|                                      i2 * input_stride[2] + |  | ||||||
|                                      i3 * input_stride[3] + Group_Size * group; |  | ||||||
|  |  | ||||||
|         // output_offset is stride for output_gm which datatype is int8_t and |  | ||||||
|         // divided by 2 is needed for int4b_t. |  | ||||||
|         const int64_t output_offset = (i1 * output_stride[1] + |  | ||||||
|                                        i2 * output_stride[2] + |  | ||||||
|                                        i3 * output_stride[3] + |  | ||||||
|                                        Group_Size * group) / 2; |  | ||||||
|         copy_in(input_offset); |  | ||||||
|  |  | ||||||
|         LocalTensor<SRC_T> input_local = input_queue.DeQue<SRC_T>(); |  | ||||||
|         LocalTensor<int4b_t> output_local = output_queue.AllocTensor<int4b_t>(); |  | ||||||
|         LocalTensor<float> cast_local = cast_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> work_local = work_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> max_local = max_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<float> min_local = min_queue.AllocTensor<float>(); |  | ||||||
|         LocalTensor<int8_t> int8_local = int8_queue.AllocTensor<int8_t>(); |  | ||||||
|         LocalTensor<half> half_local = half_queue.AllocTensor<half>(); |  | ||||||
|  |  | ||||||
|         input_to_cast(cast_local, input_local); |  | ||||||
|  |  | ||||||
|         ReduceMax(max_local, cast_local, work_local, Group_Size); |  | ||||||
|         ReduceMin(min_local, cast_local, work_local, Group_Size); |  | ||||||
|         const float max_value = max_local.GetValue(0); |  | ||||||
|         const float min_value = min_local.GetValue(0); |  | ||||||
|         float d = max_value; |  | ||||||
|         if (min_value < 0 && (-1 * min_value) > max_value) { |  | ||||||
|             d = min_value; |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         d = d / (-8); |  | ||||||
|         if (d != 0) { |  | ||||||
|             Muls(cast_local, cast_local, 1.0f / d, Group_Size); |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         // range: [-8,8] -> [0.5,16.5] -> [0,16] -> [0,15] -> [-8,7] |  | ||||||
|         float scalar = 8.5f; |  | ||||||
|         Adds(cast_local, cast_local, scalar, Group_Size); |  | ||||||
|         Cast(cast_local, cast_local, RoundMode::CAST_FLOOR, Group_Size); |  | ||||||
|         scalar = 15.0f; |  | ||||||
|         Mins(cast_local, cast_local, scalar, Group_Size); |  | ||||||
|         scalar = -8.0f; |  | ||||||
|         Adds(cast_local, cast_local, scalar, Group_Size); |  | ||||||
|  |  | ||||||
|         // float->half->int4b |  | ||||||
|         Cast(half_local, cast_local, RoundMode::CAST_NONE, Group_Size); |  | ||||||
|         Cast(output_local, half_local, RoundMode::CAST_NONE, Group_Size); |  | ||||||
|  |  | ||||||
|         output_queue.EnQue(output_local); |  | ||||||
|         copy_out(output_offset); |  | ||||||
|  |  | ||||||
|         input_queue.FreeTensor(input_local); |  | ||||||
|         work_queue.FreeTensor(work_local); |  | ||||||
|         max_queue.FreeTensor(max_local); |  | ||||||
|         min_queue.FreeTensor(min_local); |  | ||||||
|         int8_queue.FreeTensor(int8_local); |  | ||||||
|         half_queue.FreeTensor(half_local); |  | ||||||
|         cast_queue.FreeTensor(cast_local); |  | ||||||
|         return (half)d; |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|     __aicore__ inline void calculate() { |  | ||||||
|         LocalTensor<half> scale_local = scale_queue.AllocTensor<half>(); |  | ||||||
|         uint32_t scale_local_offset = 0; |  | ||||||
|         uint32_t scale_global_offset = 0; |  | ||||||
|         for (int64_t i = ir; i < ir + dr; i++) { |  | ||||||
|             for (int64_t j = 0; j < group_size_in_row; j++) { |  | ||||||
|                 half scale = calculate_group(i, j); |  | ||||||
|                 scale_local.SetValue(scale_local_offset++, scale); |  | ||||||
|                 // Copy Group_Size/2 length data each time. |  | ||||||
|                 if (scale_local_offset == Group_Size / 2) { |  | ||||||
|                     scale_local_offset = 0; |  | ||||||
|                     // TODO: OPTIMIZE ME |  | ||||||
|                     pipe_barrier(PIPE_ALL); |  | ||||||
|                     DataCopy(scale_gm[scale_global_offset], scale_local, |  | ||||||
|                                       Group_Size / 2); |  | ||||||
|                     pipe_barrier(PIPE_ALL); |  | ||||||
|                     scale_global_offset += Group_Size / 2; |  | ||||||
|                 } |  | ||||||
|             } |  | ||||||
|         } |  | ||||||
|  |  | ||||||
|         if (scale_local_offset != 0) { |  | ||||||
|             pipe_barrier(PIPE_ALL); |  | ||||||
|             DataCopyExtParams dataCopyParams; |  | ||||||
|             dataCopyParams.blockCount = 1; |  | ||||||
|             dataCopyParams.blockLen = scale_local_offset * sizeof(half); |  | ||||||
|             DataCopyPad(scale_gm[scale_global_offset], scale_local, |  | ||||||
|                         dataCopyParams); |  | ||||||
|             pipe_barrier(PIPE_ALL); |  | ||||||
|         } |  | ||||||
|         scale_queue.FreeTensor(scale_local); |  | ||||||
|     } |  | ||||||
|  |  | ||||||
|    private: |  | ||||||
|     int64_t input_ne[4]; |  | ||||||
|     size_t input_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t *scale_ne; |  | ||||||
|     size_t scale_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t output_ne[4]; |  | ||||||
|     size_t output_stride[4]; |  | ||||||
|  |  | ||||||
|     int64_t group_size_in_row; |  | ||||||
|  |  | ||||||
|     int64_t ir; |  | ||||||
|     int64_t dr; |  | ||||||
|  |  | ||||||
|     TPipe pipe; |  | ||||||
|     GlobalTensor<SRC_T> input_gm; |  | ||||||
|     GlobalTensor<half> scale_gm; |  | ||||||
|     GlobalTensor<int8_t> output_gm; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> input_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> output_queue; |  | ||||||
|     TQue<QuePosition::VECIN, BUFFER_NUM> work_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> max_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> min_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> scale_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> cast_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> int8_queue; |  | ||||||
|     TQue<QuePosition::VECOUT, BUFFER_NUM> half_queue; |  | ||||||
| }; |  | ||||||
|  |  | ||||||
| template <typename T> |  | ||||||
| __aicore__ inline void copy_to_ub(GM_ADDR gm, T *ub, size_t size) { |  | ||||||
|     auto gm_ptr = (__gm__ uint8_t *)gm; |  | ||||||
|     auto ub_ptr = (uint8_t *)(ub); |  | ||||||
|     for (int32_t i = 0; i < size; ++i, ++ub_ptr, ++gm_ptr) { |  | ||||||
|         *ub_ptr = *gm_ptr; |  | ||||||
|     } |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_quantize_f16_to_q4_0( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|     GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|  |  | ||||||
|     QUANTIZE_FLOAT_TO_Q4_0<half> op; |  | ||||||
|     op.init(input_gm, output_gm, input_ne_ub, input_nb_ub, output_ne_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| extern "C" __global__ __aicore__ void ascendc_quantize_f32_to_q4_0( |  | ||||||
|     GM_ADDR input_gm, GM_ADDR output_gm, GM_ADDR input_ne_gm, |  | ||||||
|     GM_ADDR input_nb_gm, GM_ADDR output_ne_gm) { |  | ||||||
|     int64_t input_ne_ub[4]; |  | ||||||
|     size_t input_nb_ub[4]; |  | ||||||
|     int64_t output_ne_ub[4]; |  | ||||||
|  |  | ||||||
|     copy_to_ub(input_ne_gm, input_ne_ub, 32); |  | ||||||
|     copy_to_ub(input_nb_gm, input_nb_ub, 32); |  | ||||||
|     copy_to_ub(output_ne_gm, output_ne_ub, 32); |  | ||||||
|  |  | ||||||
|     QUANTIZE_FLOAT_TO_Q4_0<float> op; |  | ||||||
|     op.init(input_gm, output_gm, input_ne_ub, input_nb_ub, output_ne_ub); |  | ||||||
|     op.calculate(); |  | ||||||
| } |  | ||||||
|  |  | ||||||
| #endif // #ifdef ASCEND_310P |  | ||||||
		Reference in New Issue
	
	Block a user
	 Chenguang Li
					Chenguang Li