Enabled updating internal IO info structures for TVM model. Compute static fields related to the model I/O.
Signed-off-by: Srikanth Yalavarthi <syalavar...@marvell.com> --- drivers/ml/cnxk/mvtvm_ml_model.c | 105 +++++++++++++++++++++++++++++++ drivers/ml/cnxk/mvtvm_ml_model.h | 1 + drivers/ml/cnxk/mvtvm_ml_ops.c | 3 + 3 files changed, 109 insertions(+) diff --git a/drivers/ml/cnxk/mvtvm_ml_model.c b/drivers/ml/cnxk/mvtvm_ml_model.c index 425a682209..86f465a645 100644 --- a/drivers/ml/cnxk/mvtvm_ml_model.c +++ b/drivers/ml/cnxk/mvtvm_ml_model.c @@ -7,10 +7,14 @@ #include <rte_mldev.h> +#include <mldev_utils.h> + #include <roc_api.h> #include "mvtvm_ml_model.h" +#include "cnxk_ml_model.h" + /* Objects list */ char mvtvm_object_list[ML_MVTVM_MODEL_OBJECT_MAX][RTE_ML_STR_MAX] = {"mod.so", "mod.json", "mod.params"}; @@ -71,3 +75,104 @@ mvtvm_ml_model_blob_parse(struct rte_ml_model_params *params, struct mvtvm_ml_mo return -EINVAL; } + +static enum rte_ml_io_type +mvtvm_ml_io_type_map(uint8_t type) +{ + switch (type) { + case kDLInt: + return RTE_ML_IO_TYPE_INT32; + case kDLUInt: + return RTE_ML_IO_TYPE_UINT32; + case kDLFloat: + return RTE_ML_IO_TYPE_FP32; + case kDLBfloat: + return RTE_ML_IO_TYPE_BFLOAT16; + } + + return RTE_ML_IO_TYPE_UNKNOWN; +} + +void +mvtvm_ml_model_io_info_update(struct cnxk_ml_model *model) +{ + struct tvmdp_model_metadata *metadata; + int32_t i; + int32_t j; + + if (model->subtype == ML_CNXK_MODEL_SUBTYPE_TVM_MRVL) + goto tvm_mrvl_model; + + metadata = &model->mvtvm.metadata; + + /* Inputs, set for layer_id = 0 */ + model->mvtvm.info.nb_inputs = metadata->model.num_input; + model->mvtvm.info.total_input_sz_d = 0; + model->mvtvm.info.total_input_sz_q = 0; + for (i = 0; i < metadata->model.num_input; i++) { + strncpy(model->mvtvm.info.input[i].name, metadata->input[i].name, + TVMDP_NAME_STRLEN); + model->mvtvm.info.input[i].dtype = + mvtvm_ml_io_type_map(metadata->input[i].datatype.code); + model->mvtvm.info.input[i].qtype = + mvtvm_ml_io_type_map(metadata->input[i].model_datatype.code); + model->mvtvm.info.input[i].nb_dims = metadata->input[i].ndim; + + model->mvtvm.info.input[i].nb_elements = 1; + for (j = 0; j < metadata->input[i].ndim; j++) { + model->mvtvm.info.input[i].shape[j] = metadata->input[i].shape[j]; + model->mvtvm.info.input[i].nb_elements *= metadata->input[i].shape[j]; + } + + model->mvtvm.info.input[i].sz_d = + model->mvtvm.info.input[i].nb_elements * + rte_ml_io_type_size_get(model->mvtvm.info.input[i].dtype); + model->mvtvm.info.input[i].sz_q = + model->mvtvm.info.input[i].nb_elements * + rte_ml_io_type_size_get(model->mvtvm.info.input[i].qtype); + + model->mvtvm.info.total_input_sz_d += model->mvtvm.info.input[i].sz_d; + model->mvtvm.info.total_input_sz_q += model->mvtvm.info.input[i].sz_q; + + plt_ml_dbg("model_id = %u, input[%u] - sz_d = %u sz_q = %u", model->model_id, i, + model->mvtvm.info.input[i].sz_d, model->mvtvm.info.input[i].sz_q); + } + + /* Outputs, set for nb_layers - 1 */ + model->mvtvm.info.nb_outputs = metadata->model.num_output; + model->mvtvm.info.total_output_sz_d = 0; + model->mvtvm.info.total_output_sz_q = 0; + for (i = 0; i < metadata->model.num_output; i++) { + strncpy(model->mvtvm.info.output[i].name, metadata->output[i].name, + TVMDP_NAME_STRLEN); + model->mvtvm.info.output[i].dtype = + mvtvm_ml_io_type_map(metadata->output[i].datatype.code); + model->mvtvm.info.output[i].qtype = + mvtvm_ml_io_type_map(metadata->output[i].model_datatype.code); + model->mvtvm.info.output[i].nb_dims = metadata->output[i].ndim; + + model->mvtvm.info.output[i].nb_elements = 1; + for (j = 0; j < metadata->output[i].ndim; j++) { + model->mvtvm.info.output[i].shape[j] = metadata->output[i].shape[j]; + model->mvtvm.info.output[i].nb_elements *= metadata->output[i].shape[j]; + } + + model->mvtvm.info.output[i].sz_d = + model->mvtvm.info.output[i].nb_elements * + rte_ml_io_type_size_get(model->mvtvm.info.output[i].dtype); + model->mvtvm.info.output[i].sz_q = + model->mvtvm.info.output[i].nb_elements * + rte_ml_io_type_size_get(model->mvtvm.info.output[i].qtype); + + model->mvtvm.info.total_output_sz_d += model->mvtvm.info.output[i].sz_d; + model->mvtvm.info.total_output_sz_q += model->mvtvm.info.output[i].sz_q; + + plt_ml_dbg("model_id = %u, output[%u] - sz_d = %u sz_q = %u", model->model_id, i, + model->mvtvm.info.output[i].sz_d, model->mvtvm.info.output[i].sz_q); + } + + return; + +tvm_mrvl_model: + cn10k_ml_layer_io_info_update(&model->mvtvm.info, &model->layer[0].glow.metadata); +} diff --git a/drivers/ml/cnxk/mvtvm_ml_model.h b/drivers/ml/cnxk/mvtvm_ml_model.h index 6c38217c15..2b25a7b568 100644 --- a/drivers/ml/cnxk/mvtvm_ml_model.h +++ b/drivers/ml/cnxk/mvtvm_ml_model.h @@ -47,5 +47,6 @@ struct mvtvm_ml_model_data { int mvtvm_ml_model_blob_parse(struct rte_ml_model_params *params, struct mvtvm_ml_model_object *object); +void mvtvm_ml_model_io_info_update(struct cnxk_ml_model *model); #endif /* _MVTVM_ML_MODEL_H_ */ diff --git a/drivers/ml/cnxk/mvtvm_ml_ops.c b/drivers/ml/cnxk/mvtvm_ml_ops.c index d9ec411385..1d585a57ff 100644 --- a/drivers/ml/cnxk/mvtvm_ml_ops.c +++ b/drivers/ml/cnxk/mvtvm_ml_ops.c @@ -179,6 +179,9 @@ mvtvm_ml_model_load(struct cnxk_ml_dev *cnxk_mldev, struct rte_ml_model_params * goto error; } + /* Update model I/O data */ + mvtvm_ml_model_io_info_update(model); + return 0; error: -- 2.41.0