diff --git a/.circleci/config.yml b/.circleci/config.yml index 4d33a6bbb..49d60e504 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -112,7 +112,7 @@ commands: - save_cache: paths: - deps - key: build-dependencies-{{ checksum "get_deps.sh" }} + key: v1-dependencies-{{ checksum "get_deps.sh" }} - run: name: Build command: make -C opt all SHOW=1 @@ -151,7 +151,7 @@ commands: - checkout-all - restore_cache: keys: - - build-dependencies-{{ checksum "get_deps.sh" }} + - v1-dependencies-{{ checksum "get_deps.sh" }} # If no exact match is found will get dependencies from source - setup-build-system - run: @@ -260,7 +260,7 @@ jobs: - checkout-all - restore_cache: keys: - - build-dependencies-{{ checksum "get_deps.sh" }} + - v1-dependencies-{{ checksum "get_deps.sh" }} # If no exact match is found will get dependencies from source - setup-build-system - run: @@ -294,7 +294,7 @@ jobs: - checkout-all - restore_cache: keys: - - build-dependencies-{{ checksum "get_deps.sh" }} + - v1-dependencies-{{ checksum "get_deps.sh" }} # If no exact match is found will get dependencies from source - setup-build-system - run: diff --git a/docs/commands.md b/docs/commands.md index 0e16a28dc..c2bd801be 100644 --- a/docs/commands.md +++ b/docs/commands.md @@ -27,7 +27,7 @@ AI.TENSORSET _Arguments_ * **key**: the tensor's key name -* **type**: the tensor's data type can be one of: `FLOAT`, `DOUBLE`, `INT8`, `INT16`, `INT32`, `INT64`, `UINT8` or `UINT16` +* **type**: the tensor's data type can be one of: `FLOAT`, `DOUBLE`, `INT8`, `INT16`, `INT32`, `INT64`, `UINT8`, `UINT16`, `BOOL` or `STRING` * **shape**: one or more dimensions, or the number of elements per axis, for the tensor * **BLOB**: indicates that data is in binary format and is provided via the subsequent `data` argument * **VALUES**: indicates that data is numeric and is provided by one or more subsequent `val` arguments @@ -45,6 +45,8 @@ This will set the key 'mytensor' to the 2x2 RedisAI tensor: ``` redis> AI.TENSORSET mytensor FLOAT 2 2 VALUES 1 2 3 4 OK +redis> AI.TENSORSET mytensor STRING 1 2 VALUES first second +OK ``` !!! note "Uninitialized Tensor Values" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index ddcb1340d..dc0971bc6 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -34,6 +34,7 @@ ADD_LIBRARY(redisai_obj OBJECT execution/command_parser.c execution/parsing/deprecated.c execution/parsing/dag_parser.c + execution/parsing/tensor_commands_parsing.c execution/parsing/model_commands_parser.c execution/parsing/script_commands_parser.c execution/parsing/parse_utils.c diff --git a/src/backends/libtorch_c/torch_c.cpp b/src/backends/libtorch_c/torch_c.cpp index b483fe750..4c3a4e6cb 100644 --- a/src/backends/libtorch_c/torch_c.cpp +++ b/src/backends/libtorch_c/torch_c.cpp @@ -198,7 +198,7 @@ void deleter(DLManagedTensor *arg) { } DLManagedTensor *toManagedDLPack(const torch::Tensor &src_) { - ATenDLMTensor *atDLMTensor(new ATenDLMTensor); + auto *atDLMTensor(new ATenDLMTensor); atDLMTensor->handle = src_; auto &src = atDLMTensor->handle; atDLMTensor->tensor.manager_ctx = atDLMTensor; @@ -213,6 +213,7 @@ DLManagedTensor *toManagedDLPack(const torch::Tensor &src_) { atDLMTensor->tensor.dl_tensor.dtype = getDLDataType(src); atDLMTensor->tensor.dl_tensor.shape = const_cast(src.sizes().data()); atDLMTensor->tensor.dl_tensor.strides = const_cast(src.strides().data()); + atDLMTensor->tensor.dl_tensor.elements_length = nullptr; atDLMTensor->tensor.dl_tensor.byte_offset = 0; return &(atDLMTensor->tensor); } diff --git a/src/backends/onnxruntime.c b/src/backends/onnxruntime.c index aee7a42d6..f6145ebef 100644 --- a/src/backends/onnxruntime.c +++ b/src/backends/onnxruntime.c @@ -170,7 +170,13 @@ ONNXTensorElementDataType RAI_GetOrtDataTypeFromDL(DLDataType dtype) { switch (dtype.bits) { case 8: return ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; - break; + default: + return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; + } + } else if (dtype.code == kDLString) { + switch (dtype.bits) { + case 8: + return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; default: return ONNX_TENSOR_ELEMENT_DATA_TYPE_UNDEFINED; } @@ -198,20 +204,22 @@ DLDataType RAI_GetDLDataTypeFromORT(ONNXTensorElementDataType dtype) { return (DLDataType){.code = kDLUInt, .bits = 16, .lanes = 1}; case ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL: return (DLDataType){.code = kDLBool, .bits = 8, .lanes = 1}; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING: + return (DLDataType){.code = kDLString, .bits = 8, .lanes = 1}; default: return (DLDataType){.bits = 0}; } } -int RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, OrtValue **input, +int RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t batches_count, OrtValue **input, OrtStatus **status_ptr) { OrtStatus *status = NULL; const OrtApi *ort = OrtGetApiBase()->GetApi(1); - size_t batch_size = 0; + int64_t batch_size = 0; size_t batch_byte_size = 0; - for (size_t i = 0; i < count; i++) { + for (size_t i = 0; i < batches_count; i++) { batch_size += ts[i]->tensor.dl_tensor.shape[0]; batch_byte_size += RAI_TensorByteSize(ts[i]); } @@ -219,21 +227,38 @@ int RAI_OrtValueFromTensors(RAI_Tensor **ts, size_t count, OrtValue **input, RAI_Tensor *t0 = ts[0]; const int ndim = t0->tensor.dl_tensor.ndim; int64_t batched_shape[ndim]; + size_t batched_tensor_len = 1; for (size_t i = 1; i < ndim; i++) { batched_shape[i] = t0->tensor.dl_tensor.shape[i]; + batched_tensor_len *= batched_shape[i]; } batched_shape[0] = batch_size; OrtValue *out; - if (count > 1) { + if (t0->tensor.dl_tensor.dtype.code == kDLString) { + ONNX_VALIDATE_STATUS(ort->CreateTensorAsOrtValue( + global_allocator, batched_shape, t0->tensor.dl_tensor.ndim, + RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out)); + size_t element_index = 0; + for (size_t i = 0; i < batches_count; i++) { + // go over all strings stored in the tensors' data from all tensors and set them in the + // ORT tensor. + size_t tensor_len = RAI_TensorLength(ts[i]); + uint64_t *offsets = RAI_TensorStringElementsOffsets(ts[i]); + for (size_t j = 0; j < tensor_len; j++) { + // send a pointer to the position of a string elements (null-terminated) in the blob + ONNX_VALIDATE_STATUS(ort->FillStringTensorElement( + out, RAI_TensorData(ts[i]) + offsets[j], element_index++)); + } + } + } else if (batches_count > 1) { ONNX_VALIDATE_STATUS( ort->CreateTensorAsOrtValue(global_allocator, batched_shape, t0->tensor.dl_tensor.ndim, RAI_GetOrtDataTypeFromDL(t0->tensor.dl_tensor.dtype), &out)) - char *ort_data; ONNX_VALIDATE_STATUS(ort->GetTensorMutableData(out, (void **)&ort_data)) size_t offset = 0; - for (size_t i = 0; i < count; i++) { + for (size_t i = 0; i < batches_count; i++) { memcpy(ort_data + offset, RAI_TensorData(ts[i]), RAI_TensorByteSize(ts[i])); offset += RAI_TensorByteSize(ts[i]); } @@ -255,99 +280,94 @@ RAI_Tensor *RAI_TensorCreateFromOrtValue(OrtValue *v, size_t batch_offset, long RAI_Error *error) { OrtStatus *status = NULL; const OrtApi *ort = OrtGetApiBase()->GetApi(1); - RAI_Tensor *ret = NULL; - int64_t *shape = NULL; - int64_t *strides = NULL; + RAI_Tensor *output_tensor = NULL; OrtTensorTypeAndShapeInfo *info = NULL; int is_tensor; ONNX_VALIDATE_STATUS(ort->IsTensor(v, &is_tensor)) if (!is_tensor) { // TODO: if not tensor, flatten the data structure (sequence or map) and store it in a - // tensor. If return value is string, emit warning. + // tensor. return NULL; } - ret = RAI_TensorNew(); - // Default device is CPU (id -1 is default, means 'no index') - DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = -1}; - + // Create an empty RAI_Tensor based on the tensor's meat-data ONNX_VALIDATE_STATUS(ort->GetTensorTypeAndShape(v, &info)) - + ONNXTensorElementDataType ort_dtype; + ONNX_VALIDATE_STATUS(ort->GetTensorElementType(info, &ort_dtype)) + DLDataType data_type = RAI_GetDLDataTypeFromORT(ort_dtype); + size_t data_type_size = data_type.bits / 8; + int64_t total_batch_size; { - size_t ndims; - ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &ndims)) - int64_t dims[ndims]; - ONNX_VALIDATE_STATUS(ort->GetDimensions(info, dims, ndims)) - enum ONNXTensorElementDataType ort_dtype; - ONNX_VALIDATE_STATUS(ort->GetTensorElementType(info, &ort_dtype)) - int64_t total_batch_size = dims[0]; - total_batch_size = total_batch_size > 0 ? total_batch_size : 1; - - shape = RedisModule_Calloc(ndims, sizeof(*shape)); - strides = RedisModule_Calloc(ndims, sizeof(*strides)); - for (int64_t i = 0; i < ndims; ++i) { - shape[i] = dims[i]; - strides[i] = 1; - } - if (batch_size != -1) { - shape[0] = batch_size; + size_t n_dims; + ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &n_dims)) + int64_t dims[n_dims]; + ONNX_VALIDATE_STATUS(ort->GetDimensions(info, dims, n_dims)) + total_batch_size = dims[0]; + + // if batch size is non-negative (meaning this tensor was a part of a larger batch), + // we restore the original first dimension of this particular tensor. + if (batch_size > 0) { + dims[0] = batch_size; } else { - batch_size = total_batch_size; - } - for (int64_t i = ndims - 2; i >= 0; --i) { - strides[i] *= strides[i + 1] * shape[i + 1]; + batch_size = total_batch_size; // = dims[0] } + output_tensor = RAI_TensorNew(data_type, (const size_t *)dims, (int)n_dims); + } - DLDataType dtype = RAI_GetDLDataTypeFromORT(ort_dtype); -#ifdef RAI_COPY_RUN_OUTPUT + size_t output_tensor_len = RAI_TensorLength(output_tensor); + size_t elem_count; + ONNX_VALIDATE_STATUS(ort->GetTensorShapeElementCount(info, &elem_count)) + RedisModule_Assert(elem_count / total_batch_size * batch_size == output_tensor_len); + + if (data_type.code != kDLString) { char *ort_data; ONNX_VALIDATE_STATUS(ort->GetTensorMutableData(v, (void **)&ort_data)) - size_t elem_count; - ONNX_VALIDATE_STATUS(ort->GetTensorShapeElementCount(info, &elem_count)) - - const size_t len = dtype.bits * elem_count / 8; - const size_t total_bytesize = len * sizeof(char); - const size_t sample_bytesize = total_bytesize / total_batch_size; - const size_t batch_bytesize = sample_bytesize * batch_size; - - char *data = RedisModule_Calloc(batch_bytesize, sizeof(*data)); - memcpy(data, ort_data + batch_offset * sample_bytesize, batch_bytesize); -#endif - - ort->ReleaseTensorTypeAndShapeInfo(info); - - // TODO: use manager_ctx to ensure ORT tensor doesn't get deallocated - // This applies to outputs - - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, -#ifdef RAI_COPY_RUN_OUTPUT - .data = data, -#else -#error zero-copy passing output memory from ORT not currently supported -#endif - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = 0}, - .manager_ctx = NULL, - .deleter = NULL}; - - return ret; + output_tensor->tensor.dl_tensor.data = RedisModule_Alloc(RAI_TensorByteSize(output_tensor)); + size_t total_byte_size = elem_count * data_type_size; + size_t sample_byte_size = total_byte_size / total_batch_size; + memcpy(RAI_TensorData(output_tensor), ort_data + batch_offset * sample_byte_size, + RAI_TensorByteSize(output_tensor)); + } else { + // Calculate the blob size for this tensor and allocate space for it + size_t element_index = batch_offset * (output_tensor_len / batch_size); + size_t blob_len = 0; + for (size_t i = 0; i < output_tensor_len; i++) { + size_t str_element_len; + ONNX_VALIDATE_STATUS( + ort->GetStringTensorElementLength(v, element_index++, &str_element_len)); + blob_len += + str_element_len + 1; // Add space for null character at the end of every string + } + output_tensor->blobSize = blob_len; + output_tensor->tensor.dl_tensor.data = RedisModule_Alloc(blob_len); + char *tensor_blob = RAI_TensorData(output_tensor); + + // Go over again and set tensor data elements one by one + element_index = batch_offset * (output_tensor_len / batch_size); + uint64_t *offsets = RAI_TensorStringElementsOffsets(output_tensor); + offsets[0] = 0; + for (size_t i = 0; i < output_tensor_len; i++) { + size_t str_element_len; + ONNX_VALIDATE_STATUS( + ort->GetStringTensorElementLength(v, element_index, &str_element_len)); + ONNX_VALIDATE_STATUS(ort->GetStringTensorElement(v, str_element_len, element_index++, + tensor_blob + offsets[i])); + *(tensor_blob + offsets[i] + str_element_len) = '\0'; + if (i < output_tensor_len - 1) { + offsets[i + 1] = offsets[i] + str_element_len + 1; + } + } } + ort->ReleaseTensorTypeAndShapeInfo(info); + return output_tensor; + error: RAI_SetError(error, RAI_EMODELCREATE, ort->GetErrorMessage(status)); ort->ReleaseStatus(status); - if (shape != NULL) { - RedisModule_Free(shape); - } - if (strides != NULL) { - RedisModule_Free(shape); - } - if (ret != NULL) { - RedisModule_Free(ret); + if (output_tensor != NULL) { + RAI_TensorFree(output_tensor); } if (info != NULL) { ort->ReleaseTensorTypeAndShapeInfo(info); @@ -505,24 +525,24 @@ int RAI_ModelRunORT(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error return REDISMODULE_ERR; } - const size_t nbatches = array_len(ectxs); - if (nbatches == 0) { + const size_t n_batches = array_len(ectxs); + if (n_batches == 0) { RAI_SetError(error, RAI_EMODELRUN, "ERR No batches to run"); return REDISMODULE_ERR; } - size_t batch_sizes[nbatches]; - size_t batch_offsets[nbatches]; + size_t batch_sizes[n_batches]; + size_t batch_offsets[n_batches]; size_t total_batch_size = 0; - const size_t ninputs = RAI_ExecutionCtx_NumInputs(ectxs[0]); - const size_t noutputs = RAI_ExecutionCtx_NumOutputs(ectxs[0]); - if (ninputs > 0) { - for (size_t b = 0; b < nbatches; ++b) { + const size_t n_inputs = RAI_ExecutionCtx_NumInputs(ectxs[0]); + const size_t n_outputs = RAI_ExecutionCtx_NumOutputs(ectxs[0]); + if (n_batches > 0) { + for (size_t b = 0; b < n_batches; ++b) { batch_sizes[b] = RAI_TensorDim(RAI_ExecutionCtx_GetInput(ectxs[b], 0), 0); total_batch_size += batch_sizes[b]; } batch_offsets[0] = 0; - for (size_t b = 1; b < nbatches; ++b) { + for (size_t b = 1; b < n_batches; ++b) { batch_offsets[b] = batch_offsets[b - 1] + batch_sizes[b - 1]; } } @@ -541,16 +561,16 @@ int RAI_ModelRunORT(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error ONNX_VALIDATE_STATUS(ort->SessionGetInputCount(session, &n_input_nodes)) ONNX_VALIDATE_STATUS(ort->SessionGetOutputCount(session, &n_output_nodes)) - if (ninputs != n_input_nodes) { + if (n_inputs != n_input_nodes) { char msg[70]; - sprintf(msg, "ERR Expected %li inputs but got %li", n_input_nodes, ninputs); + sprintf(msg, "ERR Expected %li inputs but got %li", n_input_nodes, n_inputs); RAI_SetError(error, RAI_EMODELRUN, msg); return REDISMODULE_ERR; } - if (noutputs != n_output_nodes) { + if (n_outputs != n_output_nodes) { char msg[70]; - sprintf(msg, "ERR Expected %li outputs but got %li", n_output_nodes, noutputs); + sprintf(msg, "ERR Expected %li outputs but got %li", n_output_nodes, n_outputs); RAI_SetError(error, RAI_EMODELRUN, msg); return REDISMODULE_ERR; } @@ -561,12 +581,12 @@ int RAI_ModelRunORT(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error ort->SessionGetInputName(session, i, global_allocator, &input_name)) input_names = array_append(input_names, input_name); - RAI_Tensor *batched_input_tensors[nbatches]; - for (size_t b = 0; b < nbatches; b++) { + RAI_Tensor *batched_input_tensors[n_batches]; + for (size_t b = 0; b < n_batches; b++) { batched_input_tensors[b] = RAI_ExecutionCtx_GetInput(ectxs[b], i); } OrtValue *input; - if (RAI_OrtValueFromTensors(batched_input_tensors, nbatches, &input, &status) != + if (RAI_OrtValueFromTensors(batched_input_tensors, n_batches, &input, &status) != REDISMODULE_OK) { goto error; } @@ -599,17 +619,17 @@ int RAI_ModelRunORT(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error RAI_ResetRunSessionCtxORT(run_session_index); run_options = NULL; - for (uint32_t i = 0; i < ninputs; i++) { + for (uint32_t i = 0; i < n_inputs; i++) { status = ort->AllocatorFree(global_allocator, (void *)input_names[i]); } array_free(input_names); - for (uint32_t i = 0; i < noutputs; i++) { + for (uint32_t i = 0; i < n_outputs; i++) { status = ort->AllocatorFree(global_allocator, (void *)output_names[i]); } array_free(output_names); for (size_t i = 0; i < n_output_nodes; i++) { - if (nbatches > 1) { + if (n_batches > 1) { ONNX_VALIDATE_STATUS(ort->GetTensorTypeAndShape(outputs[i], &info)) size_t ndims; ONNX_VALIDATE_STATUS(ort->GetDimensionsCount(info, &ndims)) @@ -622,7 +642,7 @@ int RAI_ModelRunORT(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error goto error; } - for (size_t b = 0; b < nbatches; b++) { + for (size_t b = 0; b < n_batches; b++) { RAI_Tensor *output_tensor = RAI_TensorCreateFromOrtValue( outputs[i], batch_offsets[b], batch_sizes[b], error); if (error->code != RAI_OK) { @@ -638,7 +658,8 @@ int RAI_ModelRunORT(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error } } else { RAI_Tensor *output_tensor = RAI_TensorCreateFromOrtValue(outputs[i], 0, -1, error); - if (error->code != RAI_OK) { + if (RAI_GetErrorCode(error) != RAI_OK) { + ort->ReleaseValue(outputs[i]); goto error; } if (output_tensor) { diff --git a/src/backends/tensorflow.c b/src/backends/tensorflow.c index e8881334b..33813bbc9 100644 --- a/src/backends/tensorflow.c +++ b/src/backends/tensorflow.c @@ -58,6 +58,13 @@ TF_DataType RAI_GetTFDataTypeFromDL(DLDataType dtype) { default: return 0; } + } else if (dtype.code == kDLString) { + switch (dtype.bits) { + case 8: + return TF_STRING; + default: + return 0; + } } return 0; } @@ -82,70 +89,73 @@ DLDataType RAI_GetDLDataTypeFromTF(TF_DataType dtype) { return (DLDataType){.code = kDLUInt, .bits = 16, .lanes = 1}; case TF_BOOL: return (DLDataType){.code = kDLBool, .bits = 8, .lanes = 1}; + case TF_STRING: + return (DLDataType){.code = kDLString, .bits = 8, .lanes = 1}; default: return (DLDataType){.bits = 0}; } - return (DLDataType){.bits = 0}; } RAI_Tensor *RAI_TensorCreateFromTFTensor(TF_Tensor *tensor, size_t batch_offset, long long batch_size) { - RAI_Tensor *ret = RAI_TensorNew(); - // Default device is CPU (id -1 is default, means 'no index') - DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = -1}; - - const size_t ndims = TF_NumDims(tensor); + int n_dims = TF_NumDims(tensor); int64_t total_batch_size = TF_Dim(tensor, 0); total_batch_size = total_batch_size > 0 ? total_batch_size : 1; - int64_t *shape = RedisModule_Calloc(ndims, sizeof(*shape)); - int64_t *strides = RedisModule_Calloc(ndims, sizeof(*strides)); - for (int64_t i = 0; i < ndims; ++i) { + size_t shape[n_dims]; + for (int i = 0; i < n_dims; ++i) { shape[i] = TF_Dim(tensor, i); - strides[i] = 1; } if (batch_size != -1) { - shape[0] = batch_size; + shape[0] = batch_size; // the TF tensor was batched } else { - batch_size = total_batch_size; - } - for (int64_t i = ndims - 2; i >= 0; --i) { - strides[i] *= strides[i + 1] * shape[i + 1]; - } - - const size_t sample_bytesize = TF_TensorByteSize(tensor) / total_batch_size; - - // FIXME: In TF, RunSession allocates memory for output tensors - // This means that we either memcpy the tensor data and let - // Redis be responsible for the memory, or we reuse the TF - // allocated memory, which might not be optimal down the road - // Note: on YOLO this has no impact on perf -#ifdef RAI_COPY_RUN_OUTPUT - const size_t len = sample_bytesize * batch_size; - char *data = RedisModule_Calloc(len, sizeof(*data)); - memcpy(data, TF_TensorData(tensor) + sample_bytesize * batch_offset, len); -#endif - - // TODO: use manager_ctx to ensure TF tensor doesn't get deallocated - // This applies to outputs - - ret->tensor = (DLManagedTensor){ - .dl_tensor = (DLTensor){.device = device, -#ifdef RAI_COPY_RUN_OUTPUT - .data = data, -#else - .data = TF_TensorData(tensor), -#endif - .ndim = ndims, - .dtype = RAI_GetDLDataTypeFromTF(TF_TensorType(tensor)), - .shape = shape, - .strides = strides, - .byte_offset = 0}, - .manager_ctx = NULL, - .deleter = NULL}; - - return ret; + batch_size = total_batch_size; // the TF tensor wasn't batched + } + + DLDataType data_type = RAI_GetDLDataTypeFromTF(TF_TensorType(tensor)); + RAI_Tensor *out = RAI_TensorNew(data_type, shape, n_dims); + size_t out_tensor_len = RAI_TensorLength(out); + + if (data_type.code == kDLString) { + TF_TString *tensor_data = TF_TensorData(tensor); + const char *strings_data[out_tensor_len]; + size_t strings_lengths[out_tensor_len]; + + // Calculate the blob size for this tensor and allocate space for it + size_t element_index = batch_offset * (out_tensor_len / batch_size); + size_t blob_len = 0; + uint64_t *offsets = RAI_TensorStringElementsOffsets(out); + offsets[0] = 0; + for (size_t i = 0; i < out_tensor_len; i++) { + size_t str_element_len = TF_TString_GetSize(tensor_data + element_index); + strings_lengths[i] = str_element_len; + strings_data[i] = TF_TString_GetDataPointer(tensor_data + element_index++); + offsets[i] = blob_len; + blob_len += str_element_len; + if (strings_data[i][str_element_len - 1] != '\0') { + blob_len++; // Add space for null character at the end of every string + } + } + out->blobSize = blob_len; + out->tensor.dl_tensor.data = RedisModule_Calloc(1, blob_len); + + // Go over again and set tensor data elements one by one + element_index = batch_offset * (out_tensor_len / batch_size); + char *tensor_blob = RAI_TensorData(out); + for (size_t i = 0; i < out_tensor_len; i++) { + memcpy(tensor_blob + offsets[i], strings_data[i], strings_lengths[i]); + TF_TString_Dealloc(tensor_data + element_index++); + } + } else { + size_t non_batched_tensor_size = TF_TensorByteSize(tensor) / total_batch_size; + size_t blob_len = non_batched_tensor_size * batch_size; + out->tensor.dl_tensor.data = RedisModule_Alloc(blob_len); + out->blobSize = blob_len; + memcpy(RAI_TensorData(out), TF_TensorData(tensor) + non_batched_tensor_size * batch_offset, + blob_len); + } + return out; } void RAI_TFDeallocator(void *data, size_t len, void *arg) { @@ -153,64 +163,60 @@ void RAI_TFDeallocator(void *data, size_t len, void *arg) { // do nothing, memory is managed by Redis } -TF_Tensor *RAI_TFTensorFromTensor(RAI_Tensor *t) { -#ifdef RAI_COPY_RUN_INPUT - TF_Tensor *out = TF_AllocateTensor(RAI_GetTFDataTypeFromDL(t->tensor.dl_tensor.dtype), - t->tensor.dl_tensor.shape, t->tensor.dl_tensor.ndim, - RAI_TensorByteSize(t)); - memcpy(TF_TensorData(out), t->tensor.dl_tensor.data, TF_TensorByteSize(out)); - return out; -#else - return TF_NewTensor(RAI_GetTFDataTypeFromDL(t->tensor.dl_tensor.dtype), - t->tensor.dl_tensor.shape, t->tensor.dl_tensor.ndim, - t->tensor.dl_tensor.data, RAI_TensorByteSize(t), &RAI_TFDeallocator, NULL); -#endif /* RAI_COPY_RUN_INPUT */ -} - -TF_Tensor *RAI_TFTensorFromTensors(RAI_Tensor **ts, size_t count) { - - if (count == 0) { - return NULL; - } +TF_Tensor *RAI_TFTensorFromTensors(RAI_Tensor **tensors, size_t count) { + RedisModule_Assert(count > 0); - size_t batch_size = 0; + int64_t batch_size = 0; size_t batch_byte_size = 0; for (size_t i = 0; i < count; i++) { - batch_size += ts[i]->tensor.dl_tensor.shape[0]; - batch_byte_size += RAI_TensorByteSize(ts[i]); - } - - RAI_Tensor *t0 = ts[0]; - - const int ndim = t0->tensor.dl_tensor.ndim; - int64_t batched_shape[ndim]; - - for (size_t i = 0; i < ndim; i++) { - batched_shape[i] = t0->tensor.dl_tensor.shape[i]; + batch_size += RAI_TensorDim(tensors[i], 0); + batch_byte_size += RAI_TensorByteSize(tensors[i]); } + // get the shapes of the batched tensor: all inner dims should be the same, + // so we go over t0 dims. + RAI_Tensor *t0 = tensors[0]; + int n_dim = RAI_TensorNumDims(t0); + int64_t batched_shape[n_dim]; batched_shape[0] = batch_size; + size_t batched_tensor_len = batch_size; + for (int i = 1; i < n_dim; i++) { + batched_shape[i] = RAI_TensorDim(t0, i); + batched_tensor_len *= batched_shape[i]; + } TF_Tensor *out = NULL; - - if (count > 1) { - out = TF_AllocateTensor(RAI_GetTFDataTypeFromDL(t0->tensor.dl_tensor.dtype), batched_shape, - t0->tensor.dl_tensor.ndim, batch_byte_size); - + if (RAI_TensorDataType(t0).code == kDLString) { + out = TF_AllocateTensor(TF_STRING, batched_shape, RAI_TensorNumDims(t0), + sizeof(TF_TString) * batched_tensor_len); + // go over the string elements and copy the data to the TF tensor + TF_TString *tf_str = (TF_TString *)TF_TensorData(out); + size_t element_ind = 0; + for (size_t i = 0; i < count; i++) { + RAI_Tensor *t = tensors[i]; + uint64_t *offsets = RAI_TensorStringElementsOffsets(t); + for (size_t j = 0; j < RAI_TensorLength(t); j++) { + TF_TString_Init(&tf_str[element_ind]); + size_t str_len = j < RAI_TensorLength(t) - 1 ? offsets[j + 1] - offsets[j] + : RAI_TensorByteSize(t) - offsets[j]; + TF_TString_Copy(&tf_str[element_ind++], RAI_TensorData(t) + offsets[j], str_len); + } + } + } else if (count > 1) { + out = TF_AllocateTensor(RAI_GetTFDataTypeFromDL(RAI_TensorDataType(t0)), batched_shape, + RAI_TensorNumDims(t0), batch_byte_size); size_t offset = 0; for (size_t i = 0; i < count; i++) { - size_t tbytesize = RAI_TensorByteSize(ts[i]); - memcpy(TF_TensorData(out) + offset, ts[i]->tensor.dl_tensor.data, tbytesize); - offset += tbytesize; + size_t tensor_byte_size = RAI_TensorByteSize(tensors[i]); + memcpy(TF_TensorData(out) + offset, RAI_TensorData(tensors[i]), tensor_byte_size); + offset += tensor_byte_size; } } else { - out = TF_NewTensor(RAI_GetTFDataTypeFromDL(t0->tensor.dl_tensor.dtype), - t0->tensor.dl_tensor.shape, t0->tensor.dl_tensor.ndim, - t0->tensor.dl_tensor.data, RAI_TensorByteSize(t0), &RAI_TFDeallocator, - NULL); + out = TF_NewTensor(RAI_GetTFDataTypeFromDL(RAI_TensorDataType(t0)), RAI_TensorShape(t0), + RAI_TensorNumDims(t0), RAI_TensorData(t0), RAI_TensorByteSize(t0), + &RAI_TFDeallocator, NULL); } - return out; } @@ -463,20 +469,22 @@ void RAI_ModelFreeTF(RAI_Model *model, RAI_Error *error) { } int RAI_ModelRunTF(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error) { - TF_Status *status = TF_NewStatus(); - + int res = REDISMODULE_ERR; const size_t nbatches = array_len(ectxs); if (nbatches == 0) { RAI_SetError(error, RAI_EMODELRUN, "ERR No batches to run"); - return 1; + return res; } - const size_t ninputs = RAI_ExecutionCtx_NumInputs(ectxs[0]); - const size_t noutputs = RAI_ExecutionCtx_NumOutputs(ectxs[0]); + TF_Status *status = TF_NewStatus(); + size_t ninputs = RAI_ExecutionCtx_NumInputs(ectxs[0]); + size_t noutputs = RAI_ExecutionCtx_NumOutputs(ectxs[0]); + + TF_Input inputs[ninputs]; + TF_Output outputs[noutputs]; + TF_Tensor *inputTensorsValues[ninputs]; - TF_Output inputs[ninputs]; TF_Tensor *outputTensorsValues[noutputs]; - TF_Output outputs[noutputs]; size_t batch_sizes[nbatches]; size_t batch_offsets[nbatches]; @@ -497,27 +505,24 @@ int RAI_ModelRunTF(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error) for (size_t i = 0; i < ninputs; ++i) { RAI_Tensor *batched_input_tensors[nbatches]; - for (size_t b = 0; b < nbatches; ++b) { batched_input_tensors[b] = RAI_ExecutionCtx_GetInput(ectxs[b], i); } inputTensorsValues[i] = RAI_TFTensorFromTensors(batched_input_tensors, nbatches); - TF_Output port; + TF_Input port; port.oper = TF_GraphOperationByName(tfGraph, RAI_ModelGetInputName(model, i)); + // this operation must exist in the graph (verified when creating the model) + RedisModule_Assert(port.oper); port.index = 0; - if (port.oper == NULL) { - return 1; - } inputs[i] = port; } for (size_t i = 0; i < noutputs; ++i) { TF_Output port; port.oper = TF_GraphOperationByName(tfGraph, RAI_ModelGetOutputName(model, i)); + // this operation must exist in the graph (verified when creating the model) + RedisModule_Assert(port.oper); port.index = 0; - if (port.oper == NULL) { - return 1; - } outputs[i] = port; } @@ -525,16 +530,9 @@ int RAI_ModelRunTF(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error) outputTensorsValues, noutputs, NULL /* target_opers */, 0 /* ntargets */, NULL /* run_Metadata */, status); - for (size_t i = 0; i < ninputs; ++i) { - TF_DeleteTensor(inputTensorsValues[i]); - } - if (TF_GetCode(status) != TF_OK) { - char *errorMessage = RedisModule_Strdup(TF_Message(status)); - RAI_SetError(error, RAI_EMODELRUN, errorMessage); - TF_DeleteStatus(status); - RedisModule_Free(errorMessage); - return 1; + RAI_SetError(error, RAI_EMODELRUN, TF_Message(status)); + goto cleanup; } for (size_t i = 0; i < noutputs; ++i) { @@ -544,10 +542,10 @@ int RAI_ModelRunTF(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error) } if (TF_Dim(outputTensorsValues[i], 0) != total_batch_size) { TF_DeleteTensor(outputTensorsValues[i]); - TF_DeleteStatus(status); RAI_SetError(error, RAI_EMODELRUN, - "ERR Model did not generate the expected batch size"); - return 1; + "ERR Tensorflow batching error: model did not generate the expected " + "batch size"); + goto cleanup; } for (size_t b = 0; b < nbatches; b++) { @@ -561,12 +559,26 @@ int RAI_ModelRunTF(RAI_Model *model, RAI_ExecutionCtx **ectxs, RAI_Error *error) RAI_ExecutionCtx_SetOutput( ectxs[0], RAI_TensorCreateFromTFTensor(outputTensorsValues[i], 0, -1), i); } - TF_DeleteTensor(outputTensorsValues[i]); } + res = REDISMODULE_OK; +cleanup: TF_DeleteStatus(status); - - return 0; + for (size_t i = 0; i < ninputs; ++i) { + RAI_Tensor *t = RAI_ExecutionCtx_GetInput(ectxs[0], i); + // free the underline string buffer for every element + if (RAI_TensorDataType(t).code == kDLString) { + TF_TString *tf_strings = TF_TensorData(inputTensorsValues[i]); + for (size_t j = 0; j < RAI_TensorLength(t); j++) { + TF_TString_Dealloc(tf_strings + j); + } + } + TF_DeleteTensor(inputTensorsValues[i]); + } + for (size_t i = 0; i < noutputs; i++) { + TF_DeleteTensor(outputTensorsValues[i]); + } + return res; } int RAI_ModelSerializeTF(RAI_Model *model, char **buffer, size_t *len, RAI_Error *error) { diff --git a/src/execution/DAG/dag.c b/src/execution/DAG/dag.c index aac10a04f..fd912e88c 100644 --- a/src/execution/DAG/dag.c +++ b/src/execution/DAG/dag.c @@ -90,8 +90,8 @@ static int _StoreTensorInKeySpace(RedisModuleCtx *ctx, RAI_Tensor *tensor, RedisModuleString *persist_key_name, RAI_Error *err) { RedisModuleKey *key; - const int status = - RAI_OpenKey_Tensor(ctx, persist_key_name, &key, REDISMODULE_READ | REDISMODULE_WRITE, err); + int status = + RAI_TensorOpenKey(ctx, persist_key_name, &key, REDISMODULE_READ | REDISMODULE_WRITE, err); if (status == REDISMODULE_ERR) { return REDISMODULE_ERR; } @@ -101,7 +101,7 @@ static int _StoreTensorInKeySpace(RedisModuleCtx *ctx, RAI_Tensor *tensor, return REDISMODULE_ERR; } // Only if we got until here, tensor is saved in keyspace. - RedisAI_ReplicateTensorSet(ctx, persist_key_name, tensor); + RAI_TensorReplicate(ctx, persist_key_name, tensor); RedisModule_CloseKey(key); return REDISMODULE_OK; } @@ -569,7 +569,7 @@ int RedisAI_DagRun_Reply(RedisModuleCtx *ctx, RedisModuleString **argv, int argc if (t == NULL) { RedisModule_ReplyWithSimpleString(ctx, "NA"); } else { - ReplyWithTensor(ctx, currentOp->fmt, t); + RAI_TensorReply(ctx, currentOp->fmt, t); } break; } diff --git a/src/execution/execution_contexts/modelRun_ctx.c b/src/execution/execution_contexts/modelRun_ctx.c index 1f0b6c224..aa4b6bd96 100644 --- a/src/execution/execution_contexts/modelRun_ctx.c +++ b/src/execution/execution_contexts/modelRun_ctx.c @@ -61,7 +61,7 @@ int ModelRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkeys, size_t ninputs = array_len(inkeys), noutputs = array_len(outkeys); for (size_t i = 0; i < ninputs; i++) { const int status = - RAI_GetTensorFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ, err); + RAI_TensorGetFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ, err); if (status == REDISMODULE_ERR) { return REDISMODULE_ERR; } diff --git a/src/execution/execution_contexts/scriptRun_ctx.c b/src/execution/execution_contexts/scriptRun_ctx.c index 6bf1eab68..fda275a80 100644 --- a/src/execution/execution_contexts/scriptRun_ctx.c +++ b/src/execution/execution_contexts/scriptRun_ctx.c @@ -139,8 +139,7 @@ int ScriptRunCtx_SetParams(RedisModuleCtx *ctx, RedisModuleString **inkeys, RedisModuleKey *key; size_t ninputs = array_len(inkeys), noutputs = array_len(outkeys); for (size_t i = 0; i < ninputs; i++) { - const int status = - RAI_GetTensorFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ, err); + int status = RAI_TensorGetFromKeyspace(ctx, inkeys[i], &key, &t, REDISMODULE_READ, err); if (status == REDISMODULE_ERR) { RedisModule_Log(ctx, "warning", "could not load input tensor %s from keyspace", RedisModule_StringPtrLen(inkeys[i], NULL)); diff --git a/src/execution/parsing/dag_parser.c b/src/execution/parsing/dag_parser.c index 0f1c584ca..791feb8a0 100644 --- a/src/execution/parsing/dag_parser.c +++ b/src/execution/parsing/dag_parser.c @@ -3,12 +3,12 @@ #include "redismodule.h" #include "util/dict.h" #include "util/string_utils.h" -#include "redis_ai_objects/tensor.h" #include "execution/execution_contexts/modelRun_ctx.h" #include "execution/command_parser.h" #include "execution/DAG/dag.h" #include "execution/DAG/dag_execute.h" #include "execution/parsing/deprecated.h" +#include "execution/parsing/tensor_commands_parsing.h" #include "execution/utils.h" #include "model_commands_parser.h" #include "script_commands_parser.h" @@ -52,8 +52,7 @@ static int _ParseDAGLoadArgs(RedisModuleCtx *ctx, RedisModuleString **argv, int const char *arg_string = RedisModule_StringPtrLen(key_name, &arg_len); RAI_Tensor *t; RedisModuleKey *key; - const int status = - RAI_GetTensorFromKeyspace(ctx, key_name, &key, &t, REDISMODULE_READ, err); + int status = RAI_TensorGetFromKeyspace(ctx, key_name, &key, &t, REDISMODULE_READ, err); if (status == REDISMODULE_ERR) { RedisModule_Log(ctx, "warning", "Could not LOAD tensor %s from keyspace into DAG", arg_string); @@ -172,24 +171,24 @@ int ParseDAGExecuteOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops, bool ro) { if (!strcasecmp(arg_string, "AI.TENSORGET")) { currentOp->commandType = REDISAI_DAG_CMD_TENSORGET; - currentOp->devicestr = "CPU"; - RAI_HoldString(currentOp->argv[1]); - currentOp->inkeys = array_append(currentOp->inkeys, currentOp->argv[1]); - currentOp->fmt = ParseTensorGetArgs(rinfo->err, currentOp->argv, currentOp->argc); + currentOp->fmt = ParseTensorGetFormat(rinfo->err, currentOp->argv, currentOp->argc); if (currentOp->fmt == TENSOR_NONE) { return REDISMODULE_ERR; } + currentOp->devicestr = "CPU"; + RAI_HoldString(currentOp->argv[1]); + currentOp->inkeys = array_append(currentOp->inkeys, currentOp->argv[1]); continue; } if (!strcasecmp(arg_string, "AI.TENSORSET")) { currentOp->commandType = REDISAI_DAG_CMD_TENSORSET; + if (ParseTensorSetArgs(currentOp->argv, currentOp->argc, ¤tOp->outTensor, + rinfo->err) != REDISMODULE_OK) { + return REDISMODULE_ERR; + } currentOp->devicestr = "CPU"; RAI_HoldString(currentOp->argv[1]); currentOp->outkeys = array_append(currentOp->outkeys, currentOp->argv[1]); - if (RAI_parseTensorSetArgs(currentOp->argv, currentOp->argc, ¤tOp->outTensor, 0, - rinfo->err) == -1) { - return REDISMODULE_ERR; - } currentOp->result = REDISMODULE_OK; continue; } diff --git a/src/execution/parsing/deprecated.c b/src/execution/parsing/deprecated.c index f09d19e8e..6817ce589 100644 --- a/src/execution/parsing/deprecated.c +++ b/src/execution/parsing/deprecated.c @@ -10,6 +10,7 @@ #include "execution/DAG/dag_execute.h" #include "execution/parsing/dag_parser.h" #include "execution/parsing/parse_utils.h" +#include "execution/parsing/tensor_commands_parsing.h" #include "execution/execution_contexts/modelRun_ctx.h" #include "execution/execution_contexts/scriptRun_ctx.h" @@ -586,22 +587,22 @@ int ParseDAGRunOps(RedisAI_RunInfo *rinfo, RAI_DagOp **ops) { if (!strcasecmp(arg_string, "AI.TENSORGET")) { currentOp->commandType = REDISAI_DAG_CMD_TENSORGET; currentOp->devicestr = "CPU"; - RAI_HoldString(currentOp->argv[1]); - currentOp->inkeys = array_append(currentOp->inkeys, currentOp->argv[1]); - currentOp->fmt = ParseTensorGetArgs(rinfo->err, currentOp->argv, currentOp->argc); + currentOp->fmt = ParseTensorGetFormat(rinfo->err, currentOp->argv, currentOp->argc); if (currentOp->fmt == TENSOR_NONE) goto cleanup; + RAI_HoldString(currentOp->argv[1]); + currentOp->inkeys = array_append(currentOp->inkeys, currentOp->argv[1]); continue; } if (!strcasecmp(arg_string, "AI.TENSORSET")) { currentOp->commandType = REDISAI_DAG_CMD_TENSORSET; currentOp->devicestr = "CPU"; - RAI_HoldString(currentOp->argv[1]); - currentOp->outkeys = array_append(currentOp->outkeys, currentOp->argv[1]); - if (RAI_parseTensorSetArgs(currentOp->argv, currentOp->argc, ¤tOp->outTensor, 0, - rinfo->err) == -1) { + if (ParseTensorSetArgs(currentOp->argv, currentOp->argc, ¤tOp->outTensor, + rinfo->err) != REDISMODULE_OK) { goto cleanup; } + RAI_HoldString(currentOp->argv[1]); + currentOp->outkeys = array_append(currentOp->outkeys, currentOp->argv[1]); currentOp->result = REDISMODULE_OK; continue; } diff --git a/src/execution/parsing/tensor_commands_parsing.c b/src/execution/parsing/tensor_commands_parsing.c new file mode 100644 index 000000000..ca634d878 --- /dev/null +++ b/src/execution/parsing/tensor_commands_parsing.c @@ -0,0 +1,119 @@ +#include +#include "tensor_commands_parsing.h" +#include "redis_ai_objects/tensor.h" +#include "util/arr.h" + +int ParseTensorSetArgs(RedisModuleString **argv, int argc, RAI_Tensor **t, RAI_Error *error) { + if (argc < 4) { + RAI_SetError(error, RAI_ETENSORSET, "wrong number of arguments for 'AI.TENSORSET' command"); + return REDISMODULE_ERR; + } + + // get the tensor data type + const char *type_str = RedisModule_StringPtrLen(argv[2], NULL); + DLDataType data_type = RAI_TensorDataTypeFromString(type_str); + if (data_type.bits == 0) { + RAI_SetError(error, RAI_ETENSORSET, "ERR invalid data type"); + return REDISMODULE_ERR; + } + + int data_fmt = TENSOR_NONE; + int n_dims = 0; + long long tensor_len = 1; + size_t *dims = array_new(size_t, 1); + int arg_pos = 3; + + // go over remaining tensor args (shapes and data) after parsing its type. + for (; arg_pos < argc; arg_pos++) { + const char *opt = RedisModule_StringPtrLen(argv[arg_pos], NULL); + if (!strcasecmp(opt, "BLOB")) { + data_fmt = TENSOR_BLOB; + // if we've found the data format, then there are no more dimensions + // check right away if the arity is correct + size_t remaining_args = argc - 1 - arg_pos; + if (remaining_args != 1) { + array_free(dims); + RAI_SetError(error, RAI_ETENSORSET, + "ERR a single binary string should come after the BLOB argument in " + "'AI.TENSORSET' command"); + return REDISMODULE_ERR; + } + arg_pos++; + break; + } else if (!strcasecmp(opt, "VALUES")) { + data_fmt = TENSOR_VALUES; + // if we've found the data_format, then there are no more dimensions + // check right away if the arity is correct + size_t remaining_args = argc - 1 - arg_pos; + if (remaining_args != tensor_len) { + array_free(dims); + RAI_SetError(error, RAI_ETENSORSET, + "ERR wrong number of values was given in 'AI.TENSORSET' command"); + return REDISMODULE_ERR; + } + arg_pos++; + break; + } else { + // Otherwise, parse the next tensor shape and append it to its dims. + long long dimension; + int ret_val = RedisModule_StringToLongLong(argv[arg_pos], &dimension); + if (ret_val != REDISMODULE_OK || dimension <= 0) { + array_free(dims); + RAI_SetError(error, RAI_ETENSORSET, + "ERR invalid or negative value found in tensor shape"); + return REDISMODULE_ERR; + } + n_dims++; + dims = array_append(dims, dimension); + tensor_len *= dimension; + } + } + + if (data_fmt == TENSOR_BLOB) { + RedisModuleString *tensor_blob_rs = argv[arg_pos]; + size_t blob_len; + const char *tensor_blob = RedisModule_StringPtrLen(tensor_blob_rs, &blob_len); + *t = RAI_TensorCreateFromBlob(data_type, dims, n_dims, tensor_blob, blob_len, error); + } else { + // Parse the rest of the arguments (tensor values) and set the values in the tensor. + // Note that it is possible that no values were given - create empty tensor in that case. + *t = RAI_TensorCreateFromValues(data_type, dims, n_dims, argc - arg_pos, &argv[arg_pos], + error); + } + array_free(dims); + + if (*t == NULL) { + return REDISMODULE_ERR; + } + return REDISMODULE_OK; +} + +uint ParseTensorGetFormat(RAI_Error *err, RedisModuleString **argv, int argc) { + uint fmt = TENSOR_NONE; + if (argc < 2 || argc > 4) { + RAI_SetError(err, RAI_EDAGBUILDER, "wrong number of arguments for 'AI.TENSORGET' command"); + return fmt; + } + if (argc == 2) { + return TENSOR_BLOB | TENSOR_META; + } + for (int i = 2; i < argc; i++) { + const char *fmtstr = RedisModule_StringPtrLen(argv[i], NULL); + if (!strcasecmp(fmtstr, "BLOB")) { + fmt |= TENSOR_BLOB; + } else if (!strcasecmp(fmtstr, "VALUES")) { + fmt |= TENSOR_VALUES; + } else if (!strcasecmp(fmtstr, "META")) { + fmt |= TENSOR_META; + } else { + RAI_SetError(err, RAI_EDAGBUILDER, "ERR unsupported data format"); + return TENSOR_NONE; + } + } + + if (fmt == TENSOR_ILLEGAL_VALUES_BLOB) { + RAI_SetError(err, RAI_EDAGBUILDER, "ERR both BLOB and VALUES specified"); + return TENSOR_NONE; + } + return fmt; +} diff --git a/src/execution/parsing/tensor_commands_parsing.h b/src/execution/parsing/tensor_commands_parsing.h new file mode 100644 index 000000000..2d0790be7 --- /dev/null +++ b/src/execution/parsing/tensor_commands_parsing.h @@ -0,0 +1,29 @@ +#pragma once + +#include "redismodule.h" +#include "redis_ai_objects/err.h" +#include "redis_ai_objects/tensor.h" + +/** + * Helper method to parse AI.TENSORGET arguments + * + * @param argv Redis command arguments, as an array of strings + * @param argc Redis command number of arguments + * @param t Destination tensor to store the parsed data + * @param error error data structure to store error message in the case of + * parsing failures + * @return REDISMODULE_OK on success, or REDISMODULE_ERR if the parsing failed + */ +int ParseTensorSetArgs(RedisModuleString **argv, int argc, RAI_Tensor **t, RAI_Error *error); + +/** + * Helper method to parse AI.TENSORGET arguments + * + * @param error error data structure to store error message in the case of + * parsing failures + * @param argv Redis command arguments, as an array of strings + * @param argc Redis command number of arguments + * @return The format in which tensor is returned. + */ + +uint ParseTensorGetFormat(RAI_Error *error, RedisModuleString **argv, int argc); diff --git a/src/redis_ai_objects/tensor.c b/src/redis_ai_objects/tensor.c index 8a8643223..a8d066ded 100644 --- a/src/redis_ai_objects/tensor.c +++ b/src/redis_ai_objects/tensor.c @@ -24,8 +24,12 @@ extern RedisModuleType *RedisAI_TensorType; +RedisModuleType *RAI_TensorRedisType(void) { return RedisAI_TensorType; } + +//************************************** static helper functions ******************* + // Check if the given value is in the range of the tensor type. -bool _ValOverflow(long long val, RAI_Tensor *t) { +static bool _ValOverflow(long long val, RAI_Tensor *t) { DLDataType dtype = t->tensor.dl_tensor.dtype; if (dtype.code == kDLInt) { unsigned long long max_abs_val = ((unsigned long long)1 << (uint)(dtype.bits - 1)); @@ -45,335 +49,473 @@ bool _ValOverflow(long long val, RAI_Tensor *t) { return false; } -DLDataType RAI_TensorDataTypeFromString(const char *typestr) { - if (strcasecmp(typestr, RAI_DATATYPE_STR_FLOAT) == 0) { - return (DLDataType){.code = kDLFloat, .bits = 32, .lanes = 1}; - } - if (strcasecmp(typestr, RAI_DATATYPE_STR_DOUBLE) == 0) { - return (DLDataType){.code = kDLFloat, .bits = 64, .lanes = 1}; - } - if (strncasecmp(typestr, "INT", 3) == 0) { - const char *bitstr = typestr + 3; - if (strcmp(bitstr, "8") == 0) { - return (DLDataType){.code = kDLInt, .bits = 8, .lanes = 1}; - } - if (strcmp(bitstr, "16") == 0) { - return (DLDataType){.code = kDLInt, .bits = 16, .lanes = 1}; - } - if (strcmp(bitstr, "32") == 0) { - return (DLDataType){.code = kDLInt, .bits = 32, .lanes = 1}; +static int _RAI_TensorParseStringValues(int argc, RedisModuleString **argv, RAI_Tensor *tensor, + RAI_Error *err) { + size_t total_len = 0; + const char *strings_data[argc]; + size_t strings_lengths[argc]; + uint64_t *strings_offsets = RAI_TensorStringElementsOffsets(tensor); + + // go over the strings and save the offset every string element in the blob. + for (int i = 0; i < argc; i++) { + strings_offsets[i] = total_len; + size_t str_len; + const char *str_data = RedisModule_StringPtrLen(argv[i], &str_len); + // Validate that the given string is a C-string (doesn't contain \0 except for the end) + for (size_t j = 0; j < str_len - 1; j++) { + if (str_data[j] == '\0') { + RAI_SetError(err, RAI_ETENSORSET, "ERR C-string value contains null character"); + return REDISMODULE_ERR; + } } - if (strcmp(bitstr, "64") == 0) { - return (DLDataType){.code = kDLInt, .bits = 64, .lanes = 1}; + strings_lengths[i] = str_len; + total_len += str_len; + if (str_data[str_len - 1] != '\0') { + total_len++; } + strings_data[i] = str_data; } - if (strncasecmp(typestr, "UINT", 4) == 0) { - const char *bitstr = typestr + 4; - if (strcmp(bitstr, "8") == 0) { - return (DLDataType){.code = kDLUInt, .bits = 8, .lanes = 1}; + tensor->tensor.dl_tensor.data = RedisModule_Calloc(1, total_len); + tensor->tensor.dl_tensor.elements_length = strings_offsets; + tensor->blobSize = total_len; + + // Copy the strings one by one to the data ptr. + for (int i = 0; i < argc; i++) { + memcpy(RAI_TensorData(tensor) + strings_offsets[i], strings_data[i], strings_lengths[i]); + } + + return REDISMODULE_OK; +} + +static int _RAI_TensorFillWithValues(int argc, RedisModuleString **argv, RAI_Tensor *t, + DLDataType data_type, RAI_Error *error) { + t->blobSize = RAI_TensorLength(t) * RAI_TensorDataSize(t); + t->tensor.dl_tensor.data = RedisModule_Alloc(t->blobSize); + + for (long long i = 0; i < argc; i++) { + if (data_type.code == kDLFloat) { + double val; + int ret_val = RedisModule_StringToDouble(argv[i], &val); + if (ret_val != REDISMODULE_OK) { + RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); + return REDISMODULE_ERR; + } + int ret_set = RAI_TensorSetValueFromDouble(t, i, val); + if (ret_set != 1) { + RAI_SetError(error, RAI_ETENSORSET, "ERR cannot specify values for this data type"); + return REDISMODULE_ERR; + } + } else { + long long val; + const int ret_val = RedisModule_StringToLongLong(argv[i], &val); + if (ret_val != REDISMODULE_OK || _ValOverflow(val, t)) { + RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); + return REDISMODULE_ERR; + } + const int ret_set = RAI_TensorSetValueFromLongLong(t, i, val); + if (ret_set != 1) { + RAI_SetError(error, RAI_ETENSORSET, "ERR cannot specify values for this data type"); + return REDISMODULE_ERR; + } } - if (strcmp(bitstr, "16") == 0) { - return (DLDataType){.code = kDLUInt, .bits = 16, .lanes = 1}; + } + return REDISMODULE_OK; +} + +// This will populate the offsets array with the start position of every string element in the blob +static int _RAI_TensorParseStringsBlob(const char *tensor_blob, size_t blob_len, size_t tensor_len, + uint64_t *offsets, RAI_Error *err) { + + size_t elements_counter = 0; + offsets[elements_counter++] = 0; + + // if we encounter null-character, we set the next element offset to the next position + for (size_t i = 0; i < blob_len - 1; i++) { + if (tensor_blob[i] == '\0') { + offsets[elements_counter++] = i + 1; } } - if (strcasecmp(typestr, "BOOL") == 0) { - return (DLDataType){.code = kDLBool, .bits = 8, .lanes = 1}; + if (tensor_blob[blob_len - 1] != '\0' || elements_counter != tensor_len) { + if (err) { + RAI_SetError(err, RAI_ETENSORSET, + "ERR Number of string elements in data blob does not match tensor length"); + } + return REDISMODULE_ERR; } - return (DLDataType){.bits = 0}; + return REDISMODULE_OK; } -static size_t Tensor_DataTypeSize(DLDataType dtype) { return dtype.bits / 8; } +static int _RAI_TensorReplyWithValues(RedisModuleCtx *ctx, RAI_Tensor *t) { + int ndims = RAI_TensorNumDims(t); + long long len = 1; + for (int i = 0; i < ndims; i++) { + len *= RAI_TensorDim(t, i); + } -int Tensor_DataTypeStr(DLDataType dtype, char *dtypestr) { - int result = REDISMODULE_ERR; + DLDataType dtype = RAI_TensorDataType(t); + RedisModule_ReplyWithArray(ctx, len); - if (dtype.code == kDLFloat) { - if (dtype.bits == 32) { - strcpy(dtypestr, RAI_DATATYPE_STR_FLOAT); - result = REDISMODULE_OK; - } else if (dtype.bits == 64) { - strcpy(dtypestr, RAI_DATATYPE_STR_DOUBLE); - result = REDISMODULE_OK; + if (dtype.code == kDLString) { + for (long long i = 0; i < len; i++) { + const char *val; + int ret = RAI_TensorGetValueAsCString(t, i, &val); + if (!ret) { + RedisModule_ReplyWithError(ctx, "ERR cannot get values for this data type"); + return REDISMODULE_ERR; + } + RedisModule_ReplyWithCString(ctx, val); } - } else if (dtype.code == kDLInt) { - if (dtype.bits == 8) { - strcpy(dtypestr, RAI_DATATYPE_STR_INT8); - result = REDISMODULE_OK; - } else if (dtype.bits == 16) { - strcpy(dtypestr, RAI_DATATYPE_STR_INT16); - result = REDISMODULE_OK; - } else if (dtype.bits == 32) { - strcpy(dtypestr, RAI_DATATYPE_STR_INT32); - result = REDISMODULE_OK; - } else if (dtype.bits == 64) { - strcpy(dtypestr, RAI_DATATYPE_STR_INT64); - result = REDISMODULE_OK; + } else if (dtype.code == kDLFloat) { + double val; + for (long long i = 0; i < len; i++) { + int ret = RAI_TensorGetValueAsDouble(t, i, &val); + if (!ret) { + RedisModule_ReplyWithError(ctx, "ERR cannot get values for this data type"); + return REDISMODULE_ERR; + } + RedisModule_ReplyWithDouble(ctx, val); } - } else if (dtype.code == kDLUInt) { - if (dtype.bits == 8) { - strcpy(dtypestr, RAI_DATATYPE_STR_UINT8); - result = REDISMODULE_OK; - } else if (dtype.bits == 16) { - strcpy(dtypestr, RAI_DATATYPE_STR_UINT16); - result = REDISMODULE_OK; + } else { + long long val; + for (long long i = 0; i < len; i++) { + int ret = RAI_TensorGetValueAsLongLong(t, i, &val); + if (!ret) { + RedisModule_ReplyWithError(ctx, "ERR cannot get values for this data type"); + return REDISMODULE_ERR; + } + RedisModule_ReplyWithLongLong(ctx, val); } - } else if (dtype.code == kDLBool && dtype.bits == 8) { - strcpy(dtypestr, RAI_DATATYPE_STR_BOOL); - result = REDISMODULE_OK; } - return result; + return REDISMODULE_OK; } -RAI_Tensor *RAI_TensorNew(void) { - RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(*ret)); - ret->refCount = 1; - ret->len = LEN_UNKNOWN; - return ret; -} +//***************** methods for creating a tensor ************************************ -RAI_Tensor *RAI_TensorCreateWithDLDataType(DLDataType dtype, long long *dims, int ndims, - bool empty) { +RAI_Tensor *RAI_TensorNew(DLDataType data_type, const size_t *dims, int n_dims) { - size_t dtypeSize = Tensor_DataTypeSize(dtype); - if (dtypeSize == 0) { - return NULL; - } + RAI_Tensor *new_tensor = RedisModule_Alloc(sizeof(RAI_Tensor)); + new_tensor->refCount = 1; + new_tensor->blobSize = 0; - RAI_Tensor *ret = RAI_TensorNew(); - int64_t *shape = RedisModule_Alloc(ndims * sizeof(*shape)); - int64_t *strides = RedisModule_Alloc(ndims * sizeof(*strides)); + // Note that n_dim can be zero (i.e., tensor is a scalar) + int64_t *shape = RedisModule_Calloc(n_dims, sizeof(*shape)); + int64_t *strides = NULL; + uint64_t *offsets = NULL; - size_t len = 1; - for (int64_t i = 0; i < ndims; ++i) { - shape[i] = dims[i]; - strides[i] = 1; - len *= dims[i]; - } - for (int64_t i = ndims - 2; i >= 0; --i) { - strides[i] *= strides[i + 1] * shape[i + 1]; + size_t tensor_len = 1; + for (int64_t i = 0; i < n_dims; ++i) { + shape[i] = (int64_t)dims[i]; + tensor_len *= dims[i]; } - // Default device is CPU (id -1 is default, means 'no index') - DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = -1}; + new_tensor->len = tensor_len; - // If we return an empty tensor, we initialize the data with zeros to avoid security - // issues. Otherwise, we only allocate without initializing (for better performance). - void *data; - if (empty) { - data = RedisModule_Calloc(len, dtypeSize); + if (data_type.code != kDLString) { + strides = RedisModule_Calloc(n_dims, sizeof(*strides)); + if (n_dims > 0) { + strides[n_dims - 1] = 1; + } + for (int i = n_dims - 2; i >= 0; --i) { + strides[i] = strides[i + 1] * shape[i + 1]; + } } else { - data = RedisModule_Alloc(len * dtypeSize); + offsets = RedisModule_Alloc(tensor_len * sizeof(*offsets)); } - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, - .data = data, - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = 0}, - .manager_ctx = NULL, - .deleter = NULL}; + // Default device is CPU (id -1 is default, means 'no index') + DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = -1}; - return ret; + new_tensor->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, + .data = NULL, + .ndim = n_dims, + .dtype = data_type, + .shape = shape, + .strides = strides, + .byte_offset = 0, + .elements_length = offsets}, + .manager_ctx = NULL, + .deleter = NULL}; + + return new_tensor; } -void RAI_RStringDataTensorDeleter(DLManagedTensor *arg) { - if (arg->dl_tensor.shape) { - RedisModule_Free(arg->dl_tensor.shape); - } - if (arg->dl_tensor.strides) { - RedisModule_Free(arg->dl_tensor.strides); +RAI_Tensor *RAI_TensorCreateFromValues(DLDataType data_type, const size_t *dims, int n_dims, + int argc, RedisModuleString **argv, RAI_Error *err) { + + RAI_Tensor *new_tensor = RAI_TensorNew(data_type, dims, n_dims); + size_t tensor_len = RAI_TensorLength(new_tensor); + + // If no values were given, we consider the empty tensor as a sequence of zeros + // (null-character in case of string tensor) + if (argc == 0) { + new_tensor->blobSize = RAI_TensorDataSize(new_tensor) * tensor_len; + new_tensor->tensor.dl_tensor.data = RedisModule_Calloc(1, new_tensor->blobSize); + if (data_type.code == kDLString) { + uint64_t *strings_offsets = RAI_TensorStringElementsOffsets(new_tensor); + for (size_t i = 0; i < tensor_len; i++) { + strings_offsets[i] = i; + } + } + return new_tensor; } - if (arg->manager_ctx) { - RedisModuleString *rstr = (RedisModuleString *)arg->manager_ctx; - RedisModule_FreeString(NULL, rstr); + + if (data_type.code == kDLString) { + // Allocate and populate the data blob and save every string element's offset within it. + if (_RAI_TensorParseStringValues(argc, argv, new_tensor, err) != REDISMODULE_OK) { + RAI_TensorFree(new_tensor); + return NULL; + } + } else { + // Parse and fill tensor data with numeric values + if (_RAI_TensorFillWithValues(argc, argv, new_tensor, data_type, err) != REDISMODULE_OK) { + RAI_TensorFree(new_tensor); + return NULL; + } } - RedisModule_Free(arg->dl_tensor.data); - RedisModule_Free(arg); + return new_tensor; } -RAI_Tensor *_TensorCreateWithDLDataTypeAndRString(DLDataType dtype, size_t dtypeSize, - long long *dims, int ndims, - RedisModuleString *rstr, RAI_Error *err) { +RAI_Tensor *RAI_TensorCreateFromBlob(DLDataType data_type, const size_t *dims, int n_dims, + const char *tensor_blob, size_t blob_len, RAI_Error *err) { + + RAI_Tensor *new_tensor = RAI_TensorNew(data_type, dims, n_dims); + size_t tensor_len = RAI_TensorLength(new_tensor); + size_t data_type_size = RAI_TensorDataSize(new_tensor); + + new_tensor->blobSize = blob_len; + new_tensor->tensor.dl_tensor.data = RedisModule_Alloc(blob_len); + if (data_type.code == kDLString) { + // find and save the offset of every individual string in the blob. Set an error if + // the number of strings in the given blob doesn't match the tensor length. + uint64_t *offsets = RAI_TensorStringElementsOffsets(new_tensor); + if (_RAI_TensorParseStringsBlob(tensor_blob, blob_len, tensor_len, offsets, err) != + REDISMODULE_OK) { + RAI_TensorFree(new_tensor); + return NULL; + } + } else { + size_t expected_n_bytes = tensor_len * data_type_size; + if (blob_len != expected_n_bytes) { + RAI_TensorFree(new_tensor); + RAI_SetError(err, RAI_ETENSORSET, + "ERR data length does not match tensor shape and type"); + return NULL; + } + } - int64_t *shape = RedisModule_Alloc(ndims * sizeof(*shape)); - int64_t *strides = RedisModule_Alloc(ndims * sizeof(*strides)); + // Copy the blob. We must copy instead of increasing the ref count since we don't have + // a way of taking ownership on the underline data pointer (this will require introducing + // a designated RedisModule API, might be optional in the future) + memcpy(RAI_TensorData(new_tensor), tensor_blob, blob_len); + return new_tensor; +} - size_t len = 1; - for (int64_t i = 0; i < ndims; ++i) { - shape[i] = dims[i]; - strides[i] = 1; - len *= dims[i]; - } - for (int64_t i = ndims - 2; i >= 0; --i) { - strides[i] *= strides[i + 1] * shape[i + 1]; - } - // Default device is CPU (id -1 is default, means 'no index') - DLDevice device = (DLDevice){.device_type = kDLCPU, .device_id = -1}; - size_t nbytes = len * dtypeSize; - - size_t blob_len; - const char *blob = RedisModule_StringPtrLen(rstr, &blob_len); - if (blob_len != nbytes) { - RedisModule_Free(shape); - RedisModule_Free(strides); - RAI_SetError(err, RAI_ETENSORSET, "ERR data length does not match tensor shape and type"); +RAI_Tensor *RAI_TensorCreate(const char *data_type_str, const long long *dims, int n_dims) { + DLDataType data_type = RAI_TensorDataTypeFromString(data_type_str); + if (data_type.bits == 0) { return NULL; } - char *data = RedisModule_Alloc(nbytes); - memcpy(data, blob, nbytes); - RAI_HoldString(rstr); - - RAI_Tensor *ret = RAI_TensorNew(); - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, - .data = data, - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = 0}, - .manager_ctx = rstr, - .deleter = RAI_RStringDataTensorDeleter}; - - return ret; -} -// Important note: the tensor data must be initialized after the creation. -RAI_Tensor *RAI_TensorCreate(const char *dataType, long long *dims, int ndims) { - DLDataType dtype = RAI_TensorDataTypeFromString(dataType); - return RAI_TensorCreateWithDLDataType(dtype, dims, ndims, false); + // return an empty tensor (with empty list of values). + return RAI_TensorCreateFromValues(data_type, (const size_t *)dims, n_dims, 0, NULL, NULL); } -RAI_Tensor *RAI_TensorCreateByConcatenatingTensors(RAI_Tensor **ts, long long n) { - - if (n == 0) { - return NULL; - } +RAI_Tensor *RAI_TensorCreateByConcatenatingTensors(RAI_Tensor **tensors, long long n) { + RedisModule_Assert(n > 0); long long total_batch_size = 0; long long batch_sizes[n]; - long long batch_offsets[n]; - - const long long ndims = RAI_TensorNumDims(ts[0]); - long long dims[ndims]; + size_t batch_offsets[n]; - // TODO check that all tensors have compatible dims + int n_dims = RAI_TensorNumDims(tensors[0]); + size_t dims[n_dims]; + // TODO: check that all tensors have compatible dims + // Calculate the total batch size after concatenation, this will be the first dim. for (long long i = 0; i < n; i++) { - batch_sizes[i] = RAI_TensorDim(ts[i], 0); + batch_sizes[i] = RAI_TensorDim(tensors[i], 0); total_batch_size += batch_sizes[i]; } - - batch_offsets[0] = 0; - for (long long i = 1; i < n; i++) { - batch_offsets[i] = batch_offsets[i - 1] + batch_sizes[i - 1]; - } - - long long sample_size = 1; - - for (long long i = 1; i < ndims; i++) { - dims[i] = RAI_TensorDim(ts[0], i); - sample_size *= dims[i]; - } dims[0] = total_batch_size; - const long long dtype_size = RAI_TensorDataSize(ts[0]); - - DLDataType dtype = RAI_TensorDataType(ts[0]); + // Get the rest of the tensor's dimensions. + for (int i = 1; i < n_dims; i++) { + dims[i] = RAI_TensorDim(tensors[0], i); + } - RAI_Tensor *ret = RAI_TensorCreateWithDLDataType(dtype, dims, ndims, false); + // Create a new tensor to store the concatenated tensor in it. + DLDataType data_type = RAI_TensorDataType(tensors[0]); + RAI_Tensor *ret = RAI_TensorNew(data_type, dims, n_dims); - for (long long i = 0; i < n; i++) { - memcpy(RAI_TensorData(ret) + batch_offsets[i] * sample_size * dtype_size, - RAI_TensorData(ts[i]), RAI_TensorByteSize(ts[i])); + // get the offsets for every tensor's blob in the batched tensor blob + batch_offsets[0] = 0; + for (long long i = 1; i < n; i++) { + batch_offsets[i] = batch_offsets[i - 1] + RAI_TensorByteSize(tensors[i - 1]); + } + ret->blobSize = batch_offsets[n - 1] + RAI_TensorByteSize(tensors[n - 1]); + ret->tensor.dl_tensor.data = RedisModule_Alloc(ret->blobSize); + + // Copy the input tensors data to the new tensor. + size_t element_ind = 0; + uint64_t *ret_offsets = RAI_TensorStringElementsOffsets(ret); + for (size_t i = 0; i < n; i++) { + memcpy(RAI_TensorData(ret) + batch_offsets[i], RAI_TensorData(tensors[i]), + RAI_TensorByteSize(tensors[i])); + if (data_type.code == kDLString) { + uint64_t *curr_tensor_offsets = RAI_TensorStringElementsOffsets(tensors[i]); + for (size_t j = 0; j < RAI_TensorLength(tensors[i]); j++) { + ret_offsets[element_ind++] = curr_tensor_offsets[j] + batch_offsets[i]; + } + } } - return ret; } RAI_Tensor *RAI_TensorCreateBySlicingTensor(RAI_Tensor *t, long long offset, long long len) { - const long long ndims = RAI_TensorNumDims(t); - long long dims[ndims]; + RedisModule_Assert(offset >= 0 || len >= 0 || offset + len <= RAI_TensorLength(t)); + int n_dims = RAI_TensorNumDims(t); + size_t dims[n_dims]; - const long long dtype_size = RAI_TensorDataSize(t); - long long sample_size = 1; - - for (long long i = 1; i < ndims; i++) { + size_t basic_tensor_len = 1; // the len of every "non-batched" tensor + for (int i = 1; i < n_dims; i++) { dims[i] = RAI_TensorDim(t, i); - sample_size *= dims[i]; + basic_tensor_len *= dims[i]; } - dims[0] = len; - DLDataType dtype = RAI_TensorDataType(t); - - RAI_Tensor *ret = RAI_TensorCreateWithDLDataType(dtype, dims, ndims, false); - - memcpy(RAI_TensorData(ret), RAI_TensorData(t) + offset * sample_size * dtype_size, - len * sample_size * dtype_size); - + // Create a new tensor to store the sliced tensor in it. + DLDataType data_type = RAI_TensorDataType(t); + size_t data_type_size = RAI_TensorDataSize(t); + RAI_Tensor *ret = RAI_TensorNew(data_type, dims, n_dims); + + // Copy the input tensor sliced data to the new tensor. + if (data_type.code == kDLString) { + size_t first_element_pos_in_slice = offset * basic_tensor_len; + size_t first_element_pos_after_slice = (offset + len) * basic_tensor_len; + uint64_t *strings_offsets = RAI_TensorStringElementsOffsets(t); + size_t blob_size = + first_element_pos_after_slice < RAI_TensorLength(t) + ? strings_offsets[first_element_pos_after_slice] - + strings_offsets[first_element_pos_in_slice] + : RAI_TensorByteSize(t) - strings_offsets[first_element_pos_in_slice]; + ret->blobSize = blob_size; + ret->tensor.dl_tensor.data = RedisModule_Alloc(blob_size); + memcpy(RAI_TensorData(ret), RAI_TensorData(t) + strings_offsets[offset * basic_tensor_len], + blob_size); + memcpy(RAI_TensorStringElementsOffsets(ret), strings_offsets + offset, + len * sizeof(uint64_t)); + } else { + size_t blob_size = len * basic_tensor_len * data_type_size; + ret->tensor.dl_tensor.data = RedisModule_Alloc(blob_size); + memcpy(RAI_TensorData(ret), RAI_TensorData(t) + offset * basic_tensor_len * data_type_size, + blob_size); + ret->blobSize = blob_size; + } return ret; } -/** - * Allocate the memory for a new Tensor and copy data fom a tensor to it. - * @param t Source tensor to copy. - * @param result Destination tensor to copy. - * @return 0 on success, or 1 if the copy failed - * failed. - */ -int RAI_TensorDeepCopy(RAI_Tensor *t, RAI_Tensor **dest) { - const long long ndims = RAI_TensorNumDims(t); - long long dims[ndims]; - - const long long dtype_size = RAI_TensorDataSize(t); - long long sample_size = 1; - - for (size_t i = 0; i < ndims; i++) { - dims[i] = RAI_TensorDim(t, i); - sample_size *= dims[i]; +DLDataType RAI_TensorDataTypeFromString(const char *type_str) { + if (strcasecmp(type_str, RAI_DATATYPE_STR_FLOAT) == 0) { + return (DLDataType){.code = kDLFloat, .bits = 32, .lanes = 1}; } - - DLDataType dtype = RAI_TensorDataType(t); - - RAI_Tensor *ret = RAI_TensorCreateWithDLDataType(dtype, dims, ndims, false); - - memcpy(RAI_TensorData(ret), RAI_TensorData(t), sample_size * dtype_size); - *dest = ret; - return 0; + if (strcasecmp(type_str, RAI_DATATYPE_STR_DOUBLE) == 0) { + return (DLDataType){.code = kDLFloat, .bits = 64, .lanes = 1}; + } + if (strncasecmp(type_str, "INT", 3) == 0) { + const char *bit_str = type_str + 3; + if (strcmp(bit_str, "8") == 0) { + return (DLDataType){.code = kDLInt, .bits = 8, .lanes = 1}; + } + if (strcmp(bit_str, "16") == 0) { + return (DLDataType){.code = kDLInt, .bits = 16, .lanes = 1}; + } + if (strcmp(bit_str, "32") == 0) { + return (DLDataType){.code = kDLInt, .bits = 32, .lanes = 1}; + } + if (strcmp(bit_str, "64") == 0) { + return (DLDataType){.code = kDLInt, .bits = 64, .lanes = 1}; + } + } + if (strncasecmp(type_str, "UINT", 4) == 0) { + const char *bit_str = type_str + 4; + if (strcmp(bit_str, "8") == 0) { + return (DLDataType){.code = kDLUInt, .bits = 8, .lanes = 1}; + } + if (strcmp(bit_str, "16") == 0) { + return (DLDataType){.code = kDLUInt, .bits = 16, .lanes = 1}; + } + } + if (strcasecmp(type_str, "BOOL") == 0) { + return (DLDataType){.code = kDLBool, .bits = 8, .lanes = 1}; + } + if (strcasecmp(type_str, "STRING") == 0) { + return (DLDataType){.code = kDLString, .bits = 8, .lanes = 1}; + } + // Invalid data type + return (DLDataType){.bits = 0}; } -// Beware: this will take ownership of dltensor. RAI_Tensor *RAI_TensorCreateFromDLTensor(DLManagedTensor *dl_tensor) { - RAI_Tensor *ret = RAI_TensorNew(); - ret->tensor = - (DLManagedTensor){.dl_tensor = (DLTensor){.device = dl_tensor->dl_tensor.device, - .data = dl_tensor->dl_tensor.data, - .ndim = dl_tensor->dl_tensor.ndim, - .dtype = dl_tensor->dl_tensor.dtype, - .shape = dl_tensor->dl_tensor.shape, - .strides = dl_tensor->dl_tensor.strides, - .byte_offset = dl_tensor->dl_tensor.byte_offset}, - .manager_ctx = dl_tensor->manager_ctx, - .deleter = dl_tensor->deleter}; - RAI_TensorLength(ret); // This will set the ret->len field. + RAI_Tensor *ret = RedisModule_Calloc(1, sizeof(RAI_Tensor)); + ret->refCount = 1; + ret->tensor = *dl_tensor; // shallow copy, takes ownership on the dl_tensor memory. + ret->len = RAI_TensorLength(ret); // compute and set the length + ret->blobSize = RAI_TensorByteSize(ret); return ret; } -DLTensor *RAI_TensorGetDLTensor(RAI_Tensor *tensor) { return &tensor->tensor.dl_tensor; } +//******************** getters and setters **************************** -DLDataType RAI_TensorDataType(RAI_Tensor *t) { return t->tensor.dl_tensor.dtype; } +int RAI_TensorGetDataTypeStr(DLDataType data_type, char *data_type_str) { + int result = REDISMODULE_ERR; -int RAI_TensorIsDataTypeEqual(RAI_Tensor *t1, RAI_Tensor *t2) { - return t1->tensor.dl_tensor.dtype.bits == t2->tensor.dl_tensor.dtype.bits && - t1->tensor.dl_tensor.dtype.code == t2->tensor.dl_tensor.dtype.code && - t1->tensor.dl_tensor.dtype.lanes == t2->tensor.dl_tensor.dtype.lanes; + if (data_type.code == kDLFloat) { + if (data_type.bits == 32) { + strcpy(data_type_str, RAI_DATATYPE_STR_FLOAT); + result = REDISMODULE_OK; + } else if (data_type.bits == 64) { + strcpy(data_type_str, RAI_DATATYPE_STR_DOUBLE); + result = REDISMODULE_OK; + } + } else if (data_type.code == kDLInt) { + if (data_type.bits == 8) { + strcpy(data_type_str, RAI_DATATYPE_STR_INT8); + result = REDISMODULE_OK; + } else if (data_type.bits == 16) { + strcpy(data_type_str, RAI_DATATYPE_STR_INT16); + result = REDISMODULE_OK; + } else if (data_type.bits == 32) { + strcpy(data_type_str, RAI_DATATYPE_STR_INT32); + result = REDISMODULE_OK; + } else if (data_type.bits == 64) { + strcpy(data_type_str, RAI_DATATYPE_STR_INT64); + result = REDISMODULE_OK; + } + } else if (data_type.code == kDLUInt) { + if (data_type.bits == 8) { + strcpy(data_type_str, RAI_DATATYPE_STR_UINT8); + result = REDISMODULE_OK; + } else if (data_type.bits == 16) { + strcpy(data_type_str, RAI_DATATYPE_STR_UINT16); + result = REDISMODULE_OK; + } + } else if (data_type.code == kDLBool && data_type.bits == 8) { + strcpy(data_type_str, RAI_DATATYPE_STR_BOOL); + result = REDISMODULE_OK; + } else if (data_type.code == kDLString) { + strcpy(data_type_str, RAI_DATATYPE_STR_STRING); + result = REDISMODULE_OK; + } + return result; } +DLTensor *RAI_TensorGetDLTensor(RAI_Tensor *tensor) { return &tensor->tensor.dl_tensor; } + +DLDataType RAI_TensorDataType(RAI_Tensor *t) { return t->tensor.dl_tensor.dtype; } + size_t RAI_TensorLength(RAI_Tensor *t) { - if (t->len == LEN_UNKNOWN) { + if (t->len == 0) { int64_t *shape = t->tensor.dl_tensor.shape; size_t len = 1; for (size_t i = 0; i < t->tensor.dl_tensor.ndim; ++i) { @@ -384,58 +526,78 @@ size_t RAI_TensorLength(RAI_Tensor *t) { return t->len; } -size_t RAI_TensorDataSize(RAI_Tensor *t) { return Tensor_DataTypeSize(RAI_TensorDataType(t)); } +size_t RAI_TensorDataSize(RAI_Tensor *t) { return RAI_TensorDataType(t).bits / 8; } -size_t RAI_TensorDataSizeFromString(const char *dataTypeStr) { - DLDataType dtype = RAI_TensorDataTypeFromString(dataTypeStr); - return Tensor_DataTypeSize(dtype); -} +int RAI_TensorNumDims(RAI_Tensor *t) { return t->tensor.dl_tensor.ndim; } -size_t RAI_TensorDataSizeFromDLDataType(DLDataType dtype) { return Tensor_DataTypeSize(dtype); } +long long RAI_TensorDim(RAI_Tensor *t, int i) { return t->tensor.dl_tensor.shape[i]; } -void RAI_TensorFree(RAI_Tensor *t) { - if (t) { - if (__atomic_sub_fetch(&t->refCount, 1, __ATOMIC_RELAXED) <= 0) { - if (t->tensor.deleter) { - t->tensor.deleter(&t->tensor); - } else { - if (t->tensor.dl_tensor.shape) { - RedisModule_Free(t->tensor.dl_tensor.shape); - } - if (t->tensor.dl_tensor.strides) { - RedisModule_Free(t->tensor.dl_tensor.strides); - } - if (t->tensor.dl_tensor.data) { - RedisModule_Free(t->tensor.dl_tensor.data); - } - RedisModule_Free(t); - } - } +size_t RAI_TensorByteSize(RAI_Tensor *t) { + + if (t->tensor.dl_tensor.dtype.code == kDLString) { + return t->blobSize; + } + if (t->blobSize == 0) { + t->blobSize = RAI_TensorLength(t) * RAI_TensorDataSize(t); } + return t->blobSize; } -int RAI_TensorSetData(RAI_Tensor *t, const char *data, size_t len) { - memcpy(t->tensor.dl_tensor.data, data, len); +char *RAI_TensorData(RAI_Tensor *t) { return t->tensor.dl_tensor.data; } + +uint64_t *RAI_TensorStringElementsOffsets(RAI_Tensor *tensor) { + return tensor->tensor.dl_tensor.elements_length; +} + +int64_t *RAI_TensorShape(RAI_Tensor *tensor) { return tensor->tensor.dl_tensor.shape; } + +int RAI_TensorGetValueAsDouble(RAI_Tensor *t, long long i, double *val) { + // Validate that i is in bound + if (i < 0 || i > RAI_TensorLength(t)) { + return 0; + } + DLDataType dtype = RAI_TensorDataType(t); + void *data = RAI_TensorData(t); + + if (dtype.code == kDLFloat) { + switch (dtype.bits) { + case 32: + *val = ((float *)data)[i]; + break; + case 64: + *val = ((double *)data)[i]; + break; + default: + return 0; + } + } else { + return 0; + } return 1; } -int RAI_TensorSetValueFromLongLong(RAI_Tensor *t, long long i, long long val) { - DLDataType dtype = t->tensor.dl_tensor.dtype; - void *data = t->tensor.dl_tensor.data; +int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { + + // Validate that i is in bound + if (i < 0 || i > RAI_TensorLength(t)) { + return 0; + } + DLDataType dtype = RAI_TensorDataType(t); + void *data = RAI_TensorData(t); if (dtype.code == kDLInt) { switch (dtype.bits) { case 8: - ((int8_t *)data)[i] = val; + *val = ((int8_t *)data)[i]; break; case 16: - ((int16_t *)data)[i] = val; + *val = ((int16_t *)data)[i]; break; case 32: - ((int32_t *)data)[i] = val; + *val = ((int32_t *)data)[i]; break; case 64: - ((int64_t *)data)[i] = val; + *val = ((int64_t *)data)[i]; break; default: return 0; @@ -443,23 +605,23 @@ int RAI_TensorSetValueFromLongLong(RAI_Tensor *t, long long i, long long val) { } else if (dtype.code == kDLUInt) { switch (dtype.bits) { case 8: - ((uint8_t *)data)[i] = val; + *val = ((uint8_t *)data)[i]; break; case 16: - ((uint16_t *)data)[i] = val; + *val = ((uint16_t *)data)[i]; break; case 32: - ((uint32_t *)data)[i] = val; + *val = ((uint32_t *)data)[i]; break; case 64: - ((uint64_t *)data)[i] = val; + *val = ((uint64_t *)data)[i]; break; default: return 0; } } else if (dtype.code == kDLBool) { if (dtype.bits == 8) { - ((uint8_t *)data)[i] = val; + *val = ((uint8_t *)data)[i]; } else { return 0; } @@ -467,68 +629,56 @@ int RAI_TensorSetValueFromLongLong(RAI_Tensor *t, long long i, long long val) { return 1; } -int RAI_TensorSetValueFromDouble(RAI_Tensor *t, long long i, double val) { - DLDataType dtype = t->tensor.dl_tensor.dtype; - void *data = t->tensor.dl_tensor.data; +int RAI_TensorGetValueAsCString(RAI_Tensor *t, long long i, const char **val) { + // Validate that i is in bound + if (i < 0 || i > RAI_TensorLength(t)) { + return 0; + } + DLDataType dtype = RAI_TensorDataType(t); + char *data = RAI_TensorData(t); - if (dtype.code == kDLFloat) { - switch (dtype.bits) { - case 32: - ((float *)data)[i] = val; - break; - case 64: - ((double *)data)[i] = val; - break; - default: - return 0; - } - } else { + if (dtype.code != kDLString) { return 0; } + uint64_t *offsets = RAI_TensorStringElementsOffsets(t); + char *str_val = data + offsets[i]; + *val = str_val; return 1; } -int RAI_TensorGetValueAsDouble(RAI_Tensor *t, long long i, double *val) { - DLDataType dtype = t->tensor.dl_tensor.dtype; - void *data = t->tensor.dl_tensor.data; - - // TODO: check i is in bound - if (dtype.code == kDLFloat) { - switch (dtype.bits) { - case 32: - *val = ((float *)data)[i]; - break; - case 64: - *val = ((double *)data)[i]; - break; - default: +int RAI_TensorSetData(RAI_Tensor *t, const char *data, size_t len) { + DLDataType data_type = RAI_TensorDataType(t); + if (data_type.code == kDLString) { + if (_RAI_TensorParseStringsBlob(data, len, RAI_TensorLength(t), + RAI_TensorStringElementsOffsets(t), + NULL) != REDISMODULE_OK) { return 0; } - } else { - return 0; + RedisModule_Free(RAI_TensorData(t)); + t->tensor.dl_tensor.data = RedisModule_Alloc(len); } + memcpy(RAI_TensorData(t), data, len); + t->blobSize = len; return 1; } -int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { - DLDataType dtype = t->tensor.dl_tensor.dtype; - void *data = t->tensor.dl_tensor.data; - - // TODO: check i is in bound +int RAI_TensorSetValueFromLongLong(RAI_Tensor *t, long long i, long long val) { + DLDataType dtype = RAI_TensorDataType(t); + void *data = RAI_TensorData(t); if (dtype.code == kDLInt) { switch (dtype.bits) { case 8: - *val = ((int8_t *)data)[i]; + ((int8_t *)data)[i] = val; break; case 16: - *val = ((int16_t *)data)[i]; + ((int16_t *)data)[i] = val; break; case 32: - *val = ((int32_t *)data)[i]; + ((int32_t *)data)[i] = val; break; case 64: - *val = ((int64_t *)data)[i]; + ((int64_t *)data)[i] = val; break; default: return 0; @@ -536,23 +686,23 @@ int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { } else if (dtype.code == kDLUInt) { switch (dtype.bits) { case 8: - *val = ((uint8_t *)data)[i]; + ((uint8_t *)data)[i] = val; break; case 16: - *val = ((uint16_t *)data)[i]; + ((uint16_t *)data)[i] = val; break; case 32: - *val = ((uint32_t *)data)[i]; + ((uint32_t *)data)[i] = val; break; case 64: - *val = ((uint64_t *)data)[i]; + ((uint64_t *)data)[i] = val; break; default: return 0; } } else if (dtype.code == kDLBool) { if (dtype.bits == 8) { - *val = ((uint8_t *)data)[i]; + ((uint8_t *)data)[i] = val; } else { return 0; } @@ -560,25 +710,66 @@ int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val) { return 1; } +int RAI_TensorSetValueFromDouble(RAI_Tensor *t, long long i, double val) { + DLDataType dtype = RAI_TensorDataType(t); + void *data = RAI_TensorData(t); + + if (dtype.code == kDLFloat) { + switch (dtype.bits) { + case 32: + ((float *)data)[i] = val; + break; + case 64: + ((double *)data)[i] = val; + break; + default: + return 0; + } + } else { + return 0; + } + return 1; +} + +//******************** tensor memory management ********************************* + RAI_Tensor *RAI_TensorGetShallowCopy(RAI_Tensor *t) { __atomic_fetch_add(&t->refCount, 1, __ATOMIC_RELAXED); return t; } -int RAI_TensorNumDims(RAI_Tensor *t) { return t->tensor.dl_tensor.ndim; } - -long long RAI_TensorDim(RAI_Tensor *t, int i) { return t->tensor.dl_tensor.shape[i]; } - -size_t RAI_TensorByteSize(RAI_Tensor *t) { - // TODO: as per dlpack it should be - // size *= (t->dtype.bits * t->dtype.lanes + 7) / 8; - return Tensor_DataTypeSize(RAI_TensorDataType(t)) * RAI_TensorLength(t); +void RAI_TensorFree(RAI_Tensor *t) { + if (t == NULL) { + return; + } + long long ref_count = __atomic_sub_fetch(&t->refCount, 1, __ATOMIC_RELAXED); + RedisModule_Assert(ref_count >= 0); + if (ref_count > 0) { + return; + } + if (t->tensor.deleter) { + t->tensor.deleter(&t->tensor); + } else { + if (t->tensor.dl_tensor.shape) { + RedisModule_Free(t->tensor.dl_tensor.shape); + } + if (t->tensor.dl_tensor.strides) { + RedisModule_Free(t->tensor.dl_tensor.strides); + } + if (t->tensor.dl_tensor.data) { + RedisModule_Free(t->tensor.dl_tensor.data); + } + if (t->tensor.dl_tensor.elements_length) { + RedisModule_Free(t->tensor.dl_tensor.elements_length); + } + RedisModule_Free(t); + } } -char *RAI_TensorData(RAI_Tensor *t) { return t->tensor.dl_tensor.data; } +//***************************** retrieve tensor from keyspace ********************* -int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - int mode, RAI_Error *err) { +int RAI_TensorOpenKey(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, + int mode, RAI_Error *err) { *key = RedisModule_OpenKey(ctx, keyName, mode); if (RedisModule_KeyType(*key) == REDISMODULE_KEYTYPE_EMPTY) { return REDISMODULE_OK; @@ -591,7 +782,7 @@ int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisMod return REDISMODULE_OK; } -int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, +int RAI_TensorGetFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, RAI_Tensor **tensor, int mode, RAI_Error *err) { *key = RedisModule_OpenKey(ctx, keyName, mode); if (RedisModule_KeyType(*key) == REDISMODULE_KEYTYPE_EMPTY) { @@ -616,278 +807,76 @@ int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, R return REDISMODULE_OK; } -void RedisAI_ReplicateTensorSet(RedisModuleCtx *ctx, RedisModuleString *key, RAI_Tensor *t) { - long long ndims = RAI_TensorNumDims(t); +void RAI_TensorReplicate(RedisModuleCtx *ctx, RedisModuleString *key, RAI_Tensor *t) { + long long n_dims = RAI_TensorNumDims(t); - char dtypestr[8]; - const int status = Tensor_DataTypeStr(RAI_TensorDataType(t), dtypestr); + char data_type_str[8]; + int status = RAI_TensorGetDataTypeStr(RAI_TensorDataType(t), data_type_str); RedisModule_Assert(status == REDISMODULE_OK); char *data = RAI_TensorData(t); - long long size = RAI_TensorByteSize(t); - - RedisModuleString *dims[ndims]; + long long size = (long long)RAI_TensorByteSize(t); - for (long long i = 0; i < ndims; i++) { - dims[i] = RedisModule_CreateStringFromLongLong(ctx, RAI_TensorDim(t, i)); + RedisModuleString *dims[n_dims]; + for (int i = 0; i < n_dims; i++) { + dims[i] = RedisModule_CreateStringFromLongLong(NULL, RAI_TensorDim(t, i)); } - RedisModule_Replicate(ctx, "AI.TENSORSET", "scvcb", key, dtypestr, dims, ndims, "BLOB", data, - size); + RedisModule_Replicate(ctx, "AI.TENSORSET", "scvcb", key, data_type_str, dims, n_dims, "BLOB", + data, size); - for (long long i = 0; i < ndims; i++) { - RedisModule_FreeString(ctx, dims[i]); + for (long long i = 0; i < n_dims; i++) { + RedisModule_FreeString(NULL, dims[i]); } } -int RAI_parseTensorSetArgs(RedisModuleString **argv, int argc, RAI_Tensor **t, int enforceArity, - RAI_Error *error) { - if (argc < 4) { - RAI_SetError(error, RAI_ETENSORSET, "wrong number of arguments for 'AI.TENSORSET' command"); - return -1; - } - - // get the tensor datatype - const char *typestr = RedisModule_StringPtrLen(argv[2], NULL); - DLDataType datatype = RAI_TensorDataTypeFromString(typestr); - size_t datasize = Tensor_DataTypeSize(datatype); - if (datasize == 0) { - RAI_SetError(error, RAI_ETENSORSET, "ERR invalid data type"); - return -1; - } - - const char *fmtstr; - int datafmt = TENSOR_NONE; - size_t ndims = 0; - long long len = 1; - long long *dims = (long long *)array_new(long long, 1); - size_t argpos = 3; - long long remaining_args = argc - 1; - size_t expected_nvalues = 0; - size_t current_nvalues = 0; - - for (; argpos <= argc - 1; argpos++) { - const char *opt = RedisModule_StringPtrLen(argv[argpos], NULL); - remaining_args = argc - 1 - argpos; - if (!strcasecmp(opt, "BLOB")) { - datafmt = TENSOR_BLOB; - // if we've found the dataformat there are no more dimensions - // check right away if the arity is correct - if (remaining_args != 1 && enforceArity == 1) { - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, - "ERR wrong number of arguments for 'AI.TENSORSET' command"); - return -1; - } - argpos++; - break; - } else if (!strcasecmp(opt, "VALUES")) { - datafmt = TENSOR_VALUES; - // if we've found the dataformat there are no more dimensions - // check right away if the arity is correct - if (remaining_args != len && enforceArity == 1) { - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, - "ERR wrong number of arguments for 'AI.TENSORSET' command"); - return -1; - } - argpos++; - break; - } else { - long long dimension; - const int retval = RedisModule_StringToLongLong(argv[argpos], &dimension); - if (retval != REDISMODULE_OK || dimension <= 0) { - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, - "ERR invalid or negative value found in tensor shape"); - return -1; - } - ndims++; - dims = array_append(dims, dimension); - len *= dimension; - } - } - - if (datafmt == TENSOR_BLOB) { - RedisModuleString *rstr = argv[argpos]; - *t = _TensorCreateWithDLDataTypeAndRString(datatype, datasize, dims, ndims, rstr, error); - } else { - bool is_empty = (datafmt == TENSOR_NONE); - *t = RAI_TensorCreateWithDLDataType(datatype, dims, ndims, is_empty); - } - if (!(*t)) { - array_free(dims); - return -1; - } - - long i = 0; - if (datafmt == TENSOR_VALUES) { - for (; (argpos <= argc - 1) && (i < len); argpos++) { - if (datatype.code == kDLFloat) { - double val; - const int retval = RedisModule_StringToDouble(argv[argpos], &val); - if (retval != REDISMODULE_OK) { - RAI_TensorFree(*t); - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); - return -1; - } - const int retset = RAI_TensorSetValueFromDouble(*t, i, val); - if (retset == -1) { - RAI_TensorFree(*t); - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, - "ERR cannot specify values for this datatype"); - return -1; - } - } else { - long long val; - const int retval = RedisModule_StringToLongLong(argv[argpos], &val); - if (retval != REDISMODULE_OK || _ValOverflow(val, *t)) { - RAI_TensorFree(*t); - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, "ERR invalid value"); - return -1; - } - const int retset = RAI_TensorSetValueFromLongLong(*t, i, val); - if (retset == -1) { - RAI_TensorFree(*t); - array_free(dims); - RAI_SetError(error, RAI_ETENSORSET, - "ERR cannot specify values for this datatype"); - return -1; - } - } - i++; - } - } - array_free(dims); - return argpos; -} - -int RAI_TensorReplyWithValues(RedisModuleCtx *ctx, RAI_Tensor *t) { - long long ndims = RAI_TensorNumDims(t); - long long len = 1; - long long i; - for (i = 0; i < ndims; i++) { - len *= RAI_TensorDim(t, i); - } - - DLDataType dtype = RAI_TensorDataType(t); - - RedisModule_ReplyWithArray(ctx, len); - - if (dtype.code == kDLFloat) { - double val; - for (i = 0; i < len; i++) { - int ret = RAI_TensorGetValueAsDouble(t, i, &val); - if (!ret) { - RedisModule_ReplyWithError(ctx, "ERR cannot get values for this datatype"); - return -1; - } - RedisModule_ReplyWithDouble(ctx, val); - } - } else { - long long val; - for (i = 0; i < len; i++) { - int ret = RAI_TensorGetValueAsLongLong(t, i, &val); - if (!ret) { - RedisModule_ReplyWithError(ctx, "ERR cannot get values for this datatype"); - return -1; - } - RedisModule_ReplyWithLongLong(ctx, val); - } - } - - return 0; -} - -uint ParseTensorGetArgs(RAI_Error *err, RedisModuleString **argv, int argc) { - uint fmt = TENSOR_NONE; - if (argc < 2 || argc > 4) { - RAI_SetError(err, RAI_EDAGBUILDER, "wrong number of arguments for 'AI.TENSORGET' command"); - return fmt; - } - if (argc == 2) { - return TENSOR_BLOB | TENSOR_META; - } - for (int i = 2; i < argc; i++) { - const char *fmtstr = RedisModule_StringPtrLen(argv[i], NULL); - if (!strcasecmp(fmtstr, "BLOB")) { - fmt |= TENSOR_BLOB; - } else if (!strcasecmp(fmtstr, "VALUES")) { - fmt |= TENSOR_VALUES; - } else if (!strcasecmp(fmtstr, "META")) { - fmt |= TENSOR_META; - } else { - RAI_SetError(err, RAI_EDAGBUILDER, "ERR unsupported data format"); - return TENSOR_NONE; - } - } - - if (fmt == TENSOR_ILLEGAL_VALUES_BLOB) { - RAI_SetError(err, RAI_EDAGBUILDER, "ERR both BLOB and VALUES specified"); - return TENSOR_NONE; - } - return fmt; -} - -int ReplyWithTensor(RedisModuleCtx *ctx, uint fmt, RAI_Tensor *t) { +int RAI_TensorReply(RedisModuleCtx *ctx, uint fmt, RAI_Tensor *t) { if (!(fmt & TENSOR_META)) { if (fmt & TENSOR_BLOB) { - long long size = RAI_TensorByteSize(t); + size_t size = RAI_TensorByteSize(t); char *data = RAI_TensorData(t); RedisModule_ReplyWithStringBuffer(ctx, data, size); return REDISMODULE_OK; } if (fmt & TENSOR_VALUES) { - int ret = RAI_TensorReplyWithValues(ctx, t); - if (ret == -1) { - return REDISMODULE_ERR; - } - return REDISMODULE_OK; + return _RAI_TensorReplyWithValues(ctx, t); } } - long long resplen = 4; + long long resp_len = 4; if (fmt & (TENSOR_BLOB | TENSOR_VALUES)) - resplen += 2; + resp_len += 2; - const long long ndims = RAI_TensorNumDims(t); + const long long n_dims = RAI_TensorNumDims(t); - char dtypestr[8]; - const int dtypestr_result = Tensor_DataTypeStr(RAI_TensorDataType(t), dtypestr); - if (dtypestr_result == REDISMODULE_ERR) { + char data_type_str[8]; + const int data_type_str_result = RAI_TensorGetDataTypeStr(RAI_TensorDataType(t), data_type_str); + if (data_type_str_result == REDISMODULE_ERR) { RedisModule_ReplyWithError(ctx, "ERR unsupported dtype"); return REDISMODULE_ERR; } - RedisModule_ReplyWithArray(ctx, resplen); + RedisModule_ReplyWithArray(ctx, resp_len); RedisModule_ReplyWithCString(ctx, "dtype"); - RedisModule_ReplyWithCString(ctx, dtypestr); + RedisModule_ReplyWithCString(ctx, data_type_str); RedisModule_ReplyWithCString(ctx, "shape"); - RedisModule_ReplyWithArray(ctx, ndims); - for (int i = 0; i < ndims; i++) { + RedisModule_ReplyWithArray(ctx, n_dims); + for (int i = 0; i < n_dims; i++) { const long long dim = RAI_TensorDim(t, i); RedisModule_ReplyWithLongLong(ctx, dim); } if (fmt & TENSOR_BLOB) { - long long size = RAI_TensorByteSize(t); + size_t size = RAI_TensorByteSize(t); char *data = RAI_TensorData(t); RedisModule_ReplyWithCString(ctx, "blob"); RedisModule_ReplyWithStringBuffer(ctx, data, size); } else if (fmt & TENSOR_VALUES) { RedisModule_ReplyWithCString(ctx, "values"); - int ret = RAI_TensorReplyWithValues(ctx, t); - if (ret != REDISMODULE_OK) { - return REDISMODULE_ERR; - } + return _RAI_TensorReplyWithValues(ctx, t); } return REDISMODULE_OK; } - -RedisModuleType *RAI_TensorRedisType(void) { return RedisAI_TensorType; } diff --git a/src/redis_ai_objects/tensor.h b/src/redis_ai_objects/tensor.h index 2e8e219d8..994223322 100644 --- a/src/redis_ai_objects/tensor.h +++ b/src/redis_ai_objects/tensor.h @@ -32,6 +32,7 @@ static const char *RAI_DATATYPE_STR_INT64 = "INT64"; static const char *RAI_DATATYPE_STR_UINT8 = "UINT8"; static const char *RAI_DATATYPE_STR_UINT16 = "UINT16"; static const char *RAI_DATATYPE_STR_BOOL = "BOOL"; +static const char *RAI_DATATYPE_STR_STRING = "STRING"; #define TENSOR_NONE 0 #define TENSOR_VALUES (1 << 0) @@ -40,64 +41,73 @@ static const char *RAI_DATATYPE_STR_BOOL = "BOOL"; #define TENSOR_ILLEGAL_VALUES_BLOB (TENSOR_VALUES | TENSOR_BLOB) /** - * Helper method to register the tensor type exported by the module. - * - * @param ctx Context in which Redis modules operate - * @return + * @brief Returns the redis module type representing a tensor. + * @return redis module type representing a tensor. */ -int RAI_TensorInit(RedisModuleCtx *ctx); +RedisModuleType *RAI_TensorRedisType(void); +//********************* methods for creating a tensor *************************** /** - * @brief Allocate an empty tensor with no data. - * @note The new tensor ref coutn is 1. + * Allocate the memory for a new tensor and set RAI_Tensor meta-data based on + * the passed 'data type`, the specified number of dimensions + * `ndims`, and n-dimension array `dims`. Note that tensor data is not initialized! * - * @return RAI_Tensor* - a pointer to the new tensor. + * @param data_type represents the data type of tensor elements + * @param dims n-dimensional array ( the dimension values are copied ) + * @param n_dims number of dimensions + * @return allocated RAI_Tensor on success, or NULL if the allocation + * failed. */ -RAI_Tensor *RAI_TensorNew(void); +RAI_Tensor *RAI_TensorNew(DLDataType data_type, const size_t *dims, int n_dims); /** - * Allocate the memory and initialise the RAI_Tensor. Creates a tensor based on - * the passed 'dataType` string and with the specified number of dimensions - * `ndims`, and n-dimension array `dims`. + * Allocate the memory and set RAI_Tensor meta-data based on + * the passed 'dataType` string, the specified number of dimensions + * `ndims`, and n-dimension array `dims`. Note that tensor data is not initialized! * - * @param dataType string containing the numeric data type of tensor elements + * @param data_type string containing the data type of tensor elements * @param dims n-dimensional array ( the dimension values are copied ) - * @param ndims number of dimensions + * @param n_dims number of dimensions * @return allocated RAI_Tensor on success, or NULL if the allocation * failed. */ -RAI_Tensor *RAI_TensorCreate(const char *dataType, long long *dims, int ndims); +RAI_Tensor *RAI_TensorCreate(const char *data_type, const long long *dims, int n_dims); /** * Allocate the memory and initialise the RAI_Tensor. Creates a tensor based on - * the passed 'DLDataType` and with the specified number of dimensions `ndims`, - * and n-dimension array `dims`. Depending on the passed `tensorAllocMode`, the - * DLTensor data will be either allocated or no allocation is performed thus - * enabling shallow copy of data (no alloc) + * the passed data type and with the specified number of dimensions `ndims`, + * and n-dimension array `dims`. The tensor will be populated with the given values. * - * @param dtype DLDataType - * @param dims n-dimensional array ( the dimension values are copied ) + * @param data_type DLDataType that represents the tensor elements data type. + * @param dims array of size ndims, contains the tensor shapes (the dimension values are copied) * @param ndims number of dimensions - * @param empty True if creating an empty tensor (need to be initialized) - * @return allocated RAI_Tensor on success, or NULL if the allocation - * failed. + * @param argc number of values to store in the vector (the size of argv array) + * @param argv array of redis module strings contains the tensor values + * @param err used to store error status if one occurs + * @return allocated RAI_Tensor on success, or NULL if operation failed. */ -RAI_Tensor *RAI_TensorCreateWithDLDataType(DLDataType dtype, long long *dims, int ndims, - bool empty); +RAI_Tensor *RAI_TensorCreateFromValues(DLDataType data_type, const size_t *dims, int n_dims, + int argc, RedisModuleString **argv, RAI_Error *err); /** - * Allocate the memory for a new Tensor and copy data fom a tensor to it. + * Allocate the memory and initialise the RAI_Tensor. Creates a tensor based on + * the passed data type and with the specified number of dimensions `ndims`, + * and n-dimension array `dims`. The tensor will be populated with the given data blob. * - * @param t Source tensor to copy. - * @param result Destination tensor to copy. - * @return 0 on success, or 1 if the copy failed - * failed. + * @param data_type DLDataType that represents the tensor elements data type. + * @param dims array of size ndims, contains the tensor shapes (the dimension values are copied) + * @param ndims number of dimensions + * @param tensor_blob a buffer contains the tensor data blob (binary string) + * @param blob_len number of bytes in tensor blob + * @param err used to store error status if one occurs + * @return allocated RAI_Tensor on success, or NULL if operation failed. */ -int RAI_TensorDeepCopy(RAI_Tensor *t, RAI_Tensor **dest); +RAI_Tensor *RAI_TensorCreateFromBlob(DLDataType data_type, const size_t *dims, int n_dims, + const char *tensor_blob, size_t blob_len, RAI_Error *err); /** * Allocate the memory and initialise the RAI_Tensor, performing a shallow copy - * of dl_tensor. Beware, this will take ownership of dltensor, and only allocate + * of dl_tensor. Beware, this will take ownership of dl_tensor, and only allocate * the data for the RAI_Tensor, meaning that the freeing the data of the input * dl_tensor, will also free the data of the returned tensor. * @@ -108,66 +118,56 @@ int RAI_TensorDeepCopy(RAI_Tensor *t, RAI_Tensor **dest); RAI_Tensor *RAI_TensorCreateFromDLTensor(DLManagedTensor *dl_tensor); /** - * @param tensor - * @return A pointer to the inner DLTensor field (do not copy). - */ -DLTensor *RAI_TensorGetDLTensor(RAI_Tensor *tensor); - -/** - * Allocate the memory and initialise the RAI_Tensor, performing a deep copy of - * the passed array of tensors. + * Allocate the memory and initialise an RAI_Tensor which is a concatenation of the input tensors. * - * @param ts input array of tensors + * @param tensors input array of tensors to concatenate * @param n number of input tensors - * @return allocated RAI_Tensor on success, or NULL if the allocation and deep - * copy failed failed. + * @return allocated RAI_Tensor on success, or NULL if the operation failed. */ RAI_Tensor *RAI_TensorCreateByConcatenatingTensors(RAI_Tensor **ts, long long n); /** - * Allocate the memory and initialise the RAI_Tensor, performing a deep copy of - * the passed tensor, at the given data offset and length. + * Allocate the memory and initialise an RAI_Tensor which is a slice of + * the passed tensor having the given length, starting at the given offset. * * @param t input tensor - * @param offset - * @param len - * @return allocated RAI_Tensor on success, or NULL if the allocation and deep - * copy failed failed. + * @param offset the index of the tensor's first dimension to start copy the data from + * @param len the length of the sliced tensor (slice ends at offset+len index) + * @return allocated RAI_Tensor on success, or NULL if the operation failed. */ RAI_Tensor *RAI_TensorCreateBySlicingTensor(RAI_Tensor *t, long long offset, long long len); /** - * Returns the length of the input tensor + * Helper method for creating the DLDataType represented by the input string * - * @param t input tensor - * @return the length of the input tensor + * @param data_type + * @return the DLDataType represented by the input string */ -size_t RAI_TensorLength(RAI_Tensor *t); +DLDataType RAI_TensorDataTypeFromString(const char *data_type); + +//*************** getters and setters ****************************** /** - * Returns the size in bytes of each element of the tensor - * - * @param t input tensor - * @return size in bytes of each the underlying tensor data type + * @param tensor + * @return A pointer to the inner DLTensor field (do not copy). */ -size_t RAI_TensorDataSize(RAI_Tensor *t); +DLTensor *RAI_TensorGetDLTensor(RAI_Tensor *tensor); /** - * Returns the size in bytes of the given DLDataType + * Returns the length (i.e., number of elements) of the input tensor * - * @param dtype DLDataType - * @return size in bytes of each the given DLDataType + * @param t input tensor + * @return the length of the input tensor */ -size_t RAI_TensorDataSizeFromDLDataType(DLDataType dtype); +size_t RAI_TensorLength(RAI_Tensor *t); /** - * Returns the size in bytes of the associated DLDataType represented by the - * input string + * Returns the size in bytes of each element of the tensor * - * @param dataType + * @param t input tensor * @return size in bytes of each the underlying tensor data type */ -size_t RAI_TensorDataSizeFromString(const char *dataType); +size_t RAI_TensorDataSize(RAI_Tensor *t); /** * Get the associated `DLDataType` for the given input tensor @@ -178,72 +178,54 @@ size_t RAI_TensorDataSizeFromString(const char *dataType); DLDataType RAI_TensorDataType(RAI_Tensor *t); /** - * Check whether two tensors have the same data type - * - * @param t1 input tensor - * @param t2 input tensor - * @return 1 if data types match, 0 otherwise - */ -int RAI_TensorIsDataTypeEqual(RAI_Tensor *t1, RAI_Tensor *t2); - -/** - * Returns the DLDataType represented by the input string + * Returns the number of dimensions for the given input tensor * - * @param dataType - * @return the DLDataType represented by the input string + * @param t input tensor + * @return number of dimensions for the given input tensor */ -DLDataType RAI_TensorDataTypeFromString(const char *dataType); +int RAI_TensorNumDims(RAI_Tensor *t); /** - * sets in dtypestr the string representing the associated DLDataType + * Returns the dimension length for the given input tensor and dimension * - * @param dtype DLDataType - * @param dtypestr output string to store the associated string representing the - * DLDataType - * @return REDISMODULE_OK on success, or REDISMODULE_ERR if failed + * @param t input tensor + * @param dim dimension + * @return the dimension length */ -int Tensor_DataTypeStr(DLDataType dtype, char *dtypestr); +long long RAI_TensorDim(RAI_Tensor *t, int dim); /** - * Frees the memory of the RAI_Tensor when the tensor reference count reaches 0. - * It is safe to call this function with a NULL input tensor. + * Returns the size in bytes of the underlying tensor data * - * @param t tensor + * @param t input tensor + * @return the size in bytes of the underlying deep learning tensor data */ -void RAI_TensorFree(RAI_Tensor *t); +size_t RAI_TensorByteSize(RAI_Tensor *t); /** - * Sets the associated data to the deep learning tensor via deep copying the - * passed data. + * Return the pointer to the tensor data blob (do not copy) * - * @param t tensor to set the data - * @param data input data - * @param len input data length - * @return 1 on success + * @param t input tensor + * @return direct access to the array data pointer */ -int RAI_TensorSetData(RAI_Tensor *t, const char *data, size_t len); +char *RAI_TensorData(RAI_Tensor *t); /** - * Sets the long value for the given tensor, at the given array data pointer - * position + * Return the pointer to the array containing the offset of every string element + * in the data array * - * @param t tensor to set the data - * @param i dl_tensor data pointer position - * @param val value to set the data from - * @return 0 on success, or 1 if the setting failed + * @param t input tensor + * @return direct access to the offsets array pointer */ -int RAI_TensorSetValueFromLongLong(RAI_Tensor *t, long long i, long long val); +uint64_t *RAI_TensorStringElementsOffsets(RAI_Tensor *tensor); /** - * Sets the double value for the given tensor, at the given array data pointer - * position + * Return the pointer to the array containing the tensors dimensions. * - * @param t tensor to set the data - * @param i dl_tensor data pointer position - * @param val value to set the data from - * @return 1 on success, or 0 if the setting failed + * @param t input tensor + * @return direct access to the shpae array pointer */ -int RAI_TensorSetValueFromDouble(RAI_Tensor *t, long long i, double val); +int64_t *RAI_TensorShape(RAI_Tensor *tensor); /** * Gets the double value from the given input tensor, at the given array data @@ -268,48 +250,81 @@ int RAI_TensorGetValueAsDouble(RAI_Tensor *t, long long i, double *val); int RAI_TensorGetValueAsLongLong(RAI_Tensor *t, long long i, long long *val); /** - * Every call to this function, will make the RAI_Tensor 't' requiring an - * additional call to RAI_TensorFree() in order to really free the tensor. - * Returns a shallow copy of the tensor. + * Gets the a const pointer to the string value from the given input tensor, at the given array data + * pointer position (do not copy the string). * - * @param t input tensor - * @return shallow copy of the tensor + * @param t tensor to get the data + * @param i dl_tensor data pointer position + * @param val value to set the data to + * @return 1 on success, or 0 if getting the data failed */ -RAI_Tensor *RAI_TensorGetShallowCopy(RAI_Tensor *t); +int RAI_TensorGetValueAsCString(RAI_Tensor *t, long long i, const char **val); /** - * Returns the number of dimensions for the given input tensor + * sets in data_type_str the string representing the associated DLDataType * - * @param t input tensor - * @return number of dimensions for the given input tensor + * @param data_type DLDataType + * @param data_type_str output string to store the associated string representing the + * DLDataType + * @return REDISMODULE_OK on success, or REDISMODULE_ERR if failed (unsupported data type) */ -int RAI_TensorNumDims(RAI_Tensor *t); +int RAI_TensorGetDataTypeStr(DLDataType data_type, char *data_type_str); /** - * Returns the dimension length for the given input tensor and dimension + * Sets the associated data to the tensor via deep copying the + * passed data. * - * @param t input tensor - * @param dim dimension - * @return the dimension length + * @param t tensor to set the data + * @param data input data + * @param len input data length + * @return 1 on success */ -long long RAI_TensorDim(RAI_Tensor *t, int dim); + +int RAI_TensorSetData(RAI_Tensor *t, const char *data, size_t len); /** - * Returns the size in bytes of the underlying deep learning tensor data + * Sets the long value for the given tensor, at the given array data pointer + * position * - * @param t input tensor - * @return the size in bytes of the underlying deep learning tensor data + * @param t tensor to set the data + * @param i dl_tensor data pointer position + * @param val value to set the data from + * @return 0 on success, or 1 if the setting failed */ -size_t RAI_TensorByteSize(RAI_Tensor *t); +int RAI_TensorSetValueFromLongLong(RAI_Tensor *t, long long i, long long val); /** - * Return the pointer the the deep learning tensor data + * Sets the double value for the given tensor, at the given array data pointer + * position + * + * @param t tensor to set the data + * @param i dl_tensor data pointer position + * @param val value to set the data from + * @return 1 on success, or 0 if the setting failed + */ +int RAI_TensorSetValueFromDouble(RAI_Tensor *t, long long i, double val); + +//************************************** tensor memory management +//*********************************** +/** + * Every call to this function, will make the RAI_Tensor 't' requiring an + * additional call to RAI_TensorFree() in order to really free the tensor. + * Returns a shallow copy of the tensor. * * @param t input tensor - * @return direct access to the array data pointer + * @return shallow copy of the tensor */ -char *RAI_TensorData(RAI_Tensor *t); +RAI_Tensor *RAI_TensorGetShallowCopy(RAI_Tensor *t); +/** + * Frees the memory of the RAI_Tensor when the tensor reference count reaches 0. + * It is safe to call this function with a NULL input tensor. + * + * @param t tensor + */ +void RAI_TensorFree(RAI_Tensor *t); + +//*************** methods for retrieval and replicating tensor from keyspace *************** /** * Helper method to open a key handler for the tensor data type * @@ -318,12 +333,13 @@ char *RAI_TensorData(RAI_Tensor *t); * @param key tensor's key handle. On success it contains an handle representing * a Redis key with the requested access mode * @param mode key access mode + * @param err used to store error status if one occurs * @return REDISMODULE_OK if it's possible to store at the specified key handle * the tensor type, or REDISMODULE_ERR if is the key not associated with a * tensor type. */ -int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, - int mode, RAI_Error *err); +int RAI_TensorOpenKey(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, + int mode, RAI_Error *err); /** * Helper method to get Tensor from keyspace. In case of a failure an @@ -335,66 +351,32 @@ int RAI_OpenKey_Tensor(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisMod * a Redis key with the requested access mode * @param tensor destination tensor structure * @param mode key access mode + * @param err used to store error status if one occurs. * @return REDISMODULE_OK if the tensor value stored at key was correctly * returned and available at *tensor variable, or REDISMODULE_ERR if there was * an error getting the Tensor */ -int RAI_GetTensorFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, +int RAI_TensorGetFromKeyspace(RedisModuleCtx *ctx, RedisModuleString *keyName, RedisModuleKey **key, RAI_Tensor **tensor, int mode, RAI_Error *err); /** * Helper method to replicate a tensor via an AI.TENSORSET command to the * replicas. This is used on MODELRUN, SCRIPTRUN, DAGRUN as a way to ensure that - * the results present on replicas match the results present on master ( since - * multiple modelruns/scripts are not ensured to have the same output values - * (non-deterministic) ). + * the results present on replicas match the results present on master * * @param ctx Context in which Redis modules operate * @param key Destination key name * @param t source tensor */ -void RedisAI_ReplicateTensorSet(RedisModuleCtx *ctx, RedisModuleString *key, RAI_Tensor *t); - -/** - * Helper method to parse AI.TENSORGET arguments - * - * @param argv Redis command arguments, as an array of strings - * @param argc Redis command number of arguments - * @param t Destination tensor to store the parsed data - * @param enforceArity flag whether to enforce arity checking - * @param error error data structure to store error message in the case of - * parsing failures - * @return processed number of arguments on success, or -1 if the parsing failed - */ -int RAI_parseTensorSetArgs(RedisModuleString **argv, int argc, RAI_Tensor **t, int enforceArity, - RAI_Error *error); - -/** - * Helper method to parse AI.TENSORGET arguments - * - * @param error error data structure to store error message in the case of - * parsing failures - * @param argv Redis command arguments, as an array of strings - * @param argc Redis command number of arguments - * @return The format in which tensor is returned. - */ - -uint ParseTensorGetArgs(RAI_Error *error, RedisModuleString **argv, int argc); +void RAI_TensorReplicate(RedisModuleCtx *ctx, RedisModuleString *key, RAI_Tensor *t); /** * Helper method to return a tensor to the client in a response to AI.TENSORGET * * @param ctx Context in which Redis modules operate. - * @param fmt The format in which tensor is returned. + * @param fmt The format in which tensor is returned (META and/or VALUES/BLOB). * @param t The tensor to reply with. * @return REDISMODULE_OK in case of success, REDISMODULE_ERR otherwise. */ - -int ReplyWithTensor(RedisModuleCtx *ctx, uint fmt, RAI_Tensor *t); - -/** - * @brief Returns the redis module type representing a tensor. - * @return redis module type representing a tensor. - */ -RedisModuleType *RAI_TensorRedisType(void); +int RAI_TensorReply(RedisModuleCtx *ctx, uint fmt, RAI_Tensor *t); diff --git a/src/redis_ai_objects/tensor_struct.h b/src/redis_ai_objects/tensor_struct.h index 4ed1d7b98..71f98da25 100644 --- a/src/redis_ai_objects/tensor_struct.h +++ b/src/redis_ai_objects/tensor_struct.h @@ -9,4 +9,5 @@ typedef struct RAI_Tensor { DLManagedTensor tensor; size_t len; long long refCount; + size_t blobSize; } RAI_Tensor; diff --git a/src/redisai.c b/src/redisai.c index de0b70362..ee71483ed 100644 --- a/src/redisai.c +++ b/src/redisai.c @@ -13,6 +13,7 @@ #include "execution/DAG/dag_execute.h" #include "execution/utils.h" #include "execution/parsing/deprecated.h" +#include "execution/parsing/tensor_commands_parsing.h" #include "execution/execution_contexts/modelRun_ctx.h" #include "execution/execution_contexts/scriptRun_ctx.h" #include "redis_ai_objects/model.h" @@ -74,8 +75,7 @@ int RedisAI_TensorSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv RedisModuleKey *key; RAI_Error err = {0}; - const int status = - RAI_OpenKey_Tensor(ctx, argv[1], &key, REDISMODULE_READ | REDISMODULE_WRITE, &err); + int status = RAI_TensorOpenKey(ctx, argv[1], &key, REDISMODULE_READ | REDISMODULE_WRITE, &err); if (status == REDISMODULE_ERR) { RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); RAI_ClearError(&err); @@ -83,10 +83,8 @@ int RedisAI_TensorSet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv } RAI_Tensor *t = NULL; - const int parse_result = RAI_parseTensorSetArgs(argv, argc, &t, 1, &err); - - // if the number of parsed args is negative something went wrong - if (parse_result < 0) { + int parse_result = ParseTensorSetArgs(argv, argc, &t, &err); + if (parse_result != REDISMODULE_OK) { RedisModule_CloseKey(key); RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); RAI_ClearError(&err); @@ -114,20 +112,21 @@ int RedisAI_TensorGet_RedisCommand(RedisModuleCtx *ctx, RedisModuleString **argv RAI_Tensor *t; RedisModuleKey *key; RAI_Error err = {0}; - const int status = RAI_GetTensorFromKeyspace(ctx, argv[1], &key, &t, REDISMODULE_READ, &err); + int status = RAI_TensorGetFromKeyspace(ctx, argv[1], &key, &t, REDISMODULE_READ, &err); if (status == REDISMODULE_ERR) { RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); RAI_ClearError(&err); return REDISMODULE_ERR; } - uint fmt = ParseTensorGetArgs(&err, argv, argc); + uint fmt = ParseTensorGetFormat(&err, argv, argc); + + // TENSOR_NONE is returned in case that args are invalid. if (fmt == TENSOR_NONE) { RedisModule_ReplyWithError(ctx, RAI_GetErrorOneLine(&err)); RAI_ClearError(&err); - // This means that args are invalid. return REDISMODULE_ERR; } - ReplyWithTensor(ctx, fmt, t); + RAI_TensorReply(ctx, fmt, t); return REDISMODULE_OK; } diff --git a/src/serialization/AOF/rai_aof_rewrite.c b/src/serialization/AOF/rai_aof_rewrite.c index 6a81f3830..5aedb307a 100644 --- a/src/serialization/AOF/rai_aof_rewrite.c +++ b/src/serialization/AOF/rai_aof_rewrite.c @@ -4,7 +4,7 @@ void RAI_AOFRewriteTensor(RedisModuleIO *aof, RedisModuleString *key, void *valu RAI_Tensor *tensor = (RAI_Tensor *)value; char dtypestr[64]; - Tensor_DataTypeStr(RAI_TensorDataType(tensor), dtypestr); + RAI_TensorGetDataTypeStr(RAI_TensorDataType(tensor), dtypestr); char *data = RAI_TensorData(tensor); long long size = RAI_TensorByteSize(tensor); diff --git a/src/serialization/RDB/decoder/current/v4/decode_v4.c b/src/serialization/RDB/decoder/current/v4/decode_v4.c new file mode 100644 index 000000000..ae4f749cb --- /dev/null +++ b/src/serialization/RDB/decoder/current/v4/decode_v4.c @@ -0,0 +1,53 @@ +#include "decode_v4.h" +#include "../../previous/v3/decode_v3.h" +#include "assert.h" + +/** + * In case of IO errors, the default return values are: + * numbers - 0 + * strings - null + * So only when it is necessary check for IO errors. + */ + +void *RAI_RDBLoadTensor_v4(RedisModuleIO *io) { + DLDataTypeCode code = RedisModule_LoadUnsigned(io); + uint8_t bits = RedisModule_LoadUnsigned(io); + DLDataType data_type = (DLDataType){.code = code, .bits = bits, .lanes = 1}; + + int ndims = (int)RedisModule_LoadSigned(io); + size_t shape[ndims]; + for (size_t i = 0; i < ndims; ++i) { + shape[i] = RedisModule_LoadSigned(io); + } + + RAI_Tensor *tensor = RAI_TensorNew(data_type, shape, ndims); + + size_t blob_len; + char *data = RedisModule_LoadStringBuffer(io, &blob_len); + if (RedisModule_IsIOError(io)) + goto error; + + tensor->blobSize = blob_len; + tensor->tensor.dl_tensor.data = data; + + if (data_type.code == kDLString) { + for (size_t i = 0; i < RAI_TensorLength(tensor); i++) { + tensor->tensor.dl_tensor.elements_length[i] = RedisModule_LoadUnsigned(io); + } + } + if (RedisModule_IsIOError(io)) + goto error; + return tensor; + +error: + RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB"); + RAI_TensorFree(tensor); + if (data) { + RedisModule_Free(data); + } + return NULL; +} + +void *RAI_RDBLoadModel_v4(RedisModuleIO *io) { return RAI_RDBLoadModel_v3(io); } + +void *RAI_RDBLoadScript_v4(RedisModuleIO *io) { return RAI_RDBLoadScript_v3(io); } diff --git a/src/serialization/RDB/decoder/current/v4/decode_v4.h b/src/serialization/RDB/decoder/current/v4/decode_v4.h new file mode 100644 index 000000000..6bf9f4181 --- /dev/null +++ b/src/serialization/RDB/decoder/current/v4/decode_v4.h @@ -0,0 +1,8 @@ +#pragma once +#include "serialization/serialization_include.h" + +void *RAI_RDBLoadTensor_v4(RedisModuleIO *io); + +void *RAI_RDBLoadModel_v4(RedisModuleIO *io); + +void *RAI_RDBLoadScript_v4(RedisModuleIO *io); diff --git a/src/serialization/RDB/decoder/decode_previous.c b/src/serialization/RDB/decoder/decode_previous.c index 69773bff6..15e29664f 100644 --- a/src/serialization/RDB/decoder/decode_previous.c +++ b/src/serialization/RDB/decoder/decode_previous.c @@ -2,6 +2,7 @@ #include "previous/v0/decode_v0.h" #include "previous/v1/decode_v1.h" #include "previous/v2/decode_v2.h" +#include "previous/v3/decode_v3.h" void *Decode_PreviousTensor(RedisModuleIO *rdb, int encver) { switch (encver) { @@ -11,6 +12,8 @@ void *Decode_PreviousTensor(RedisModuleIO *rdb, int encver) { return RAI_RDBLoadTensor_v1(rdb); case 2: return RAI_RDBLoadTensor_v2(rdb); + case 3: + return RAI_RDBLoadTensor_v3(rdb); default: assert(false && "Invalid encoding version"); } @@ -25,6 +28,8 @@ void *Decode_PreviousModel(RedisModuleIO *rdb, int encver) { return RAI_RDBLoadModel_v1(rdb); case 2: return RAI_RDBLoadModel_v2(rdb); + case 3: + return RAI_RDBLoadModel_v3(rdb); default: assert(false && "Invalid encoding version"); } @@ -39,6 +44,8 @@ void *Decode_PreviousScript(RedisModuleIO *rdb, int encver) { return RAI_RDBLoadScript_v1(rdb); case 2: return RAI_RDBLoadScript_v2(rdb); + case 3: + return RAI_RDBLoadScript_v3(rdb); default: assert(false && "Invalid encoding version"); } diff --git a/src/serialization/RDB/decoder/previous/v0/decode_v0.c b/src/serialization/RDB/decoder/previous/v0/decode_v0.c index f1aef8ead..057ea3244 100644 --- a/src/serialization/RDB/decoder/previous/v0/decode_v0.c +++ b/src/serialization/RDB/decoder/previous/v0/decode_v0.c @@ -2,18 +2,11 @@ #include "assert.h" void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) { - int64_t *shape = NULL; - int64_t *strides = NULL; - DLDevice device; device.device_type = RedisModule_LoadUnsigned(io); device.device_id = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - // For now we only support CPU tensors (except during model and script run) - // device_id for default CPU should be -1 (in previous versions we might saved - // it as 0). + // For now, we only support CPU tensors (except during model and script run) RedisModule_Assert(device.device_type == kDLCPU); if (device.device_id != -1) { device.device_id = -1; @@ -24,46 +17,31 @@ void *RAI_RDBLoadTensor_v0(RedisModuleIO *io) { dtype.code = RedisModule_LoadUnsigned(io); dtype.lanes = RedisModule_LoadUnsigned(io); - size_t ndims = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - shape = RedisModule_Calloc(ndims, sizeof(*shape)); + int ndims = RedisModule_LoadUnsigned(io); + size_t shape[ndims]; for (size_t i = 0; i < ndims; ++i) { shape[i] = RedisModule_LoadUnsigned(io); } - strides = RedisModule_Calloc(ndims, sizeof(*strides)); + int64_t strides[ndims]; for (size_t i = 0; i < ndims; ++i) { strides[i] = RedisModule_LoadUnsigned(io); } - size_t byte_offset = RedisModule_LoadUnsigned(io); - size_t len; - char *data = RedisModule_LoadStringBuffer(io, &len); + size_t blob_len; + char *data = RedisModule_LoadStringBuffer(io, &blob_len); if (RedisModule_IsIOError(io)) - goto cleanup; + goto error; - RAI_Tensor *ret = RAI_TensorNew(); - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, - .data = data, - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = byte_offset}, - .manager_ctx = NULL, - .deleter = NULL}; - return ret; + RAI_Tensor *tensor = RAI_TensorNew(dtype, shape, ndims); + tensor->blobSize = blob_len; + tensor->tensor.dl_tensor.data = data; -cleanup: - if (shape) - RedisModule_Free(shape); - if (strides) - RedisModule_Free(strides); - RedisModule_LogIOError(io, "warning", - "Experienced a short read while reading a tensor from RDB"); + return tensor; + +error: + RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB"); return NULL; } diff --git a/src/serialization/RDB/decoder/previous/v1/decode_v1.c b/src/serialization/RDB/decoder/previous/v1/decode_v1.c index c31e54c5c..33bdce165 100644 --- a/src/serialization/RDB/decoder/previous/v1/decode_v1.c +++ b/src/serialization/RDB/decoder/previous/v1/decode_v1.c @@ -1,5 +1,5 @@ #include "decode_v1.h" -#include "assert.h" +#include "../v0/decode_v0.h" /** * In case of IO errors, the default return values are: @@ -8,70 +8,7 @@ * So only when it is necessary check for IO errors. */ -void *RAI_RDBLoadTensor_v1(RedisModuleIO *io) { - int64_t *shape = NULL; - int64_t *strides = NULL; - - DLDevice device; - device.device_type = RedisModule_LoadUnsigned(io); - device.device_id = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - // For now we only support CPU tensors (except during model and script run) - // device_id for default CPU should be -1 (in previous versions we might saved - // it as 0). - RedisModule_Assert(device.device_type == kDLCPU); - if (device.device_id != -1) { - device.device_id = -1; - } - - DLDataType dtype; - dtype.bits = RedisModule_LoadUnsigned(io); - dtype.code = RedisModule_LoadUnsigned(io); - dtype.lanes = RedisModule_LoadUnsigned(io); - - size_t ndims = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - shape = RedisModule_Calloc(ndims, sizeof(*shape)); - for (size_t i = 0; i < ndims; ++i) { - shape[i] = RedisModule_LoadUnsigned(io); - } - - strides = RedisModule_Calloc(ndims, sizeof(*strides)); - for (size_t i = 0; i < ndims; ++i) { - strides[i] = RedisModule_LoadUnsigned(io); - } - - size_t byte_offset = RedisModule_LoadUnsigned(io); - - size_t len; - char *data = RedisModule_LoadStringBuffer(io, &len); - if (RedisModule_IsIOError(io)) - goto cleanup; - - RAI_Tensor *ret = RAI_TensorNew(); - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, - .data = data, - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = byte_offset}, - .manager_ctx = NULL, - .deleter = NULL}; - return ret; - -cleanup: - if (shape) - RedisModule_Free(shape); - if (strides) - RedisModule_Free(strides); - RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB"); - return NULL; -} +void *RAI_RDBLoadTensor_v1(RedisModuleIO *io) { return RAI_RDBLoadTensor_v0(io); } void *RAI_RDBLoadModel_v1(RedisModuleIO *io) { diff --git a/src/serialization/RDB/decoder/previous/v2/decode_v2.c b/src/serialization/RDB/decoder/previous/v2/decode_v2.c index 7b08f5ca1..8403c534c 100644 --- a/src/serialization/RDB/decoder/previous/v2/decode_v2.c +++ b/src/serialization/RDB/decoder/previous/v2/decode_v2.c @@ -1,5 +1,5 @@ #include "decode_v2.h" -#include "assert.h" +#include "../v0/decode_v0.h" /** * In case of IO errors, the default return values are: @@ -8,70 +8,7 @@ * So only when it is necessary check for IO errors. */ -void *RAI_RDBLoadTensor_v2(RedisModuleIO *io) { - int64_t *shape = NULL; - int64_t *strides = NULL; - - DLDevice device; - device.device_type = RedisModule_LoadUnsigned(io); - device.device_id = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - // For now we only support CPU tensors (except during model and script run) - // device_id for default CPU should be -1 (in previous versions we might saved - // it as 0). - RedisModule_Assert(device.device_type == kDLCPU); - if (device.device_id != -1) { - device.device_id = -1; - } - - DLDataType dtype; - dtype.bits = RedisModule_LoadUnsigned(io); - dtype.code = RedisModule_LoadUnsigned(io); - dtype.lanes = RedisModule_LoadUnsigned(io); - - size_t ndims = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - shape = RedisModule_Calloc(ndims, sizeof(*shape)); - for (size_t i = 0; i < ndims; ++i) { - shape[i] = RedisModule_LoadUnsigned(io); - } - - strides = RedisModule_Calloc(ndims, sizeof(*strides)); - for (size_t i = 0; i < ndims; ++i) { - strides[i] = RedisModule_LoadUnsigned(io); - } - - size_t byte_offset = RedisModule_LoadUnsigned(io); - - size_t len; - char *data = RedisModule_LoadStringBuffer(io, &len); - if (RedisModule_IsIOError(io)) - goto cleanup; - - RAI_Tensor *ret = RAI_TensorNew(); - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, - .data = data, - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = byte_offset}, - .manager_ctx = NULL, - .deleter = NULL}; - return ret; - -cleanup: - if (shape) - RedisModule_Free(shape); - if (strides) - RedisModule_Free(strides); - RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB"); - return NULL; -} +void *RAI_RDBLoadTensor_v2(RedisModuleIO *io) { return RAI_RDBLoadTensor_v0(io); } void *RAI_RDBLoadModel_v2(RedisModuleIO *io) { diff --git a/src/serialization/RDB/decoder/current/v3/decode_v3.c b/src/serialization/RDB/decoder/previous/v3/decode_v3.c similarity index 76% rename from src/serialization/RDB/decoder/current/v3/decode_v3.c rename to src/serialization/RDB/decoder/previous/v3/decode_v3.c index 6c5de9f81..05db265cd 100644 --- a/src/serialization/RDB/decoder/current/v3/decode_v3.c +++ b/src/serialization/RDB/decoder/previous/v3/decode_v3.c @@ -1,5 +1,5 @@ #include "decode_v3.h" -#include "assert.h" +#include "../v0/decode_v0.h" /** * In case of IO errors, the default return values are: @@ -8,68 +8,7 @@ * So only when it is necessary check for IO errors. */ -void *RAI_RDBLoadTensor_v3(RedisModuleIO *io) { - int64_t *shape = NULL; - int64_t *strides = NULL; - - DLDevice device; - device.device_type = RedisModule_LoadUnsigned(io); - device.device_id = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - // For now we only support CPU tensors (except during model and script run) - assert(device.device_type == kDLCPU); - if (device.device_id != -1) { - device.device_id = -1; - } - - DLDataType dtype; - dtype.bits = RedisModule_LoadUnsigned(io); - dtype.code = RedisModule_LoadUnsigned(io); - dtype.lanes = RedisModule_LoadUnsigned(io); - - size_t ndims = RedisModule_LoadUnsigned(io); - if (RedisModule_IsIOError(io)) - goto cleanup; - - shape = RedisModule_Calloc(ndims, sizeof(*shape)); - for (size_t i = 0; i < ndims; ++i) { - shape[i] = RedisModule_LoadUnsigned(io); - } - - strides = RedisModule_Calloc(ndims, sizeof(*strides)); - for (size_t i = 0; i < ndims; ++i) { - strides[i] = RedisModule_LoadUnsigned(io); - } - - size_t byte_offset = RedisModule_LoadUnsigned(io); - - size_t len; - char *data = RedisModule_LoadStringBuffer(io, &len); - if (RedisModule_IsIOError(io)) - goto cleanup; - - RAI_Tensor *ret = RAI_TensorNew(); - ret->tensor = (DLManagedTensor){.dl_tensor = (DLTensor){.device = device, - .data = data, - .ndim = ndims, - .dtype = dtype, - .shape = shape, - .strides = strides, - .byte_offset = byte_offset}, - .manager_ctx = NULL, - .deleter = NULL}; - return ret; - -cleanup: - if (shape) - RedisModule_Free(shape); - if (strides) - RedisModule_Free(strides); - RedisModule_LogIOError(io, "error", "Experienced a short read while reading a tensor from RDB"); - return NULL; -} +void *RAI_RDBLoadTensor_v3(RedisModuleIO *io) { return RAI_RDBLoadTensor_v0(io); } void *RAI_RDBLoadModel_v3(RedisModuleIO *io) { diff --git a/src/serialization/RDB/decoder/current/v3/decode_v3.h b/src/serialization/RDB/decoder/previous/v3/decode_v3.h similarity index 100% rename from src/serialization/RDB/decoder/current/v3/decode_v3.h rename to src/serialization/RDB/decoder/previous/v3/decode_v3.h diff --git a/src/serialization/RDB/decoder/rai_rdb_decoder.c b/src/serialization/RDB/decoder/rai_rdb_decoder.c index 344d8a88d..20294bf32 100644 --- a/src/serialization/RDB/decoder/rai_rdb_decoder.c +++ b/src/serialization/RDB/decoder/rai_rdb_decoder.c @@ -1,8 +1,8 @@ #include "rai_rdb_decoder.h" -#include "current/v3/decode_v3.h" +#include "current/v4/decode_v4.h" -void *RAI_RDBLoadTensor(RedisModuleIO *io) { return RAI_RDBLoadTensor_v3(io); } +void *RAI_RDBLoadTensor(RedisModuleIO *io) { return RAI_RDBLoadTensor_v4(io); } -void *RAI_RDBLoadModel(RedisModuleIO *io) { return RAI_RDBLoadModel_v3(io); } +void *RAI_RDBLoadModel(RedisModuleIO *io) { return RAI_RDBLoadModel_v4(io); } -void *RAI_RDBLoadScript(RedisModuleIO *io) { return RAI_RDBLoadScript_v3(io); } +void *RAI_RDBLoadScript(RedisModuleIO *io) { return RAI_RDBLoadScript_v4(io); } diff --git a/src/serialization/RDB/encoder/rai_rdb_encode.c b/src/serialization/RDB/encoder/rai_rdb_encode.c index 4a018f53a..4b5d98a28 100644 --- a/src/serialization/RDB/encoder/rai_rdb_encode.c +++ b/src/serialization/RDB/encoder/rai_rdb_encode.c @@ -1,8 +1,8 @@ #include "rai_rdb_encode.h" -#include "v3/encode_v3.h" +#include "v4/encode_v4.h" -void RAI_RDBSaveTensor(RedisModuleIO *io, void *value) { RAI_RDBSaveTensor_v3(io, value); } +void RAI_RDBSaveTensor(RedisModuleIO *io, void *value) { RAI_RDBSaveTensor_v4(io, value); } -void RAI_RDBSaveModel(RedisModuleIO *io, void *value) { RAI_RDBSaveModel_v3(io, value); } +void RAI_RDBSaveModel(RedisModuleIO *io, void *value) { RAI_RDBSaveModel_v4(io, value); } -void RAI_RDBSaveScript(RedisModuleIO *io, void *value) { RAI_RDBSaveScript_v3(io, value); } +void RAI_RDBSaveScript(RedisModuleIO *io, void *value) { RAI_RDBSaveScript_v4(io, value); } diff --git a/src/serialization/RDB/encoder/v3/encode_v3.h b/src/serialization/RDB/encoder/v3/encode_v3.h deleted file mode 100644 index 4d7939d33..000000000 --- a/src/serialization/RDB/encoder/v3/encode_v3.h +++ /dev/null @@ -1,8 +0,0 @@ -#pragma once -#include "../../../serialization_include.h" - -void RAI_RDBSaveTensor_v3(RedisModuleIO *io, void *value); - -void RAI_RDBSaveModel_v3(RedisModuleIO *io, void *value); - -void RAI_RDBSaveScript_v3(RedisModuleIO *io, void *value); diff --git a/src/serialization/RDB/encoder/v3/encode_v3.c b/src/serialization/RDB/encoder/v4/encode_v4.c similarity index 78% rename from src/serialization/RDB/encoder/v3/encode_v3.c rename to src/serialization/RDB/encoder/v4/encode_v4.c index 343134033..65036b244 100644 --- a/src/serialization/RDB/encoder/v3/encode_v3.c +++ b/src/serialization/RDB/encoder/v4/encode_v4.c @@ -1,29 +1,28 @@ -#include "encode_v3.h" +#include "encode_v4.h" -void RAI_RDBSaveTensor_v3(RedisModuleIO *io, void *value) { +void RAI_RDBSaveTensor_v4(RedisModuleIO *io, void *value) { RAI_Tensor *tensor = (RAI_Tensor *)value; - size_t ndim = tensor->tensor.dl_tensor.ndim; - - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.device.device_type); - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.device.device_id); - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.bits); RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.code); - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.lanes); - RedisModule_SaveUnsigned(io, ndim); - for (size_t i = 0; i < ndim; i++) { - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.shape[i]); - } + RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.dtype.bits); + + size_t ndim = tensor->tensor.dl_tensor.ndim; + RedisModule_SaveSigned(io, ndim); for (size_t i = 0; i < ndim; i++) { - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.strides[i]); + RedisModule_SaveSigned(io, tensor->tensor.dl_tensor.shape[i]); } - RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.byte_offset); - size_t size = RAI_TensorByteSize(tensor); + size_t size = RAI_TensorByteSize(tensor); RedisModule_SaveStringBuffer(io, tensor->tensor.dl_tensor.data, size); + + if (tensor->tensor.dl_tensor.dtype.code == kDLString) { + for (size_t i = 0; i < RAI_TensorLength(tensor); i++) { + RedisModule_SaveUnsigned(io, tensor->tensor.dl_tensor.elements_length[i]); + } + } } -void RAI_RDBSaveModel_v3(RedisModuleIO *io, void *value) { +void RAI_RDBSaveModel_v4(RedisModuleIO *io, void *value) { RAI_Model *model = (RAI_Model *)value; char *buffer = NULL; size_t len = 0; @@ -69,7 +68,7 @@ void RAI_RDBSaveModel_v3(RedisModuleIO *io, void *value) { } } -void RAI_RDBSaveScript_v3(RedisModuleIO *io, void *value) { +void RAI_RDBSaveScript_v4(RedisModuleIO *io, void *value) { RAI_Script *script = (RAI_Script *)value; RedisModule_SaveStringBuffer(io, script->devicestr, strlen(script->devicestr) + 1); diff --git a/src/serialization/RDB/encoder/v4/encode_v4.h b/src/serialization/RDB/encoder/v4/encode_v4.h new file mode 100644 index 000000000..dc0016d0e --- /dev/null +++ b/src/serialization/RDB/encoder/v4/encode_v4.h @@ -0,0 +1,8 @@ +#pragma once +#include "../../../serialization_include.h" + +void RAI_RDBSaveTensor_v4(RedisModuleIO *io, void *value); + +void RAI_RDBSaveModel_v4(RedisModuleIO *io, void *value); + +void RAI_RDBSaveScript_v4(RedisModuleIO *io, void *value); diff --git a/src/version.h b/src/version.h index ac57685e3..5662f0a37 100644 --- a/src/version.h +++ b/src/version.h @@ -10,4 +10,4 @@ /* API versions. */ #define REDISAI_LLAPI_VERSION 1 -static const long long REDISAI_ENC_VER = 3; +static const long long REDISAI_ENC_VER = 4; diff --git a/tests/flow/test_data/identity_string.onnx b/tests/flow/test_data/identity_string.onnx new file mode 100644 index 000000000..ac7f4e80d --- /dev/null +++ b/tests/flow/test_data/identity_string.onnx @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:0c2dd3190846d88672dd870052d8ed7c7fb49df965b386f1db0b0fc3a849b280 +size 132 diff --git a/tests/flow/test_data/identity_string.pb b/tests/flow/test_data/identity_string.pb new file mode 100644 index 000000000..49d05b12f --- /dev/null +++ b/tests/flow/test_data/identity_string.pb @@ -0,0 +1,3 @@ +version https://git-lfs.github.com/spec/v1 +oid sha256:e7c501a876b112fe84e926292cf8aced8960bd6ca591bcec512d75701985e0f9 +size 93 diff --git a/tests/flow/test_data/onnx_string_identity.py b/tests/flow/test_data/onnx_string_identity.py new file mode 100644 index 000000000..e9d82128b --- /dev/null +++ b/tests/flow/test_data/onnx_string_identity.py @@ -0,0 +1,18 @@ +import onnx +from onnx import helper, TensorProto, ModelProto, defs + +# Create string_identity op +string_input = helper.make_tensor_value_info('input:0', TensorProto.STRING, [2, 2]) +string_output = helper.make_tensor_value_info('output:0', TensorProto.STRING, [2, 2]) +string_identity_def = helper.make_node('Identity', + inputs=['input:0'], + outputs=['output:0'], + name='StringIdentity') + +# Make a graph that contains this single op +graph_def = helper.make_graph(nodes=[string_identity_def], name='optional_input_graph', inputs=[string_input], + outputs=[string_output]) + +model_def = helper.make_model(graph_def) +onnx.checker.check_model(model_def) +onnx.save(model_def, "identity_string.onnx") diff --git a/tests/flow/test_data/tf_string_identity.py b/tests/flow/test_data/tf_string_identity.py new file mode 100644 index 000000000..49bbd008c --- /dev/null +++ b/tests/flow/test_data/tf_string_identity.py @@ -0,0 +1,10 @@ +import tensorflow as tf +from tensorflow.python.framework.graph_util import convert_variables_to_constants + +with tf.compat.v1.Session() as sess: + + a = tf.compat.v1.placeholder(tf.string, name='input_str') + c = tf.identity(a, name='c') + + frozen_graph = convert_variables_to_constants(sess, sess.graph_def, ['c']) + tf.compat.v1.train.write_graph(frozen_graph, './', 'identity_string.pb', as_text=False) diff --git a/tests/flow/test_serializations.py b/tests/flow/test_serializations.py index 5ab4f6a1f..d6e9a383d 100644 --- a/tests/flow/test_serializations.py +++ b/tests/flow/test_serializations.py @@ -273,6 +273,72 @@ def test_v3_tensor(self): self.env.assertEqual(values, [1, 2]) +class test_v4_rdb_load: + + def __init__(self): + self.env = Env() + + def test_v4_tf_model(self): + key_name = "tf_graph{1}" + con = get_connection(self.env, key_name) + model_rdb = b"\a\x81\x00\x8f\xff0\xe0\xc4,\x04\x02\x00\x05\x04CPU\x00\x05\x0bTF_GRAPH_V4\x02\x04\x02\x02\x02C\xe8\x02\x02\x05\x02a\x00\x05\x02b\x00\x02\x01\x05\x04mul\x00\x02@\x96\x02\x01\x05\xc3@j@\x96\x1f\n,\n\x01a\x12\x0bPlaceholder*\x0b\n\x05dtype\x12\x020\x01*\x05\r\n\x05sha \x0c\x05\x04:\x02\x18\x01\n -\x00b\xe0!-\x15\x19\n\x03mul\x12\x03Mul\x1a\x01a\x1a\x01b*\a\n\x01T@W\x0f\n\x1b\n\x01c\x12\bIdentity\x1a@'\xe0\x00\x1c\x01\x12\x00\x00\t\x00\aW\xce\xb9\xdc+\x8f\x97" + self.env.assertEqual(con.execute_command('FLUSHALL'), True) + con.restore(key_name, 0, model_rdb, True) + _, backend, _, device, _, tag, _, batchsize, _, minbatchsize, _, inputs, _, outputs, _, minbatchtimeout = \ + con.execute_command("AI.MODELGET", key_name, "META") + self.env.assertEqual([backend, device, tag, batchsize, minbatchsize, minbatchtimeout, inputs, outputs], + [b"TF", b"CPU", b"TF_GRAPH_V4", 4, 2, 1000, [b"a", b"b"], [b"mul"]]) + tf_model_run(self.env, key_name) + + def test_v4_torch_model(self): + key_name = "pt_minimal{1}" + con = get_connection(self.env, key_name) + model_rdb = b"\a\x81\x00\x8f\xff0\xe0\xc4,\x04\x02\x02\x05\x04CPU\x00\x05\rPT_MINIMAL_V4\x02\x04\x02\x02\x02C\xe8\x02\x02\x05\x02a\x00\x05\x02b\x00\x02\x01\x05\x01\x00\x02EH\x02\x01\x05\xc3C\x0eEH\x0ePK\x03\x04\x00\x00\b\b\x00\x00\x86\xb0zO\x00\xe0\x02\x00\x1a\x12\x00\x10\x00pt-minimal/versionFB\x0c\x00Z\xe0\x02\x00\n1\nPK\a\bS\xfcQg\x02 ;@\x03\x00P Q\x00\x14 Q\x00\b\xe0\bQ\x02\x1c\x004\xe0\x03Q\x13code/__torch__.pyFB0\xe0\x04[\xe0\x1b\x00\x1f5LK\n\x830\x10\xdd{\x8a\xb7T\xb0\x82\xdb\x80\xbd\x81\xbb\xeeJ\t\xa3\x19\xab\x90fd\x12[z\xfb\x1f\x06\xad\xab\xf7\x7f\xb2\xda7k\\$\xd8\xc8\t\x1d\xdab\xf4\x14#\xfao/n\xf3\\\x1eP\x99\x02\xb0v\x1f%\xa5\x17\xa7\xbc\xb06\x97\xef\x8f\xec&\xa5%,\xe1\t\x83A\xc4g\xc7\xf1\x84I\xf4C\xea\xca\xc8~2\x1fy\x99D\xc7\xd9\xda\xe6\xfc\xads\x0f \x83\x1b\x87(z\xc8\xe1\x94\x15.\xd7?5{\xa2\x9c6\r\xd8_\r\x1ar\xae\xa4\x1aC\r\xf2\xebL][\x15?A\x0b\x04a\xc1#I\x8e!\x0b\x00\xc8 \x03\xe1\x11\x0b\x02&\x00\x1e\xe1\x14\x0b\x0c.debug_pklFB\x1a\xe1\x12\x15\x1f5\x8eA\n\xc20\x10EcU\x90\x82+/0\xcb\x8a%\ap\xe5V\x06t\xdb\x9d\xa4mB\"m\xd3\x1f\xa4\x11q\xe7\xca\x1e\xc7S\xa8\xd72)\xe4m\x06\xde\x87?\xff\x99d\xb8)\x88\xc7\x10I\x90\x8cf\x86\xe1\x1f\xbc\xf0]\x9c\xbd\x05\xcf\xc1i[IzU\x8e\x0e\x95U\xbd\xbb\xb4\xdcI]\xa7!\xac\xb9\x00\xa1\xed\x9d\xd9\x1f:\x1bx#r`9\x94\xdb\xfd\x14\x06,w7\xdb\x01\x83\x1d\x94\xa9I\x8a\xb5o\r\x15\xaaS-kh\x15\xff0s\\\x8df\x81G<\xf9\xb7\x1f\x19\a|et\xbf\xe8\x9cY\xd2a\b\x04\xa7\x94\x1a\x02\x97!\x04\x00\xb2 \x03A\b\xe2\rf\x02\x18\x00#\xe1\x05\b\anstants.`\xfa\x00\x1f\xe0\x12\xfa`\x00\x03\x80\x02).Au\x03m/\tW a\x00\x00@\x03\xe0\x11l\x02\x13\x00;\xe0\x03l\x03data\x80g\x007\xe0\x17g\xe0\x0f\x00\x02\x80\x02c\xe2\x00\xc2\x10\nMyModule\nq\x00)\x81}(X#U\x0f\x00trainingq\x01\x88ubq\x02`\xac\x04z\xb8\x18\x811 \x1b@\x03\x02PK\x01C5#0\x83\x82\xe3\x03J\x00\x12 \x17\xe0\x05\x00\xe3\t\x90\x80?\xe3\x01p\xe2\x03~\x00\x1c\xe0\x04<\x00R \r\xe0\x02?\xe3\b~\xe0\aI\xe1\x03\xbf\x00& ;\xe0\x01\x00\x01^\x01\xe0\x15I\xe2\x01\xbc\x80S\xe0\x01\xdd\xe1\x03\xa6\x00\x18 \x17\xe0\x01\x00D?\xe0\x04\x9d\xe2\x02\a\xe0\aE\xe1\x03?\x00\x13\xe0\x01B \x00\x00\xd4!K\xe0\x02E\xc1\xe0\x04PK\x06\x06, \x1e@\x00\x02\x1e\x03-@\x06`\x00\x00\x05`\x05\xe0\x01\a\x00e \xd8@\x00\x01\x81\x03@\x05A\x9c\x01\x06\a \x06\x01\x00\xe6BV \x00@\x1e\x03PK\x05\x06 \n ;\x00\x05`/@+\x01\x00\x00\x00\t\x00K\x1a\xed\xeabG\xbdV" + self.env.assertEqual(con.execute_command('FLUSHALL'), True) + con.restore(key_name, 0, model_rdb, True) + _, backend, _, device, _, tag, _, batchsize, _, minbatchsize, _ , inputs, _, outputs, _, minbatchtimeout = con.execute_command("AI.MODELGET", key_name, "META") + self.env.assertEqual([backend, device, tag, batchsize, minbatchsize, minbatchtimeout, inputs, outputs], [b"TORCH", b"CPU", b"PT_MINIMAL_V4", 4, 2, 1000, [b'a', b'b'], [b'']]) + torch_model_run(self.env, key_name) + + def test_v4_troch_script(self): + key_name = "torch_script{1}" + con = get_connection(self.env, key_name) + script_rdb = b"\a\x81\x00\x8f\xd2\t\x12\x0fL\x04\x05\x04CPU\x00\x05\x0fTORCH_SCRIPT_V4\x05\xc3@|A\x01\x16def bar(tensors: List[T`\r\x05], key\xc0\x13\x01st@\x10\x02arg\xe0\x03\x10\x03):\n \x00\x03a = \xa0A\x02[0]`\x12\x00b\xe0\x02\x12\x001\x80\x12\x05return ,\x05+ b\n\nd\x80y\b_variadic\xe0I\x82\x00l\xe0\x03\x82\x00:\xe0\b\x83\x00l \xaa\x01\n\x00\x02\x02\x05\x04bar\x00\x05\rbar_variadic\x00\x00\t\x00\xad\xbe\xd5\xb3d\x05\xb8\xe8" + con.restore(key_name, 0, script_rdb, True) + _, device, _, tag, _, entry_points = con.execute_command("AI.SCRIPTGET", key_name, "META") + self.env.assertEqual([device, tag, entry_points], [b"CPU", b"TORCH_SCRIPT_V4", [b'bar', b'bar_variadic']]) + torch_script_run(self.env, key_name) + + def test_v4_onnx_model(self): + key_name = "linear_iris{1}" + con = get_connection(self.env, key_name) + model_rdb = b"\a\x81\x00\x8f\xff0\xe0\xc4,\x04\x02\x03\x05\x04CPU\x00\x05\x13ONNX_LINEAR_IRIS_V4\x02\x04\x02\x02\x02C\xe8\x02\x01\x05\x0cfloat_input\x00\x02\x01\x05\tvariable\x00\x02A\x15\x02\x01\x05\xc3@\xe6A\x15\x17\b\x05\x12\bskl2onnx\x1a\x051.4.9\"\aai.@\x0f\x1f(\x002\x00:\xe2\x01\n\x82\x01\n\x0bfloat_input\x12\bvariabl\x12e\x1a\x0fLinearRegressor\"\xe0\a\x10\x1f*%\n\x0ccoefficients=K\xfe\xc2\xbd=\xf7\xbe\x1c\xbd=/ii>=\x12\xe81\x1a?\xa0\x01\x06*\x14\n\nintercep $\x03\xa8\x1d\xb7= \x15\x01:\n\xa0\x88\x1f.ml\x12 2d76caf265cd4138a74199640a1\x06fc408Z\x1d\xe0\x05\xa5\n\x0e\n\x0c\b\x01\x12\b\n\x02\b\x01 \x03\x03\x04b\x1a\n\xe0\x00\xb7\xe0\x06\x1b\x03\x01B\x0e\n\xe0\x02j\x01\x10\x01\x00\t\x00\t'\xd6\x1b-\xad\xf2\x0f" + self.env.assertEqual(con.execute_command('FLUSHALL'), True) + con.restore(key_name, 0, model_rdb, True) + _, backend, _, device, _, tag, _, batchsize, _, minbatchsize, _ , inputs, _, outputs, _, minbatchtimeout = con.execute_command("AI.MODELGET", key_name, "META") + self.env.assertEqual([backend, device, tag, batchsize, minbatchsize, minbatchtimeout, inputs, outputs], [b"ONNX", b"CPU", b"ONNX_LINEAR_IRIS_V4", 4, 2, 1000, [b'float_input'], [b'variable']]) + onnx_model_run(self.env, key_name) + + def test_v4_tensor(self): + key_name = "tensor{1}" + con = get_connection(self.env, key_name) + tensor_rdb = b'\x07\x81\x00\x8f\xd3\x10\xd4\x8eD\x04\x02\x00\x02 \x02\x02\x02\x02\x02\x01\x05\x08\x01\x00\x00\x00\x02\x00\x00\x00\x00\t\x00Viy\xab4\xbe\xdd\x82' + self.env.assertEqual(con.execute_command('FLUSHALL'), True) + con.restore(key_name, 0, tensor_rdb, True) + _, tensor_type, _, tensor_shape = con.execute_command('AI.TENSORGET', key_name, 'META') + self.env.assertEqual([tensor_type, tensor_shape], [b"INT32", [2, 1]]) + values = con.execute_command('AI.TENSORGET', key_name, 'VALUES') + self.env.assertEqual(values, [1, 2]) + + # test RDB load of string tensor + str_tensor_rdb = b'\x07\x81\x00\x8f\xd3\x10\xd4\x8eD\x04\x02\x07\x02\x08\x02\x01\x02\x02\x05\x12str_val1\x00str_val2\x00\x02\x00\x02\t\x00\t\x00\x8b\x05Z\x0f:\x877O' + con.restore('string_tensor{1}', 0, str_tensor_rdb, True) + _, tensor_type, _, tensor_shape = con.execute_command('AI.TENSORGET', 'string_tensor{1}', 'META') + self.env.assertEqual([tensor_type, tensor_shape], [b"STRING", [2]]) + values = con.execute_command('AI.TENSORGET', 'string_tensor{1}', 'VALUES') + self.env.assertEqual(values, [b'str_val1', b'str_val2']) + + class TestAofRewrite: def __init__(self): diff --git a/tests/flow/tests_common.py b/tests/flow/tests_common.py index cd09f1017..178b819b2 100644 --- a/tests/flow/tests_common.py +++ b/tests/flow/tests_common.py @@ -37,147 +37,87 @@ def test_common_tensorset(env): def test_common_tensorset_error_replies(env): con = get_connection(env, '{0}') - test_data_path = os.path.join(os.path.dirname(__file__), 'test_data') - sample_filename = os.path.join(test_data_path, 'one.raw') - - with open(sample_filename, 'rb') as f: - sample_raw = f.read() + sample_raw = load_file_content('one.raw') ret = con.execute_command('AI.TENSORSET', 'sample_raw_ok{0}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw) env.assertEqual(ret, b'OK') # WRONGTYPE Operation against a key holding the wrong kind of value - try: - con.execute_command('SET','non-tensor{0}','value') - con.execute_command('AI.TENSORSET', 'non-tensor{0}', 'INT32', 2, 'unsupported', 2, 3) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual(exception.__str__(), "WRONGTYPE Operation against a key holding the wrong kind of value") + con.execute_command('SET', 'non-tensor{0}', 'value') + check_error_message(env, con, "WRONGTYPE Operation against a key holding the wrong kind of value", + 'AI.TENSORSET', 'non-tensor{0}', 'INT32', 2, 'unsupported', 2, 3) - # ERR invalid data type - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'INT128', 2, 'VALUES', 2, 3) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual(exception.__str__(), "invalid data type") + # ERR unsupported tensor type + check_error_message(env, con, "invalid data type", + 'AI.TENSORSET', 'z{0}', 'INT128', 2, 'VALUES', 2, 3) - # ERR invalid or negative value found in tensor shape - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'INT32', -1, 'VALUES', 2, 3) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("invalid or negative value found in tensor shape",exception.__str__()) + # ERR negative value found in tensor shape + check_error_message(env, con, "invalid or negative value found in tensor shape", + 'AI.TENSORSET', 'z{0}', 'INT32', -1, 'VALUES', 2, 3) # ERR invalid argument found in tensor shape - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'INT32', 2, 'unsupported', 2, 3) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("invalid or negative value found in tensor shape",exception.__str__()) + check_error_message(env, con, "invalid or negative value found in tensor shape", + 'AI.TENSORSET', 'z{0}', 'INT32', 2, 'not_a_number', 'VALUES', 2, 3) # ERR invalid value - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'FLOAT', 2, 'VALUES', 2, 'A') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("invalid value",exception.__str__()) + check_error_message(env, con, "invalid value", + 'AI.TENSORSET', 'z{0}', 'FLOAT', 2, 'VALUES', 2, 'A') # ERR invalid value - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'INT32', 2, 'VALUES', 2, 'A') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual(exception.__str__(), "invalid value") + check_error_message(env, con, "invalid value", + 'AI.TENSORSET', 'z{0}', 'INT32', 2, 'VALUES', 2, '1.87') # ERR invalid value - overflow - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'BOOL', 2, 'VALUES', 1, 2) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual(exception.__str__(), "invalid value") + check_error_message(env, con, "invalid value", + 'AI.TENSORSET', 'z{0}', 'BOOL', 2, 'VALUES', 1, 2) + + # ERR invalid string value + check_error_message(env, con, "C-string value contains null character", + 'AI.TENSORSET', 'z{0}', 'STRING', 2, 'VALUES', 'valid C-string', 'invalid\0C-string') # ERR invalid value - overflow - try: - con.execute_command('AI.TENSORSET', 'z{0}', 'INT8', 2, 'VALUES', -1, -128) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual(exception.__str__(), "invalid value") + check_error_message(env, con, "invalid value", + 'AI.TENSORSET', 'z{0}', 'INT8', 2, 'VALUES', -1, -128) - try: - con.execute_command('AI.TENSORSET', '1{0}') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) + # ERR insufficient number of args + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORSET' command", + 'AI.TENSORSET') - try: - con.execute_command('AI.TENSORSET', 'y{0}', 'FLOAT') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) + # ERR insufficient number of args + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORSET' command", + 'AI.TENSORSET', 'z{0}') - try: - con.execute_command('AI.TENSORSET', 'y{0}', 'FLOAT', 2, 'VALUES') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) + # ERR insufficient number of args + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORSET' command", + 'AI.TENSORSET', 'z{0}', 'FLOAT') - try: - con.execute_command('AI.TENSORSET', 'y{0}', 'FLOAT', 2, 'VALUES', 1) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) + # ERR insufficient number of values (expected 2) + check_error_message(env, con, "wrong number of values was given in 'AI.TENSORSET' command", + 'AI.TENSORSET', 'z{0}', 'FLOAT', 2, 'VALUES', 1) - try: - con.execute_command('AI.TENSORSET', 'y{0}', 'FLOAT', 2, 'VALUES', '1') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) + # ERR too many values (expected 2) + check_error_message(env, con, "wrong number of values was given in 'AI.TENSORSET' command", + 'AI.TENSORSET', 'z{0}', 'FLOAT', 2, 'VALUES', 1, 2, 3) - try: - con.execute_command('AI.TENSORSET', 'blob_tensor_moreargs{0}', 'FLOAT', 2, 'BLOB', '\x00', 'extra-argument') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("wrong number of arguments for 'AI.TENSORSET' command", exception.__str__()) + # ERR in blob - extra argument + check_error_message(env, con, "a single binary string should come after the BLOB argument in 'AI.TENSORSET' command", + 'AI.TENSORSET', 'blob_tensor_more_args{0}', 'FLOAT', 2, 'BLOB', '\x00', 'extra-argument') - try: - con.execute_command('AI.TENSORSET', 'blob_tensor_lessargs{0}', 'FLOAT', 2, 'BLOB') - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("wrong number of arguments for 'AI.TENSORSET' command", exception.__str__()) + # ERR in blob - missing argument + check_error_message(env, con, "a single binary string should come after the BLOB argument in 'AI.TENSORSET' command", + 'AI.TENSORSET', 'blob_tensor_less_args{0}', 'FLOAT', 2, 'BLOB') - # ERR data length does not match tensor shape and type - try: - con.execute_command('AI.TENSORSET', 'sample_raw_wrong_blob_for_dim{0}', 'FLOAT', 1, 1, 28, 280, 'BLOB', sample_raw) - env.assertFalse(True) - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("data length does not match tensor shape and type", exception.__str__()) + # ERR in blob - blob size is not compatible with tensor meta data + check_error_message(env, con, "data length does not match tensor shape and type", + 'AI.TENSORSET', 'sample_raw_wrong_blob_for_dim{0}', 'FLOAT', 1, 1, 28, 280, 'BLOB', sample_raw) + + # ERR in string tensor blob - number of strings is not compatible with tensor meta data + check_error_message(env, con, "Number of string elements in data blob does not match tensor length", + 'AI.TENSORSET', 'z{0}', 'STRING', 2, 'BLOB', 'single c-string\0') + + # ERR in string tensor blob - number of strings is not compatible with tensor meta data + check_error_message(env, con, "Number of string elements in data blob does not match tensor length", + 'AI.TENSORSET', 'z{0}', 'STRING', 2, 'BLOB', 'C-string\0followed by a non C-string') def test_common_tensorget(env): @@ -245,32 +185,33 @@ def test_common_tensorget(env): def test_common_tensorget_error_replies(env): con = get_connection(env, '{0}') + # ERR insufficient args + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORGET' command", + 'AI.TENSORGET') + + # ERR too many arguments + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORGET' command", + 'AI.TENSORGET', 'T_FLOAT{0}', 'META', 'VALUES', 'extra-arg') + # ERR tensor key is empty - try: - con.execute_command('AI.TENSORGET', 'empty{0}', 'unsupported') - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("tensor key is empty or in a different shard",exception.__str__()) + check_error_message(env, con, "tensor key is empty or in a different shard", + 'AI.TENSORGET', 'empty{0}') # WRONGTYPE Operation against a key holding the wrong kind of value - try: - con.execute_command('SET', 'non-tensor{0}', 'value') - con.execute_command('AI.TENSORGET', 'non-tensor{0}', 'unsupported') - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("WRONGTYPE Operation against a key holding the wrong kind of value",exception.__str__()) + con.execute_command('SET', 'non-tensor{0}', 'value') + check_error_message(env, con, "WRONGTYPE Operation against a key holding the wrong kind of value", + 'AI.TENSORGET', 'non-tensor{0}', 'unsupported') - # ERR unsupported data format ret = con.execute_command('AI.TENSORSET', "T_FLOAT{0}", "FLOAT", 2, 'VALUES', 1, 1) env.assertEqual(ret, b'OK') - try: - con.execute_command('AI.TENSORGET', 'T_FLOAT{0}', 'unsupported') - except Exception as e: - exception = e - env.assertEqual(type(exception), redis.exceptions.ResponseError) - env.assertEqual("unsupported data format",exception.__str__()) + + # ERR unsupported data format + check_error_message(env, con, "unsupported data format", + 'AI.TENSORGET', 'T_FLOAT{0}', 'unsupported_format') + + # ERR invalid data format + check_error_message(env, con, "both BLOB and VALUES specified", + 'AI.TENSORGET', 'T_FLOAT{0}', 'BLOB', 'VALUES') def test_common_tensorset_multiproc(env): @@ -366,3 +307,52 @@ def test_info_command(env): env.assertTrue('ai_children_used_cpu_sys' in cpu.keys()) env.assertTrue('ai_children_used_cpu_user' in cpu.keys()) env.assertTrue('ai_queue_CPU_bthread_n1_used_cpu_total' in cpu.keys()) + + +def test_string_tensor(env): + con = get_connection(env, '{0}') + + # test creation of string tensor from values - every C-string is valid (null-terminated) + # we also allow omitting the final null-character from the string + ret = con.execute_command('AI.TENSORSET', 'string_tensor_from_val{0}', 'STRING', 2, 'VALUES', 'str_val1', 'str_val2\0') + env.assertEqual(ret, b'OK') + + tensor_reply_values = con.execute_command('AI.TENSORGET', 'string_tensor_from_val{0}', 'VALUES') + tensor_reply_blob = con.execute_command('AI.TENSORGET', 'string_tensor_from_val{0}', 'BLOB') + env.assertEqual(tensor_reply_values, [b'str_val1', b'str_val2']) + env.assertEqual(tensor_reply_blob, b'str_val1\0str_val2\0') + + # test creation of string tensor from blob + ret = con.execute_command('AI.TENSORSET', 'string_tensor_from_blob{0}', 'STRING', 2, 'BLOB', 'str_blob1\0str_blob2\0') + env.assertEqual(ret, b'OK') + + tensor_reply_values = con.execute_command('AI.TENSORGET', 'string_tensor_from_blob{0}', 'VALUES') + tensor_reply_blob = con.execute_command('AI.TENSORGET', 'string_tensor_from_blob{0}', 'BLOB') + env.assertEqual(tensor_reply_values, [b'str_blob1', b'str_blob2']) + env.assertEqual(tensor_reply_blob, b'str_blob1\0str_blob2\0') + + # test for empty string tensor + ret = con.execute_command('AI.TENSORSET', 'empty_string_tensor{0}', 'STRING', 2) + env.assertEqual(ret, b'OK') + _, tensor_dtype, _, tensor_dim, _, tensor_values = con.execute_command('AI.TENSORGET', 'empty_string_tensor{0}', 'META', 'VALUES') + env.assertEqual(tensor_dtype, b'STRING') + env.assertEqual(tensor_dim, [2]) + env.assertEqual(tensor_values, [b'', b'']) + + # advanced - tensor with more than one dimension + ret = con.execute_command('AI.TENSORSET', 'string_tensor_few_dims{0}', 'STRING', 2, 2, 'BLOB', + 'str11\0str12\0str21\0str22\0') + env.assertEqual(ret, b'OK') + + _, tensor_dtype, _, tensor_dim, _, tensor_values = con.execute_command('AI.TENSORGET', 'string_tensor_few_dims{0}', 'META', 'VALUES') + env.assertEqual(tensor_dtype, b'STRING') + env.assertEqual(tensor_dim, [2, 2]) + env.assertEqual(tensor_values, [b'str11', b'str12', b'str21', b'str22']) + + # advanced - tensor with non ascii characters + ret = con.execute_command('AI.TENSORSET', 'string_tensor_non-ascii{0}', 'STRING', 2, 'BLOB', + 'english\0עברית\0') + env.assertEqual(ret, b'OK') + + tensor_reply_values = con.execute_command('AI.TENSORGET', 'string_tensor_non-ascii{0}', 'VALUES') + env.assertEqual(tensor_reply_values[1].decode('utf-8'), 'עברית') diff --git a/tests/flow/tests_dag_basic.py b/tests/flow/tests_dag_basic.py index 2345515ad..df03e4d22 100644 --- a/tests/flow/tests_dag_basic.py +++ b/tests/flow/tests_dag_basic.py @@ -231,30 +231,21 @@ def test_dag_with_timeout(env): env.assertEqual(b'TIMEDOUT', res) -def test_dag_with_error(env): - if not TEST_TF: +def test_dag_with_string_tensor(env): + if not TEST_ONNX: + env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True) return con = get_connection(env, '{1}') - tf_model = load_file_content('graph.pb') - ret = con.execute_command('AI.MODELSTORE', 'tf_model{1}', 'TF', DEVICE, - 'INPUTS', 2, 'a', 'b', - 'OUTPUTS', 1, 'mul', - 'BLOB', tf_model) - env.assertEqual(b'OK', ret) - - # Run the model from DAG context, where MODELEXECUTE op fails due to dim mismatch in one of the tensors inputs: - # the input tensor 'b' is considered as tensor with dim 2X2X3 initialized with zeros, while the model expects that - # both inputs to node 'mul' will be with dim 2. - ret = con.execute_command('AI.DAGEXECUTE_RO', 'ROUTING', '{1}', - '|>', 'AI.TENSORSET', 'a', 'FLOAT', 2, 'VALUES', 2, 3, - '|>', 'AI.TENSORSET', 'b', 'FLOAT', 2, 2, 3, - '|>', 'AI.MODELEXECUTE', 'tf_model{1}', 'INPUTS', 2, 'a', 'b', 'OUTPUTS', 1, 'tD', - '|>', 'AI.TENSORGET', 'tD', 'VALUES') - - # Expect that the MODELEXECUTE op will raise an error, and the last TENSORGET op will not be executed - env.assertEqual(ret[0], b'OK') - env.assertEqual(ret[1], b'OK') - env.assertEqual(ret[3], b'NA') - env.assertEqual(type(ret[2]), redis.exceptions.ResponseError) - env.assertTrue(str(ret[2]).find('Incompatible shapes: [2] vs. [2,2,3] \t [[{{node mul}}]]') >= 0) + model_pb = load_file_content('identity_string.onnx') + ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + + # Execute onnx model whose input is string tensor with shape [2,2], that outputs the input + string_tensor_blob = b'input11\0input12\0input21\0input22\0' + ret = con.execute_command('AI.DAGEXECUTE', 'ROUTING', '{1}', + '|>', 'AI.TENSORSET', 'in_tensor{1}', 'STRING', 2, 2, 'BLOB', string_tensor_blob, + '|>', 'AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'in_tensor{1}', 'OUTPUTS', 1, 'out_tensor{1}', + '|>', 'AI.TENSORGET', 'out_tensor{1}', 'VALUES') + + env.assertEqual(ret, [b'OK', b'OK', [b'input11', b'input12', b'input21', b'input22']]) diff --git a/tests/flow/tests_dag_errors.py b/tests/flow/tests_dag_errors.py index acb8bd07b..77e16f99d 100644 --- a/tests/flow/tests_dag_errors.py +++ b/tests/flow/tests_dag_errors.py @@ -236,3 +236,48 @@ def test_dag_crossslot_violation_errors(env): 'OUTPUTS', 1, 'resultTensor:{1}', ) check_error_message(env, con, "CROSSSLOT Keys in request don't hash to the same slot", *command) + + +def test_dag_tensorget_tensorset_errors(env): + con = get_connection(env, '{1}') + + # ERR insufficient args for tensor get + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORGET' command", + "AI.DAGEXECUTE ROUTING a{1} |> AI.TENSORSET a{1} FLOAT 1 |> AI.TENSORGET") + + # ERR too many args for tensor get + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORGET' command", + "AI.DAGEXECUTE ROUTING a{1} |> AI.TENSORSET a{1} FLOAT 1 |> AI.TENSORGET a{1} META BLOB extra") + + # ERR insufficient args for tensor set + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORSET' command", + "AI.DAGEXECUTE ROUTING a{1} |> AI.TENSORSET a{1} FLOAT 1 |> AI.TENSORSET") + + +def test_dag_error_before_tensorget_op(env): + if not TEST_TF: + return + + con = get_connection(env, '{1}') + tf_model = load_file_content('graph.pb') + ret = con.execute_command('AI.MODELSTORE', 'tf_model{1}', 'TF', DEVICE, + 'INPUTS', 2, 'a', 'b', + 'OUTPUTS', 1, 'mul', + 'BLOB', tf_model) + env.assertEqual(b'OK', ret) + + # Run the model from DAG context, where MODELEXECUTE op fails due to dim mismatch in one of the tensors inputs: + # the input tensor 'b' is considered as tensor with dim 2X2X3 initialized with zeros, while the model expects that + # both inputs to node 'mul' will be with dim 2. + ret = con.execute_command('AI.DAGEXECUTE_RO', 'ROUTING', '{1}', + '|>', 'AI.TENSORSET', 'a', 'FLOAT', 2, 'VALUES', 2, 3, + '|>', 'AI.TENSORSET', 'b', 'FLOAT', 2, 2, 3, + '|>', 'AI.MODELEXECUTE', 'tf_model{1}', 'INPUTS', 2, 'a', 'b', 'OUTPUTS', 1, 'tD', + '|>', 'AI.TENSORGET', 'tD', 'VALUES') + + # Expect that the MODELEXECUTE op will raise an error, and the last TENSORGET op will not be executed + env.assertEqual(ret[0], b'OK') + env.assertEqual(ret[1], b'OK') + env.assertEqual(ret[3], b'NA') + env.assertEqual(type(ret[2]), redis.exceptions.ResponseError) + env.assertTrue('Incompatible shapes: [2] vs. [2,2,3] \t [[{{node mul}}]]' in str(ret[2])) diff --git a/tests/flow/tests_deprecated_commands.py b/tests/flow/tests_deprecated_commands.py index 708afa07b..d7210297f 100644 --- a/tests/flow/tests_deprecated_commands.py +++ b/tests/flow/tests_deprecated_commands.py @@ -474,6 +474,18 @@ def test_dagrun_common_errors(env): check_error_message(env, con, "PERSIST cannot be specified in a read-only DAG", "AI.DAGRUN_RO PERSIST 1 tensor1{1} |> AI.TENSORSET tensor1{1} FLOAT 1 2 VALUES 5 10") + # ERR insufficient args for tensor get + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORGET' command", + "AI.DAGRUN |> AI.TENSORSET a{1} FLOAT 1 |> AI.TENSORGET") + + # ERR too many args for tensor get + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORGET' command", + "AI.DAGRUN |> AI.TENSORSET a{1} FLOAT 1 |> AI.TENSORGET a{1} META BLOB extra") + + # ERR insufficient args for tensor set + check_error_message(env, con, "wrong number of arguments for 'AI.TENSORSET' command", + "AI.DAGRUN |> AI.TENSORSET a{1} FLOAT 1 |> AI.TENSORSET") + def test_dagrun_modelrun_multidevice_resnet_ensemble_alias(env): if (not TEST_TF or not TEST_PT): diff --git a/tests/flow/tests_llapi.py b/tests/flow/tests_llapi.py index ae2878fc1..b7ef44938 100644 --- a/tests/flow/tests_llapi.py +++ b/tests/flow/tests_llapi.py @@ -154,3 +154,14 @@ def test_dagrun_multidevice_resnet(env): ret = con.execute_command("RAI_llapi.DAG_resnet") env.assertEqual(ret, b'DAG resnet success') + + +@with_test_module +def test_tensor_create(env): + con = get_connection(env, '{1}') + ret = con.execute_command("RAI_llapi.CreateTensor") + env.assertEqual(ret, b'create tensor test success') + ret = con.execute_command("RAI_llapi.ConcatenateTensors") + env.assertEqual(ret, b'concatenate tensors test success') + ret = con.execute_command("RAI_llapi.SliceTensor") + env.assertEqual(ret, b'slice tensor test success') diff --git a/tests/flow/tests_onnx.py b/tests/flow/tests_onnx.py index aa2ee836b..1fd83d27b 100644 --- a/tests/flow/tests_onnx.py +++ b/tests/flow/tests_onnx.py @@ -52,7 +52,7 @@ def test_onnx_modelrun_mnist(env): con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 1, 1, 28, 28, 'BLOB', sample_raw) check_error_message(env, con, "Number of keys given as INPUTS here does not match model definition", - 'AI.MODELEXECUTE', 'm{1}', 'INPUTS', 3, 'a{1}', 'b{1}', 'c{1}', 'OTUPUTS', 'c{1}') + 'AI.MODELEXECUTE', 'm{1}', 'INPUTS', 3, 'a{1}', 'b{1}', 'c{1}', 'OUTPUTS', 'c{1}') con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'a{1}', 'OUTPUTS', 1, 'b{1}') @@ -68,6 +68,63 @@ def test_onnx_modelrun_mnist(env): env.assertEqual(values2, values) +def test_onnx_string_tensors(env): + if not TEST_ONNX: + env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True) + return + + con = get_connection(env, '{1}') + model_pb = load_file_content('identity_string.onnx') + ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'ONNX', DEVICE, 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + + # Execute onnx model whose input is string tensor with shape [2,2], that outputs the input + string_tensor_blob = b'input11\0input12\0input21\0input22\0' + con.execute_command('AI.TENSORSET', 'in_tensor{1}', 'STRING', 2, 2, 'BLOB', string_tensor_blob) + ret = con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'in_tensor{1}', 'OUTPUTS', 1, 'out_tensor{1}') + env.assertEqual(ret, b'OK') + + _, tensor_dtype, _, tensor_dim, _, tensor_values = con.execute_command('AI.TENSORGET', 'out_tensor{1}', 'META', 'VALUES') + env.assertEqual(tensor_dtype, b'STRING') + env.assertEqual(tensor_dim, [2, 2]) + env.assertEqual(tensor_values, [b'input11', b'input12', b'input21', b'input22']) + + if env.useSlaves: + ensureSlaveSynced(con, env) + slave_con = env.getSlaveConnection() + slave_tensor_values = slave_con.execute_command('AI.TENSORGET', 'out_tensor{1}', 'VALUES') + env.assertEqual(tensor_values, slave_tensor_values) + + +def test_onnx_string_tensors_batching(env): + if not TEST_ONNX: + env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True) + return + + con = get_connection(env, '{1}') + model_pb = load_file_content('identity_string.onnx') + ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'ONNX', DEVICE, 'BATCHSIZE', 2, 'MINBATCHSIZE', 2, + 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + con.execute_command('AI.TENSORSET', 'first_batch{1}', 'STRING', 1, 2, 'VALUES', 'this is\0', 'the first batch\0') + con.execute_command('AI.TENSORSET', 'second_batch{1}', 'STRING', 1, 2, 'VALUES', 'that is\0', 'the second batch\0') + + def run(): + con2 = get_connection(env, '{1}') + con2.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'first_batch{1}', 'OUTPUTS', 1, 'first_output{1}') + + t = threading.Thread(target=run) + t.start() + + con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'second_batch{1}', 'OUTPUTS', 1, 'second_output{1}') + t.join() + + out_values = con.execute_command('AI.TENSORGET', 'first_batch{1}', 'VALUES') + env.assertEqual(out_values, [b'this is', b'the first batch']) + out_values = con.execute_command('AI.TENSORGET', 'second_batch{1}', 'VALUES') + env.assertEqual(out_values, [b'that is', b'the second batch']) + + def test_onnx_modelrun_batchdim_mismatch(env): if not TEST_ONNX: env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True) @@ -260,36 +317,6 @@ def test_onnx_modelrun_disconnect(env): env.assertEqual(ret, None) -def test_onnx_model_rdb_save_load(env): - env.skipOnCluster() - if env.useAof or not TEST_ONNX: - env.debugPrint("skipping {}".format(sys._getframe().f_code.co_name), force=True) - return - - linear_model = load_file_content('linear_iris.onnx') - - con = get_connection(env, '{1}') - ret = con.execute_command('AI.MODELSTORE', 'linear{1}', 'ONNX', DEVICE, 'BLOB', linear_model) - env.assertEqual(ret, b'OK') - - model_serialized_memory = con.execute_command('AI.MODELGET', 'linear{1}', 'BLOB') - - ensureSlaveSynced(con, env) - ret = con.execute_command('SAVE') - env.assertEqual(ret, True) - - env.stop() - env.start() - con = get_connection(env, '{1}') - model_serialized_after_rdbload = con.execute_command('AI.MODELGET', 'linear{1}', 'BLOB') - env.assertEqual(len(model_serialized_memory), len(model_serialized_after_rdbload)) - env.assertEqual(len(linear_model), len(model_serialized_after_rdbload)) - # Assert in memory model binary is equal to loaded model binary - env.assertTrue(model_serialized_memory == model_serialized_after_rdbload) - # Assert input model binary is equal to loaded model binary - env.assertTrue(linear_model == model_serialized_after_rdbload) - - def tests_onnx_info(env): if not TEST_ONNX: env.debugPrint("skipping {} since TEST_ONNX=0".format(sys._getframe().f_code.co_name), force=True) diff --git a/tests/flow/tests_pytorch.py b/tests/flow/tests_pytorch.py index 5cdd0be92..e4052b0ad 100644 --- a/tests/flow/tests_pytorch.py +++ b/tests/flow/tests_pytorch.py @@ -562,48 +562,6 @@ def test_pytorch_modelscan_scriptscan(env): env.assertEqual(2, len(ret[1])) -def test_pytorch_model_rdb_save_load(env): - env.skipOnCluster() - if env.useAof or not TEST_PT: - env.debugPrint("skipping {}".format(sys._getframe().f_code.co_name), force=True) - return - if DEVICE == "GPU": - env.debugPrint("skipping {} since it's hanging CI".format(sys._getframe().f_code.co_name), force=True) - return - - model_pb = load_file_content('pt-minimal.pt') - - con = get_connection(env, '{1}') - - ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'TORCH', DEVICE, 'BLOB', model_pb) - env.assertEqual(ret, b'OK') - - model_serialized_memory = con.execute_command('AI.MODELGET', 'm{1}', 'BLOB') - - con.execute_command('AI.TENSORSET', 'a{1}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3) - con.execute_command('AI.TENSORSET', 'b{1}', 'FLOAT', 2, 2, 'VALUES', 2, 3, 2, 3) - con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 2, 'a{1}', 'b{1}', 'OUTPUTS', 1, 'c{1}') - _, dtype_memory, _, shape_memory, _, data_memory = con.execute_command('AI.TENSORGET', 'c{1}', 'META', 'VALUES') - - ensureSlaveSynced(con, env) - ret = con.execute_command('SAVE') - env.assertEqual(ret, True) - - env.stop() - env.start() - con = get_connection(env, '{1}') - model_serialized_after_rdbload = con.execute_command('AI.MODELGET', 'm{1}', 'BLOB') - con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 2, 'a{1}', 'b{1}', 'OUTPUTS', 1, 'c{1}') - _, dtype_after_rdbload, _, shape_after_rdbload, _, data_after_rdbload = con.execute_command('AI.TENSORGET', 'c{1}', 'META', 'VALUES') - - # Assert in memory model metadata is equal to loaded model metadata - env.assertTrue(model_serialized_memory[1:6] == model_serialized_after_rdbload[1:6]) - # Assert in memory tensor data is equal to loaded tensor data - env.assertTrue(dtype_memory == dtype_after_rdbload) - env.assertTrue(shape_memory == shape_after_rdbload) - env.assertTrue(data_memory == data_after_rdbload) - - def test_parallelism(): env = Env(moduleArgs='INTRA_OP_PARALLELISM 1 INTER_OP_PARALLELISM 1') if not TEST_PT: diff --git a/tests/flow/tests_tensorflow.py b/tests/flow/tests_tensorflow.py index f9ec102a7..c590ae6f6 100644 --- a/tests/flow/tests_tensorflow.py +++ b/tests/flow/tests_tensorflow.py @@ -698,3 +698,64 @@ def test_tf_info(env): backends_info = get_info_section(con, 'backends_info') env.assertTrue('ai_TensorFlow_version' in backends_info) + + +@skip_if_no_TF +def test_tf_string_tensors(env): + con = get_connection(env, '{1}') + + model_pb = load_file_content('identity_string.pb') + ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'TF', DEVICE, 'INPUTS', 1, 'input_str', + 'OUTPUTS', 1, 'c', 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + + # Execute tf model whose input is string tensor (with any shape), that outputs the input + # TF will copy each string and use a designated structure (TF_TString). This structure handles + # differently for short and long strings (24 chars is the threshold) - short strings are kept in an array + # that TF_TString holds, while long strings are allocated such that TF_TString only holds a pointer to the data. + string_tensor_blob = b'input11\0input12\0input21\0here we want to test a string longer than 24 chars,' \ + b' to force heap alloc in tf\0' + con.execute_command('AI.TENSORSET', 'in_tensor{1}', 'STRING', 2, 2, 'BLOB', string_tensor_blob) + con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'in_tensor{1}', 'OUTPUTS', 1, 'out_tensor{1}') + env.assertEqual(ret, b'OK') + + _, tensor_dtype, _, tensor_dim, _, tensor_values = con.execute_command('AI.TENSORGET', 'out_tensor{1}', 'META', 'VALUES') + env.assertEqual(tensor_dtype, b'STRING') + env.assertEqual(tensor_dim, [2, 2]) + env.assertEqual(tensor_values, [b'input11', b'input12', b'input21', + b'here we want to test a string longer than 24 chars, to force heap alloc in tf']) + + if env.useSlaves: + ensureSlaveSynced(con, env) + slave_con = env.getSlaveConnection() + slave_tensor_values = slave_con.execute_command('AI.TENSORGET', 'out_tensor{1}', 'VALUES') + env.assertEqual(tensor_values, slave_tensor_values) + + +@skip_if_no_TF +def test_tf_string_tensors_batching(env): + con = get_connection(env, '{1}') + + model_pb = load_file_content('identity_string.pb') + ret = con.execute_command('AI.MODELSTORE', 'm{1}', 'TF', DEVICE, 'BATCHSIZE', 2, 'MINBATCHSIZE', 2, + 'INPUTS', 1, 'input_str', 'OUTPUTS', 1, 'c', 'BLOB', model_pb) + env.assertEqual(ret, b'OK') + con.execute_command('AI.TENSORSET', 'first_batch{1}', 'STRING', 1, 2, 'VALUES', 'this is', 'the first batch') + con.execute_command('AI.TENSORSET', 'second_batch{1}', 'STRING', 1, 2, 'VALUES', 'that is', 'the second batch') + + def run(): + con2 = get_connection(env, '{1}') + con2.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'first_batch{1}', + 'OUTPUTS', 1, 'first_output{1}') + + t = threading.Thread(target=run) + t.start() + + con.execute_command('AI.MODELEXECUTE', 'm{1}', 'INPUTS', 1, 'second_batch{1}', + 'OUTPUTS', 1, 'second_output{1}') + t.join() + + out_values = con.execute_command('AI.TENSORGET', 'first_batch{1}', 'VALUES') + env.assertEqual(out_values, [b'this is', b'the first batch']) + out_values = con.execute_command('AI.TENSORGET', 'second_batch{1}', 'VALUES') + env.assertEqual(out_values, [b'that is', b'the second batch']) diff --git a/tests/module/LLAPI.c b/tests/module/LLAPI.c index 685313536..99d9bbc61 100644 --- a/tests/module/LLAPI.c +++ b/tests/module/LLAPI.c @@ -294,6 +294,121 @@ int RAI_llapi_DAG_resnet(RedisModuleCtx *ctx, RedisModuleString **argv, int argc return RedisModule_ReplyWithSimpleString(ctx, "DAG resnet success"); } +int RAI_llapi_CreateTensor(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + + if (argc > 1) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + int n_dims = 2; + long long dims[] = {1, 4}; + // Try to create a tensor with a non-supported data type. + RAI_Tensor *t = RedisAI_TensorCreate("INVALID", dims, n_dims); + if (t != NULL) { + return RedisModule_ReplyWithSimpleString(ctx, + "invalid data type tensor create test failed"); + } + + // create an empty tensor and validate that in contains zeros + t = RedisAI_TensorCreate("INT8", dims, n_dims); + int8_t expected_blob[8] = {0}; + if (t == NULL || RedisAI_TensorLength(t) != dims[0] * dims[1] || + memcmp(RedisAI_TensorData(t), expected_blob, 4) != 0) { + return RedisModule_ReplyWithSimpleString(ctx, "empty tensor create test failed"); + } + RedisAI_TensorFree(t); + + // This should fail since the blob contains only one null-terminated string, while the tensor's + // len should be 4. + RAI_Tensor *t1 = RedisAI_TensorCreate("STRING", dims, n_dims); + const char *data_blob1 = "only one string\0"; + if (RedisAI_TensorSetData(t1, data_blob1, strlen(data_blob1)) != 0) { + return RedisModule_ReplyWithSimpleString(ctx, "invalid string tensor data set test failed"); + } + RedisAI_TensorFree(t1); + return RedisModule_ReplyWithSimpleString(ctx, "create tensor test success"); +} + +int RAI_llapi_ConcatenateTensors(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + if (argc > 1) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + + int n_dims = 2; + long long dims[] = {1, 4}; + + // test concatenation of string tensors + RAI_Tensor *t1 = RedisAI_TensorCreate("STRING", dims, n_dims); + const char *data_blob1 = "first\0second\0third\0forth\0"; + size_t len_data_blob1 = 25; + if (RedisAI_TensorSetData(t1, data_blob1, len_data_blob1) != 1) { + return RedisModule_ReplyWithSimpleString(ctx, "string tensor data set test failed"); + } + + // the second tensor's shape is [2,4], while the previous shape was [1,4] + dims[0] = 2; + const char *data_blob2 = "A\0B\0C\0D\0E\0F\0G\0H\0"; + size_t len_data_blob2 = 16; + RAI_Tensor *t2 = RedisAI_TensorCreate("STRING", dims, n_dims); + if (RedisAI_TensorSetData(t2, data_blob2, len_data_blob2) != 1) { + return RedisModule_ReplyWithSimpleString(ctx, "string tensor data set test failed"); + } + + RAI_Tensor *tensors[] = {t1, t2}; + RAI_Tensor *batched_tensor = RedisAI_TensorCreateByConcatenatingTensors(tensors, 2); + RedisAI_TensorFree(t1); + RedisAI_TensorFree(t2); + const char *expected_batched_data = "first\0second\0third\0forth\0A\0B\0C\0D\0E\0F\0G\0H\0"; + size_t expected_batched_data_len = len_data_blob1 + len_data_blob2; + if (batched_tensor == NULL || RedisAI_TensorNumDims(batched_tensor) != 2 || + RedisAI_TensorDim(batched_tensor, 0) != 3 || RedisAI_TensorDim(batched_tensor, 1) != 4 || + memcmp(expected_batched_data, RedisAI_TensorData(batched_tensor), + expected_batched_data_len) != 0) { + return RedisModule_ReplyWithSimpleString(ctx, "string tensor concatenation test failed"); + } + RedisAI_TensorFree(batched_tensor); + return RedisModule_ReplyWithSimpleString(ctx, "concatenate tensors test success"); +} + +int RAI_llapi_SliceTensor(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { + REDISMODULE_NOT_USED(argv); + if (argc > 1) { + RedisModule_WrongArity(ctx); + return REDISMODULE_OK; + } + int n_dims = 2; + long long dims[] = {3, 4}; + RAI_Tensor *batched_tensor = RedisAI_TensorCreate("STRING", dims, n_dims); + const char *batched_data = "first\0second\0third\0forth\0A\0B\0C\0D\0E\0F\0G\0H\0"; + size_t len_data_batch1 = 25; + size_t len_data_batch2 = 16; + RedisAI_TensorSetData(batched_tensor, batched_data, len_data_batch1 + len_data_batch2); + + // test slicing string tensors + RAI_Tensor *t1 = RedisAI_TensorCreateBySlicingTensor(batched_tensor, 0, 1); + RAI_Tensor *t2 = RedisAI_TensorCreateBySlicingTensor(batched_tensor, 1, 2); + RedisAI_TensorFree(batched_tensor); + + if (t1 == NULL || RedisAI_TensorNumDims(t1) != 2 || RedisAI_TensorDim(t1, 0) != 1 || + RedisAI_TensorDim(t1, 1) != 4 || + memcmp(batched_data, RedisAI_TensorData(t1), len_data_batch1) != 0) { + return RedisModule_ReplyWithSimpleString(ctx, "string tensor slicing test failed"); + } + + if (t2 == NULL || RedisAI_TensorNumDims(t2) != 2 || RedisAI_TensorDim(t2, 0) != 2 || + RedisAI_TensorDim(t2, 1) != 4 || + memcmp(batched_data + len_data_batch1, RedisAI_TensorData(t2), len_data_batch2) != 0) { + return RedisModule_ReplyWithSimpleString(ctx, "string tensor slicing test failed"); + } + RedisAI_TensorFree(t1); + RedisAI_TensorFree(t2); + return RedisModule_ReplyWithSimpleString(ctx, "slice tensor test success"); +} + int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) { REDISMODULE_NOT_USED(argv); REDISMODULE_NOT_USED(argc); @@ -327,5 +442,20 @@ int RedisModule_OnLoad(RedisModuleCtx *ctx, RedisModuleString **argv, int argc) return REDISMODULE_ERR; } + if (RedisModule_CreateCommand(ctx, "RAI_llapi.CreateTensor", RAI_llapi_CreateTensor, "", 0, 0, + 0) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } + + if (RedisModule_CreateCommand(ctx, "RAI_llapi.ConcatenateTensors", RAI_llapi_ConcatenateTensors, + "", 0, 0, 0) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } + + if (RedisModule_CreateCommand(ctx, "RAI_llapi.SliceTensor", RAI_llapi_SliceTensor, "", 0, 0, + 0) == REDISMODULE_ERR) { + return REDISMODULE_ERR; + } + return REDISMODULE_OK; }