ExplicitKokkosSurfaceSolver.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
//
// Copyright 2026 Caleb Buahin
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

 
#ifndef OPENSWMM_ENGINE_2D_GPU_EXPLICIT_KOKKOS_SURFACE_SOLVER_HPP
#define OPENSWMM_ENGINE_2D_GPU_EXPLICIT_KOKKOS_SURFACE_SOLVER_HPP
 
#include <array>
#include <string>
#include <utility>
#include <vector>
 
#include "../solver/ISurfaceSolver.hpp"
#include "../solver/InertialEdges.hpp"
#include "KokkosTypes.hpp"
#include "KokkosPerfCounters.hpp"
 
namespace openswmm::twoD::gpu {
 
class ExplicitKokkosSurfaceSolver final : public ISurfaceSolver {
public:
    ExplicitKokkosSurfaceSolver() = default;
    ~ExplicitKokkosSurfaceSolver() override { finalize(); }

    void initialize(MeshData& mesh, SurfaceStateData& state, SolverOptions2D& opts) override;
    double advance(double t_current, double t_target) override;
    void reinitialize(double t0) override;
    void resyncFromVolumes(double t0) override;
    void finalize() override;

    long last_num_steps() const noexcept override { return last_steps_; }
    double last_step_size() const noexcept override { return last_dt_; }
    const std::vector<double>& last_coupling_exchange() const noexcept override {
        return exch_host_;
    }

    RunStats run_stats() const noexcept override;
    bool is_initialized() const noexcept override { return initialized_; }
 
private:
    friend struct ExplicitKokkosSurfaceSolverTestAccess;
    kperf::Counters perf_counters_;
    // Host-side handles (owned by SurfaceRouter2D; outlive the solver).
    MeshData* mesh_ = nullptr;
    SurfaceStateData* state_ = nullptr;
    SolverOptions2D* opts_ = nullptr;
    void refreshConveyance(); 
    InertialEdges edges_; 
    DView d_conveyance_;  
 
    // ---- device geometry (const after initialize) -------------------------
    DView d_tri_area_, d_tri_cz_, d_tri_cx_, d_tri_cy_, d_vz_, d_vx_, d_vy_;
    IView d_cell_v_;            
    DView d_vs_wt_;             
    IView d_vs_ptr_, d_vs_idx_; 
    DView d_edge_length_, d_mannings_n_;
    // InertialEdges mirror
    IView d_cL_, d_cR_, d_slotL_, d_slotR_, d_cell_ptr_, d_cell_edge_;
    DView d_xi_, d_inv_dx_, d_zface_, d_ze_lo_, d_ze_hi_, d_nx_, d_ny_, d_mx_, d_my_;
    DView d_n2_, d_sign_, d_lchar_;
 
    // ---- device state -----------------------------------------------------
    DView d_volume_, d_head_, d_depth_;
    DView d_q_, d_faccL_, d_faccR_, d_qcx_, d_qcy_;
    DView d_rain_, d_coup_, d_evap_, d_infil_;

    DView d_infil_applied_;
    DView d_evap_applied_;
    std::vector<double> evap_applied_host_;
    DView d_coupling_applied_;
    DView d_node_fullvol_;
    std::vector<double> infil_applied_host_;    
    std::vector<double> coupling_applied_host_; 
    DView d_edge_flux_;                         
    IView d_active_, d_pin_t0_, d_tier_, d_face_tier_;
    IView d_cells_compact_, d_edges_compact_; 
    IView d_scratch_;                         
    DView d_dtcell_;                          
 
    // Boundary edges (non-WALL): built host-side at initialize.
    IView d_bc_cell_, d_bc_slot_, d_bc_type_;
    DView d_bc_accum_, d_bc_slope_, d_bc_head_, d_bc_flow_;
    DView d_bc_in_, d_bc_out_; 
    std::vector<double> bc_in_host_, bc_out_host_;
    DView d_bc_requested_, d_bc_last_flow_, d_bc_last_dt_;
    std::vector<double> bc_requested_host_, bc_last_flow_host_, bc_last_dt_host_;
    DView d_bc_q_; 
    std::vector<int> bc_cell_host_, bc_slot_host_;
 
    // Live junction exchange (windowless coupling).
    IView d_cp_cell_, d_cp_vertex_, d_cp_node_;
    DView d_cp_cd_, d_cp_area_;
    // WEIR_ORIFICE per point: law, structure {z_crest, l_weir, a_orif} and
    // coefficients {c_w, c_sw, c_oi, c_oo, h_eps}; COUPLING_HEAD cell runs.
    IView d_cp_law_, d_cp_head_b_, d_cp_head_e_, d_head_cells_, d_cp_eq_;
    DView d_cp_geo_, d_cp_coef_;
    DView d_node_pond_area_; 
    DView d_exch_, d_node_drawn_;
    DView d_node_head_, d_node_depth_, d_node_volume_;
    DView d_node_invert_, d_node_fulldepth_;
    std::vector<double> exch_host_;
 
    // ---- host mirrors / control -------------------------------------------
    std::array<int, 9> tier_off_{};  
    std::array<int, 9> ftier_off_{}; 
    int n_active_ = 0;
    double dt0_ = 0.0;
    int K_ = 1;
    bool have_perot_ = false;
    double t_last_sync_ = 0.0;
    int cycles_since_rebuild_ = 1000;
    double last_dt_ = 0.0;
    long last_steps_ = 0, substeps_run_ = 0, face_passes_ = 0;
    bool initialized_ = false;
 
    std::array<long, 8> tier_occupancy_{};
    std::vector<std::pair<double, int>> telemetry_;
    std::string telemetry_path_;
 
    // ---- internals (public for nvcc extended-lambda access) ----------------
public:
    void reconstructAllDev();
    void settleAccumulatorsDev();
    void lazySourcesDev(double t);
    void syncAndRebuild(double t);
    void refreshDt0();         
    void collapseToGlobalDt(); 
    void fireFaces(int tier, double dt_f, bool global_step = false);
    void fireCells(int tier, double dt_c, bool global_step = false);
    void runMacroCycle(double dt0, int nsub);
    void pushForcings();
    void pushNodeState();
    void publishAndCopyBack(double t_current, double t_target);
};
 
} // namespace openswmm::twoD::gpu
 
#endif // OPENSWMM_ENGINE_2D_GPU_EXPLICIT_KOKKOS_SURFACE_SOLVER_HPP