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 425a682209f..86f465a645f 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 6c38217c158..2b25a7b568e 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 5c30bbf6b89..a783e16e6eb 100644
--- a/drivers/ml/cnxk/mvtvm_ml_ops.c
+++ b/drivers/ml/cnxk/mvtvm_ml_ops.c
@@ -181,6 +181,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

Reply via email to