Repository: ggerganov/llama.cpp
File: ggml/src/ggml-cann/aclnn_ops.cpp (C++)

Code before:
        src->ne[1] * src->ne[2] * src->ne[3] * ggml_element_size(src);
    ggml_cann_pool_alloc zero_tensor_allocator(ctx.pool(), zero_tensor_n_bytes);
    aclTensor* acl_rstd =
        aclnn_zero(ctx, zero_tensor_allocator.get(), zero_tensor_n_bytes,
                   src->ne, GGML_MAX_DIMS, ggml_cann_type_mapping(src->type),
                   ggml_element_size(src));
    GGML_CANN_CALL_ACLNN_OP(ctx, RmsNorm, acl_src, acl_gamma, eps, acl_dst, acl_rstd);
    ggml_cann_release_resources(ctx, acl_src, acl_dst, acl_gamma, acl_rstd);
}

// TODO: performace is low.
void ggml_cann_diag_mask(ggml_backend_cann_context& ctx, ggml_tensor* dst,
                         float value) {
    ggml_tensor* src = dst->src[0];

    aclTensor* acl_src = ggml_cann_create_tensor(src);
    aclTensor* acl_dst = ggml_cann_create_tensor(dst);

    const int n_past = ((int32_t*)dst->op_params)[0];

    size_t one_tensor_n_bytes = src->ne[0] * src->ne[1] * src->ne[2] *
                                src->ne[3] * ggml_element_size(src);
    ggml_cann_pool_alloc one_tensor_allocator(ctx.pool(), one_tensor_n_bytes);

    aclTensor* mask_tensor =
        aclnn_values(ctx, one_tensor_allocator.get(), one_tensor_n_bytes,
                     src->ne, GGML_MAX_DIMS, ggml_cann_type_mapping(src->type),
                     ggml_element_size(src), value);

    aclScalar* alpha = nullptr;
    float alphaValue = 1.0f;
    alpha = aclCreateScalar(&alphaValue, aclDataType::ACL_FLOAT);

    GGML_CANN_CALL_ACLNN_OP(ctx, InplaceTriu, mask_tensor, n_past + 1);
    GGML_CANN_CALL_ACLNN_OP(ctx, Tril, acl_src, n_past + 1, acl_dst);
    GGML_CANN_CALL_ACLNN_OP(ctx, InplaceAdd, acl_dst, mask_tensor, alpha);
    ggml_cann_release_resources(ctx, alpha, acl_src, acl_dst, mask_tensor);
}

/**
 * @brief Permutes the dimensions of a tensor according to a specified order.
 *
 * This function permutes the dimensions of the source tensor `acl_src`
 * according to the order specified in the `new_dim` array and stores the result
 * in the destination tensor `acl_dst`.
 *
 * @param ctx The context for the CANN backend operations.
 * @param acl_src The source tensor whose dimensions will be permuted.
 * @param acl_dst The destination tensor where the permuted result will be
 * stored.
 * @param new_dim An array specifying the new order of dimensions for the

Review comment (performance):
I think cache_element is not need to define. Make it 0 as default. the first time this function is called, it will go to cache expansion to malloc memory.

Code after:

/**
 * @brief Get or expand a cached float32 tensor filled with a scalar value.
 *
 * This function manages cached device memory for float32 tensors. If the current
 * cache size is insufficient for the requested tensor shape, the old memory will
 * be released and new memory will be allocated. The allocated buffer is then
 * initialized either with zeros (when @p value == 0.0f) or with the given scalar
 * value using CANN operations. Finally, an aclTensor object is created from the
 * cached memory and returned.
 *
 * @param ctx           The CANN backend context that manages device memory.
 * @param buffer        A pointer to the cached device buffer (will be allocated
 *                      or reallocated if necessary).
 * @param cache_element The current number of cached elements. This will be
 *                      updated when the cache is expanded.
 * @param ne            The tensor shape array (number of elements in each dimension).
 * @param nb            The stride size for each dimension.
 * @param dims          The number of tensor dimensions.
 * @param value         The scalar value used to fill the tensor (supports zero
 *                      initialization via memset or arbitrary values via fill_scalar).
 * @return              An aclTensor pointer created from the cached buffer.
 */
static aclTensor* get_f32_cache_acl_tensor(
    ggml_backend_cann_context& ctx,
    void** buffer,
    int64_t &cache_element,
    int64_t* ne,
    size_t* nb,
    int64_t dims,
    float value) {
    // Calculate total number of elements
    int64_t n_element = 1;
    for (int i = 0; i < dims; i++) {
        n_element *= ne[i];
    }
    size_t size = n_element * sizeof(float);

    // Allocate or expand cache if needed
    if (cache_element < n_element) {
        if (*buffer != nullptr) {
            aclrtFree(*buffer);
            *buffer = nullptr;
        }

        ACL_CHECK(aclrtMalloc(buffer, size, ACL_MEM_MALLOC_HUGE_FIRST));
        cache_element = n_element;

        // Initialize cache
        if (value == 0.0f) {
            ACL_CHECK(aclrtMemsetAsync(*buffer, size, 0, size, ctx.stream()));