implement tensors in binsparse

.gitignore

@@ -6,3 +6,4 @@ scripts
 venv
 build
 ._*
+tensor_test_files

CMakeLists.txt

@@ -13,7 +13,7 @@ set(CMAKE_CXX_STANDARD 20)
 
 set(CMAKE_C_FLAGS "-O3 -march=native")
 
-add_library(binsparse-rc STATIC)
+add_library(binsparse-rc SHARED)
 
 add_subdirectory(include)
 add_subdirectory(src)

compile_flags.txt

@@ -0,0 +1,4 @@
+-I./include
+-I./build/include
+-DBSP_USE_HDF5
+-I/usr/include/hdf5/serial

examples/CMakeLists.txt

@@ -19,5 +19,6 @@ add_example(bsp-ls)
 add_example(benchmark_read)
 add_example(benchmark_read_parallel)
 add_example(benchmark_write)
+add_example(tensor_test)
 
 add_subdirectory(cpp)

examples/cpp/CMakeLists.txt

@@ -15,5 +15,6 @@ add_example(mtx2bsp)
 add_example(bsp2mtx)
 add_example(check_equivalence)
 add_example(bsp-ls)
+add_example(tensor_test)
 add_example(benchmark_read)
 add_example(benchmark_write)

examples/cpp/tensor_test.cpp

@@ -0,0 +1 @@
+#include "../tensor_test.c"

examples/simple_read.c

@@ -7,7 +7,7 @@
 #include <binsparse/binsparse.h>
 
 int main(int argc, char** argv) {
-  char* file_name = "test.hdf5";
+  char* file_name = (char*) "test.hdf5";
 
   hid_t f = H5Fopen(file_name, H5F_ACC_RDWR, H5P_DEFAULT);
 

examples/tensor_test.c

@@ -0,0 +1,22 @@
+#include <binsparse/read_tensor.h>
+#include <binsparse/tensor.h>
+#include <binsparse/write_tensor.h>
+
+int main(int argc, char** argv) {
+  if (argc < 3) {
+    fprintf(stderr,
+            "usage: ./tensor_test [file_name.h5] [output_file_name.h5]\n");
+    return 1;
+  }
+  char* file_name = argv[1];
+  bsp_tensor_t tensor = bsp_read_tensor(argv[1], NULL);
+  printf("rank: %d\n", tensor.rank);
+  printf("dims:");
+  for (int i = 0; i < tensor.rank; i++) {
+    printf("%ld, ", tensor.dims[i]);
+  }
+  printf("\n");
+  bsp_write_tensor(argv[2], tensor, NULL, NULL, 9);
+  bsp_destroy_tensor_t(tensor);
+  return 0;
+}

include/binsparse/read_tensor.h

@@ -0,0 +1,19 @@
+#pragma once
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+char* key_with_index(const char* key, size_t index);
+
+#ifdef BSP_USE_HDF5
+#include <hdf5.h>
+
+bsp_tensor_t bsp_read_tensor_from_group(hid_t f);
+#endif
+
+bsp_tensor_t bsp_read_tensor(const char* file_name, const char* group);
+
+#ifdef __cplusplus
+}
+#endif

include/binsparse/tensor.h

