CWG-E158 Dataset generation This reverts commit fbf0837e61a71aed6e30bfbfe655fe9b942dc20f.
diff --git a/CMakeLists.txt b/CMakeLists.txt index 1b3da62..868599b 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt
@@ -569,6 +569,18 @@ twopass_encoder) endif() + if(CONFIG_ML_PART_SPLIT) + include(FetchContent) + message(STATUS "Building CNPY ...") + + set(ZLIB_USE_STATIC_LIBS "ON") + include_directories("${AOM_ROOT}/third_party/cnpy") + add_subdirectory(third_party/cnpy) + + find_package(ZLIB REQUIRED) + target_link_libraries(aom PRIVATE ${ZLIB_LIBRARIES}) + endif() + if(ENABLE_TOOLS) if(CONFIG_ENTROPY_STATS AND NOT BUILD_SHARED_LIBS) @@ -970,6 +982,15 @@ endif() endif() +if(CONFIG_ML_PART_SPLIT AND CONFIG_AV1_ENCODER) + if(ENABLE_EXAMPLES OR ENABLE_TOOLS) + foreach(aom_app ${AOM_APP_TARGETS}) + add_dependencies(${aom_app} cnpy) + target_link_libraries(${aom_app} PRIVATE cnpy) + endforeach() + endif() +endif() + if(ENABLE_EXAMPLES AND "${CMAKE_GENERATOR}" MATCHES "Makefiles$") # For historical purposes place the example binaries in the example directory.
diff --git a/aom/aomcx.h b/aom/aomcx.h index f8f5b9a..b83e06d 100644 --- a/aom/aomcx.h +++ b/aom/aomcx.h
@@ -1264,6 +1264,9 @@ AV1E_SET_FILM_GRAIN_BLOCK_SIZE = 174, #endif +#if CONFIG_ML_PART_SPLIT + AV1E_SET_PY_DATAFILE_NAME = 175, +#endif }; /*!\brief aom 1-D scaling mode @@ -1701,6 +1704,8 @@ AOM_CTRL_USE_TYPE(AV1E_SET_FILM_GRAIN_BLOCK_SIZE, int) #define AOM_CTRL_AV1E_SET_FILM_GRAIN_BLOCK_SIZE #endif +AOM_CTRL_USE_TYPE(AV1E_SET_PY_DATAFILE_NAME, const char *) +#define AOM_CTRL_AV1E_SET_PY_DATAFILE_NAME AOM_CTRL_USE_TYPE(AV1E_SET_CDF_UPDATE_MODE, unsigned int) #define AOM_CTRL_AV1E_SET_CDF_UPDATE_MODE
diff --git a/apps/aomenc.c b/apps/aomenc.c index da49301..b1fa81c 100644 --- a/apps/aomenc.c +++ b/apps/aomenc.c
@@ -228,6 +228,9 @@ #endif AV1E_SET_SUBGOP_CONFIG_STR, AV1E_SET_SUBGOP_CONFIG_PATH, +#if CONFIG_ML_PART_SPLIT + AV1E_SET_PY_DATAFILE_NAME, +#endif // CONFIG_ML_PART_SPLIT AV1E_SET_FRAME_OUTPUT_ORDER_DERIVATION, AV1E_SET_ENABLE_CDF_AVERAGING, AV1E_SET_ENABLE_BRU, @@ -429,6 +432,9 @@ #endif &g_av1_codec_arg_defs.subgop_config_str, &g_av1_codec_arg_defs.subgop_config_path, +#if CONFIG_ML_PART_SPLIT + &g_av1_codec_arg_defs.py_datafile_name, +#endif &g_av1_codec_arg_defs.frame_hash_metadata, &g_av1_codec_arg_defs.frame_hash_per_plane, NULL, @@ -1010,6 +1016,11 @@ } #endif + if (key == AV1E_SET_PY_DATAFILE_NAME) { + config->py_datafile_name = arg->val; + return; + } + if (key == AV1E_SET_SUBGOP_CONFIG_STR) { config->subgop_config_str = arg->val; return; @@ -1808,6 +1819,10 @@ AV1E_SET_FILM_GRAIN_BLOCK_SIZE, stream->config.film_grain_block_size); #endif + if (stream->config.py_datafile_name) { + AOM_CODEC_CONTROL_TYPECHECKED(&stream->encoder, AV1E_SET_PY_DATAFILE_NAME, + stream->config.py_datafile_name); + } AOM_CODEC_CONTROL_TYPECHECKED(&stream->encoder, AV1E_SET_COLOR_RANGE, stream->config.color_range);
diff --git a/av1/arg_defs.c b/av1/arg_defs.c index 32d1679..e92f9f4 100644 --- a/av1/arg_defs.c +++ b/av1/arg_defs.c
@@ -892,4 +892,9 @@ .enable_short_refresh_frame_flags = ARG_DEF(NULL, "enable-short-refresh-frame-flags", 1, "Signal refresh frame flags with N bits. (0: N = 8, 1 : N = 3)"), +#if CONFIG_ML_PART_SPLIT + .py_datafile_name = + ARG_DEF(NULL, "py-datafile", 1, + "Where to store debug data to be later analyzed from colab."), +#endif };
diff --git a/av1/arg_defs.h b/av1/arg_defs.h index 3046182..cf58db4 100644 --- a/av1/arg_defs.h +++ b/av1/arg_defs.h
@@ -302,6 +302,9 @@ arg_def_t crop_win_top_offset; arg_def_t crop_win_bottom_offset; #endif // CONFIG_CROP_WIN_CWG_F220 +#if CONFIG_ML_PART_SPLIT + arg_def_t py_datafile_name; +#endif } av1_codec_arg_definitions_t; extern const av1_codec_arg_definitions_t g_av1_codec_arg_defs;
diff --git a/av1/av1.cmake b/av1/av1.cmake index b715be1..d7736d1 100644 --- a/av1/av1.cmake +++ b/av1/av1.cmake
@@ -318,7 +318,9 @@ "${AOM_ROOT}/av1/encoder/simple_intrapred_tflite_model_64x64.h" "${AOM_ROOT}/av1/encoder/simple_intrapred_tflite_model_32x32.h" "${AOM_ROOT}/av1/encoder/simple_intrapred_tflite_model_16x16.h" - "${AOM_ROOT}/av1/encoder/sms_part_split_prune_tflite_model.h") + "${AOM_ROOT}/av1/encoder/sms_part_split_prune_tflite_model.h" + "${AOM_ROOT}/tools/ml/py_bridge.h" + "${AOM_ROOT}/tools/ml/py_bridge.cc") endif() if(CONFIG_DIP_EXT_PRUNING)
diff --git a/av1/av1_cx_iface.c b/av1/av1_cx_iface.c index efe446a..226dd21 100644 --- a/av1/av1_cx_iface.c +++ b/av1/av1_cx_iface.c
@@ -37,6 +37,10 @@ #include "aom_dsp/psnr.h" #include "aom_ports/aom_timer.h" +#if CONFIG_ML_PART_SPLIT +#include "tools/ml/py_bridge.h" +#endif + #define MAG_SIZE (4) struct av1_extracfg { @@ -251,6 +255,9 @@ int crop_win_top_offset; // cropping window top offset int crop_win_bottom_offset; // cropping window bottom offset #endif // CONFIG_CROP_WIN_CWG_F220 +#if CONFIG_ML_PART_SPLIT + char *py_datafile_name; +#endif }; // Example subgop configs. Currently not used by default. @@ -587,6 +594,9 @@ 0, // crop_win_top_offset 0, // crop_win_bottom_offset #endif // CONFIG_CROP_WIN_CWG_F220 +#if CONFIG_ML_PART_SPLIT + NULL, // py_datafile_name +#endif // CONFIG_ML_PART_SPLIT }; // clang-format on @@ -1760,6 +1770,9 @@ part_cfg->enable_ext_partitions = extra_cfg->enable_ext_partitions; part_cfg->min_partition_size = extra_cfg->min_partition_size; part_cfg->max_partition_size = extra_cfg->max_partition_size; +#if CONFIG_ML_PART_SPLIT + part_cfg->py_datafile_name = extra_cfg->py_datafile_name; +#endif // Set intra mode configuration. intra_mode_cfg->enable_angle_delta = extra_cfg->enable_angle_delta; @@ -2682,6 +2695,13 @@ return update_extra_cfg(ctx, &extra_cfg); } +static aom_codec_err_t ctrl_set_py_datafile(aom_codec_alg_priv_t *ctx, + va_list args) { + struct av1_extracfg extra_cfg = ctx->extra_cfg; + extra_cfg.py_datafile_name = CAST(AV1E_SET_PY_DATAFILE_NAME, args); + return update_extra_cfg(ctx, &extra_cfg); +} + static aom_codec_err_t ctrl_set_film_grain_table(aom_codec_alg_priv_t *ctx, va_list args) { struct av1_extracfg extra_cfg = ctx->extra_cfg; @@ -3147,6 +3167,17 @@ fprintf(stdout, "]\n"); #endif // CONFIG_CWG_F317_TEST_PATTERN } +#if CONFIG_ML_PART_SPLIT + MultiThreadInfo *const mt_info = &cpi->mt_info; + for (int t = 0; t < mt_info->num_workers; ++t) { + EncWorkerData *const thread_data = &mt_info->tile_thr_data[t]; + py_datafile_fold(&cpi->td.py_bridge, &thread_data->td->py_bridge); + } + if (cpi->oxcf.part_cfg.py_datafile_name) { + printf("Storing PY data to: %s\n", cpi->oxcf.part_cfg.py_datafile_name); + py_datafile_close(&cpi->td.py_bridge, cpi->oxcf.part_cfg.py_datafile_name); + } +#endif } // TODO(Mufaddal): Check feasibility of abstracting functions related to LAP @@ -4071,6 +4102,9 @@ } else if (arg_match_helper(&arg, &g_av1_codec_arg_defs.film_grain_test, argv, err_string)) { extra_cfg.film_grain_test_vector = arg_parse_int_helper(&arg, err_string); + } else if (arg_match_helper(&arg, &g_av1_codec_arg_defs.py_datafile_name, + argv, err_string)) { + extra_cfg.py_datafile_name = value; } else if (arg_match_helper(&arg, &g_av1_codec_arg_defs.film_grain_table, argv, err_string)) { extra_cfg.film_grain_table_filename = value; @@ -4614,6 +4648,9 @@ #if CONFIG_FGS_BLOCK_SIZE { AV1E_SET_FILM_GRAIN_BLOCK_SIZE, ctrl_set_film_grain_block_size }, #endif +#if CONFIG_ML_PART_SPLIT + { AV1E_SET_PY_DATAFILE_NAME, ctrl_set_py_datafile }, +#endif { AV1E_SET_DENOISE_NOISE_LEVEL, ctrl_set_denoise_noise_level }, { AV1E_SET_DENOISE_BLOCK_SIZE, ctrl_set_denoise_block_size }, { AV1E_ENABLE_MOTION_VECTOR_UNIT_TEST, ctrl_enable_motion_vector_unit_test },
diff --git a/av1/encoder/context_tree.c b/av1/encoder/context_tree.c index 9e7767f..7fc3a94 100644 --- a/av1/encoder/context_tree.c +++ b/av1/encoder/context_tree.c
@@ -876,3 +876,28 @@ // Search from the highest common ancestor return counterpart_from_different_partition(current, pc_tree); } +void av1_setup_sms_pred_buf(struct AV1Common *cm, struct ThreadData *td) { + unsigned int width = ((cm->width + MAX_SB_SIZE - 1) >> MAX_SB_SIZE_LOG2) + << MAX_SB_SIZE_LOG2; + unsigned int height = ((cm->height + MAX_SB_SIZE - 1) >> MAX_SB_SIZE_LOG2) + << MAX_SB_SIZE_LOG2; + uint16_t *buf = + (uint16_t *)aom_memalign(16, width * height * sizeof(uint16_t)); + struct buf_2d tmp = { .buf = buf, + .buf0 = buf, + .width = cm->width, + .height = cm->height, + .stride = width }; + td->sms_pred_buf = tmp; + printf("SMS frame w:%d h:%d w:%d h:%d\n", cm->width, cm->height, width, + height); +} + +void av1_free_sms_pred_buf(struct ThreadData *td) { + if (td->sms_pred_buf.buf0 != NULL) { + printf("SMS frame free\n"); + aom_free(td->sms_pred_buf.buf0); + td->sms_pred_buf.buf0 = NULL; + td->sms_pred_buf.buf = NULL; + } +}
diff --git a/av1/encoder/context_tree.h b/av1/encoder/context_tree.h index e2bf79a..450a7e7 100644 --- a/av1/encoder/context_tree.h +++ b/av1/encoder/context_tree.h
@@ -144,6 +144,8 @@ void av1_free_sms_tree(struct ThreadData *td); void av1_setup_sms_bufs(struct AV1Common *cm, struct ThreadData *td); void av1_free_sms_bufs(struct ThreadData *td); +void av1_setup_sms_pred_buf(struct AV1Common *cm, struct ThreadData *td); +void av1_free_sms_pred_buf(struct ThreadData *td); #ifdef __cplusplus } // extern "C"
diff --git a/av1/encoder/encoder.c b/av1/encoder/encoder.c index 1a3094f..7ef3159 100644 --- a/av1/encoder/encoder.c +++ b/av1/encoder/encoder.c
@@ -87,6 +87,7 @@ #if CONFIG_ML_PART_SPLIT #include "av1/encoder/part_split_prune_tflite.h" +#include "tools/ml/py_bridge.h" #endif // CONFIG_ML_PART_SPLIT #if CONFIG_DIP_EXT_PRUNING @@ -1128,6 +1129,7 @@ #if CONFIG_DIP_EXT_PRUNING intra_dip_mode_prune_close(&(cpi->td.dip_pruning_model)); #endif // CONFIG_DIP_EXT_PRUNING + av1_free_sms_pred_buf(&cpi->td); av1_free_pmc(cpi->td.firstpass_ctx, av1_num_planes(cm)); cpi->td.firstpass_ctx = NULL; alloc_compressor_data(cpi); @@ -1210,6 +1212,13 @@ } cpi->alloc_pyramid = oxcf->tool_cfg.enable_global_motion; +#if CONFIG_ML_PART_SPLIT + if (cpi->oxcf.part_cfg.py_datafile_name) { + size_t count; + py_datafile_delete_old_files(cpi->oxcf.part_cfg.py_datafile_name, &count); + printf("\x1b[92mDeleted %d old stats\x1b[0m\n", count); + } +#endif } static INLINE void init_frame_info(FRAME_INFO *frame_info, @@ -1568,6 +1577,7 @@ av1_free_shared_coeff_buffer(&thread_data->td->shared_coeff_buf); av1_free_sms_tree(thread_data->td); av1_free_sms_bufs(thread_data->td); + av1_free_sms_pred_buf(&cpi->td); #if CONFIG_ML_PART_SPLIT av2_part_prune_tflite_close(&(thread_data->td->partition_model)); #endif // CONFIG_ML_PART_SPLIT
diff --git a/av1/encoder/encoder.h b/av1/encoder/encoder.h index c105776..344bf7a 100644 --- a/av1/encoder/encoder.h +++ b/av1/encoder/encoder.h
@@ -254,6 +254,10 @@ * Indicates the maximum aspect ratio of allowed partition block sizes. */ unsigned int max_partition_aspect_ratio; + +#if CONFIG_ML_PART_SPLIT + const char *py_datafile_name; +#endif } PartitionCfg; /*! @@ -1722,6 +1726,12 @@ int warped_used[2]; } RD_COUNTS; +typedef struct { + unsigned int valid : 1; + unsigned int low_test : 1; + unsigned int high_test : 1; +} part_ml_inf; + typedef struct ThreadData { MACROBLOCK mb; RD_COUNTS rd_counts; @@ -1754,9 +1764,23 @@ PICK_MODE_CONTEXT *firstpass_ctx; #if CONFIG_ML_PART_SPLIT void *partition_model; + void *py_bridge; #endif // CONFIG_ML_PART_SPLIT void *dip_pruning_model; struct buf_2d sms_pred_buf; + + // unsigned int prune_non[32]; + // unsigned int prune_spl[32]; + // unsigned int prune_hor[32]; + // unsigned int prune_ver[32]; + // unsigned int force_non[32]; + // unsigned int force_spl[32]; + // unsigned int force_hor[32]; + // unsigned int force_ver[32]; + // unsigned int prune_tot[32]; + // uint64_t ftr_time[32]; + // uint64_t inf_time[32]; + // uint64_t non_time[32]; } ThreadData; struct EncWorkerData;
diff --git a/av1/encoder/encoder_alloc.h b/av1/encoder/encoder_alloc.h index e6d5336..58ee80c 100644 --- a/av1/encoder/encoder_alloc.h +++ b/av1/encoder/encoder_alloc.h
@@ -15,6 +15,9 @@ #include "av1/encoder/encoder.h" #include "av1/encoder/encodetxb.h" +#if CONFIG_ML_PART_SPLIT +#include "tools/ml/py_bridge.h" +#endif #if CONFIG_ML_PART_SPLIT #include "av1/encoder/part_split_prune_tflite.h" @@ -86,6 +89,7 @@ av1_setup_shared_coeff_buffer(&cpi->common, &cpi->td.shared_coeff_buf); av1_setup_sms_tree(cpi, &cpi->td); av1_setup_sms_bufs(&cpi->common, &cpi->td); + av1_setup_sms_pred_buf(&cpi->common, &cpi->td); cpi->td.firstpass_ctx = av1_alloc_pmc(cm, SHARED_PART, 0, 0, BLOCK_16X16, NULL, PARTITION_NONE, 0, cm->seq_params.subsampling_x, cm->seq_params.subsampling_y, @@ -278,6 +282,41 @@ #if CONFIG_DIP_EXT_PRUNING intra_dip_mode_prune_close(&(cpi->td.dip_pruning_model)); #endif // CONFIG_DIP_EXT_PRUNING + av1_free_sms_pred_buf(&cpi->td); + + // for (int i = 0; i < 32; i++) { + // cm->prune_non[i] += cpi->td.prune_non[i]; + // cm->prune_spl[i] += cpi->td.prune_spl[i]; + // cm->prune_ver[i] += cpi->td.prune_ver[i]; + // cm->prune_hor[i] += cpi->td.prune_hor[i]; + // cm->force_non[i] += cpi->td.force_non[i]; + // cm->force_spl[i] += cpi->td.force_spl[i]; + // cm->force_ver[i] += cpi->td.force_ver[i]; + // cm->force_hor[i] += cpi->td.force_hor[i]; + // cm->prune_tot[i] += cpi->td.prune_tot[i]; + // cm->ftr_time[i] += cpi->td.ftr_time[i]; + // cm->inf_time[i] += cpi->td.inf_time[i]; + // cm->non_time[i] += cpi->td.non_time[i]; + + // if (cm->prune_non[i] || cm->prune_spl[i] || cm->prune_ver[i] || + // cm->prune_hor[i]) { + // printf( + // "%d: pnon=%d%%,pspl=%d%%,pver=%d%%,phor=%d%% " + // "fnon=%d%%,fspl=%d%%,fver=%d%%,fhor=%d%%\n", + // i, 100 * cm->prune_non[i] / cm->prune_tot[i], + // 100 * cm->prune_spl[i] / cm->prune_tot[i], + // 100 * cm->prune_ver[i] / cm->prune_tot[i], + // 100 * cm->prune_hor[i] / cm->prune_tot[i], + // 100 * cm->force_non[i] / cm->prune_tot[i], + // 100 * cm->force_spl[i] / cm->prune_tot[i], + // 100 * cm->force_ver[i] / cm->prune_tot[i], + // 100 * cm->force_hor[i] / cm->prune_tot[i]); + // } + // if (cm->ftr_time[i] || cm->inf_time[i]) { + // printf("%d: ftr:%" PRIu64 ",inf:%" PRIu64 ",non:%" PRIu64 "\n", i, + // cm->ftr_time[i], cm->inf_time[i], cm->non_time[i]); + // } + //} aom_free(cpi->td.mb.palette_buffer); release_compound_type_rd_buffers(&cpi->td.mb.comp_rd_buffer); aom_free(cpi->td.mb.tmp_conv_dst);
diff --git a/av1/encoder/encoder_utils.c b/av1/encoder/encoder_utils.c index 35bcf3f..938a19f 100644 --- a/av1/encoder/encoder_utils.c +++ b/av1/encoder/encoder_utils.c
@@ -849,6 +849,7 @@ #if CONFIG_DIP_EXT_PRUNING intra_dip_mode_prune_close(&(cpi->td.dip_pruning_model)); #endif // CONFIG_DIP_EXT_PRUNING + av1_free_sms_pred_buf(&cpi->td); av1_free_pmc(cpi->td.firstpass_ctx, num_planes); cpi->td.firstpass_ctx = NULL; alloc_compressor_data(cpi); @@ -936,6 +937,7 @@ #if CONFIG_DIP_EXT_PRUNING intra_dip_mode_prune_close(&(cpi->td.dip_pruning_model)); #endif // CONFIG_DIP_EXT_PRUNING + av1_free_sms_pred_buf(&cpi->td); av1_free_pmc(cpi->td.firstpass_ctx, av1_num_planes(cm)); cpi->td.firstpass_ctx = NULL; alloc_compressor_data(cpi);
diff --git a/av1/encoder/ethread.c b/av1/encoder/ethread.c index 99b645a..c5f9e8e 100644 --- a/av1/encoder/ethread.c +++ b/av1/encoder/ethread.c
@@ -544,6 +544,7 @@ // Set up sms_tree. av1_setup_sms_tree(cpi, thread_data->td); av1_setup_sms_bufs(cm, thread_data->td); + av1_setup_sms_pred_buf(cm, thread_data->td); for (int x = 0; x < 2; x++) for (int y = 0; y < 2; y++)
diff --git a/av1/encoder/motion_search_facade.c b/av1/encoder/motion_search_facade.c index c311fc1..0dcd5d0 100644 --- a/av1/encoder/motion_search_facade.c +++ b/av1/encoder/motion_search_facade.c
@@ -1450,7 +1450,7 @@ /*!\brief A simplified version for motion search used for speed features. Same * as \ref av1_simple_motion_search, but it is used with ERP. */ -int_mv av1_simple_motion_search_ext(AV1_COMP *const cpi, +int_mv av1_simple_motion_search_ext(AV1_COMP *const cpi, ThreadData *td, const TileInfo *const tile, MACROBLOCK *x, int mi_row, int mi_col, BLOCK_SIZE bsize, int ref, FULLPEL_MV start_mv, @@ -1569,8 +1569,27 @@ mbmi->mv[0] = best_mv; // Get a copy of the prediction output + struct buf_2d dst_old = xd->plane[AOM_PLANE_Y].dst; + uint32_t buf_off = dst_old.buf - dst_old.buf0; + uint32_t tx = buf_off % dst_old.stride; + uint32_t ty = buf_off / dst_old.stride; + td->sms_pred_buf.buf = + td->sms_pred_buf.buf0 + tx + ty * td->sms_pred_buf.stride; + xd->plane[AOM_PLANE_Y].dst = td->sms_pred_buf; av1_enc_build_inter_predictor(cm, xd, mi_row, mi_col, NULL, bsize, AOM_PLANE_Y, AOM_PLANE_Y); + xd->plane[AOM_PLANE_Y].dst = dst_old; + // copy the data back + int blk_w = block_size_wide[bsize]; + int blk_h = block_size_high[bsize]; + uint16_t *ptr0 = dst_old.buf, *ptr1 = td->sms_pred_buf.buf; + for (int i = 0; i < blk_h; i++) { + for (int j = 0; j < blk_w; j++) { + ptr0[j] = ptr1[j]; + } + ptr0 += dst_old.stride; + ptr1 += td->sms_pred_buf.stride; + } aom_clear_system_state();
diff --git a/av1/encoder/motion_search_facade.h b/av1/encoder/motion_search_facade.h index 67d410d..545526b 100644 --- a/av1/encoder/motion_search_facade.h +++ b/av1/encoder/motion_search_facade.h
@@ -82,7 +82,7 @@ int mi_row, int mi_col, BLOCK_SIZE bsize, const FULLPEL_MV start_mv, int use_subpixel, unsigned int *sse, unsigned int *var); -int_mv av1_simple_motion_search_ext(AV1_COMP *const cpi, +int_mv av1_simple_motion_search_ext(AV1_COMP *const cpi, ThreadData *td, const TileInfo *const tile, MACROBLOCK *x, int mi_row, int mi_col, BLOCK_SIZE bsize, int ref, FULLPEL_MV start_mv,
diff --git a/av1/encoder/partition_ml.c b/av1/encoder/partition_ml.c index 2f6d4fb..e121473 100644 --- a/av1/encoder/partition_ml.c +++ b/av1/encoder/partition_ml.c
@@ -365,8 +365,8 @@ }; #define USE_MODEL(tgt_intra, tgt_level, tgt_bsize, model_type, low, high, qp_low, qp_high) {\ - if (tgt_intra == intra && (tgt_level == harsh_level) && tgt_bsize == bsize && *num_models < 12) { \ - if (!model_in_list(model_type, out, *num_models)) { \ + if (tgt_intra == intra && (tgt_level == harsh_level || ds_mode) && tgt_bsize == bsize && *num_models < 12) { \ + if (!ds_mode || !model_in_list(model_type, out, *num_models)) { \ out[*num_models] = model_type; \ struct ModelParams tmp = {low, high, qp_low, qp_high}; \ params[*num_models] = tmp; \ @@ -376,7 +376,7 @@ } static void get_model_type(bool intra, BLOCK_SIZE bsize, int harsh_level, - MODEL_TYPE* out, struct ModelParams* params, int* num_models) { + bool ds_mode, MODEL_TYPE* out, struct ModelParams* params, int* num_models) { // TODO: de-dup the guys *num_models = 0; // lvl blk model low high qp_l qp_h @@ -657,7 +657,9 @@ int mi_col, BLOCK_SIZE bsize, const TileInfo *tile_info, ThreadData *td, bool search_none_after_rect, - bool* prune_list) { + bool* prune_list, MLResult* ml_results, + size_t* ml_count) { + *ml_count = 0; for (int i = 0; i < 4; i++) { prune_list[i] = false; } @@ -682,8 +684,8 @@ MODEL_TYPE model_types[12]; struct ModelParams model_params[12]; int num_models = 0; - get_model_type(key_frame, bsize, harsh_level, model_types, model_params, - &num_models); + get_model_type(key_frame, bsize, harsh_level, !!ml_results, model_types, + model_params, &num_models); float ml_input[FEATURE_INTER_MAX] = { 0.0f }; float ml_output[1] = { 0.0f }; bool has_features = false; @@ -715,6 +717,34 @@ assert(!had_error); + // When in dataset mode we still run ML but we don't do anything + // about the result. + // ml_results being NULL indicates, not need to collect + if (ml_results) { + params.thresh_low = 0; + params.thresh_high = 1; + MLResult* result = &ml_results[(*ml_count)++]; + for (int i = 0; i < FEATURE_INTER_MAX; i++) + result->x[i] = ml_input[i]; + result->y = ml_output[0]; + result->model = model_type; + result->x_shape = get_model_n_features(model_type); + + // also check the correctness of the inverse + struct InputSpec input_spec; + av2_model_input_spec(model_type, &input_spec); + if (input_spec.valid) { + for (int i = 0; i < result->x_shape; i++) { + float diff = 1 / input_spec.std[i] - input_spec.invstd[i]; + if (diff >= 0.00005) { + printf("[%d] %f != 1/%f, %f\n", model_type, input_spec.invstd[i], + input_spec.std[i], diff); + break; + } + } + } + } + if (had_error) continue; else {
diff --git a/av1/encoder/partition_ml.h b/av1/encoder/partition_ml.h index 8b1ddd2..8d65677 100644 --- a/av1/encoder/partition_ml.h +++ b/av1/encoder/partition_ml.h
@@ -77,5 +77,8 @@ int mi_col, BLOCK_SIZE bsize, const TileInfo *tile_info, ThreadData *td, bool search_none_after_rect, - bool* prune_list); + bool* prune_list, MLResult* mls_results, + size_t* ml_count); + + #endif // AV1_ENCODER_PARTITION_ML_H_
diff --git a/av1/encoder/partition_search.c b/av1/encoder/partition_search.c index 20660d4..d659c0f 100644 --- a/av1/encoder/partition_search.c +++ b/av1/encoder/partition_search.c
@@ -45,6 +45,7 @@ #if CONFIG_ML_PART_SPLIT #include "av1/encoder/partition_ml.h" +#include "tools/ml/py_bridge.h" #endif static void update_partition_cdfs_and_counts(MACROBLOCKD *xd, int blk_col, @@ -3758,6 +3759,9 @@ return; } + struct aom_usec_timer non_timer; + aom_usec_timer_start(&non_timer); + int pt_cost = 0; RD_STATS best_remain_rdcost; @@ -3819,6 +3823,8 @@ partition_timer_on = 0; } #endif + aom_usec_timer_mark(&non_timer); + // td->non_time[bsize] += aom_usec_timer_elapsed(&non_timer); *pb_source_variance = x->source_variance; if (none_rd) *none_rd = this_rdc->rdcost; @@ -6145,6 +6151,8 @@ // Partition block source pixel variance. unsigned int pb_source_variance = UINT_MAX; #if CONFIG_ML_PART_SPLIT + MLResult ml_results[8]; + size_t ml_count = 0; int next_force_prune_flags[2][3] = { { 0, 0, 0 }, { 0, 0, 0 } }; // Don't use ML pruning if this is the second attempt to find a valid // partition. @@ -6168,9 +6176,10 @@ // prune either one or both. if (!force_prune_flags[PRUNE_OTHER]) { bool prune_list[4]; + MLResult *ml_arg = cpi->oxcf.part_cfg.py_datafile_name ? ml_results : NULL; int ml_result = av1_ml_part_split_infer(cpi, x, mi_row, mi_col, bsize, tile_info, td, - search_none_after_rect, prune_list); + search_none_after_rect, prune_list, ml_arg, &ml_count); if (ml_result == ML_PART_FORCE_NONE || ml_result == ML_PART_FORCE_SPLIT || ml_result == ML_PART_FORCE_VERT || ml_result == ML_PART_FORCE_HORZ) { part_search_state.prune_partition_3[0] = 1; @@ -6553,6 +6562,48 @@ // Reset the PC_TREE deallocation flag. int pc_tree_dealloc = 0; +#if CONFIG_ML_PART_SPLIT + // Only storing dataset if the datafile is provided, and only if: + // * evaluated this branch and not bailed out; + // * partition picked on its merit and not based on its context; + if (part_search_state.found_best_partition && + part_search_state.forced_partition == PARTITION_INVALID && + part_search_state.partition_none_allowed && + part_search_state.partition_rect_allowed[HORZ] && + part_search_state.partition_rect_allowed[VERT] && + cpi->oxcf.part_cfg.py_datafile_name) { + + float label = pc_tree->partitioning; + float true_label = pc_tree->partitioning; + if (pc_tree->partitioning == PARTITION_VERT && + pc_tree->vertical[pc_tree->region_type][0]->partitioning == + PARTITION_HORZ && + pc_tree->vertical[pc_tree->region_type][1]->partitioning == + PARTITION_HORZ) + label = PARTITION_SPLIT; + if (pc_tree->partitioning == PARTITION_HORZ && + pc_tree->horizontal[pc_tree->region_type][0]->partitioning == + PARTITION_VERT && + pc_tree->horizontal[pc_tree->region_type][1]->partitioning == + PARTITION_VERT) + label = PARTITION_SPLIT; + for (size_t i = 0; i < ml_count; i++) { + char data_file[64]; + char var_name[16]; + sprintf(data_file, "model_%d_%s", ml_results[i].model, + get_model_name(ml_results[i].model)); + sprintf(var_name, "x_%d", cm->cur_frame->absolute_poc); + py_datafile_add_1d(&td->py_bridge, data_file, var_name, ml_results[i].x, + ml_results[i].x_shape); + sprintf(var_name, "y_%d", cm->cur_frame->absolute_poc); + py_datafile_add_0d(&td->py_bridge, data_file, var_name, &ml_results[i].y); + sprintf(var_name, "label_%d", cm->cur_frame->absolute_poc); + py_datafile_add_0d(&td->py_bridge, data_file, var_name, &label); + sprintf(var_name, "true_%d", cm->cur_frame->absolute_poc); + py_datafile_add_0d(&td->py_bridge, data_file, var_name, &true_label); + } + } +#endif // If a valid partition is found and reconstruction is required for future // sub-blocks in the same group. if (part_search_state.found_best_partition && pc_tree->index != 3) {
diff --git a/av1/encoder/partition_strategy.c b/av1/encoder/partition_strategy.c index 9677441..4998078 100644 --- a/av1/encoder/partition_strategy.c +++ b/av1/encoder/partition_strategy.c
@@ -1259,7 +1259,7 @@ for (int idx = 0; idx < sms_data->num_start_mvs; idx++) { const MV start_mv = sms_data->start_mv_list[idx]; const FULLPEL_MV start_mv_full = get_fullmv_from_mv(&start_mv); - av1_simple_motion_search_ext(cpi, tile, x, mi_row, mi_col, bsize, + av1_simple_motion_search_ext(cpi, td, tile, x, mi_row, mi_col, bsize, ref_frame, start_mv_full, 1, 1, sms_data); sms_data->var = cpi->fn_ptr[bsize].vf(src_buf, src_stride, dst_buf, dst_stride, &sms_data->sse);
diff --git a/third_party/cnpy/CMakeLists.txt b/third_party/cnpy/CMakeLists.txt new file mode 100644 index 0000000..d56e666 --- /dev/null +++ b/third_party/cnpy/CMakeLists.txt
@@ -0,0 +1,13 @@ +CMAKE_MINIMUM_REQUIRED(VERSION 3.0 FATAL_ERROR) +if(COMMAND cmake_policy) + cmake_policy(SET CMP0003 NEW) +endif(COMMAND cmake_policy) + +project(CNPY) + +set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -std=c++11") + +find_package(ZLIB REQUIRED) +include_directories(${ZLIB_INCLUDE_DIRS}) + +add_library(cnpy STATIC "cnpy.cpp")
diff --git a/third_party/cnpy/cnpy.cpp b/third_party/cnpy/cnpy.cpp new file mode 100644 index 0000000..2d28578 --- /dev/null +++ b/third_party/cnpy/cnpy.cpp
@@ -0,0 +1,340 @@ +//Copyright (C) 2011 Carl Rogers +//Released under MIT License +//license available in LICENSE file, or at http://www.opensource.org/licenses/mit-license.php + +#include"cnpy.h" +#include<complex> +#include<cstdlib> +#include<algorithm> +#include<cstring> +#include<iomanip> +#include<stdint.h> +#include<stdexcept> +#include <regex> + +char cnpy::BigEndianTest() { + int x = 1; + return (((char *)&x)[0]) ? '<' : '>'; +} + +char cnpy::map_type(const std::type_info& t) +{ + if(t == typeid(float) ) return 'f'; + if(t == typeid(double) ) return 'f'; + if(t == typeid(long double) ) return 'f'; + + if(t == typeid(int) ) return 'i'; + if(t == typeid(char) ) return 'i'; + if(t == typeid(short) ) return 'i'; + if(t == typeid(long) ) return 'i'; + if(t == typeid(long long) ) return 'i'; + + if(t == typeid(unsigned char) ) return 'u'; + if(t == typeid(unsigned short) ) return 'u'; + if(t == typeid(unsigned long) ) return 'u'; + if(t == typeid(unsigned long long) ) return 'u'; + if(t == typeid(unsigned int) ) return 'u'; + + if(t == typeid(bool) ) return 'b'; + + if(t == typeid(std::complex<float>) ) return 'c'; + if(t == typeid(std::complex<double>) ) return 'c'; + if(t == typeid(std::complex<long double>) ) return 'c'; + + else return '?'; +} + +template<> std::vector<char>& cnpy::operator+=(std::vector<char>& lhs, const std::string rhs) { + lhs.insert(lhs.end(),rhs.begin(),rhs.end()); + return lhs; +} + +template<> std::vector<char>& cnpy::operator+=(std::vector<char>& lhs, const char* rhs) { + //write in little endian + size_t len = strlen(rhs); + lhs.reserve(len); + for(size_t byte = 0; byte < len; byte++) { + lhs.push_back(rhs[byte]); + } + return lhs; +} + +void cnpy::parse_npy_header(unsigned char* buffer,size_t& word_size, std::vector<size_t>& shape, bool& fortran_order) { + //std::string magic_string(buffer,6); + uint8_t major_version = *reinterpret_cast<uint8_t*>(buffer+6); + uint8_t minor_version = *reinterpret_cast<uint8_t*>(buffer+7); + uint16_t header_len = *reinterpret_cast<uint16_t*>(buffer+8); + std::string header(reinterpret_cast<char*>(buffer+9),header_len); + + size_t loc1, loc2; + + //fortran order + loc1 = header.find("fortran_order")+16; + fortran_order = (header.substr(loc1,4) == "True" ? true : false); + + //shape + loc1 = header.find("("); + loc2 = header.find(")"); + + std::regex num_regex("[0-9][0-9]*"); + std::smatch sm; + shape.clear(); + + std::string str_shape = header.substr(loc1+1,loc2-loc1-1); + while(std::regex_search(str_shape, sm, num_regex)) { + shape.push_back(std::stoi(sm[0].str())); + str_shape = sm.suffix().str(); + } + + //endian, word size, data type + //byte order code | stands for not applicable. + //not sure when this applies except for byte array + loc1 = header.find("descr")+9; + bool littleEndian = (header[loc1] == '<' || header[loc1] == '|' ? true : false); + assert(littleEndian); + + //char type = header[loc1+1]; + //assert(type == map_type(T)); + + std::string str_ws = header.substr(loc1+2); + loc2 = str_ws.find("'"); + word_size = atoi(str_ws.substr(0,loc2).c_str()); +} + +void cnpy::parse_npy_header(FILE* fp, size_t& word_size, std::vector<size_t>& shape, bool& fortran_order) { + char buffer[256]; + size_t res = fread(buffer,sizeof(char),11,fp); + if(res != 11) + throw std::runtime_error("parse_npy_header: failed fread"); + std::string header = fgets(buffer,256,fp); + assert(header[header.size()-1] == '\n'); + + size_t loc1, loc2; + + //fortran order + loc1 = header.find("fortran_order"); + if (loc1 == std::string::npos) + throw std::runtime_error("parse_npy_header: failed to find header keyword: 'fortran_order'"); + loc1 += 16; + fortran_order = (header.substr(loc1,4) == "True" ? true : false); + + //shape + loc1 = header.find("("); + loc2 = header.find(")"); + if (loc1 == std::string::npos || loc2 == std::string::npos) + throw std::runtime_error("parse_npy_header: failed to find header keyword: '(' or ')'"); + + std::regex num_regex("[0-9][0-9]*"); + std::smatch sm; + shape.clear(); + + std::string str_shape = header.substr(loc1+1,loc2-loc1-1); + while(std::regex_search(str_shape, sm, num_regex)) { + shape.push_back(std::stoi(sm[0].str())); + str_shape = sm.suffix().str(); + } + + //endian, word size, data type + //byte order code | stands for not applicable. + //not sure when this applies except for byte array + loc1 = header.find("descr"); + if (loc1 == std::string::npos) + throw std::runtime_error("parse_npy_header: failed to find header keyword: 'descr'"); + loc1 += 9; + bool littleEndian = (header[loc1] == '<' || header[loc1] == '|' ? true : false); + assert(littleEndian); + + //char type = header[loc1+1]; + //assert(type == map_type(T)); + + std::string str_ws = header.substr(loc1+2); + loc2 = str_ws.find("'"); + word_size = atoi(str_ws.substr(0,loc2).c_str()); +} + +void cnpy::parse_zip_footer(FILE* fp, uint16_t& nrecs, size_t& global_header_size, size_t& global_header_offset) +{ + std::vector<char> footer(22); + fseek(fp,-22,SEEK_END); + size_t res = fread(&footer[0],sizeof(char),22,fp); + if(res != 22) + throw std::runtime_error("parse_zip_footer: failed fread"); + + uint16_t disk_no, disk_start, nrecs_on_disk, comment_len; + disk_no = *(uint16_t*) &footer[4]; + disk_start = *(uint16_t*) &footer[6]; + nrecs_on_disk = *(uint16_t*) &footer[8]; + nrecs = *(uint16_t*) &footer[10]; + global_header_size = *(uint32_t*) &footer[12]; + global_header_offset = *(uint32_t*) &footer[16]; + comment_len = *(uint16_t*) &footer[20]; + + assert(disk_no == 0); + assert(disk_start == 0); + assert(nrecs_on_disk == nrecs); + assert(comment_len == 0); +} + +cnpy::NpyArray load_the_npy_file(FILE* fp) { + std::vector<size_t> shape; + size_t word_size; + bool fortran_order; + cnpy::parse_npy_header(fp,word_size,shape,fortran_order); + + cnpy::NpyArray arr(shape, word_size, fortran_order); + size_t nread = fread(arr.data<char>(),1,arr.num_bytes(),fp); + if(nread != arr.num_bytes()) + throw std::runtime_error("load_the_npy_file: failed fread"); + return arr; +} + +cnpy::NpyArray load_the_npz_array(FILE* fp, uint32_t compr_bytes, uint32_t uncompr_bytes) { + + std::vector<unsigned char> buffer_compr(compr_bytes); + std::vector<unsigned char> buffer_uncompr(uncompr_bytes); + size_t nread = fread(&buffer_compr[0],1,compr_bytes,fp); + if(nread != compr_bytes) + throw std::runtime_error("load_the_npy_file: failed fread"); + + int err; + z_stream d_stream; + + d_stream.zalloc = Z_NULL; + d_stream.zfree = Z_NULL; + d_stream.opaque = Z_NULL; + d_stream.avail_in = 0; + d_stream.next_in = Z_NULL; + err = inflateInit2(&d_stream, -MAX_WBITS); + + d_stream.avail_in = compr_bytes; + d_stream.next_in = &buffer_compr[0]; + d_stream.avail_out = uncompr_bytes; + d_stream.next_out = &buffer_uncompr[0]; + + err = inflate(&d_stream, Z_FINISH); + err = inflateEnd(&d_stream); + + std::vector<size_t> shape; + size_t word_size; + bool fortran_order; + cnpy::parse_npy_header(&buffer_uncompr[0],word_size,shape,fortran_order); + + cnpy::NpyArray array(shape, word_size, fortran_order); + + size_t offset = uncompr_bytes - array.num_bytes(); + memcpy(array.data<unsigned char>(),&buffer_uncompr[0]+offset,array.num_bytes()); + + return array; +} + +cnpy::npz_t cnpy::npz_load(std::string fname) { + FILE* fp = fopen(fname.c_str(),"rb"); + + if(!fp) { + throw std::runtime_error("npz_load: Error! Unable to open file "+fname+"!"); + } + + cnpy::npz_t arrays; + + while(1) { + std::vector<char> local_header(30); + size_t headerres = fread(&local_header[0],sizeof(char),30,fp); + if(headerres != 30) + throw std::runtime_error("npz_load: failed fread"); + + //if we've reached the global header, stop reading + if(local_header[2] != 0x03 || local_header[3] != 0x04) break; + + //read in the variable name + uint16_t name_len = *(uint16_t*) &local_header[26]; + std::string varname(name_len,' '); + size_t vname_res = fread(&varname[0],sizeof(char),name_len,fp); + if(vname_res != name_len) + throw std::runtime_error("npz_load: failed fread"); + + //erase the lagging .npy + varname.erase(varname.end()-4,varname.end()); + + //read in the extra field + uint16_t extra_field_len = *(uint16_t*) &local_header[28]; + if(extra_field_len > 0) { + std::vector<char> buff(extra_field_len); + size_t efield_res = fread(&buff[0],sizeof(char),extra_field_len,fp); + if(efield_res != extra_field_len) + throw std::runtime_error("npz_load: failed fread"); + } + + uint16_t compr_method = *reinterpret_cast<uint16_t*>(&local_header[0]+8); + uint32_t compr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+18); + uint32_t uncompr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+22); + + if(compr_method == 0) {arrays[varname] = load_the_npy_file(fp);} + else {arrays[varname] = load_the_npz_array(fp,compr_bytes,uncompr_bytes);} + } + + fclose(fp); + return arrays; +} + +cnpy::NpyArray cnpy::npz_load(std::string fname, std::string varname) { + FILE* fp = fopen(fname.c_str(),"rb"); + + if(!fp) throw std::runtime_error("npz_load: Unable to open file "+fname); + + while(1) { + std::vector<char> local_header(30); + size_t header_res = fread(&local_header[0],sizeof(char),30,fp); + if(header_res != 30) + throw std::runtime_error("npz_load: failed fread"); + + //if we've reached the global header, stop reading + if(local_header[2] != 0x03 || local_header[3] != 0x04) break; + + //read in the variable name + uint16_t name_len = *(uint16_t*) &local_header[26]; + std::string vname(name_len,' '); + size_t vname_res = fread(&vname[0],sizeof(char),name_len,fp); + if(vname_res != name_len) + throw std::runtime_error("npz_load: failed fread"); + vname.erase(vname.end()-4,vname.end()); //erase the lagging .npy + + //read in the extra field + uint16_t extra_field_len = *(uint16_t*) &local_header[28]; + fseek(fp,extra_field_len,SEEK_CUR); //skip past the extra field + + uint16_t compr_method = *reinterpret_cast<uint16_t*>(&local_header[0]+8); + uint32_t compr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+18); + uint32_t uncompr_bytes = *reinterpret_cast<uint32_t*>(&local_header[0]+22); + + if(vname == varname) { + NpyArray array = (compr_method == 0) ? load_the_npy_file(fp) : load_the_npz_array(fp,compr_bytes,uncompr_bytes); + fclose(fp); + return array; + } + else { + //skip past the data + uint32_t size = *(uint32_t*) &local_header[22]; + fseek(fp,size,SEEK_CUR); + } + } + + fclose(fp); + + //if we get here, we haven't found the variable in the file + throw std::runtime_error("npz_load: Variable name "+varname+" not found in "+fname); +} + +cnpy::NpyArray cnpy::npy_load(std::string fname) { + + FILE* fp = fopen(fname.c_str(), "rb"); + + if(!fp) throw std::runtime_error("npy_load: Unable to open file "+fname); + + NpyArray arr = load_the_npy_file(fp); + + fclose(fp); + return arr; +} + + +
diff --git a/third_party/cnpy/cnpy.h b/third_party/cnpy/cnpy.h new file mode 100644 index 0000000..0d3bb4c --- /dev/null +++ b/third_party/cnpy/cnpy.h
@@ -0,0 +1,269 @@ +//Copyright (C) 2011 Carl Rogers +//Released under MIT License +//license available in LICENSE file, or at http://www.opensource.org/licenses/mit-license.php + +#ifndef LIBCNPY_H_ +#define LIBCNPY_H_ + +#include<string> +#include<stdexcept> +#include<sstream> +#include<vector> +#include<cstdio> +#include<typeinfo> +#include<iostream> +#include<cassert> +#include<zlib.h> +#include<map> +#include<memory> +#include<stdint.h> +#include<numeric> + +namespace cnpy { + + struct NpyArray { + NpyArray(const std::vector<size_t>& _shape, size_t _word_size, bool _fortran_order) : + shape(_shape), word_size(_word_size), fortran_order(_fortran_order) + { + num_vals = 1; + for(size_t i = 0;i < shape.size();i++) num_vals *= shape[i]; + data_holder = std::shared_ptr<std::vector<char>>( + new std::vector<char>(num_vals * word_size)); + } + + NpyArray() : shape(0), word_size(0), fortran_order(0), num_vals(0) { } + + template<typename T> + T* data() { + return reinterpret_cast<T*>(&(*data_holder)[0]); + } + + template<typename T> + const T* data() const { + return reinterpret_cast<T*>(&(*data_holder)[0]); + } + + template<typename T> + std::vector<T> as_vec() const { + const T* p = data<T>(); + return std::vector<T>(p, p+num_vals); + } + + size_t num_bytes() const { + return data_holder->size(); + } + + std::shared_ptr<std::vector<char>> data_holder; + std::vector<size_t> shape; + size_t word_size; + bool fortran_order; + size_t num_vals; + }; + + using npz_t = std::map<std::string, NpyArray>; + + char BigEndianTest(); + char map_type(const std::type_info& t); + template<typename T> std::vector<char> create_npy_header(const std::vector<size_t>& shape); + void parse_npy_header(FILE* fp,size_t& word_size, std::vector<size_t>& shape, bool& fortran_order); + void parse_npy_header(unsigned char* buffer,size_t& word_size, std::vector<size_t>& shape, bool& fortran_order); + void parse_zip_footer(FILE* fp, uint16_t& nrecs, size_t& global_header_size, size_t& global_header_offset); + npz_t npz_load(std::string fname); + NpyArray npz_load(std::string fname, std::string varname); + NpyArray npy_load(std::string fname); + + template<typename T> std::vector<char>& operator+=(std::vector<char>& lhs, const T rhs) { + //write in little endian + for(size_t byte = 0; byte < sizeof(T); byte++) { + char val = *((char*)&rhs+byte); + lhs.push_back(val); + } + return lhs; + } + + template<> std::vector<char>& operator+=(std::vector<char>& lhs, const std::string rhs); + template<> std::vector<char>& operator+=(std::vector<char>& lhs, const char* rhs); + + + template<typename T> void npy_save(std::string fname, const T* data, const std::vector<size_t> shape, std::string mode = "w") { + FILE* fp = NULL; + std::vector<size_t> true_data_shape; //if appending, the shape of existing + new data + + if(mode == "a") fp = fopen(fname.c_str(),"r+b"); + + if(fp) { + //file exists. we need to append to it. read the header, modify the array size + size_t word_size; + bool fortran_order; + parse_npy_header(fp,word_size,true_data_shape,fortran_order); + assert(!fortran_order); + + if(word_size != sizeof(T)) { + std::cout<<"libnpy error: "<<fname<<" has word size "<<word_size<<" but npy_save appending data sized "<<sizeof(T)<<"\n"; + assert( word_size == sizeof(T) ); + } + if(true_data_shape.size() != shape.size()) { + std::cout<<"libnpy error: npy_save attempting to append misdimensioned data to "<<fname<<"\n"; + assert(true_data_shape.size() != shape.size()); + } + + for(size_t i = 1; i < shape.size(); i++) { + if(shape[i] != true_data_shape[i]) { + std::cout<<"libnpy error: npy_save attempting to append misshaped data to "<<fname<<"\n"; + assert(shape[i] == true_data_shape[i]); + } + } + true_data_shape[0] += shape[0]; + } + else { + fp = fopen(fname.c_str(),"wb"); + true_data_shape = shape; + } + + std::vector<char> header = create_npy_header<T>(true_data_shape); + size_t nels = std::accumulate(shape.begin(),shape.end(),1,std::multiplies<size_t>()); + + fseek(fp,0,SEEK_SET); + fwrite(&header[0],sizeof(char),header.size(),fp); + fseek(fp,0,SEEK_END); + fwrite(data,sizeof(T),nels,fp); + fclose(fp); + } + + template<typename T> void npz_save(std::string zipname, std::string fname, const T* data, const std::vector<size_t>& shape, std::string mode = "w") + { + //first, append a .npy to the fname + fname += ".npy"; + + //now, on with the show + FILE* fp = NULL; + uint16_t nrecs = 0; + size_t global_header_offset = 0; + std::vector<char> global_header; + + if(mode == "a") fp = fopen(zipname.c_str(),"r+b"); + + if(fp) { + //zip file exists. we need to add a new npy file to it. + //first read the footer. this gives us the offset and size of the global header + //then read and store the global header. + //below, we will write the the new data at the start of the global header then append the global header and footer below it + size_t global_header_size; + parse_zip_footer(fp,nrecs,global_header_size,global_header_offset); + fseek(fp,global_header_offset,SEEK_SET); + global_header.resize(global_header_size); + size_t res = fread(&global_header[0],sizeof(char),global_header_size,fp); + if(res != global_header_size){ + throw std::runtime_error("npz_save: header read error while adding to existing zip"); + } + fseek(fp,global_header_offset,SEEK_SET); + } + else { + fp = fopen(zipname.c_str(),"wb"); + } + + std::vector<char> npy_header = create_npy_header<T>(shape); + + size_t nels = std::accumulate(shape.begin(),shape.end(),1,std::multiplies<size_t>()); + size_t nbytes = nels*sizeof(T) + npy_header.size(); + + //get the CRC of the data to be added + uint32_t crc = crc32(0L,(uint8_t*)&npy_header[0],npy_header.size()); + crc = crc32(crc,(uint8_t*)data,nels*sizeof(T)); + + //build the local header + std::vector<char> local_header; + local_header += "PK"; //first part of sig + local_header += (uint16_t) 0x0403; //second part of sig + local_header += (uint16_t) 20; //min version to extract + local_header += (uint16_t) 0; //general purpose bit flag + local_header += (uint16_t) 0; //compression method + local_header += (uint16_t) 0; //file last mod time + local_header += (uint16_t) 0; //file last mod date + local_header += (uint32_t) crc; //crc + local_header += (uint32_t) nbytes; //compressed size + local_header += (uint32_t) nbytes; //uncompressed size + local_header += (uint16_t) fname.size(); //fname length + local_header += (uint16_t) 0; //extra field length + local_header += fname; + + //build global header + global_header += "PK"; //first part of sig + global_header += (uint16_t) 0x0201; //second part of sig + global_header += (uint16_t) 20; //version made by + global_header.insert(global_header.end(),local_header.begin()+4,local_header.begin()+30); + global_header += (uint16_t) 0; //file comment length + global_header += (uint16_t) 0; //disk number where file starts + global_header += (uint16_t) 0; //internal file attributes + global_header += (uint32_t) 0; //external file attributes + global_header += (uint32_t) global_header_offset; //relative offset of local file header, since it begins where the global header used to begin + global_header += fname; + + //build footer + std::vector<char> footer; + footer += "PK"; //first part of sig + footer += (uint16_t) 0x0605; //second part of sig + footer += (uint16_t) 0; //number of this disk + footer += (uint16_t) 0; //disk where footer starts + footer += (uint16_t) (nrecs+1); //number of records on this disk + footer += (uint16_t) (nrecs+1); //total number of records + footer += (uint32_t) global_header.size(); //nbytes of global headers + footer += (uint32_t) (global_header_offset + nbytes + local_header.size()); //offset of start of global headers, since global header now starts after newly written array + footer += (uint16_t) 0; //zip file comment length + + //write everything + fwrite(&local_header[0],sizeof(char),local_header.size(),fp); + fwrite(&npy_header[0],sizeof(char),npy_header.size(),fp); + fwrite(data,sizeof(T),nels,fp); + fwrite(&global_header[0],sizeof(char),global_header.size(),fp); + fwrite(&footer[0],sizeof(char),footer.size(),fp); + fclose(fp); + } + + template<typename T> void npy_save(std::string fname, const std::vector<T> data, std::string mode = "w") { + std::vector<size_t> shape; + shape.push_back(data.size()); + npy_save(fname, &data[0], shape, mode); + } + + template<typename T> void npz_save(std::string zipname, std::string fname, const std::vector<T> data, std::string mode = "w") { + std::vector<size_t> shape; + shape.push_back(data.size()); + npz_save(zipname, fname, &data[0], shape, mode); + } + + template<typename T> std::vector<char> create_npy_header(const std::vector<size_t>& shape) { + + std::vector<char> dict; + dict += "{'descr': '"; + dict += BigEndianTest(); + dict += map_type(typeid(T)); + dict += std::to_string(sizeof(T)); + dict += "', 'fortran_order': False, 'shape': ("; + dict += std::to_string(shape[0]); + for(size_t i = 1;i < shape.size();i++) { + dict += ", "; + dict += std::to_string(shape[i]); + } + if(shape.size() == 1) dict += ","; + dict += "), }"; + //pad with spaces so that preamble+dict is modulo 16 bytes. preamble is 10 bytes. dict needs to end with \n + int remainder = 16 - (10 + dict.size()) % 16; + dict.insert(dict.end(),remainder,' '); + dict.back() = '\n'; + + std::vector<char> header; + header += (char) 0x93; + header += "NUMPY"; + header += (char) 0x01; //major version of numpy format + header += (char) 0x00; //minor version of numpy format + header += (uint16_t) dict.size(); + header.insert(header.end(),dict.begin(),dict.end()); + + return header; + } + + +} + +#endif
diff --git a/tools/ml/py_bridge.cc b/tools/ml/py_bridge.cc new file mode 100644 index 0000000..36e692f --- /dev/null +++ b/tools/ml/py_bridge.cc
@@ -0,0 +1,202 @@ +#include "tools/ml/py_bridge.h" + +#include <algorithm> +#include <iterator> +#include <cstdlib> +#include "cnpy.h" + +#include <stdio.h> +#include <stdlib.h> +#include <string.h> +#include <dirent.h> +#include <fnmatch.h> // For pattern matching +#include <unistd.h> // For unlink (file deletion) +#include <filesystem> +#include <sys/stat.h> + +struct Var { + std::vector<size_t> shape; + std::vector<double> data; +}; + +struct DataFile { + std::map<std::string, Var> vars; +}; + +struct PyDataFile { + std::map<std::string, DataFile> files; +}; + +extern "C" void* py_datafile_open() { + struct PyDataFile* df = new struct PyDataFile(); + return (void*)df; +} + +static size_t get_data_length(std::vector<size_t> shape) { + int data_len = 1; + for (size_t v : shape) { + data_len *= v; + } + return data_len; +} + +static void py_datafile_add_int(void** context, const char* file, + const char* var_name, float* var_data, + std::vector<size_t> shape) { + struct PyDataFile* df = (struct PyDataFile*)*context; + if (df == nullptr) { + printf("Creating py datafile\n"); + df = new struct PyDataFile(); + *context = (void*)df; + } + if (!df->files[file].vars.count(var_name)) { + df->files[file].vars[var_name].shape = shape; + } else { + const auto& old_shape = df->files[file].vars[var_name].shape; + if (old_shape != shape) { + std::cout << "Illegal shape for var:" << var_name << ", was "; + std::copy(old_shape.begin(), old_shape.end(), std::ostream_iterator<int>(std::cout, ", ")); + std::cout << " now "; + std::copy(shape.begin(), shape.end(), std::ostream_iterator<int>(std::cout, ", ")); + std::cout << std::endl; + std::exit(1); + } + } + auto& data = df->files[file].vars[var_name].data; + for (size_t i = 0; i < get_data_length(shape); i++) { + data.push_back(var_data[i]); + } +} + +extern "C" void py_datafile_add_0d(void** context, const char* file, + const char* var_name, float* var_data) { + py_datafile_add_int(context, file, var_name, var_data, {}); +} + +extern "C" void py_datafile_add_1d(void** context, const char* file, + const char* var_name, float* var_data, + size_t shape0) { + py_datafile_add_int(context, file, var_name, var_data, {shape0}); +} + +extern "C" void py_datafile_add_2d(void** context, const char* file, + const char* var_name, float* var_data, + size_t shape0, size_t shape1) { + py_datafile_add_int(context, file, var_name, var_data, {shape0, shape1}); +} + +void py_datafile_fold(void** _to, void** _from) { + struct PyDataFile* to = (struct PyDataFile*)*_to; + struct PyDataFile* from = (struct PyDataFile*)*_from; + if (from == nullptr) { + std::cout << "\x1b[91mpy_data_fold: from is null\x1b[0m" << std::endl; + return; + } + if (to == nullptr) { + to = new struct PyDataFile(); + *_to = (void*)to; + } + for (const auto& e0 : from->files) { + for (const auto& e1 : e0.second.vars) { + const auto& from_var = e1.second; + if (!to->files[e0.first].vars.count(e1.first)) { + to->files[e0.first].vars[e1.first].shape = from_var.shape; + } else { + const auto& to_shape = to->files[e0.first].vars[e1.first].shape; + if (to_shape != from_var.shape) { + std::cout << "Illegal shape for var:" << e1.first << ", was "; + std::copy(to_shape.begin(), to_shape.end(), std::ostream_iterator<int>(std::cout, ", ")); + std::cout << " now "; + std::copy(from_var.shape.begin(), from_var.shape.end(), std::ostream_iterator<int>(std::cout, ", ")); + std::cout << std::endl; + std::exit(1); + } + } + auto& to_var = to->files[e0.first].vars[e1.first]; + to_var.data.insert(to_var.data.end(), from_var.data.begin(), from_var.data.end()); + } + } + delete from; + *_from = nullptr; +} + +extern "C" int py_datafile_delete_old_files(const char* _filename, size_t* count) { + DIR *dir; + struct dirent *entry; + char filepath[256]; + char pattern[128]; + const char* dirname; + const char* basename; + + *count = 0; + char path_copy[strlen(_filename) + 1]; + strcpy(path_copy, _filename); + char *last_slash = strrchr(path_copy, '/'); + if (last_slash != NULL) { + *last_slash = '\0'; + dirname = path_copy; + basename = last_slash + 1; + } else { + dirname = "."; + basename = path_copy; + } + + if ((dir = opendir(dirname)) == NULL) { + return 1; + } + + sprintf(pattern, "%s*.npz", basename); + + int ret = 0; + while ((entry = readdir(dir))!= NULL) { + if (fnmatch(pattern, entry->d_name, 0) == 0) { + snprintf(filepath, sizeof(filepath), "%s/%s", dirname, entry->d_name); + if (unlink(filepath)) { + ret = 1; + } else { + (*count)++; + } + } + } + closedir(dir); + return ret; +} + +bool dir_exists(const std::string& file_path) { + std::filesystem::path path = std::filesystem::path(file_path).parent_path(); + if (path.empty()) + return true; + struct stat info; + if (stat(path.c_str(), &info) != 0) { + return false; // Path does not exist or is not accessible + } + return (info.st_mode & S_IFDIR) != 0; // Check if it's a directory +} + +extern "C" void py_datafile_close(void** context, const char* _filename) { + struct PyDataFile* df = (struct PyDataFile*)*context; + if (df == nullptr) { + std::cout << "\x1b[91mpy_data_close: df is null\x1b[0m" << std::endl; + return; + } + if (!dir_exists(std::string(_filename))) { + std::cout << "\x1b[91mpy_data_close: tring to create a file " + << std::string(_filename) << " in a directory that doesn't " + << "exist.\x1b[0m" << std::endl; + return; + } + for (const auto& e0 : df->files) { + const auto filename = std::string(_filename) + "_" + e0.first + ".npz"; + for (const auto& e1 : e0.second.vars) { + const auto& var = e1.second; + const unsigned int data_length = get_data_length(var.shape); + std::vector<size_t> npz_shape; + npz_shape.push_back(var.data.size() / data_length); + npz_shape.insert(npz_shape.end(), var.shape.begin(), var.shape.end()); + cnpy::npz_save(filename, e1.first, &var.data[0], npz_shape, "a"); + } + } + delete df; + *context = nullptr; +} +
diff --git a/tools/ml/py_bridge.h b/tools/ml/py_bridge.h new file mode 100644 index 0000000..8d4e40f --- /dev/null +++ b/tools/ml/py_bridge.h
@@ -0,0 +1,23 @@ +#ifndef TOOLS_ML_PY_BRIDGE_H_ +#define TOOLS_ML_PY_BRIDGE_H_ + +#ifdef __cplusplus +extern "C" { +#endif +#include <stddef.h> + +void* py_datafile_open(); +void py_datafile_close(void** context, const char* filename); +void py_datafile_add_0d(void** context, const char* file, const char* var_name, + float* var_data); +void py_datafile_add_1d(void** context, const char* file, const char* var_name, + float* var_data, size_t shape0); +void py_datafile_add_2d(void** context, const char* file, const char* var_name, + float* var_data, size_t shape0, size_t shape1); +void py_datafile_fold(void** to, void** from); +int py_datafile_delete_old_files(const char* _filename, size_t* count); +#ifdef __cplusplus +} +#endif + +#endif // TOOLS_ML_PY_BRIDGE_H_