From 1929598052be482156348248e3a3546fc77823d5 Mon Sep 17 00:00:00 2001 From: Invisi Date: Fri, 17 Oct 2025 15:10:09 +0700 Subject: [PATCH] RKNPU2 backend is implemented --- CMakeLists.txt | 28 + ggml/CMakeLists.txt | 1 + ggml/src/CMakeLists.txt | 1 + ggml/src/ggml-backend-reg.cpp | 16 + ggml/src/ggml-rknpu2/CMakeLists.txt | 27 + ggml/src/ggml-rknpu2/dma-alloc.cpp | 54 ++ ggml/src/ggml-rknpu2/dma-alloc.h | 30 + ggml/src/ggml-rknpu2/ggml-rknpu2.cpp | 878 ++++++++++++++++++ ggml/src/ggml-rknpu2/ggml-rknpu2.h | 13 + ggml/src/ggml-rknpu2/libs/include/rknn_api.h | 804 ++++++++++++++++ .../ggml-rknpu2/libs/include/rknn_custom_op.h | 144 +++ .../libs/include/rknn_matmul_api.h | 544 +++++++++++ 12 files changed, 2540 insertions(+) create mode 100644 ggml/src/ggml-rknpu2/CMakeLists.txt create mode 100644 ggml/src/ggml-rknpu2/dma-alloc.cpp create mode 100644 ggml/src/ggml-rknpu2/dma-alloc.h create mode 100644 ggml/src/ggml-rknpu2/ggml-rknpu2.cpp create mode 100644 ggml/src/ggml-rknpu2/ggml-rknpu2.h create mode 100644 ggml/src/ggml-rknpu2/libs/include/rknn_api.h create mode 100644 ggml/src/ggml-rknpu2/libs/include/rknn_custom_op.h create mode 100644 ggml/src/ggml-rknpu2/libs/include/rknn_matmul_api.h diff --git a/CMakeLists.txt b/CMakeLists.txt index caea48c50..0a4b432b9 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -115,6 +115,7 @@ option(LLAMA_TESTS_INSTALL "llama: install tests" ON) # 3rd party libs option(LLAMA_OPENSSL "llama: use openssl to support HTTPS" ON) option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF) +option(LLAMA_RKNPU2 "llama: use Rockchip NPU" OFF) # Required for relocatable CMake package @@ -169,6 +170,33 @@ llama_option_depr(WARNING LLAMA_CURL) include("cmake/license.cmake") license_add_file("llama.cpp" "LICENSE") +if (LLAMA_RKNPU2) + set(GGML_RKNPU2 ON) +endif() + +if (NOT MSVC) + if (LLAMA_SANITIZE_THREAD) + message(STATUS "Using -fsanitize=thread") + + add_compile_options(-fsanitize=thread) + link_libraries (-fsanitize=thread) + endif() + + if (LLAMA_SANITIZE_ADDRESS) + message(STATUS "Using -fsanitize=address") + + add_compile_options(-fsanitize=address -fno-omit-frame-pointer) + link_libraries (-fsanitize=address) + endif() + + if (LLAMA_SANITIZE_UNDEFINED) + message(STATUS "Using -fsanitize=undefined") + + add_compile_options(-fsanitize=undefined) + link_libraries (-fsanitize=undefined) + endif() +endif() + # # 3rd-party # diff --git a/ggml/CMakeLists.txt b/ggml/CMakeLists.txt index 5834e544b..cd4c14425 100644 --- a/ggml/CMakeLists.txt +++ b/ggml/CMakeLists.txt @@ -259,6 +259,7 @@ set (GGML_OPENCL_TARGET_VERSION "300" CACHE STRING "ggml: OpenCL API version to target") option(GGML_HEXAGON "ggml: enable Hexagon backend" OFF) +option(GGML_RKNPU2 "ggml: use RKNPU2" OFF) set(GGML_HEXAGON_FP32_QUANTIZE_GROUP_SIZE 128 CACHE STRING "ggml: quantize group size (32, 64, or 128)") # toolchain for vulkan-shaders-gen diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 78853304d..35feeca97 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -461,6 +461,7 @@ ggml_add_backend(OpenCL) ggml_add_backend(Hexagon) ggml_add_backend(ZenDNN) ggml_add_backend(OPENVINO) +ggml_add_backend(RKNPU2) foreach (target ggml-base ggml) target_include_directories(${target} PUBLIC $ $) diff --git a/ggml/src/ggml-backend-reg.cpp b/ggml/src/ggml-backend-reg.cpp index 058710921..0c4010a38 100644 --- a/ggml/src/ggml-backend-reg.cpp +++ b/ggml/src/ggml-backend-reg.cpp @@ -86,6 +86,19 @@ #include "ggml-openvino.h" #endif +#ifdef GGML_USE_RKNPU2 +#include "ggml-rknpu2.h" +#endif + +// disable C++17 deprecation warning for std::codecvt_utf8 +#if defined(__clang__) +# pragma clang diagnostic push +# pragma clang diagnostic ignored "-Wdeprecated-declarations" +#elif defined(__GNUC__) +# pragma GCC diagnostic push +# pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + namespace fs = std::filesystem; static std::string path_str(const fs::path & path) { @@ -163,6 +176,9 @@ struct ggml_backend_registry { #endif #ifdef GGML_USE_CPU register_backend(ggml_backend_cpu_reg()); +#endif +#ifdef GGML_USE_RKNPU2 + register_backend(ggml_backend_rknpu2_reg()); #endif } diff --git a/ggml/src/ggml-rknpu2/CMakeLists.txt b/ggml/src/ggml-rknpu2/CMakeLists.txt new file mode 100644 index 000000000..407de9f7c --- /dev/null +++ b/ggml/src/ggml-rknpu2/CMakeLists.txt @@ -0,0 +1,27 @@ +# Finding OpenMP +find_package(OpenMP REQUIRED) + +# Adding backend files +ggml_add_backend_library(ggml-rknpu2 + ggml-rknpu2.cpp + dma-alloc.cpp +) + +# Linking rknpu2 header files +target_include_directories(ggml-rknpu2 + PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/libs/include + PUBLIC + ${CMAKE_CURRENT_SOURCE_DIR} +) + +# Linking rknpu2 library +target_link_libraries(ggml-rknpu2 PRIVATE + ${CMAKE_CURRENT_SOURCE_DIR}/libs/librknnrt.so +) + +# Adding OpenMP in build +if(OpenMP_FOUND) + add_definitions(-DHAVE_OPENMP=1) + list(APPEND CMAKE_CXX_FLAGS "${OpenMP_CXX_FLAGS}") +endif() \ No newline at end of file diff --git a/ggml/src/ggml-rknpu2/dma-alloc.cpp b/ggml/src/ggml-rknpu2/dma-alloc.cpp new file mode 100644 index 000000000..29ebd0b78 --- /dev/null +++ b/ggml/src/ggml-rknpu2/dma-alloc.cpp @@ -0,0 +1,54 @@ +#include "dma-alloc.h" +#include +#include +#include +#include +#include +#include +#include + +DmaBuffer dma_alloc(size_t size) { + DmaBuffer buffer; + buffer.size = size; + + const char* path = "/dev/dma_heap/system"; + int dma_heap_fd = open(path, O_RDWR); + if (dma_heap_fd < 0) { + fprintf(stderr, "DMA_ALLOC: Failed to open %s: %s\n", path, strerror(errno)); + return buffer; + } + + dma_heap_allocation_data buf_data; + memset(&buf_data, 0, sizeof(buf_data)); + buf_data.len = size; + buf_data.fd_flags = O_CLOEXEC | O_RDWR; + + if (ioctl(dma_heap_fd, DMA_HEAP_IOCTL_ALLOC, &buf_data) < 0) { + fprintf(stderr, "DMA_ALLOC: ioctl DMA_HEAP_IOCTL_ALLOC failed: %s\n", strerror(errno)); + close(dma_heap_fd); + return buffer; + } + + close(dma_heap_fd); + + buffer.fd = buf_data.fd; + buffer.virt_addr = mmap(NULL, size, PROT_READ | PROT_WRITE, MAP_SHARED, buffer.fd, 0); + + if (buffer.virt_addr == MAP_FAILED) { + fprintf(stderr, "DMA_ALLOC: mmap failed: %s\n", strerror(errno)); + close(buffer.fd); + buffer.fd = -1; + buffer.virt_addr = nullptr; + } + + return buffer; +} + +void dma_free(const DmaBuffer& buffer) { + if (buffer.virt_addr) { + munmap(buffer.virt_addr, buffer.size); + } + if (buffer.fd >= 0) { + close(buffer.fd); + } +} \ No newline at end of file diff --git a/ggml/src/ggml-rknpu2/dma-alloc.h b/ggml/src/ggml-rknpu2/dma-alloc.h new file mode 100644 index 000000000..f5efa3ba3 --- /dev/null +++ b/ggml/src/ggml-rknpu2/dma-alloc.h @@ -0,0 +1,30 @@ +#pragma once + +#include +#include + +// --- Structs --- +struct dma_heap_allocation_data { + uint64_t len; + uint32_t fd; + uint32_t fd_flags; + uint64_t heap_flags; +}; + +// --- IOCTL-commands --- + +#define DMA_HEAP_IOC_MAGIC 'H' +#define DMA_HEAP_IOCTL_ALLOC _IOWR(DMA_HEAP_IOC_MAGIC, 0x0, struct dma_heap_allocation_data) + +// Information of allocated block +struct DmaBuffer { + int fd = -1; + void* virt_addr = nullptr; + size_t size = 0; +}; + +// --- Functions --- + +DmaBuffer dma_alloc(size_t size); + +void dma_free(const DmaBuffer& buffer); diff --git a/ggml/src/ggml-rknpu2/ggml-rknpu2.cpp b/ggml/src/ggml-rknpu2/ggml-rknpu2.cpp new file mode 100644 index 000000000..3e5a5cf0a --- /dev/null +++ b/ggml/src/ggml-rknpu2/ggml-rknpu2.cpp @@ -0,0 +1,878 @@ +#include "ggml-rknpu2.h" +#include "ggml-backend-impl.h" +#include "ggml-impl.h" +#include "ggml-quants.h" + +#include "dma-alloc.h" + +#include +#include + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#define UNUSED(x) (void)(x) + +// Performance logging flag +#define RKNPU_PERF_LOG 0 + +// Backend logging flag +#define RKNPU_DEBUG 0 + +// RAII-struct for performance logging +#if RKNPU_PERF_LOG +struct RknpuPerfLogger { + std::string name; + std::chrono::high_resolution_clock::time_point start; + + RknpuPerfLogger(std::string name) : name(std::move(name)), start(std::chrono::high_resolution_clock::now()) {} + + ~RknpuPerfLogger() { + auto end = std::chrono::high_resolution_clock::now(); + auto duration = std::chrono::duration_cast(end - start).count(); + fprintf(stderr, "[PERF] %-40s: %8lld us\n", name.c_str(), (long long)duration); + } +}; +#endif + +#if RKNPU_DEBUG +#define RKNPU_LOG_INFO(...) fprintf(stderr, __VA_ARGS__) +#else +#define RKNPU_LOG_INFO(...) +#endif + +#define RKNPU_LOG_ERROR(...) fprintf(stderr, __VA_ARGS__) + +// Macro for RKNN API calls +#define RKNN_CHECK(stmt, msg) \ + do { \ + int ret = (stmt); \ + if (ret < 0) { \ + RKNPU_LOG_ERROR("RKNN error %d at %s:%d: %s\n", ret, \ + __FILE__, __LINE__, msg); \ + assert(false); \ + } \ + } while (0) + + +// --- Helper functions --- + +// Packing KxN FP16 (row-major: idx [k,n] -> k*N + n) into native RKNN: (N/16, K/32, 16, 32) +static void pack_B_segment_fp16_native_from_KN( + uint16_t * dst, const uint16_t * src, + int K, int N_total, int n_offset, int n_segment) { + + GGML_ASSERT(K % 32 == 0 && N_total > 0 && K > 0); + GGML_ASSERT(n_offset % 16 == 0 && n_segment % 16 == 0 && n_offset + n_segment <= N_total); + + const size_t s0 = (size_t)(K / 32) * 16 * 32; + const size_t s1 = 16 * 32; + const size_t s2 = 32; + + for (int i = 0; i < n_segment / 16; ++i) { + for (int j = 0; j < K / 32; ++j) { + const size_t dst_block = (size_t) i * s0 + (size_t) j * s1; + for (int ii = 0; ii < 16; ++ii) { + const size_t n_global = (size_t)n_offset + (size_t)i * 16 + (size_t)ii; + + const uint16_t * src_ptr = src + n_global * K + j * 32; + uint16_t * dst_ptr = dst + dst_block + ii * s2; + + uint16x8_t d0 = vld1q_u16(src_ptr + 0); + uint16x8_t d1 = vld1q_u16(src_ptr + 8); + uint16x8_t d2 = vld1q_u16(src_ptr + 16); + uint16x8_t d3 = vld1q_u16(src_ptr + 24); + + vst1q_u16(dst_ptr + 0, d0); + vst1q_u16(dst_ptr + 8, d1); + vst1q_u16(dst_ptr + 16, d2); + vst1q_u16(dst_ptr + 24, d3); + } + } + } +} + +// --- Hashers --- + +// Function for hash combinations +template +inline void hash_combine(std::size_t& seed, const T& v) { + std::hash hasher; + seed ^= hasher(v) + 0x9e3779b9 + (seed << 6) + (seed >> 2); +} + +// Hasher for std::pair +struct PairHasher { + template + std::size_t operator()(const std::pair& p) const { + std::size_t seed = 0; + hash_combine(seed, p.first); + hash_combine(seed, p.second); + return seed; + } +}; + +// Hasher for std::tuple +struct TupleHasher { + template + std::size_t operator()(const std::tuple& t) const { + std::size_t seed = 0; + std::apply([&](const auto&... args) { + (hash_combine(seed, args), ...); + }, t); + return seed; + } +}; + +// --- Segmenters --- + +// Matrix segment information +struct MatrixSegment { + int offset_n; // Segment offset + int size_n; // Segment size + int core_id; // Segment core ID +}; + +// Split B-matrix into segments +static std::vector compute_matrix_segments(int N, int num_cores = 3) { + std::vector segments; + + // N divisible by 16 for F16 + const int alignment = 16; + + // Basic segment size + int base_segment_size = (N / num_cores / alignment) * alignment; + int remaining = N - (base_segment_size * num_cores); + + int offset = 0; + for (int i = 0; i < num_cores; i++) { + MatrixSegment seg; + seg.offset_n = offset; + seg.size_n = base_segment_size; + seg.core_id = i; + + if (i < remaining / alignment) { + seg.size_n += alignment; + } + + offset += seg.size_n; + segments.push_back(seg); + } + + return segments; +} + +// --- Structs --- + +// RKNN buffer context +struct ggml_backend_rknpu_buffer_context { + rknn_tensor_mem * rknn_mem; + DmaBuffer dma_buf; + std::string name; +}; + +// RKNN matmul operation context +struct rknpu_matmul_context { + rknn_matmul_info info; + rknn_matmul_io_attr io_attr; + rknn_matmul_ctx ctx = 0; + + rknpu_matmul_context(int M, int K, int N, rknn_matmul_type type) { + memset(&info, 0, sizeof(info)); + info.M = M; + info.K = K; + info.N = N; + info.type = type; + info.B_layout = RKNN_MM_LAYOUT_NATIVE; + info.AC_layout = RKNN_MM_LAYOUT_NORM; + + int ret = rknn_matmul_create(&ctx, &info, &io_attr); + if (ret < 0) { + RKNPU_LOG_ERROR("rknn_matmul_create failed for %dx%dx%d\n", M, K, N); + ctx = 0; + } + } + + ~rknpu_matmul_context() { + if (ctx != 0) { + rknn_matmul_destroy(ctx); + } + } +}; + + +// Backend main context +struct ggml_backend_rknpu_context { + std::string name; + std::mutex mutex; + + // RKNN matmul contexts cache + std::unordered_map, std::shared_ptr, TupleHasher> matmul_ctx_cache; + + // B-matrices cache + std::unordered_map, std::shared_ptr, TupleHasher> b_npu_buffer_cache; + + // A- and C-matrices cache + std::unordered_map, std::shared_ptr, TupleHasher> a_buffer_cache; + std::unordered_map, std::shared_ptr, TupleHasher> c_buffer_cache; + + std::shared_ptr get_matmul_ctx(int M, int K, int N, int core_id) { + std::lock_guard lock(mutex); + auto key = std::make_tuple(M, K, N, core_id); + auto it = matmul_ctx_cache.find(key); + if (it != matmul_ctx_cache.end()) { + return it->second; + } + auto ctx = std::make_shared(M, K, N, RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT32); + if (ctx->ctx == 0) { + return nullptr; + } + + rknn_core_mask core_mask; + switch(core_id) { + case 0: core_mask = RKNN_NPU_CORE_0; break; + case 1: core_mask = RKNN_NPU_CORE_1; break; + case 2: core_mask = RKNN_NPU_CORE_2; break; + default: core_mask = RKNN_NPU_CORE_AUTO; break; + } + + int ret = rknn_matmul_set_core_mask(ctx->ctx, core_mask); + if (ret != RKNN_SUCC) { + RKNPU_LOG_ERROR("Failed to set core mask %d for core %d\n", core_mask, core_id); + } + + matmul_ctx_cache[key] = ctx; + return ctx; + } +}; + +// RKNN memory global context +struct rknpu_memory_context { + rknn_matmul_ctx mem_ctx = 0; + std::mutex mutex; + + rknpu_memory_context() { + rknn_matmul_info dummy_info; + memset(&dummy_info, 0, sizeof(dummy_info)); + dummy_info.M = 32; + dummy_info.K = 32; + dummy_info.N = 32; + dummy_info.type = RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT32; + + rknn_matmul_io_attr dummy_io_attr; + int ret = rknn_matmul_create(&mem_ctx, &dummy_info, &dummy_io_attr); + if (ret < 0) { + RKNPU_LOG_ERROR("Failed to create dummy matmul context for memory operations: %d\n", ret); + mem_ctx = 0; + } + } + + ~rknpu_memory_context() { + if (mem_ctx != 0) { + rknn_matmul_destroy(mem_ctx); + } + } + + rknn_matmul_ctx get_ctx() { + std::lock_guard lock(mutex); + return mem_ctx; + } +}; + +static rknpu_memory_context & get_rknpu_memory_context() { + static rknpu_memory_context g_mem_ctx; + return g_mem_ctx; +} + + +// +// Backend +// + +static const char * ggml_backend_rknpu_name(ggml_backend_t backend) { + UNUSED(backend); + return "RKNPU"; +} + +static void ggml_backend_rknpu_free(ggml_backend_t backend) { + ggml_backend_rknpu_context * ctx = (ggml_backend_rknpu_context *)backend->context; + delete ctx; + delete backend; +} + +// Function for getting buffer from cache or creating new one +static std::shared_ptr get_or_create_npu_buffer( + ggml_backend_rknpu_context* backend_ctx, + rknn_matmul_ctx matmul_ctx, + size_t size, + const std::tuple& key, + std::unordered_map, std::shared_ptr, TupleHasher>& cache +) { + std::lock_guard lock(backend_ctx->mutex); + auto it = cache.find(key); + if (it != cache.end()) { + return it->second; + } + + rknn_tensor_mem* mem = rknn_create_mem(matmul_ctx, size); + if (!mem) { return nullptr; } + + auto mem_ctx_for_deleter = get_rknpu_memory_context().get_ctx(); + auto deleter = [mem_ctx_for_deleter](rknn_tensor_mem* m) { + if (m && mem_ctx_for_deleter != 0) { + rknn_destroy_mem(mem_ctx_for_deleter, m); + } + }; + + std::shared_ptr mem_shared(mem, deleter); + cache[key] = mem_shared; + return mem_shared; +} + +static enum ggml_status ggml_backend_rknpu_graph_compute(ggml_backend_t backend, struct ggml_cgraph* cgraph) { + auto* backend_ctx = (ggml_backend_rknpu_context*)backend->context; + const int NUM_CORES = 3; + + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor* node = cgraph->nodes[i]; + if (node->op != GGML_OP_MUL_MAT) continue; + + #if RKNPU_PERF_LOG + std::string node_name = std::string(node->name) + " (" + std::to_string(i) + ")"; + RknpuPerfLogger perf_total(node_name + " - TOTAL"); + #endif + + const struct ggml_tensor* src0 = node->src[0]; // Weights : (K x N) + const struct ggml_tensor* src1 = node->src[1]; // Activations : (M x K) + struct ggml_tensor* dst = node; + + const int M = (int)src1->ne[1]; + const int K = (int)src0->ne[0]; + const int N = (int)src0->ne[1]; + + auto segments = compute_matrix_segments(N, NUM_CORES); + + std::vector> matmul_ctxs(NUM_CORES); + std::vector> mem_B_segments(NUM_CORES); + std::shared_ptr mem_A_shared; + std::vector> mem_C_segments(NUM_CORES); + + // ===== 1. Preraring contexts ===== + { + #if RKNPU_PERF_LOG + RknpuPerfLogger perf_ctx(node_name + " - Context preparation"); + #endif + + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + int N_segment = segments[core_id].size_n; + if (N_segment == 0) continue; + matmul_ctxs[core_id] = backend_ctx->get_matmul_ctx(M, K, N_segment, core_id); + + if (!matmul_ctxs[core_id] || matmul_ctxs[core_id]->ctx == 0) { + RKNPU_LOG_ERROR("Failed to create matmul context for core %d\n", core_id); + return GGML_STATUS_FAILED; + } + } + } + + // ===== 2. Preparing B-matrix ===== + { + #if RKNPU_PERF_LOG + RknpuPerfLogger perf_b(node_name + " - B matrix preparation"); + #endif + + std::vector need_upload(NUM_CORES, false); + std::vector> packed_segments(NUM_CORES); + + // Creating or getting memory from cache (sequentially) + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n == 0) continue; + std::lock_guard lock(backend_ctx->mutex); + auto cache_key = std::make_tuple(src0, core_id); + auto it = backend_ctx->b_npu_buffer_cache.find(cache_key); + + if (it != backend_ctx->b_npu_buffer_cache.end()) { + mem_B_segments[core_id] = it->second; + } else { + need_upload[core_id] = true; + auto& matmul_ctx = matmul_ctxs[core_id]; + rknn_tensor_mem* mem = rknn_create_mem(matmul_ctx->ctx, matmul_ctx->io_attr.B.size); + if (!mem) return GGML_STATUS_FAILED; + + auto mem_ctx_for_deleter = get_rknpu_memory_context().get_ctx(); + auto deleter = [mem_ctx_for_deleter](rknn_tensor_mem* m) { if (m) rknn_destroy_mem(mem_ctx_for_deleter, m); }; + mem_B_segments[core_id] = std::shared_ptr(mem, deleter); + backend_ctx->b_npu_buffer_cache[cache_key] = mem_B_segments[core_id]; + } + } + + // Packing segments into native format (parallel) + #pragma omp parallel for num_threads(NUM_CORES) + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (need_upload[core_id]) { + auto& matmul_ctx = matmul_ctxs[core_id]; + packed_segments[core_id].resize(matmul_ctx->io_attr.B.size / sizeof(uint16_t)); + + pack_B_segment_fp16_native_from_KN( + packed_segments[core_id].data(), + (const uint16_t*)src0->data, + K, N, segments[core_id].offset_n, segments[core_id].size_n); + } + } + + // Setting and syncronizing memory in NPU (sequentially) + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n == 0) continue; + if (need_upload[core_id]) { + memcpy(mem_B_segments[core_id]->virt_addr, packed_segments[core_id].data(), packed_segments[core_id].size() * sizeof(uint16_t)); + RKNN_CHECK(rknn_mem_sync(matmul_ctxs[core_id]->ctx, mem_B_segments[core_id].get(), RKNN_MEMORY_SYNC_TO_DEVICE), "sync B segment"); + } + RKNN_CHECK(rknn_matmul_set_io_mem(matmul_ctxs[core_id]->ctx, mem_B_segments[core_id].get(), &matmul_ctxs[core_id]->io_attr.B), "set_io_mem B segment"); + } + } + + // ===== 3. Preparing A-matrix ===== + { + #if RKNPU_PERF_LOG + RknpuPerfLogger perf_a(node_name + " - A matrix preparation"); + #endif + + auto cache_key = std::make_tuple(M, K, -1); + auto& matmul_ctx_0 = matmul_ctxs[0]; + + mem_A_shared = get_or_create_npu_buffer(backend_ctx, matmul_ctx_0->ctx, matmul_ctx_0->io_attr.A.size, cache_key, backend_ctx->a_buffer_cache); + if (!mem_A_shared) return GGML_STATUS_FAILED; + + const float* x = (const float*)src1->data; + const int row_stride = (int)(src1->nb[1] / sizeof(float)); + uint16_t* dst_base = (uint16_t*)mem_A_shared->virt_addr; + + #pragma omp parallel for + for (int m = 0; m < M; ++m) { + const float* src_row = x + (size_t)m * row_stride; + uint16_t* dst_row = dst_base + (size_t)m * K; + int k = 0; + for (; k <= K - 8; k += 8) { + float32x4_t f32_vec_0 = vld1q_f32(src_row + k); + float32x4_t f32_vec_1 = vld1q_f32(src_row + k + 4); + float16x8_t f16_vec = vcombine_f16(vcvt_f16_f32(f32_vec_0), vcvt_f16_f32(f32_vec_1)); + vst1q_u16(dst_row + k, (uint16x8_t)f16_vec); + } + for (; k < K; ++k) { + dst_row[k] = GGML_FP32_TO_FP16(src_row[k]); + } + } + + RKNN_CHECK(rknn_mem_sync(matmul_ctx_0->ctx, mem_A_shared.get(), RKNN_MEMORY_SYNC_TO_DEVICE), "sync A TO_DEVICE"); + + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n == 0) continue; + RKNN_CHECK(rknn_matmul_set_io_mem(matmul_ctxs[core_id]->ctx, mem_A_shared.get(), &matmul_ctxs[core_id]->io_attr.A), "set_io_mem A for core"); + } + } + + // ===== 4. Preparing C-matrix ===== + { + #if RKNPU_PERF_LOG + RknpuPerfLogger perf_c(node_name + " - C matrix preparation"); + #endif + + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n == 0) continue; + auto& matmul_ctx = matmul_ctxs[core_id]; + auto cache_key = std::make_tuple(M, segments[core_id].size_n, core_id); + mem_C_segments[core_id] = get_or_create_npu_buffer(backend_ctx, matmul_ctx->ctx, matmul_ctx->io_attr.C.size, cache_key, backend_ctx->c_buffer_cache); + if (!mem_C_segments[core_id]) return GGML_STATUS_FAILED; + RKNN_CHECK(rknn_matmul_set_io_mem(matmul_ctx->ctx, mem_C_segments[core_id].get(), &matmul_ctx->io_attr.C), "set_io_mem C"); + } + } + + // ===== 5. Running operation ===== + { + #if RKNPU_PERF_LOG + RknpuPerfLogger perf_run(node_name + " - NPU parallel execution"); + #endif + + #pragma omp parallel for num_threads(NUM_CORES) + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n > 0) { + int ret = rknn_matmul_run(matmul_ctxs[core_id]->ctx); + if (ret != RKNN_SUCC) { + RKNPU_LOG_ERROR("rknn_matmul_run failed for core %d with error %d\n", core_id, ret); + } + } + } + } + + // ===== 6. Collecting results ===== + { + #if RKNPU_PERF_LOG + RknpuPerfLogger perf_gather(node_name + " - Result gathering"); + #endif + + float* dst_data = (float*)dst->data; + + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n == 0) continue; + RKNN_CHECK(rknn_mem_sync(matmul_ctxs[core_id]->ctx, mem_C_segments[core_id].get(), RKNN_MEMORY_SYNC_FROM_DEVICE), "sync C FROM_DEVICE"); + } + + #pragma omp parallel for + for (int m = 0; m < M; m++) { + for (int core_id = 0; core_id < NUM_CORES; core_id++) { + if (segments[core_id].size_n == 0) continue; + + int N_offset = segments[core_id].offset_n; + int N_segment = segments[core_id].size_n; + float* src_segment_base = (float*)mem_C_segments[core_id]->virt_addr; + + memcpy(dst_data + (size_t)m * N + N_offset, + src_segment_base + (size_t)m * N_segment, + N_segment * sizeof(float)); + } + } + } + } + + return GGML_STATUS_SUCCESS; +} + + +// +// Buffer +// + +static void ggml_backend_rknpu_buffer_free_buffer(ggml_backend_buffer_t buffer) { + ggml_backend_rknpu_buffer_context * ctx = (ggml_backend_rknpu_buffer_context *)buffer->context; + rknn_matmul_ctx mem_ctx = get_rknpu_memory_context().get_ctx(); + if (mem_ctx != 0 && ctx->rknn_mem != nullptr) { + RKNN_CHECK(rknn_destroy_mem(mem_ctx, ctx->rknn_mem), "rknn_destroy_mem"); + } + dma_free(ctx->dma_buf); + delete ctx; +} + +static void * ggml_backend_rknpu_buffer_get_base(ggml_backend_buffer_t buffer) { + ggml_backend_rknpu_buffer_context * ctx = (ggml_backend_rknpu_buffer_context *)buffer->context; + return ctx->rknn_mem->virt_addr; +} + +static enum ggml_status ggml_backend_rknpu_buffer_init_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor) { + UNUSED(buffer); + UNUSED(tensor); + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_rknpu_buffer_set_tensor(ggml_backend_buffer_t buffer, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + auto * ctx = (ggml_backend_rknpu_buffer_context *) buffer->context; + + uint8_t * base = (uint8_t *) ctx->rknn_mem->virt_addr; + uint8_t * dst = (uint8_t *) tensor->data + offset; + GGML_ASSERT(dst >= base && dst + size <= base + ctx->rknn_mem->size); + + memcpy(dst, data, size); + + auto mem_ctx = get_rknpu_memory_context().get_ctx(); + RKNN_CHECK(rknn_mem_sync(mem_ctx, ctx->rknn_mem, RKNN_MEMORY_SYNC_TO_DEVICE), "mem_sync write"); +} + +static void ggml_backend_rknpu_buffer_get_tensor(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + ggml_backend_rknpu_buffer_context * ctx = (ggml_backend_rknpu_buffer_context *)buffer->context; + rknn_matmul_ctx mem_ctx = get_rknpu_memory_context().get_ctx(); + + RKNN_CHECK(rknn_mem_sync(mem_ctx, ctx->rknn_mem, RKNN_MEMORY_SYNC_FROM_DEVICE), "rknn_mem_sync from device"); + + memcpy(data, (const char *)tensor->data + offset, size); +} + +static void ggml_backend_rknpu_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + ggml_backend_rknpu_buffer_context * ctx = (ggml_backend_rknpu_buffer_context *)buffer->context; + memset(ctx->rknn_mem->virt_addr, value, ctx->rknn_mem->size); + + rknn_matmul_ctx mem_ctx = get_rknpu_memory_context().get_ctx(); + RKNN_CHECK(rknn_mem_sync(mem_ctx, ctx->rknn_mem, RKNN_MEMORY_SYNC_TO_DEVICE), "rknn_mem_sync to device after clear"); +} + + +// +// Buffer Type +// + +static const char * ggml_backend_rknpu_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + UNUSED(buft); + return "RKNPU"; +} + +static ggml_backend_buffer_t ggml_backend_rknpu_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + UNUSED(buft); + rknn_matmul_ctx mem_ctx = get_rknpu_memory_context().get_ctx(); + if (mem_ctx == 0) { + RKNPU_LOG_ERROR("RKNPU memory context not initialized, cannot allocate buffer.\n"); + return NULL; + } + + DmaBuffer dma_buf = dma_alloc(size); + if (dma_buf.fd < 0) { + RKNPU_LOG_ERROR("dma_alloc failed to allocate %zu bytes\n", size); + return NULL; + } + + rknn_tensor_mem * rknn_mem = rknn_create_mem_from_fd(mem_ctx, dma_buf.fd, dma_buf.virt_addr, size, 0); + if (rknn_mem == NULL) { + RKNPU_LOG_ERROR("rknn_create_mem_from_fd failed for size %zu\n", size); + dma_free(dma_buf); + return NULL; + } + + ggml_backend_rknpu_buffer_context * ctx = new ggml_backend_rknpu_buffer_context{rknn_mem, dma_buf, "rknpu_dma_buffer"}; + + static const ggml_backend_buffer_i rknpu_buffer_interface = { + /* .free_buffer = */ ggml_backend_rknpu_buffer_free_buffer, + /* .get_base = */ ggml_backend_rknpu_buffer_get_base, + /* .init_tensor = */ ggml_backend_rknpu_buffer_init_tensor, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_rknpu_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_rknpu_buffer_get_tensor, + /* .cpy_tensor = */ NULL, + /* .clear = */ ggml_backend_rknpu_buffer_clear, + /* .reset = */ NULL, + }; + + return ggml_backend_buffer_init(buft, rknpu_buffer_interface, ctx, size); +} + +static size_t ggml_backend_rknpu_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + UNUSED(buft); + return 64; +} + +static size_t ggml_backend_rknpu_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * tensor) { + UNUSED(buft); + return ggml_nbytes(tensor); +} + + +// +// Device +// + +static const char * ggml_backend_rknpu_device_get_name(ggml_backend_dev_t dev) { + UNUSED(dev); + return "RKNPU"; +} + +static const char * ggml_backend_rknpu_device_get_description(ggml_backend_dev_t dev) { + UNUSED(dev); + return "Rockchip NPU"; +} + +static void ggml_backend_rknpu_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { + UNUSED(dev); + *free = 0; + *total = 0; +} + +static enum ggml_backend_dev_type ggml_backend_rknpu_device_get_type(ggml_backend_dev_t dev) { + UNUSED(dev); + return GGML_BACKEND_DEVICE_TYPE_ACCEL; +} + +static void ggml_backend_rknpu_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_dev_props * props) { + props->name = ggml_backend_rknpu_device_get_name(dev); + props->description = ggml_backend_rknpu_device_get_description(dev); + props->type = ggml_backend_rknpu_device_get_type(dev); + ggml_backend_rknpu_device_get_memory(dev, &props->memory_free, &props->memory_total); + props->device_id = NULL; + + props->caps.async = false; + props->caps.host_buffer = false; + props->caps.buffer_from_host_ptr = false; + props->caps.events = false; +} + +static bool ggml_backend_rknpu_device_supports_op(ggml_backend_dev_t dev, const struct ggml_tensor * op) { + UNUSED(dev); + RKNPU_LOG_INFO("[RKNPU_SUPPORT] Checking support for op: %s, name: %s\n", ggml_op_name(op->op), op->name); + + switch (op->op) { + case GGML_OP_NONE: + RKNPU_LOG_INFO(" - ACCEPT: Op NONE is supported for RKNPU (leaf/weights).\n"); + return true; + + case GGML_OP_MUL_MAT: { + const struct ggml_tensor * src0 = op->src[0]; // Weights + const struct ggml_tensor * src1 = op->src[1]; // Activations + + RKNPU_LOG_INFO(" - src0: type=%s, dims=[%lld, %lld, %lld, %lld]\n", ggml_type_name(src0->type), src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3]); + RKNPU_LOG_INFO(" - src1: type=%s, dims=[%lld, %lld, %lld, %lld]\n", ggml_type_name(src1->type), src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3]); + + // Checking for weights in F16 and activations in F32 + if (src0->type != GGML_TYPE_F16 || src1->type != GGML_TYPE_F32) { + RKNPU_LOG_INFO(" - REJECT: Unsupported type combination. Want src0=F16, src1=F32.\n"); + return false; + } + + // Checking for K divisible 32, N divisible 16 + if (src0->ne[0] % 32 != 0 || src0->ne[1] % 16 != 0) { + RKNPU_LOG_INFO(" - REJECT: K (%lld) must be a multiple of 32 and N (%lld) a multiple of 16.\n", src0->ne[0], src0->ne[1]); + return false; + } + + // Checking for exact dimentions + if (src1->ne[0] != src0->ne[0]) { + RKNPU_LOG_INFO(" - REJECT: K dimensions do not match (%lld vs %lld).\n", src1->ne[0], src0->ne[0]); + return false; + } + + // Checking contiguous memory + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1)) { + RKNPU_LOG_INFO(" - REJECT: Tensors are not contiguous.\n"); + return false; + } + + RKNPU_LOG_INFO(" - ACCEPT: Op %s can be offloaded to RKNPU.\n", op->name); + return true; + } + default: + RKNPU_LOG_INFO(" - REJECT: Op %s is not supported by RKNPU backend.\n", ggml_op_name(op->op)); + return false; + } +} + +static ggml_backend_t ggml_backend_rknpu_device_init_backend(ggml_backend_dev_t dev, const char * params) { + UNUSED(dev); + UNUSED(params); + + ggml_backend_rknpu_context * ctx = new ggml_backend_rknpu_context{"RKNPU Backend"}; + + static const struct ggml_backend_i rknpu_backend_interface = { + /* .get_name = */ ggml_backend_rknpu_name, + /* .free = */ ggml_backend_rknpu_free, + /* .set_tensor_async = */ NULL, + /* .get_tensor_async = */ NULL, + /* .cpy_tensor_async = */ NULL, + /* .synchronize = */ NULL, + /* .graph_plan_create = */ NULL, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_rknpu_graph_compute, + /* .event_record = */ NULL, + /* .event_wait = */ NULL, + /* .graph_optimize = */ NULL, + }; + + return new ggml_backend{ + /* .guid = */ {0}, + /* .iface = */ rknpu_backend_interface, + /* .device = */ dev, + /* .context = */ ctx, + }; +} + + +// +// Registry +// + +static const char * ggml_backend_rknpu_reg_get_name(ggml_backend_reg_t reg) { + UNUSED(reg); + return "RKNPU"; +} + +static size_t ggml_backend_rknpu_reg_get_device_count(ggml_backend_reg_t reg) { + UNUSED(reg); + if (get_rknpu_memory_context().get_ctx() != 0) { + return 1; + } + return 0; +} + +static ggml_backend_dev_t ggml_backend_rknpu_reg_get_device(ggml_backend_reg_t reg, size_t index) { + if (index != 0) { + return NULL; + } + + static const struct ggml_backend_buffer_type_i rknpu_buffer_type_interface = { + /* .get_name = */ ggml_backend_rknpu_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_rknpu_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_rknpu_buffer_type_get_alignment, + /* .get_max_size = */ NULL, + /* .get_alloc_size = */ ggml_backend_rknpu_buffer_type_get_alloc_size, + /* .is_host = */ NULL, + }; + + static struct ggml_backend_buffer_type rknpu_buffer_type = { + /* .iface = */ rknpu_buffer_type_interface, + /* .device = */ NULL, + /* .context = */ NULL, + }; + + static const struct ggml_backend_device_i rknpu_device_interface = { + /* .get_name = */ ggml_backend_rknpu_device_get_name, + /* .get_description = */ ggml_backend_rknpu_device_get_description, + /* .get_memory = */ ggml_backend_rknpu_device_get_memory, + /* .get_type = */ ggml_backend_rknpu_device_get_type, + /* .get_props = */ ggml_backend_rknpu_device_get_props, + /* .init_backend = */ ggml_backend_rknpu_device_init_backend, + /* .get_buffer_type = */ [](ggml_backend_dev_t dev) { UNUSED(dev); return &rknpu_buffer_type; }, + /* .get_host_buffer_type = */ NULL, + /* .buffer_from_host_ptr = */ NULL, + /* .supports_op = */ ggml_backend_rknpu_device_supports_op, + /* .supports_buft = */ [](ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { UNUSED(dev); return buft == &rknpu_buffer_type; }, + /* .offload_op = */ NULL, + /* .event_new = */ NULL, + /* .event_free = */ NULL, + /* .event_synchronize = */ NULL, + }; + + static struct ggml_backend_device rknpu_device = { + /* .iface = */ rknpu_device_interface, + /* .reg = */ reg, + /* .context = */ NULL, + }; + + if (rknpu_buffer_type.device == NULL) { + rknpu_buffer_type.device = &rknpu_device; + } + + return &rknpu_device; +} + + +// +// Public API +// + +GGML_API ggml_backend_reg_t ggml_backend_rknpu2_reg(void) { + static const struct ggml_backend_reg_i rknpu_reg_interface = { + /* .get_name = */ ggml_backend_rknpu_reg_get_name, + /* .get_device_count = */ ggml_backend_rknpu_reg_get_device_count, + /* .get_device = */ ggml_backend_rknpu_reg_get_device, + /* .get_proc_address = */ NULL, + }; + + static struct ggml_backend_reg rknpu_backend_reg = { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ rknpu_reg_interface, + /* .context = */ NULL, + }; + + return &rknpu_backend_reg; +} + +#ifdef GGML_BACKEND_DL +GGML_BACKEND_DL_IMPL(ggml_backend_rknpu2_reg) +#endif \ No newline at end of file diff --git a/ggml/src/ggml-rknpu2/ggml-rknpu2.h b/ggml/src/ggml-rknpu2/ggml-rknpu2.h new file mode 100644 index 000000000..24de1e015 --- /dev/null +++ b/ggml/src/ggml-rknpu2/ggml-rknpu2.h @@ -0,0 +1,13 @@ +#pragma once + +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + +GGML_API ggml_backend_reg_t ggml_backend_rknpu2_reg(void); + +#ifdef __cplusplus +} +#endif \ No newline at end of file diff --git a/ggml/src/ggml-rknpu2/libs/include/rknn_api.h b/ggml/src/ggml-rknpu2/libs/include/rknn_api.h new file mode 100644 index 000000000..6b294d2fc --- /dev/null +++ b/ggml/src/ggml-rknpu2/libs/include/rknn_api.h @@ -0,0 +1,804 @@ +/**************************************************************************** +* +* Copyright (c) 2017 - 2022 by Rockchip Corp. All rights reserved. +* +* The material in this file is confidential and contains trade secrets +* of Rockchip Corporation. This is proprietary information owned by +* Rockchip Corporation. No part of this work may be disclosed, +* reproduced, copied, transmitted, or used in any way for any purpose, +* without the express written permission of Rockchip Corporation. +* +*****************************************************************************/ + + +#ifndef _RKNN_API_H +#define _RKNN_API_H + +#ifdef __cplusplus +extern "C" { +#endif + +#include + +/* + Definition of extended flag for rknn_init. +*/ +/* set high priority context. */ +#define RKNN_FLAG_PRIOR_HIGH 0x00000000 + +/* set medium priority context */ +#define RKNN_FLAG_PRIOR_MEDIUM 0x00000001 + +/* set low priority context. */ +#define RKNN_FLAG_PRIOR_LOW 0x00000002 + +/* asynchronous mode. + when enable, rknn_outputs_get will not block for too long because it directly retrieves the result of + the previous frame which can increase the frame rate on single-threaded mode, but at the cost of + rknn_outputs_get not retrieves the result of the current frame. + in multi-threaded mode you do not need to turn this mode on. */ +#define RKNN_FLAG_ASYNC_MASK 0x00000004 + +/* collect performance mode. + when enable, you can get detailed performance reports via rknn_query(ctx, RKNN_QUERY_PERF_DETAIL, ...), + but it will reduce the frame rate. */ +#define RKNN_FLAG_COLLECT_PERF_MASK 0x00000008 + +/* allocate all memory in outside, includes weight/internal/inputs/outputs */ +#define RKNN_FLAG_MEM_ALLOC_OUTSIDE 0x00000010 + +/* weight sharing with the same network structure */ +#define RKNN_FLAG_SHARE_WEIGHT_MEM 0x00000020 + +/* send fence fd from outside */ +#define RKNN_FLAG_FENCE_IN_OUTSIDE 0x00000040 + +/* get fence fd from inside */ +#define RKNN_FLAG_FENCE_OUT_OUTSIDE 0x00000080 + +/* dummy init flag: could only get total_weight_size and total_internal_size by rknn_query*/ +#define RKNN_FLAG_COLLECT_MODEL_INFO_ONLY 0x00000100 + +/* allocate internal memory in outside */ +#define RKNN_FLAG_INTERNAL_ALLOC_OUTSIDE 0x00000200 + +/* set GPU as the preferred execution backend When the operator is not supported by the NPU */ +#define RKNN_FLAG_EXECUTE_FALLBACK_PRIOR_DEVICE_GPU 0x00000400 + +/* enable allocate sram type buffers */ +#define RKNN_FLAG_ENABLE_SRAM 0x00000800 + +/* sram type buffers are shared among different contexts */ +#define RKNN_FLAG_SHARE_SRAM 0x00001000 + +/* default nice -19, this flag can disable default priority */ +#define RKNN_FLAG_DISABLE_PROC_HIGH_PRIORITY 0x00002000 + +/* don't flush input buffer cache, the user must ensure that the input tensor has flushed the cache before calling rknn_run. +!!! Don't use this flags when you call rknn_inputs_set() to set input data. */ +#define RKNN_FLAG_DISABLE_FLUSH_INPUT_MEM_CACHE 0x00004000 + +/* Don't invalid output buffer cache. + Users cannot directly access output_mem->virt_addr, + which will cause cache consistency problems. + If you want to use output_mem->virt_addr, + you must use rknn_mem_sync (ctx, mem, RKNN_MEMORY_SYNC_FROM_DEVICE) to flush the cache. + This flags is generally used when the output data of the NPU is not accessed by the CPU, + but is accessed by the GPU or RGA to reduce the time required to flush the cache. + !!! Don't use this flags when you call rknn_outputs_get() to get output data.*/ +#define RKNN_FLAG_DISABLE_FLUSH_OUTPUT_MEM_CACHE 0x00008000 + +/* This flag is used when the model data buffer is allocated by NPU, and can be accessed by NPU directly. */ +#define RKNN_FLAG_MODEL_BUFFER_ZERO_COPY 0x00010000 + +/* This flag is a memory allocation flag, which is used in rknn_create_mem2() when no context is available. */ +#define RKNN_MEM_FLAG_ALLOC_NO_CONTEXT 0x00020000 + + +/* + Error code returned by the RKNN API. +*/ +#define RKNN_SUCC 0 /* execute succeed. */ +#define RKNN_ERR_FAIL -1 /* execute failed. */ +#define RKNN_ERR_TIMEOUT -2 /* execute timeout. */ +#define RKNN_ERR_DEVICE_UNAVAILABLE -3 /* device is unavailable. */ +#define RKNN_ERR_MALLOC_FAIL -4 /* memory malloc fail. */ +#define RKNN_ERR_PARAM_INVALID -5 /* parameter is invalid. */ +#define RKNN_ERR_MODEL_INVALID -6 /* model is invalid. */ +#define RKNN_ERR_CTX_INVALID -7 /* context is invalid. */ +#define RKNN_ERR_INPUT_INVALID -8 /* input is invalid. */ +#define RKNN_ERR_OUTPUT_INVALID -9 /* output is invalid. */ +#define RKNN_ERR_DEVICE_UNMATCH -10 /* the device is unmatch, please update rknn sdk + and npu driver/firmware. */ +#define RKNN_ERR_INCOMPATILE_PRE_COMPILE_MODEL -11 /* This RKNN model use pre_compile mode, but not compatible with current driver. */ +#define RKNN_ERR_INCOMPATILE_OPTIMIZATION_LEVEL_VERSION -12 /* This RKNN model set optimization level, but not compatible with current driver. */ +#define RKNN_ERR_TARGET_PLATFORM_UNMATCH -13 /* This RKNN model set target platform, but not compatible with current platform. */ + +/* + Definition for tensor +*/ +#define RKNN_MAX_DIMS 16 /* maximum dimension of tensor. */ +#define RKNN_MAX_NUM_CHANNEL 15 /* maximum channel number of input tensor. */ +#define RKNN_MAX_NAME_LEN 256 /* maximum name lenth of tensor. */ +#define RKNN_MAX_DYNAMIC_SHAPE_NUM 512 /* maximum number of dynamic shape for each input. */ + +#ifdef __arm__ +typedef uint32_t rknn_context; +#else +typedef uint64_t rknn_context; +#endif + + +/* + The query command for rknn_query +*/ +typedef enum _rknn_query_cmd { + RKNN_QUERY_IN_OUT_NUM = 0, /* query the number of input & output tensor. */ + RKNN_QUERY_INPUT_ATTR = 1, /* query the attribute of input tensor. */ + RKNN_QUERY_OUTPUT_ATTR = 2, /* query the attribute of output tensor. */ + RKNN_QUERY_PERF_DETAIL = 3, /* query the detail performance, need set + RKNN_FLAG_COLLECT_PERF_MASK when call rknn_init, + this query needs to be valid after rknn_outputs_get. */ + RKNN_QUERY_PERF_RUN = 4, /* query the time of run, + this query needs to be valid after rknn_outputs_get. */ + RKNN_QUERY_SDK_VERSION = 5, /* query the sdk & driver version */ + + RKNN_QUERY_MEM_SIZE = 6, /* query the weight & internal memory size */ + RKNN_QUERY_CUSTOM_STRING = 7, /* query the custom string */ + + RKNN_QUERY_NATIVE_INPUT_ATTR = 8, /* query the attribute of native input tensor. */ + RKNN_QUERY_NATIVE_OUTPUT_ATTR = 9, /* query the attribute of native output tensor. */ + + RKNN_QUERY_NATIVE_NC1HWC2_INPUT_ATTR = 8, /* query the attribute of native input tensor. */ + RKNN_QUERY_NATIVE_NC1HWC2_OUTPUT_ATTR = 9, /* query the attribute of native output tensor. */ + + RKNN_QUERY_NATIVE_NHWC_INPUT_ATTR = 10, /* query the attribute of native input tensor. */ + RKNN_QUERY_NATIVE_NHWC_OUTPUT_ATTR = 11, /* query the attribute of native output tensor. */ + + RKNN_QUERY_DEVICE_MEM_INFO = 12, /* query the attribute of rknn memory information. */ + + RKNN_QUERY_INPUT_DYNAMIC_RANGE = 13, /* query the dynamic shape range of rknn input tensor. */ + RKNN_QUERY_CURRENT_INPUT_ATTR = 14, /* query the current shape of rknn input tensor, only valid for dynamic rknn model*/ + RKNN_QUERY_CURRENT_OUTPUT_ATTR = 15, /* query the current shape of rknn output tensor, only valid for dynamic rknn model*/ + + RKNN_QUERY_CURRENT_NATIVE_INPUT_ATTR = 16, /* query the current native shape of rknn input tensor, only valid for dynamic rknn model*/ + RKNN_QUERY_CURRENT_NATIVE_OUTPUT_ATTR = 17, /* query the current native shape of rknn output tensor, only valid for dynamic rknn model*/ + + + RKNN_QUERY_CMD_MAX +} rknn_query_cmd; + +/* + the tensor data type. +*/ +typedef enum _rknn_tensor_type { + RKNN_TENSOR_FLOAT32 = 0, /* data type is float32. */ + RKNN_TENSOR_FLOAT16, /* data type is float16. */ + RKNN_TENSOR_INT8, /* data type is int8. */ + RKNN_TENSOR_UINT8, /* data type is uint8. */ + RKNN_TENSOR_INT16, /* data type is int16. */ + RKNN_TENSOR_UINT16, /* data type is uint16. */ + RKNN_TENSOR_INT32, /* data type is int32. */ + RKNN_TENSOR_UINT32, /* data type is uint32. */ + RKNN_TENSOR_INT64, /* data type is int64. */ + RKNN_TENSOR_BOOL, + RKNN_TENSOR_INT4, + RKNN_TENSOR_BFLOAT16, + + RKNN_TENSOR_TYPE_MAX +} rknn_tensor_type; + +inline static const char* get_type_string(rknn_tensor_type type) +{ + switch(type) { + case RKNN_TENSOR_FLOAT32: return "FP32"; + case RKNN_TENSOR_FLOAT16: return "FP16"; + case RKNN_TENSOR_INT8: return "INT8"; + case RKNN_TENSOR_UINT8: return "UINT8"; + case RKNN_TENSOR_INT16: return "INT16"; + case RKNN_TENSOR_UINT16: return "UINT16"; + case RKNN_TENSOR_INT32: return "INT32"; + case RKNN_TENSOR_UINT32: return "UINT32"; + case RKNN_TENSOR_INT64: return "INT64"; + case RKNN_TENSOR_BOOL: return "BOOL"; + case RKNN_TENSOR_INT4: return "INT4"; + case RKNN_TENSOR_BFLOAT16: return "BF16"; + default: return "UNKNOW"; + } +} + +/* + the quantitative type. +*/ +typedef enum _rknn_tensor_qnt_type { + RKNN_TENSOR_QNT_NONE = 0, /* none. */ + RKNN_TENSOR_QNT_DFP, /* dynamic fixed point. */ + RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC, /* asymmetric affine. */ + + RKNN_TENSOR_QNT_MAX +} rknn_tensor_qnt_type; + +inline static const char* get_qnt_type_string(rknn_tensor_qnt_type type) +{ + switch(type) { + case RKNN_TENSOR_QNT_NONE: return "NONE"; + case RKNN_TENSOR_QNT_DFP: return "DFP"; + case RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC: return "AFFINE"; + default: return "UNKNOW"; + } +} + +/* + the tensor data format. +*/ +typedef enum _rknn_tensor_format { + RKNN_TENSOR_NCHW = 0, /* data format is NCHW. */ + RKNN_TENSOR_NHWC, /* data format is NHWC. */ + RKNN_TENSOR_NC1HWC2, /* data format is NC1HWC2. */ + RKNN_TENSOR_UNDEFINED, + + RKNN_TENSOR_FORMAT_MAX +} rknn_tensor_format; + +/* + the mode of running on target NPU core. +*/ +typedef enum _rknn_core_mask { + RKNN_NPU_CORE_AUTO = 0, /* default, run on NPU core randomly. */ + RKNN_NPU_CORE_0 = 1, /* run on NPU core 0. */ + RKNN_NPU_CORE_1 = 2, /* run on NPU core 1. */ + RKNN_NPU_CORE_2 = 4, /* run on NPU core 2. */ + RKNN_NPU_CORE_0_1 = RKNN_NPU_CORE_0 | RKNN_NPU_CORE_1, /* run on NPU core 0 and core 1. */ + RKNN_NPU_CORE_0_1_2 = RKNN_NPU_CORE_0_1 | RKNN_NPU_CORE_2, /* run on NPU core 0 and core 1 and core 2. */ + RKNN_NPU_CORE_ALL = 0xffff, /* auto choice, run on NPU cores depending on platform */ + + RKNN_NPU_CORE_UNDEFINED, +} rknn_core_mask; + +inline static const char* get_format_string(rknn_tensor_format fmt) +{ + switch(fmt) { + case RKNN_TENSOR_NCHW: return "NCHW"; + case RKNN_TENSOR_NHWC: return "NHWC"; + case RKNN_TENSOR_NC1HWC2: return "NC1HWC2"; + case RKNN_TENSOR_UNDEFINED: return "UNDEFINED"; + default: return "UNKNOW"; + } +} + +/* + the information for RKNN_QUERY_IN_OUT_NUM. +*/ +typedef struct _rknn_input_output_num { + uint32_t n_input; /* the number of input. */ + uint32_t n_output; /* the number of output. */ +} rknn_input_output_num; + +/* + the information for RKNN_QUERY_INPUT_ATTR / RKNN_QUERY_OUTPUT_ATTR. +*/ +typedef struct _rknn_tensor_attr { + uint32_t index; /* input parameter, the index of input/output tensor, + need set before call rknn_query. */ + + uint32_t n_dims; /* the number of dimensions. */ + uint32_t dims[RKNN_MAX_DIMS]; /* the dimensions array. */ + char name[RKNN_MAX_NAME_LEN]; /* the name of tensor. */ + + uint32_t n_elems; /* the number of elements. */ + uint32_t size; /* the bytes size of tensor. */ + + rknn_tensor_format fmt; /* the data format of tensor. */ + rknn_tensor_type type; /* the data type of tensor. */ + rknn_tensor_qnt_type qnt_type; /* the quantitative type of tensor. */ + int8_t fl; /* fractional length for RKNN_TENSOR_QNT_DFP. */ + int32_t zp; /* zero point for RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC. */ + float scale; /* scale for RKNN_TENSOR_QNT_AFFINE_ASYMMETRIC. */ + + uint32_t w_stride; /* the stride of tensor along the width dimention of input, + Note: it is read-only, 0 means equal to width. */ + uint32_t size_with_stride; /* the bytes size of tensor with stride. */ + + uint8_t pass_through; /* pass through mode, for rknn_set_io_mem interface. + if TRUE, the buf data is passed directly to the input node of the rknn model + without any conversion. the following variables do not need to be set. + if FALSE, the buf data is converted into an input consistent with the model + according to the following type and fmt. so the following variables + need to be set.*/ + uint32_t h_stride; /* the stride along the height dimention of input, + Note: it is write-only, if it was set to 0, h_stride = height. */ +} rknn_tensor_attr; + +typedef struct _rknn_input_range { + uint32_t index; /* input parameter, the index of input/output tensor, + need set before call rknn_query. */ + uint32_t shape_number; /* the number of shape. */ + rknn_tensor_format fmt; /* the data format of tensor. */ + char name[RKNN_MAX_NAME_LEN]; /* the name of tensor. */ + uint32_t dyn_range[RKNN_MAX_DYNAMIC_SHAPE_NUM][RKNN_MAX_DIMS]; /* the dynamic input dimensions range. */ + uint32_t n_dims; /* the number of dimensions. */ + +} rknn_input_range; + +/* + the information for RKNN_QUERY_PERF_DETAIL. +*/ +typedef struct _rknn_perf_detail { + char* perf_data; /* the string pointer of perf detail. don't need free it by user. */ + uint64_t data_len; /* the string length. */ +} rknn_perf_detail; + +/* + the information for RKNN_QUERY_PERF_RUN. +*/ +typedef struct _rknn_perf_run { + int64_t run_duration; /* real inference time (us) */ +} rknn_perf_run; + +/* + the information for RKNN_QUERY_SDK_VERSION. +*/ +typedef struct _rknn_sdk_version { + char api_version[256]; /* the version of rknn api. */ + char drv_version[256]; /* the version of rknn driver. */ +} rknn_sdk_version; + +/* + the information for RKNN_QUERY_MEM_SIZE. +*/ +typedef struct _rknn_mem_size { + uint32_t total_weight_size; /* the weight memory size */ + uint32_t total_internal_size; /* the internal memory size, exclude inputs/outputs */ + uint64_t total_dma_allocated_size; /* total dma memory allocated size */ + uint32_t total_sram_size; /* total system sram size reserved for rknn */ + uint32_t free_sram_size; /* free system sram size reserved for rknn */ + uint32_t reserved[10]; /* reserved */ +} rknn_mem_size; + +/* + the information for RKNN_QUERY_CUSTOM_STRING. +*/ +typedef struct _rknn_custom_string { + char string[1024]; /* the string of custom, lengths max to 1024 bytes */ +} rknn_custom_string; + +/* + The flags of rknn_tensor_mem. +*/ +typedef enum _rknn_tensor_mem_flags { + RKNN_TENSOR_MEMORY_FLAGS_ALLOC_INSIDE = 1, /*Used to mark in rknn_destroy_mem() whether it is necessary to release the "mem" pointer itself. + If the flag RKNN_TENSOR_MEMORY_FLAGS_ALLOC_INSIDE is set, rknn_destroy_mem() will call free(mem).*/ + RKNN_TENSOR_MEMORY_FLAGS_FROM_FD = 2, /*Used to mark in rknn_create_mem_from_fd() whether it is necessary to release the "mem" pointer itself. + If the flag RKNN_TENSOR_MEMORY_FLAGS_FROM_FD is set, rknn_destroy_mem() will call free(mem).*/ + RKNN_TENSOR_MEMORY_FLAGS_FROM_PHYS = 3, /*Used to mark in rknn_create_mem_from_phys() whether it is necessary to release the "mem" pointer itself. + If the flag RKNN_TENSOR_MEMORY_FLAGS_FROM_PHYS is set, rknn_destroy_mem() will call free(mem).*/ + RKNN_TENSOR_MEMORY_FLAGS_UNKNOWN +} rknn_tensor_mem_flags; + +/* + The mode to allocate rknn memory. +*/ +typedef enum _rknn_mem_alloc_flags { + RKNN_FLAG_MEMORY_FLAGS_DEFAULT = 0 << 0, /* Same with RKNN_FLAG_MEMORY_CACHEABLE */ + RKNN_FLAG_MEMORY_CACHEABLE = 1 << 0, /* Create Cacheable memory. */ + RKNN_FLAG_MEMORY_NON_CACHEABLE = 1 << 1, /* Create NON-Cacheable memory. */ + RKNN_FLAG_MEMORY_TRY_ALLOC_SRAM = 1 << 2, /* Try to allocate memory in SRAM if possible. if SRAM is not enough, allocate rest memory in DRAM. */ +} rknn_mem_alloc_flags; + +/* + The mode to sync cacheable rknn memory. +*/ +typedef enum _rknn_mem_sync_mode { + RKNN_MEMORY_SYNC_TO_DEVICE = 0x1, /* the mode used for consistency of device access after CPU accesses data. */ + RKNN_MEMORY_SYNC_FROM_DEVICE = 0x2, /* the mode used for consistency of CPU access after device accesses data. */ + RKNN_MEMORY_SYNC_BIDIRECTIONAL = RKNN_MEMORY_SYNC_TO_DEVICE | RKNN_MEMORY_SYNC_FROM_DEVICE, /* the mode used for consistency of data access + between device and CPU in both directions. */ +} rknn_mem_sync_mode; + +/* + the memory information of tensor. +*/ +typedef struct _rknn_tensor_memory { + void* virt_addr; /* the virtual address of tensor buffer. */ + uint64_t phys_addr; /* the physical address of tensor buffer. */ + int32_t fd; /* the fd of tensor buffer. */ + int32_t offset; /* indicates the offset of the memory. */ + uint32_t size; /* the size of tensor buffer. */ + uint32_t flags; /* the flags of tensor buffer, reserved */ + void * priv_data; /* the private data of tensor buffer. */ +} rknn_tensor_mem; + +/* + the input information for rknn_input_set. +*/ +typedef struct _rknn_input { + uint32_t index; /* the input index. */ + void* buf; /* the input buf for index. */ + uint32_t size; /* the size of input buf. */ + uint8_t pass_through; /* pass through mode. + if TRUE, the buf data is passed directly to the input node of the rknn model + without any conversion. the following variables do not need to be set. + if FALSE, the buf data is converted into an input consistent with the model + according to the following type and fmt. so the following variables + need to be set.*/ + rknn_tensor_type type; /* the data type of input buf. */ + rknn_tensor_format fmt; /* the data format of input buf. + currently the internal input format of NPU is NCHW by default. + so entering NCHW data can avoid the format conversion in the driver. */ +} rknn_input; + +/* + the output information for rknn_outputs_get. +*/ +typedef struct _rknn_output { + uint8_t want_float; /* want transfer output data to float */ + uint8_t is_prealloc; /* whether buf is pre-allocated. + if TRUE, the following variables need to be set. + if FALSE, the following variables do not need to be set. */ + uint32_t index; /* the output index. */ + void* buf; /* the output buf for index. + when is_prealloc = FALSE and rknn_outputs_release called, + this buf pointer will be free and don't use it anymore. */ + uint32_t size; /* the size of output buf. */ +} rknn_output; + +/* + the extend information for rknn_init. +*/ +typedef struct _rknn_init_extend { + rknn_context ctx; /* rknn context */ + int32_t real_model_offset; /* real rknn model file offset, only valid when init context with rknn file path and zero-copy model model */ + uint32_t real_model_size; /* real rknn model file size, only valid when init context with rknn file path and zero-copy model model */ + int32_t model_buffer_fd; /* the fd of model buffer. */ + uint32_t model_buffer_flags; /* the flags of model_buffer */ + uint8_t reserved[112]; /* reserved */ +} rknn_init_extend; + +/* + the extend information for rknn_run. +*/ +typedef struct _rknn_run_extend { + uint64_t frame_id; /* output parameter, indicate current frame id of run. */ + int32_t non_block; /* block flag of run, 0 is block else 1 is non block */ + int32_t timeout_ms; /* timeout for block mode, in milliseconds */ + int32_t fence_fd; /* fence fd from other unit */ +} rknn_run_extend; + +/* + the extend information for rknn_outputs_get. +*/ +typedef struct _rknn_output_extend { + uint64_t frame_id; /* output parameter, indicate the frame id of outputs, corresponds to + struct rknn_run_extend.frame_id.*/ +} rknn_output_extend; + + +/* rknn_init + + initial the context and load the rknn model. + + input: + rknn_context* context the pointer of context handle. + void* model if size > 0, pointer to the rknn model, if size = 0, filepath to the rknn model. + uint32_t size the size of rknn model. + uint32_t flag extend flag, see the define of RKNN_FLAG_XXX_XXX. + rknn_init_extend* extend the extend information of init. + return: + int error code. +*/ +int rknn_init(rknn_context* context, void* model, uint32_t size, uint32_t flag, rknn_init_extend* extend); + +/* rknn_dup_context + + initial the context and load the rknn model. + + input: + rknn_context* context_in the pointer of context in handle. + rknn_context* context_out the pointer of context out handle. + return: + int error code. +*/ +int rknn_dup_context(rknn_context* context_in, rknn_context* context_out); + +/* rknn_destroy + + unload the rknn model and destroy the context. + + input: + rknn_context context the handle of context. + return: + int error code. +*/ +int rknn_destroy(rknn_context context); + + +/* rknn_query + + query the information about model or others. see rknn_query_cmd. + + input: + rknn_context context the handle of context. + rknn_query_cmd cmd the command of query. + void* info the buffer point of information. + uint32_t size the size of information. + return: + int error code. +*/ +int rknn_query(rknn_context context, rknn_query_cmd cmd, void* info, uint32_t size); + + +/* rknn_inputs_set + + set inputs information by input index of rknn model. + inputs information see rknn_input. + + input: + rknn_context context the handle of context. + uint32_t n_inputs the number of inputs. + rknn_input inputs[] the arrays of inputs information, see rknn_input. + return: + int error code +*/ +int rknn_inputs_set(rknn_context context, uint32_t n_inputs, rknn_input inputs[]); + +/* + rknn_set_batch_core_num + + set rknn batch core_num. + + input: + rknn_context context the handle of context. + int core_num the core number. + return: + int error code. + +*/ +int rknn_set_batch_core_num(rknn_context context, int core_num); + +/* rknn_set_core_mask + + set the core mask for the model.(only supported on multi-core NPU platform) + + RKNN_NPU_CORE_AUTO: auto mode, default value + RKNN_NPU_CORE_0: core 0 mode + RKNN_NPU_CORE_1: core 1 mode + RKNN_NPU_CORE_2: core 2 mode + RKNN_NPU_CORE_0_1: combine core 0/1 mode + RKNN_NPU_CORE_0_1_2: combine core 0/1/2 mode + RKNN_NPU_CORE_ALL: auto mode, select multiple npu cores to run depending on platform + + + input: + rknn_context context the handle of context. + rknn_core_mask core_mask the core mask. + return: + int error code. +*/ +int rknn_set_core_mask(rknn_context context, rknn_core_mask core_mask); + +/* rknn_run + + run the model to execute inference. + + input: + rknn_context context the handle of context. + rknn_run_extend* extend the extend information of run. + return: + int error code. +*/ +int rknn_run(rknn_context context, rknn_run_extend* extend); + + +/* rknn_wait + + wait the model after execute inference. + + input: + rknn_context context the handle of context. + rknn_run_extend* extend the extend information of run. + return: + int error code. +*/ +int rknn_wait(rknn_context context, rknn_run_extend* extend); + + +/* rknn_outputs_get + + wait the inference to finish and get the outputs. + this function will block until inference finish. + the results will set to outputs[]. + + input: + rknn_context context the handle of context. + uint32_t n_outputs the number of outputs. + rknn_output outputs[] the arrays of output, see rknn_output. + rknn_output_extend* the extend information of output. + return: + int error code. +*/ +int rknn_outputs_get(rknn_context context, uint32_t n_outputs, rknn_output outputs[], rknn_output_extend* extend); + + +/* rknn_outputs_release + + release the outputs that get by rknn_outputs_get. + after called, the rknn_output[x].buf get from rknn_outputs_get will + also be free when rknn_output[x].is_prealloc = FALSE. + + input: + rknn_context context the handle of context. + uint32_t n_ouputs the number of outputs. + rknn_output outputs[] the arrays of output. + return: + int error code +*/ +int rknn_outputs_release(rknn_context context, uint32_t n_ouputs, rknn_output outputs[]); + + +/* new api for zero copy */ + +/* rknn_create_mem_from_phys (memory allocated outside) + + initialize tensor memory from physical address. + + input: + rknn_context ctx the handle of context. + uint64_t phys_addr physical address. + void *virt_addr virtual address. + uint32_t size the size of tensor buffer. + return: + rknn_tensor_mem the pointer of tensor memory information. +*/ +rknn_tensor_mem* rknn_create_mem_from_phys(rknn_context ctx, uint64_t phys_addr, void *virt_addr, uint32_t size); + + +/* rknn_create_mem_from_fd (memory allocated outside) + + initialize tensor memory from file description. + + input: + rknn_context ctx the handle of context. + int32_t fd file description. + void *virt_addr virtual address. + uint32_t size the size of tensor buffer. + int32_t offset indicates the offset of the memory (virt_addr without offset). + return: + rknn_tensor_mem the pointer of tensor memory information. +*/ +rknn_tensor_mem* rknn_create_mem_from_fd(rknn_context ctx, int32_t fd, void *virt_addr, uint32_t size, int32_t offset); + + +/* rknn_create_mem_from_mb_blk (memory allocated outside) + + create tensor memory from mb_blk. + + input: + rknn_context ctx the handle of context. + void *mb_blk mb_blk allocate from system api. + int32_t offset indicates the offset of the memory. + return: + rknn_tensor_mem the pointer of tensor memory information. +*/ +rknn_tensor_mem* rknn_create_mem_from_mb_blk(rknn_context ctx, void *mb_blk, int32_t offset); + + +/* rknn_create_mem (memory allocated inside) + + create tensor memory. + + input: + rknn_context ctx the handle of context. + uint32_t size the size of tensor buffer. + return: + rknn_tensor_mem the pointer of tensor memory information. +*/ +rknn_tensor_mem* rknn_create_mem(rknn_context ctx, uint32_t size); + +/* rknn_create_mem2 (memory allocated inside) + + create tensor memory. + + input: + rknn_context ctx the handle of context. + uint64_t size the size of tensor buffer. + uint64_t alloc_flags memory allocation flags. + return: + rknn_tensor_mem the pointer of tensor memory information. +*/ +rknn_tensor_mem* rknn_create_mem2(rknn_context ctx, uint64_t size, uint64_t alloc_flags); + +/* rknn_destroy_mem (support allocate inside and outside) + + destroy tensor memory. + + input: + rknn_context ctx the handle of context. + rknn_tensor_mem *mem the pointer of tensor memory information. + return: + int error code +*/ +int rknn_destroy_mem(rknn_context ctx, rknn_tensor_mem *mem); + + +/* rknn_set_weight_mem + + set the weight memory. + + input: + rknn_context ctx the handle of context. + rknn_tensor_mem *mem the array of tensor memory information + return: + int error code. +*/ +int rknn_set_weight_mem(rknn_context ctx, rknn_tensor_mem *mem); + + +/* rknn_set_internal_mem + + set the internal memory. + + input: + rknn_context ctx the handle of context. + rknn_tensor_mem *mem the array of tensor memory information + return: + int error code. +*/ +int rknn_set_internal_mem(rknn_context ctx, rknn_tensor_mem *mem); + + +/* rknn_set_io_mem + + set the input and output tensors buffer. + + input: + rknn_context ctx the handle of context. + rknn_tensor_mem *mem the array of tensor memory information. + rknn_tensor_attr *attr the attribute of input or output tensor buffer. + return: + int error code. +*/ +int rknn_set_io_mem(rknn_context ctx, rknn_tensor_mem *mem, rknn_tensor_attr *attr); + +/* rknn_set_input_shape(deprecated) + + set the input tensor shape (only valid for dynamic shape rknn model). + + input: + rknn_context ctx the handle of context. + rknn_tensor_attr *attr the attribute of input or output tensor buffer. + return: + int error code. +*/ +int rknn_set_input_shape(rknn_context ctx, rknn_tensor_attr* attr); + +/* rknn_set_input_shapes + + set all the input tensor shapes. graph will run under current set of input shapes after rknn_set_input_shapes.(only valid for dynamic shape rknn model). + + input: + rknn_context ctx the handle of context. + uint32_t n_inputs the number of inputs. + rknn_tensor_attr attr[] the attribute array of all input tensors. + return: + int error code. +*/ +int rknn_set_input_shapes(rknn_context ctx, uint32_t n_inputs, rknn_tensor_attr attr[]); + +/* rknn_mem_sync + + sync cacheable rknn memory when both cpu and device access data. + + input: + rknn_context context the handle of context. + rknn_tensor_mem *mem the pointer of tensor memory information. + rknn_mem_sync_mode mode the mode of sync cache. + return: + int error code. +*/ +int rknn_mem_sync(rknn_context context, rknn_tensor_mem* mem, rknn_mem_sync_mode mode); + +#ifdef __cplusplus +} //extern "C" +#endif + +#endif //_RKNN_API_H diff --git a/ggml/src/ggml-rknpu2/libs/include/rknn_custom_op.h b/ggml/src/ggml-rknpu2/libs/include/rknn_custom_op.h new file mode 100644 index 000000000..253263f77 --- /dev/null +++ b/ggml/src/ggml-rknpu2/libs/include/rknn_custom_op.h @@ -0,0 +1,144 @@ +/**************************************************************************** + * + * Copyright (c) 2017 - 2023 by Rockchip Corp. All rights reserved. + * + * The material in this file is confidential and contains trade secrets + * of Rockchip Corporation. This is proprietary information owned by + * Rockchip Corporation. No part of this work may be disclosed, + * reproduced, copied, transmitted, or used in any way for any purpose, + * without the express written permission of Rockchip Corporation. + * + *****************************************************************************/ + +#ifndef _RKNN_CUSTOM_OP_H +#define _RKNN_CUSTOM_OP_H + +#ifdef __cplusplus +extern "C" { +#endif + +#include "rknn_api.h" + +#include + +/* + Error code returned by the RKNN Custom Operator API. +*/ +#define RKNN_WARNING_SKIP_CUSTOM_OP_COMPUTE -14 /* if custom op init callback funtion return this code and op type is supported by RKNN, it will use RKNN implementation. */ + +#define RKNN_CUSTOM_OP_MAX_STR_LEN 64 +#define RKNN_CUSTOM_OP_MAX_VALUE_LEN 32 +#define RKNN_CUSTOM_OP_EXPORT __attribute__((visibility("default"))) + +#ifdef __arm__ +typedef uint32_t rknn_custom_op_interal_context; +#else +typedef uint64_t rknn_custom_op_interal_context; +#endif +/* + the backend execution device of custom operator. +*/ +typedef enum _rknn_target_type +{ + RKNN_TARGET_TYPE_CPU = 1, /* backend device is cpu */ + RKNN_TARGET_TYPE_GPU = 2, /* backend device is gpu */ + RKNN_TARGET_TYPE_MAX +} rknn_target_type; + +typedef struct _rknn_gpu_op_context +{ + void* cl_context; + void* cl_command_queue; + void* cl_kernel; + +} rknn_gpu_op_context; + +typedef struct _rknn_custom_op_context +{ + rknn_target_type target; /* custom op backend target */ + rknn_custom_op_interal_context internal_ctx; /* the context of custom op*/ + rknn_gpu_op_context gpu_ctx; /* the gpu context of custom op */ + void* priv_data; /* the private data managed by user */ +} rknn_custom_op_context; + +typedef struct _rknn_custom_op_tensor +{ + rknn_tensor_attr attr; /* the attribute of tensor buffer. */ + rknn_tensor_mem mem; /* the memory information of tensor. */ +} rknn_custom_op_tensor; + +typedef struct _rknn_custom_op_attr +{ + char name[RKNN_MAX_NAME_LEN]; /* the name of operator atrributes. */ + rknn_tensor_type dtype; /* the data type of operator attributes, indicate the 'array' type. */ + uint32_t n_elems; /* the number of 'array'. */ + void* data; /* the array pointer of operator attributes, the data type of each element is determined by type. */ +} rknn_custom_op_attr; + +/* + the information of custom operator to add to the rknn_context. +*/ +typedef struct _rknn_custom_op +{ + uint32_t version; /* custom op version */ + rknn_target_type target; /* custom op backend target */ + char op_type[RKNN_MAX_NAME_LEN]; /* custom op type */ + + char cl_kernel_name[RKNN_MAX_NAME_LEN]; /* the opencl kernel name used by custom op */ + char* cl_kernel_source; /* if cl_source_size > 0, pointer to the cl kernel source string, if cl_source_size = 0, + filepath to the cl kernel file. */ + uint64_t cl_source_size; /* the size of cl_kernel_source */ + char cl_build_options[RKNN_MAX_NAME_LEN]; /* the options for opencl to build clProgram used by custom op */ + + /** + * The callback function sets that the users need to code + */ + int (*init)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs, + rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [optional] custom op kernel init falllback function*/ + int (*prepare)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs, + rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [optional] custom op kernel prepare falllback function*/ + int (*compute)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs, + rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [required] custom op kernel compute falllback function */ + int (*compute_native)(rknn_custom_op_context* op_ctx, rknn_custom_op_tensor* inputs, uint32_t n_inputs, + rknn_custom_op_tensor* outputs, uint32_t n_outputs); /* [optional] custom op kernel compute with native attribute falllback function */ + int (*destroy)(rknn_custom_op_context* op_ctx); /* [optional] custom op kernel compute falllback function */ + +} rknn_custom_op; + +/** + * dlopen custom op with so required this function + */ +typedef rknn_custom_op* (*get_custom_op_func)(); + +/* rknn_register_custom_ops + + Register custom operators to rknn_context. + Steps to use a custom op: + 1. Create a rknn_custom_op structure array and fill in it. + 2. Setup prepare/compute/compute_native/destroy callback function and add them to the + rknn_custom_op.(compute is required and other function is optional, compute_native is not supported now, set it + to nullptr) + 3. Call rknn_register_custom_ops to register the op type after rknn_init. + input: + rknn_context ctx the handle of context. + rknn_custom_op* op the custom operator array, each of which contains operator information and calllback function. + uint32_t custom_op_num the length of rknn_custom_op array. + return: + int error code. +*/ +int rknn_register_custom_ops(rknn_context ctx, rknn_custom_op* op, uint32_t custom_op_num); + +/* rknn_custom_op_get_op_attr + + input: + rknn_custom_op_context* op_ctx the handle of custom op context. + const char* attr_name the attribute name of operator. + rknn_custom_op_attr* op_attr the data and information of operator attributes. +*/ +void rknn_custom_op_get_op_attr(rknn_custom_op_context* op_ctx, const char* attr_name, rknn_custom_op_attr* op_attr); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif //_RKNN_CUSTOM_OP_H diff --git a/ggml/src/ggml-rknpu2/libs/include/rknn_matmul_api.h b/ggml/src/ggml-rknpu2/libs/include/rknn_matmul_api.h new file mode 100644 index 000000000..7a321b32c --- /dev/null +++ b/ggml/src/ggml-rknpu2/libs/include/rknn_matmul_api.h @@ -0,0 +1,544 @@ +/**************************************************************************** + * + * Copyright (c) 2017 - 2018 by Rockchip Corp. All rights reserved. + * + * The material in this file is confidential and contains trade secrets + * of Rockchip Corporation. This is proprietary information owned by + * Rockchip Corporation. No part of this work may be disclosed, + * reproduced, copied, transmitted, or used in any way for any purpose, + * without the express written permission of Rockchip Corporation. + * + *****************************************************************************/ + +#ifndef _RKNN_MATMUL_API_H +#define _RKNN_MATMUL_API_H + +#ifdef __cplusplus +extern "C" { +#endif + +#include "rknn_api.h" + +typedef rknn_context rknn_matmul_ctx; + +typedef enum _rknn_matmul_quant_type +{ + RKNN_QUANT_TYPE_PER_LAYER_SYM = 0, + RKNN_QUANT_TYPE_PER_LAYER_ASYM = 1, + RKNN_QUANT_TYPE_PER_CHANNEL_SYM = 2, + RKNN_QUANT_TYPE_PER_CHANNEL_ASYM = 3, + RKNN_QUANT_TYPE_PER_GROUP_SYM = 4, + RKNN_QUANT_TYPE_PER_GROUP_ASYM = 5, +} rknn_matmul_quant_type; + +typedef struct _rknn_quant_params +{ + char name[RKNN_MAX_NAME_LEN]; + + // matmul tensor scale + float* scale; + int32_t scale_len; + + // matmul tensor zero point + int32_t* zp; + int32_t zp_len; + +} rknn_quant_params; + +typedef enum _rknn_matmul_type +{ + RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT32 = 1, + RKNN_INT8_MM_INT8_TO_INT32 = 2, + RKNN_INT8_MM_INT8_TO_INT8 = 3, + RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT16 = 4, + RKNN_FLOAT16_MM_INT8_TO_FLOAT32 = 5, + RKNN_FLOAT16_MM_INT8_TO_FLOAT16 = 6, + RKNN_FLOAT16_MM_INT4_TO_FLOAT32 = 7, + RKNN_FLOAT16_MM_INT4_TO_FLOAT16 = 8, + RKNN_INT8_MM_INT8_TO_FLOAT32 = 9, + RKNN_INT4_MM_INT4_TO_INT16 = 10, + RKNN_INT8_MM_INT4_TO_INT32 = 11, + RKNN_FLOAT16_MM_INT4_TO_BFLOAT16 = 12, + RKNN_INT8_MM_INT4_TO_FLOAT16 = 15, +} rknn_matmul_type; + +inline static const char* get_matmul_type_string(rknn_matmul_type type) +{ + switch (type) { + case RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT32: + return "RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT32"; + case RKNN_INT8_MM_INT8_TO_INT32: + return "RKNN_INT8_MM_INT8_TO_INT32"; + case RKNN_INT8_MM_INT8_TO_INT8: + return "RKNN_INT8_MM_INT8_TO_INT8"; + case RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT16: + return "RKNN_FLOAT16_MM_FLOAT16_TO_FLOAT16"; + case RKNN_FLOAT16_MM_INT8_TO_FLOAT32: + return "RKNN_FLOAT16_MM_INT8_TO_FLOAT32"; + case RKNN_FLOAT16_MM_INT8_TO_FLOAT16: + return "RKNN_FLOAT16_MM_INT8_TO_FLOAT16"; + case RKNN_INT4_MM_INT4_TO_INT16: + return "RKNN_INT4_MM_INT4_TO_INT16"; + case RKNN_FLOAT16_MM_INT4_TO_FLOAT32: + return "RKNN_FLOAT16_MM_INT4_TO_FLOAT32"; + case RKNN_FLOAT16_MM_INT4_TO_FLOAT16: + return "RKNN_FLOAT16_MM_INT4_TO_FLOAT16"; + case RKNN_INT8_MM_INT4_TO_INT32: + return "RKNN_INT8_MM_INT4_TO_INT32"; + case RKNN_INT8_MM_INT8_TO_FLOAT32: + return "RKNN_INT8_MM_INT8_TO_FLOAT32"; + case RKNN_FLOAT16_MM_INT4_TO_BFLOAT16: + return "RKNN_FLOAT16_MM_INT4_TO_BFLOAT16"; + default: + return "UNKNOW"; + } +} + +typedef struct _rknn_matmul_tensor_attr +{ + char name[RKNN_MAX_NAME_LEN]; + + // indicate A(M, K) or B(K, N) or C(M, N) + uint32_t n_dims; + uint32_t dims[RKNN_MAX_DIMS]; + + // matmul tensor size + uint32_t size; + + // matmul tensor data type + // int8 : A, B + // int32: C + rknn_tensor_type type; + +} rknn_matmul_tensor_attr; + +typedef struct _rknn_matmul_io_attr +{ + // indicate A(M, K) or B(K, N) or C(M, N) + rknn_matmul_tensor_attr A; + rknn_matmul_tensor_attr B; + rknn_matmul_tensor_attr C; +} rknn_matmul_io_attr; + +/* + matmul dynamic shape struct +*/ +typedef struct _rknn_matmul_shape +{ + int32_t M; + int32_t K; + int32_t N; +} rknn_matmul_shape; + +/* + the layout of matmul input/output tensor. +*/ +typedef enum +{ + RKNN_MM_LAYOUT_NORM = 0, + RKNN_MM_LAYOUT_NATIVE = 1, + RKNN_MM_LAYOUT_TP_NORM = 2, +} rknn_matmul_layout; + +/* + matmul information struct + */ +typedef struct rknn_matmul_info_t +{ + int32_t M; + int32_t K; // limit: RK3566/3568: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte; + // RK3562: int8 type must be aligned with 32byte, float16 type must be aligned with 32byte; + // RK3588/3576: int8 type must be aligned with 32byte, float16 type must be aligned with 32byte, + // int4 type must be aligned with 32byte; + int32_t N; // limit: RK3566/3568: int8 type must be aligned with 16byte, float16 type must be aligned with 8byte; + // RK3562: int8 type must be aligned with 16byte, float16 type must be aligned with 8byte; + // RK3588/3576: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte, + // int4 type must be aligned with 64byte; + // matmul data type + // int4: int4(A) x int4(B) -> int16(C) + // int8: int8(A) x int8(B) -> int32(C) + // float16: float16(A) x float16(B) -> float32(C) + rknn_matmul_type type; + + // matmul native layout for B + // 0: normal layout + // 1: native layout + int16_t B_layout; + + // matmul quant type for B + // A and C only support per layer + // 0: per layer + // 1: per channel + // 2: per group + int16_t B_quant_type; + + // matmul native layout for A and C + // 0: normal layout + // 1: native layout + int16_t AC_layout; + + // matmul quant type for A and C, only support 0 + int16_t AC_quant_type; + + // iommu domain id, each domain has 4GB of space + int32_t iommu_domain_id; + + // B_quant_type set 2, group size is enable + int16_t group_size; + + // reserved field + int8_t reserved[34]; +} rknn_matmul_info; + +/* rknn_matmul_create + + params: + rknn_matmul_ctx *ctx the handle of context. + rknn_matmul_info *info the matmal information. + rknn_matmul_io_attr *io_attr inputs/output attribute + return: + int error code +*/ +int rknn_matmul_create(rknn_matmul_ctx* ctx, rknn_matmul_info* info, rknn_matmul_io_attr* io_attr); + +/* rknn_matmul_create_dynamic_shape + + params: + rknn_matmul_ctx *ctx the handle of context. + rknn_matmul_info *info the matmal information. + int shape_num the supported shape number of matmul. + rknn_matmul_shape dynamic_shapes[] the supported M,K,N shape struct array. + rknn_matmul_io_attr *io_attr the array of inputs and output attribute + return: + int error code +*/ +/* + 原来的info.M, K, N无效 +*/ +int rknn_matmul_create_dynamic_shape(rknn_matmul_ctx* ctx, rknn_matmul_info* info, int shape_num, + rknn_matmul_shape dynamic_shapes[], rknn_matmul_io_attr io_attrs[]); + +/* rknn_matmul_set_io_mem + + params: + rknn_matmul_ctx ctx the handle of context. + rknn_tensor_mem *mem the pointer of tensor memory information. + rknn_matmul_tensor_attr *attr the attribute of input or output tensor buffer. + return: + int error code. + + formula: + C = A * B, + + limit: + K max: k <= 10240 + K limit: RK3566/3568: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte; + RK3562: int8 type must be aligned with 32byte, float16 type must be aligned with 32byte; + RK3588/3576: int8 type must be aligned with 32byte, float16 type must be aligned with 32byte, + int4 type must be aligned with 32byte; + N limit: RK3566/3568: int8 type must be aligned with 16byte, float16 type must be aligned with 8byte; + RK3562: int8 type must be aligned with 16byte, float16 type must be aligned with 8byte; + RK3588/3576: int8 type must be aligned with 32byte, float16 type must be aligned with 16byte, + int4 type must be aligned with 64byte; + A shape: M x K + normal layout: (M, K) + [M1K1, M1K2, ..., M1Kk, + M2K1, M2K2, ..., M2Kk, + ... + MmK1, MmK2, ..., MmKk] + for RK3566/3568： + int8: + native layout: (K / 8, M, 8) + [K1M1, K2M1, ..., K8M1, + K9M2, K10M2, ..., K16M2, + ... + K(k-7)Mm, K(k-6)Mm, ..., KkMm] + float16: + native layout: (K / 4, M, 4) + [K1M1, K2M1, ..., K4M1, + K9M2, K10M2, ..., K8M2, + ... + K(k-3)Mm, K(k-2)Mm, ..., KkMm] + for RK3562： + int8: + native layout: (K / 16, M, 16) + [K1M1, K2M1, ..., K16M1, + K17M2, K18M2, ..., K32M2, + ... + K(k-15)Mm, K(k-14)Mm, ..., KkMm] + float16: + native layout: (K / 8, M, 8) + [K1M1, K2M1, ..., K8M1, + K9M2, K10M2, ..., K16M2, + ... + K(k-7)Mm, K(k-6)Mm, ..., KkMm] + for RK3588/3576： + int4: + native layout: (K / 32, M, 32) + [K1M1, K2M1, ..., K32M1, + K33M2, K10M2, ..., K64M2, + ... + K(k-31)Mm, K(k-30)Mm, ..., KkMm] + int8: + native layout: (K / 16, M, 16) + [K1M1, K2M1, ..., K16M1, + K17M2, K18M2, ..., K32M2, + ... + K(k-15)Mm, K(k-14)Mm, ..., KkMm] + float16: + native layout: (K / 8, M, 8) + [K1M1, K2M1, ..., K8M1, + K9M2, K10M2, ..., K16M2, + ... + K(k-7)Mm, K(k-6)Mm, ..., KkMm] + B shape: K x N + normal layout: (K, N) + [K1N1, K1N2, ..., K1Nn, + K2N1, K2N2, ..., K2Nn, + ... + KkN1, KkN2, ..., KkNn] + for RK3566/3568： + int8: + native layout: (N / 16, K / 32, 16, 32) + [K1N1, K2N1, ..., K32N1, + K1N2, K2N2, ..., K32N2, + ... + K1N16, K2N16, ..., K32N16, + K33N1, K34N1, ..., K64N1, + K33N2, K34N2, ..., K64N2, + ... + K(k-31)N16, K(k-30)N16, ..., KkN16, + K1N17, K2N17, ..., K32N17, + K1N18, K2N18, ..., K32N18, + ... + K(k-31)Nn, K(k-30)Nn, ..., KkNn] + float16: + native layout: (N / 8, K / 16, 8, 16) + [K1N1, K2N1, ..., K16N1, + K1N2, K2N2, ..., K16N2, + ... + K1N8, K2N8, ..., K16N8, + K17N1, K18N1, ..., K32N1, + K17N2, K18N2, ..., K32N2, + ... + K(k-15)N8, K(k-30)N8, ..., KkN8, + K1N9, K2N9, ..., K16N9, + K1N10, K2N10, ..., K16N10, + ... + K(k-15)Nn, K(k-14)Nn, ..., KkNn] + for RK3562： + int8: + native layout: (N / 16, K / 32, 16, 32) + [K1N1, K2N1, ..., K32N1, + K1N2, K2N2, ..., K32N2, + ... + K1N16, K2N16, ..., K32N16, + K33N1, K34N1, ..., K64N1, + K33N2, K34N2, ..., K64N2, + ... + K(k-31)N16, K(k-30)N16, ..., KkN16, + K1N17, K2N17, ..., K32N17, + K1N18, K2N18, ..., K32N18, + ... + K(k-31)Nn, K(k-30)Nn, ..., KkNn] + float16: + native layout: (N / 8, K / 32, 8, 32) + [K1N1, K2N1, ..., K32N1, + K1N2, K2N2, ..., K32N2, + ... + K1N8, K2N8, ..., K32N8, + K33N1, K34N1, ..., K64N1, + K33N2, K34N2, ..., K64N2, + ... + K(k-31)N8, K(k-30)N8, ..., KkN8, + K1N9, K2N9, ..., K16N9, + K1N10, K2N10, ..., K16N10, + ... + K(k-31)Nn, K(k-30)Nn, ..., KkNn] + for RK3588： + when K > 8192, the B data will be split into T segments. + int T = std::ceil(K / 8192); + For example: normal layout -> native layout + K = 20488, N = 4096, T = 3, the data will be split into 3 segments. + subN = rknn_matmul_io_attr.B.dims[2]; + subK = rknn_matmul_io_attr.B.dims[3]; + (8196, 4096) (4096 / subN, 8196 / subK, subN, subK) + (K, N) = (20488, 4096) -> (8196, 4096) -> (4096 / subN, 8196 / subK, subN, subK) + normal layout (4096, 4096) (4096 / subN, 4096 / subK, subN, subK) + T normal layout T native layout + It is recommended to use the rknn_B_normal_layout_to_native_layout interface for direct data conversion. + for RK3576： + when K > 4096, the B data will be split into T segments. + int T = std::ceil(K / 4096); + For example: normal layout -> native layout + K = 10240, N = 2048, T = 3, the data will be split into 3 segments. + subN = rknn_matmul_io_attr.B.dims[2]; + subK = rknn_matmul_io_attr.B.dims[3]; + (4096, 2048) (2048 / subN, 4096 / subK, subN, subK) + (K, N) = (10240, 2048) -> (4096, 2048) -> (2048 / subN, 4096 / subK, subN, subK) + normal layout (2048, 2048) (2048 / subN, 2048 / subK, subN, subK) + T normal layout T native layout + It is recommended to use the rknn_B_normal_layout_to_native_layout interface for direct data conversion. + for RK3588/3576： + int4: + native layout: (N / 64, K / 32, 64, 32) + [K1N1, K2N1, ..., K32N1, + K1N2, K2N2, ..., K32N2, + ... + K1N64, K2N64, ..., K32N64, + K33N1, K34N1, ..., K64N1, + K33N2, K34N2, ..., K64N2, + ... + K(k-31)N64, K(k-30)N64, ..., KkN64, + K1N65, K2N65, ..., K32N65, + K1N66, K2N66, ..., K32N66, + ... + K(k-31)Nn, K(k-30)Nn, ..., KkNn] + int8: + native layout: (N / 32, K / 32, 32, 32) + [K1N1, K2N1, ..., K32N1, + K1N2, K2N2, ..., K32N2, + ... + K1N32, K2N32, ..., K32N32, + K33N1, K34N1, ..., K64N1, + K33N2, K34N2, ..., K64N2, + ... + K(k-31)N32, K(k-30)N32, ..., KkN32, + K1N33, K2N33, ..., K32N33, + K1N34, K2N34, ..., K32N34, + ... + K(k-31)Nn, K(k-30)Nn, ..., KkNn] + float16: + native layout: (N / 16, K / 32, 16, 32) + [K1N1, K2N1, ..., K32N1, + K1N2, K2N2, ..., K32N2, + ... + K1N16, K2N16, ..., K32N16, + K33N1, K34N1, ..., K64N1, + K33N2, K34N2, ..., K64N2, + ... + K(k-31)N16, K(k-30)N16, ..., KkN16, + K1N17, K2N17, ..., K32N17, + K1N18, K2N18, ..., K32N18, + ... + K(k-31)Nn, K(k-30)Nn, ..., KkNn] + C shape: M x N + normal layout: (M, N) + [M1N1, M1N2, ..., M1Nn, + M2N1, M2N2, ..., M2Nn, + ... + MmN1, MmN2, ..., MmNn] + native layout: (N / 4, M, 4) + [N1M1, N2M1, ..., N4M1, + N5M2, N6M2, ..., N8M2, + ... + N(n-3)Mm, N(n-2)Mm, ..., NnMm] + for RK3588： + int4: + native layout: (N / 8, M, 8) + [N1M1, N2M1, ..., N8M1, + N9M2, N10M2, ..., N16M2, + ... + N(n-7)Mm, N(n-6)Mm, ..., NnMm] + */ +int rknn_matmul_set_io_mem(rknn_matmul_ctx ctx, rknn_tensor_mem* mem, rknn_matmul_tensor_attr* attr); + +/* rknn_matmul_set_core_mask + + set rknn core mask.(only support RK3588 in current) + + RKNN_NPU_CORE_AUTO: auto mode, default value + RKNN_NPU_CORE_0: core 0 mode + RKNN_NPU_CORE_1: core 1 mode + RKNN_NPU_CORE_2: core 2 mode + RKNN_NPU_CORE_0_1: combine core 0/1 mode + RKNN_NPU_CORE_0_1_2: combine core 0/1/2 mode + + input: + rknn_matmul_ctx context the handle of context. + rknn_core_mask core_mask the core mask. + return: + int error code. +*/ +int rknn_matmul_set_core_mask(rknn_matmul_ctx context, rknn_core_mask core_mask); + +/* rknn_matmul_set_quant_params + + set quant params.(only support matmul type RKNN_INT8_MM_INT8_TO_INT8, RKNN_INT8_MM_INT8_TO_INT32) + + input: + rknn_matmul_ctx context the handle of context. + rknn_quant_params params quant params. + return: + int error code. +*/ +int rknn_matmul_set_quant_params(rknn_matmul_ctx context, rknn_quant_params* params); + +/* rknn_matmul_get_quant_params + + get per channel quant params.(only support matmul type RKNN_INT8_MM_INT8_TO_INT32) + + input: + rknn_matmul_ctx context the handle of context. + rknn_quant_params params quant params. + float scale get scale for user. + return: + int error code. +*/ +int rknn_matmul_get_quant_params(rknn_matmul_ctx ctx, rknn_quant_params* params, float* scale); + +/* rknn_matmul_set_dynamic_shape + + set the matmul input/output shape. matmul will run under current input shape after rknn_matmul_set_dynamic_shape, + only support M dynamicly now. + + input: + rknn_matmul_ctx ctx the handle of context. + rknn_matmul_shape* shape the M,K,N shape of matmul currently + return: + int error code. +*/ +int rknn_matmul_set_dynamic_shape(rknn_matmul_ctx ctx, rknn_matmul_shape* shape); + +/* rknn_matmul_run + + run the matmul in blocking mode + + params: + rknn_matmul_ctx ctx the handle of context. + return: + int error code. + */ +int rknn_matmul_run(rknn_matmul_ctx ctx); + +/* rknn_matmul_destroy + + destroy the matmul context + + params: + rknn_matmul_ctx ctx the handle of context. + return: + int error code. + */ +int rknn_matmul_destroy(rknn_matmul_ctx ctx); + +/* rknn_B_normal_layout_to_native_layout + + change the B normal layout buffer to native layout buffer + + params: + void* B_input B normal layout buffer. + void* B_output B native layout buffer. + int K K + int N N + rknn_matmul_info info matmul info + return: + int error code. + */ +int rknn_B_normal_layout_to_native_layout(void* B_input, void* B_output, int K, int N, rknn_matmul_info* info); + +#ifdef __cplusplus +} // extern "C" +#endif + +#endif // _RKNN_MATMUL_API_H \ No newline at end of file