@@ -0,0 +1,130 @@
+#pragma once
+
+#include <binsparse/array.h>
+#include <binsparse/structure.h>
+
+typedef enum {
+  BSP_TENSOR_SPARSE = 0,
+  BSP_TENSOR_DENSE = 1,
+  BSP_TENSOR_ELEMENT = 2,
+} bsp_level_kind_t;
+
+typedef struct {
+  bsp_level_kind_t kind;
+  // data here should be bsp_element_t*, bsp_sparse_t*, or bsp_dense_t*
+  void* data;
+} bsp_level_t;
+
+// corresponds to BSP_TENSOR_ELEMENT
+typedef struct {
+  bsp_array_t values;
+} bsp_element_t;
+
+// corresponds to BSP_TENSOR_DENSE
+typedef struct {
+  int rank;
+  // pointers_to, while it will only ever point to one bsp_array_t, must be kept
+  // as a pointer (rather than a struct) because there are cases where it MUST
+  // be null.
+  bsp_array_t* pointers_to;
+  // indices is supposed to be an array of bsp_array_t's.
+  bsp_array_t* indices;
+  bsp_level_t* child;
+} bsp_sparse_t;
+
+typedef struct {
+  int rank;
+  bsp_level_t* child;
+} bsp_dense_t;
+
+typedef struct {
+  int rank;
+  size_t* dims;
+  size_t* transpose;
+  size_t nnz;
+  bool is_iso;
+
+  bsp_level_t* level;
+  // don't think too much about this at the moment.
+  bsp_structure_t structure;
+} bsp_tensor_t;
+
+static inline bsp_tensor_t bsp_construct_default_tensor_t() {
+  bsp_tensor_t tensor;
+  tensor.structure = BSP_GENERAL;
+  tensor.is_iso = false;
+  tensor.nnz = tensor.rank = 0;
+  tensor.dims = NULL;
+  tensor.transpose = NULL;
+
+  tensor.level = NULL;
+  return tensor;
+}
+
+static void bsp_destroy_level_t(bsp_level_t* level) {
+  if (level == NULL)
+    return;
+  switch (level->kind) {
+  case BSP_TENSOR_ELEMENT: {
+    bsp_element_t* element = (bsp_element_t*) level->data;
+    bsp_destroy_array_t(element->values);
+    free(element);
+    break;
+  }
+  case BSP_TENSOR_DENSE: {
+    bsp_dense_t* dense = (bsp_dense_t*) level->data;
+    bsp_destroy_level_t(dense->child);
+    free(dense);
+    break;
+  }
+  case BSP_TENSOR_SPARSE: {
+    bsp_sparse_t* sparse = (bsp_sparse_t*) level->data;
+
+    if (sparse->pointers_to != NULL)
+      bsp_destroy_array_t(*sparse->pointers_to);
+    if (sparse->indices != NULL) {
+      for (int i = 0; i < sparse->rank; i++) {
+        bsp_destroy_array_t(sparse->indices[i]);
+      }
+    }
+    bsp_destroy_level_t(sparse->child);
+    free(sparse);
+    break;
+  }
+  default:;
+  }
+}
+
+static bsp_array_t bsp_get_tensor_values(bsp_tensor_t tensor) {
+  bsp_level_t* level = tensor.level;
+  while (level != NULL) {
+    switch (level->kind) {
+    case BSP_TENSOR_ELEMENT: {
+      bsp_element_t* element = (bsp_element_t*) level->data;
+      return element->values;
+      break;
+    }
+    case BSP_TENSOR_SPARSE: {
+      bsp_sparse_t* sparse = (bsp_sparse_t*) level->data;
+      level = sparse->child;
+      break;
+    }
+    case BSP_TENSOR_DENSE: {
+      bsp_dense_t* dense = (bsp_dense_t*) level->data;
+      level = dense->child;
+      break;
+    }
+    default:;
+    }
+  }
+  // this should never happen!
+  assert(false);
+}
+
+static inline void bsp_destroy_tensor_t(bsp_tensor_t tensor) {
+  bsp_destroy_level_t(tensor.level);
+  if (tensor.dims != NULL)
+    free(tensor.dims);
+  if (tensor.transpose != NULL)
+    free(tensor.transpose);
+}

include/binsparse/write_tensor.h

@@ -0,0 +1,30 @@
+/*
+ * SPDX-FileCopyrightText: 2024 Binsparse Developers
+ *
+ * SPDX-License-Identifier: BSD-3-Clause
+ */
+
+#pragma once
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+// TODO: make cJSON optional.
+
+#include <binsparse/tensor.h>
+#include <cJSON/cJSON.h>
+
+#ifdef BSP_USE_HDF5
+#include <hdf5.h>
+
+int bsp_write_tensor_to_group(hid_t f, bsp_tensor_t tensor, cJSON* user_json,
+                              int compression_level);
+#endif
+
+int bsp_write_tensor(const char* fname, bsp_tensor_t tensor, const char* group,
+                     cJSON* user_json, int compression_level);
+
+#ifdef __cplusplus
+}
+#endif

src/CMakeLists.txt

@@ -4,5 +4,7 @@
 
 target_sources(binsparse-rc PRIVATE
   src/read_matrix.c
+  src/read_tensor.c
   src/write_matrix.c
+  src/write_tensor.c
 )