From 588ca709fb41b7c3c64228fb13ff0d1b56437086 Mon Sep 17 00:00:00 2001 From: John Date: Sat, 17 Jun 2023 16:40:57 +0200 Subject: [PATCH] added the tensor size calculation routines --- ggml.c | 30 ++++++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/ggml.c b/ggml.c index 5e832b474..0bd9a0fa2 100644 --- a/ggml.c +++ b/ggml.c @@ -4177,6 +4177,36 @@ void ggml_scratch_load(struct ggml_context * ctx) { //////////////////////////////////////////////////////////////////////////////// +size_t ggml_sizeof_tensor_data_impl(enum ggml_type type, int n_dims, const int64_t* ne) { + size_t size_needed = 0; + size_needed += GGML_TYPE_SIZE[type]*(ne[0]/GGML_BLCK_SIZE[type]); + for (int i = 1; i < n_dims; i++) { + size_needed *= ne[i]; + } + // align to GGML_MEM_ALIGN + size_needed = ((size_needed + GGML_MEM_ALIGN - 1)/GGML_MEM_ALIGN)*GGML_MEM_ALIGN; + return size_needed; +} + +size_t ggml_sizeof_tensor_1d(enum ggml_type type, int64_t ne0) { + return ggml_tensor_overhead() + ggml_sizeof_tensor_data_impl(type, 1, &ne0); +} + +size_t ggml_sizeof_tensor_2d(enum ggml_type type, int64_t ne0, int64_t ne1) { + const int64_t ne[2] = {ne0, ne1}; + return ggml_tensor_overhead() + ggml_sizeof_tensor_data_impl(type, 2, &ne); +} + +size_t ggml_sizeof_tensor_3d(enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2) { + const int64_t ne[3] = {ne0, ne1, ne2}; + return ggml_tensor_overhead() + ggml_sizeof_tensor_data_impl(type, 3, &ne); +} + +size_t ggml_sizeof_tensor_4d(enum ggml_type type, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + const int64_t ne[4] = {ne0, ne1, ne2, ne3}; + return ggml_tensor_overhead() + ggml_sizeof_tensor_data_impl(type, 4, &ne); +} + struct ggml_tensor * ggml_new_tensor_impl( struct ggml_context * ctx, enum ggml_type type,