1 // SPDX-FileCopyrightText: Copyright (c) 2017-2025, HONEE contributors. 2 // SPDX-License-Identifier: Apache-2.0 OR BSD-2-Clause 3 4 #include <navierstokes.h> 5 #include <smartsim-impl.h> 6 7 typedef struct { 8 PetscInt write_data_interval; 9 size_t local_array_dims[2]; 10 PetscBool overwrite_training_data; 11 } *SmartSimSolutionData; 12 13 PetscErrorCode TSMonitor_SmartSimSolutionSetup(TS ts, PetscViewerAndFormat *ctx) { 14 SmartSimSolutionData smartsimsol; 15 Honee honee; 16 MPI_Comm comm = PetscObjectComm((PetscObject)ts); 17 18 PetscFunctionBeginUser; 19 PetscCall(TSGetApplicationContext(ts, &honee)); 20 PetscCall(PetscNew(&smartsimsol)); 21 smartsimsol->overwrite_training_data = PETSC_TRUE; 22 23 PetscOptionsBegin(comm, NULL, "SmartSim Solution Writing", NULL); 24 PetscCall(PetscOptionsBool("-ts_monitor_smartsim_solution_overwrite_data", "Overwrite old solution data in the database", NULL, 25 smartsimsol->overwrite_training_data, &smartsimsol->overwrite_training_data, NULL)); 26 PetscOptionsEnd(); 27 28 { // Get solution vector size 29 PetscSection output_section; 30 PetscInt num_dofs, num_comps; 31 DM dm, output_dm; 32 33 PetscCall(TSGetDM(ts, &dm)); 34 PetscCall(DMGetOutputDM(dm, &output_dm)); 35 36 PetscCall(DMGetGlobalVectorInfo(output_dm, &num_dofs, NULL, NULL)); 37 PetscCall(DMGetGlobalSection(output_dm, &output_section)); 38 PetscCall(PetscSectionGetFieldComponents(output_section, 0, &num_comps)); 39 smartsimsol->local_array_dims[0] = num_dofs / num_comps; 40 smartsimsol->local_array_dims[1] = num_comps; 41 } 42 ctx->data = smartsimsol; 43 PetscFunctionReturn(PETSC_SUCCESS); 44 } 45 46 PetscErrorCode TSMonitor_SmartSimSolution(TS ts, PetscInt step_num, PetscReal solution_time, Vec Q, PetscViewerAndFormat *ctx) { 47 Honee honee; 48 SmartSimData smartsim; 49 Vec Q_output; 50 PetscMPIInt rank; 51 DM dm, output_dm; 52 53 PetscFunctionBeginUser; 54 SmartSimSolutionData smartsimsol = ctx->data; 55 PetscCall(TSGetApplicationContext(ts, &honee)); 56 PetscCall(HoneeGetSmartSimData(honee, &smartsim)); 57 PetscCallMPI(MPI_Comm_rank(honee->comm, &rank)); 58 59 if (step_num % ctx->view_interval != 0) PetscFunctionReturn(PETSC_SUCCESS); 60 61 PetscCall(TSGetDM(ts, &dm)); 62 PetscCall(DMGetOutputDM(dm, &output_dm)); 63 PetscCall(DMGetGlobalVector(output_dm, &Q_output)); 64 65 PetscCall(UpdateBoundaryValues(honee, honee->Q_loc, solution_time)); 66 PetscCall(DMGlobalToLocal(honee->dm, Q, INSERT_VALUES, honee->Q_loc)); 67 PetscCall(DMLocalToGlobal(output_dm, honee->Q_loc, INSERT_VALUES, Q_output)); 68 69 { // -- Send solution data to SmartSim 70 char array_key[PETSC_MAX_PATH_LEN]; 71 size_t array_key_len; 72 73 if (smartsimsol->overwrite_training_data) { 74 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.flow_solution", smartsim->rank_id_name)); 75 } else { 76 PetscCall(PetscSNPrintf(array_key, sizeof array_key, "%s.flow_solution.%" PetscInt_FMT, smartsim->rank_id_name, step_num)); 77 } 78 PetscCall(PetscStrlen(array_key, &array_key_len)); 79 80 { 81 const PetscScalar *Q_output_array; 82 PetscCall(VecGetArrayRead(Q_output, &Q_output_array)); 83 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Write, 0, 0, 0, 0)); 84 PetscCallSmartRedis(put_tensor(smartsim->client, array_key, array_key_len, (void *)Q_output_array, smartsimsol->local_array_dims, 2, 85 SRTensorTypeDouble, SRMemLayoutContiguous)); 86 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Write, 0, 0, 0, 0)); 87 PetscCall(VecRestoreArrayRead(Q_output, &Q_output_array)); 88 } 89 } 90 91 if (rank % smartsim->collocated_database_num_ranks == 0) { 92 const char tensor_name[] = "flow_solution_step"; 93 size_t dim_2[1] = {2}; 94 PetscInt64 step_array[2] = {step_num, step_num}; 95 96 PetscCall(PetscLogEventBegin(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 97 PetscCallSmartRedis( 98 put_tensor(smartsim->client, tensor_name, strlen(tensor_name), step_array, dim_2, 1, SRTensorTypeInt64, SRMemLayoutContiguous)); 99 PetscCall(PetscLogEventEnd(HONEE_SmartRedis_Meta, 0, 0, 0, 0)); 100 } 101 102 PetscCall(DMRestoreGlobalVector(output_dm, &Q_output)); 103 PetscFunctionReturn(PETSC_SUCCESS); 104 } 105