From f85f4cb568a18fe324bf6faf23880189eb5c82b3 Mon Sep 17 00:00:00 2001 From: dungeon2/cli Date: Thu, 16 Jul 2026 00:06:09 -0400 Subject: [PATCH 1/3] wip --- CMakeLists.txt | 2 - cmake/nvshmem.cmake | 40 --- docs/user_guide/distributed.rst | 4 + src/CMakeLists.txt | 10 - src/hydro/fused_cubed_sphere_dispatch.cu | 54 +--- src/hydro/fused_recon_riemann_dispatch.hpp | 52 ++-- src/hydro/hydro_forward_fused.cpp | 277 ++++++++---------- src/hydro/hydro_options.cpp | 2 +- tests/CMakeLists.txt | 22 +- ...nc_smoke.cu => test_fused_cs_ucx_smoke.cu} | 194 +++++------- 10 files changed, 252 insertions(+), 405 deletions(-) delete mode 100644 cmake/nvshmem.cmake rename tests/{test_fused_cs_sync_smoke.cu => test_fused_cs_ucx_smoke.cu} (73%) diff --git a/CMakeLists.txt b/CMakeLists.txt index 4b96c50e..bfb6dba4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -129,8 +129,6 @@ if (CUDAToolkit_FOUND) set(CMAKE_CUDA_ARCHITECTURES ${_cuda_arch_list}) endif() -include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/nvshmem.cmake) - message(STATUS "${Blue}${PROJECT_NAME}-4. Setting up project libraries ...${ColorReset}") configure_file(${CMAKE_SOURCE_DIR}/configure.h.in configure.h @ONLY) add_subdirectory(src) diff --git a/cmake/nvshmem.cmake b/cmake/nvshmem.cmake deleted file mode 100644 index aad50fd5..00000000 --- a/cmake/nvshmem.cmake +++ /dev/null @@ -1,40 +0,0 @@ -set(NVSHMEM ${CUDA} CACHE INTERNAL - "Enable NVSHMEM-backed kernel-side communication" FORCE) - -if(NOT NVSHMEM) - return() -endif() - -get_filename_component(_torch_lib_dir "${TORCH_LIBRARY}" DIRECTORY) -get_filename_component(_torch_package_dir "${_torch_lib_dir}" DIRECTORY) -get_filename_component(_torch_site_packages_dir "${_torch_package_dir}" DIRECTORY) - -if(TARGET torch_nvshmem) - find_path(NVSHMEM_INCLUDE_DIR nvshmem.h - HINTS "${_torch_site_packages_dir}/nvidia/nvshmem" - PATH_SUFFIXES include) - set(NVSHMEM_TORCH_TARGET torch_nvshmem) - message(STATUS "NVSHMEM: using LibTorch torch_nvshmem target") -else() - find_path(NVSHMEM_INCLUDE_DIR nvshmem.h - HINTS - "${_torch_site_packages_dir}/nvidia/nvshmem" - ENV NVSHMEM_HOME - ENV NVSHMEM_ROOT - PATH_SUFFIXES include) - find_library(NVSHMEM_LIBRARY - NAMES torch_nvshmem nvshmem - HINTS - "${_torch_lib_dir}" - "${_torch_site_packages_dir}/nvidia/nvshmem" - ENV NVSHMEM_HOME - ENV NVSHMEM_ROOT - PATH_SUFFIXES lib lib64) -endif() - -if(NOT NVSHMEM_INCLUDE_DIR OR (NOT NVSHMEM_LIBRARY AND NOT NVSHMEM_TORCH_TARGET)) - message(FATAL_ERROR - "snapy requires NVSHMEM headers and either LibTorch's " - "torch_nvshmem target or a standalone libnvshmem when CUDA=ON. " - "Set NVSHMEM_HOME or NVSHMEM_ROOT if using a standalone install.") -endif() diff --git a/docs/user_guide/distributed.rst b/docs/user_guide/distributed.rst index a610b808..d9d175df 100644 --- a/docs/user_guide/distributed.rst +++ b/docs/user_guide/distributed.rst @@ -248,6 +248,10 @@ UCX transports for the cluster: export DEVICE=cuda export UCX_TLS=rc,cuda_copy,cuda_ipc +The fused cubed-sphere CUDA path also uses this UCX backend. It exchanges one +packed edge-state buffer with each neighboring process before launching the +fused seam-flux kernel; no separate symmetric-memory runtime is required. + Load Balancing -------------- diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index a3c54ca3..e95a874f 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -142,16 +142,6 @@ if (CUDAToolkit_FOUND) ${C10_CUDA_LIBRARY} ) - target_include_directories(${namel}_cuda_${buildl} - SYSTEM PUBLIC ${NVSHMEM_INCLUDE_DIR}) - if(NVSHMEM_TORCH_TARGET) - target_link_libraries(${namel}_cuda_${buildl} - PUBLIC ${NVSHMEM_TORCH_TARGET}) - else() - target_link_libraries(${namel}_cuda_${buildl} - PUBLIC ${NVSHMEM_LIBRARY}) - endif() - add_library(snapy::snap_cu ALIAS ${namel}_cuda_${buildl}) endif() diff --git a/src/hydro/fused_cubed_sphere_dispatch.cu b/src/hydro/fused_cubed_sphere_dispatch.cu index 00791653..002aebd1 100644 --- a/src/hydro/fused_cubed_sphere_dispatch.cu +++ b/src/hydro/fused_cubed_sphere_dispatch.cu @@ -2,7 +2,6 @@ #include #include #include -#include // C/C++ #include @@ -155,20 +154,6 @@ __global__ void cs_pack_kernel(T const *w, T *buf, } } -__global__ void cs_sync_kernel(uint32_t **signal_pads, int rank, - int world_size) { - c10d::symmetric_memory::sync_remote_blocks(signal_pads, rank, - world_size); -} - -__global__ void cs_release_reads_kernel(uint32_t **signal_pads, int rank, - int world_size) { - if (blockIdx.x != 1) - return; - c10d::symmetric_memory::sync_remote_blocks(signal_pads, rank, - world_size); -} - template __device__ void cs_flux_common(DeviceCubedSphereFluxParams params, int i, int edge_pos, int side) { @@ -348,7 +333,8 @@ __global__ void cs_flux_all_kernel(DeviceCubedSphereFluxParams base_params, } // namespace -void fused_cubed_sphere_pack_cuda(torch::Tensor w, torch::Tensor symm_buffer, +void fused_cubed_sphere_pack_cuda(torch::Tensor w, + torch::Tensor exchange_buffer, FusedCubedSpherePackParams const ¶ms) { at::cuda::CUDAGuard device_guard(w.device()); int nc3 = w.size(1); @@ -380,26 +366,16 @@ void fused_cubed_sphere_pack_cuda(torch::Tensor w, torch::Tensor symm_buffer, scalar_t(params.pressure_floor), params.eos_limiter}; cs_pack_kernel<<>>( - w.data_ptr(), symm_buffer.data_ptr(), + w.data_ptr(), exchange_buffer.data_ptr(), device_params); }); C10_CUDA_KERNEL_LAUNCH_CHECK(); } -void fused_cubed_sphere_sync_cuda(uint32_t **symm_signal_pads_dev, - int symm_rank, int symm_world_size, - torch::Device device) { - at::cuda::CUDAGuard device_guard(device); - auto stream = at::cuda::getCurrentCUDAStream(); - cs_sync_kernel<<<1, std::max(32, symm_world_size), 0, stream>>>( - symm_signal_pads_dev, symm_rank, symm_world_size); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, torch::Tensor flux3, - torch::Tensor symm_buffer, - void **symm_buffer_ptrs_dev, + torch::Tensor exchange_buffer, + void **exchange_buffer_ptrs_dev, FusedCubedSphereFluxParams const ¶ms) { at::cuda::CUDAGuard device_guard(w.device()); int nc3 = w.size(1); @@ -417,10 +393,10 @@ void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, AT_DISPATCH_FLOATING_TYPES(w.scalar_type(), "fused_cubed_sphere_flux", [&] { DeviceCubedSphereFluxParams device_params{ - symm_buffer.data_ptr(), + exchange_buffer.data_ptr(), flux2.defined() ? flux2.data_ptr() : nullptr, flux3.defined() ? flux3.data_ptr() : nullptr, - symm_buffer_ptrs_dev, + exchange_buffer_ptrs_dev, nvar, nc3, nc2, @@ -433,7 +409,7 @@ void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, } void fused_cubed_sphere_flux_all_cuda( - torch::Tensor symm_buffer, FusedCubedSphereFluxAllPtrs ptrs, + torch::Tensor exchange_buffer, FusedCubedSphereFluxAllPtrs ptrs, FusedCubedSphereFluxAllParams const ¶ms) { at::cuda::CUDAGuard device_guard(params.device); int edge_len = std::max(params.nc2, params.nc3); @@ -447,10 +423,10 @@ void fused_cubed_sphere_flux_all_cuda( AT_DISPATCH_FLOATING_TYPES(params.dtype, "fused_cubed_sphere_flux_all", [&] { DeviceCubedSphereFluxParams device_params{ - symm_buffer.data_ptr(), + exchange_buffer.data_ptr(), nullptr, nullptr, - ptrs.symm_buffer, + ptrs.exchange_buffer, params.nvar, params.nc3, params.nc2, @@ -465,14 +441,4 @@ void fused_cubed_sphere_flux_all_cuda( C10_CUDA_KERNEL_LAUNCH_CHECK(); } -void fused_cubed_sphere_release_cuda(uint32_t **symm_signal_pads_dev, - int symm_rank, int symm_world_size, - torch::Device device) { - at::cuda::CUDAGuard device_guard(device); - auto stream = at::cuda::getCurrentCUDAStream(); - cs_release_reads_kernel<<<2, std::max(32, symm_world_size), 0, stream>>>( - symm_signal_pads_dev, symm_rank, symm_world_size); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - } // namespace snap diff --git a/src/hydro/fused_recon_riemann_dispatch.hpp b/src/hydro/fused_recon_riemann_dispatch.hpp index 87d4114e..ae80b29f 100644 --- a/src/hydro/fused_recon_riemann_dispatch.hpp +++ b/src/hydro/fused_recon_riemann_dispatch.hpp @@ -49,15 +49,12 @@ struct FusedReconRiemannParams { void fused_recon_riemann_cuda(torch::Tensor w, torch::Tensor flux, torch::Tensor face_pressure, - FusedReconRiemannParams const ¶ms); - -// Multi-block-per-process cubed-sphere fused exchange, split into host-callable -// phases so the caller can coordinate the concurrent local-block threads: -// every local block packs its edge states into slice `local_block` of a shared -// per-process symmetric buffer, one designated block publishes cross-process -// visibility (sync), every local block computes and overwrites its cross-panel -// boundary flux (reading own + peer slices, peer selected by process rank), and -// one designated block closes the read epoch (release). + FusedReconRiemannParams const& params); + +// Multi-block-per-process cubed-sphere fused exchange. Every local block packs +// edge states into slice `local_block` of a shared per-process buffer. The host +// exchanges that buffer with remote peers, then one kernel computes and +// overwrites every local panel's cross-panel boundary flux. struct FusedCubedSpherePanelParams { torch::Tensor side_meta; int face; @@ -79,12 +76,9 @@ struct FusedCubedSpherePackParams { bool eos_limiter; }; -void fused_cubed_sphere_pack_cuda(torch::Tensor w, torch::Tensor symm_buffer, - FusedCubedSpherePackParams const ¶ms); - -void fused_cubed_sphere_sync_cuda(uint32_t **symm_signal_pads_dev, - int symm_rank, int symm_world_size, - torch::Device device); +void fused_cubed_sphere_pack_cuda(torch::Tensor w, + torch::Tensor exchange_buffer, + FusedCubedSpherePackParams const& params); struct FusedCubedSphereFluxParams { FusedCubedSpherePanelParams panel; @@ -93,9 +87,9 @@ struct FusedCubedSphereFluxParams { void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, torch::Tensor flux3, - torch::Tensor symm_buffer, - void **symm_buffer_ptrs_dev, - FusedCubedSphereFluxParams const ¶ms); + torch::Tensor exchange_buffer, + void** exchange_buffer_ptrs_dev, + FusedCubedSphereFluxParams const& params); // Process-level seam flux: one launch overwrites the cross-panel boundary flux // for ALL local panels. flux2_ptrs_dev/flux3_ptrs_dev are device arrays of the @@ -103,13 +97,13 @@ void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, // int array of each panel's face id, and side_meta_all is [bpp, 4, kStride]. // Coords are the shared equiangular angular grid (same for every panel). struct FusedCubedSphereFluxAllPtrs { - void **flux2; - void **flux3; - void **symm_buffer; - void **x2v; - void **x2f; - void **x3v; - void **x3f; + void** flux2; + void** flux3; + void** exchange_buffer; + void** x2v; + void** x2f; + void** x3v; + void** x3f; }; struct FusedCubedSphereFluxAllParams { @@ -126,11 +120,7 @@ struct FusedCubedSphereFluxAllParams { }; void fused_cubed_sphere_flux_all_cuda( - torch::Tensor symm_buffer, FusedCubedSphereFluxAllPtrs ptrs, - FusedCubedSphereFluxAllParams const ¶ms); - -void fused_cubed_sphere_release_cuda(uint32_t **symm_signal_pads_dev, - int symm_rank, int symm_world_size, - torch::Device device); + torch::Tensor exchange_buffer, FusedCubedSphereFluxAllPtrs ptrs, + FusedCubedSphereFluxAllParams const& params); } // namespace snap diff --git a/src/hydro/hydro_forward_fused.cpp b/src/hydro/hydro_forward_fused.cpp index 71e5fe7d..70a36be5 100644 --- a/src/hydro/hydro_forward_fused.cpp +++ b/src/hydro/hydro_forward_fused.cpp @@ -2,6 +2,7 @@ #include #include #include +#include #include #include #include @@ -14,7 +15,6 @@ #include #include #include -#include #include "../eos/ideal_moist.hpp" #include "../sedimentation/sedimentation.hpp" @@ -24,7 +24,7 @@ namespace snap { namespace { -FusedReconScheme fused_recon_scheme(std::string const &type, +FusedReconScheme fused_recon_scheme(std::string const& type, bool velocity_group, bool shock) { if (type == "weno5") { return velocity_group && !shock ? FusedReconScheme::CP5 @@ -42,8 +42,8 @@ FusedReconScheme fused_recon_scheme(std::string const &type, type); } -bool fused_combo_supported(std::string const &eos_type, - std::string const &riemann_type) { +bool fused_combo_supported(std::string const& eos_type, + std::string const& riemann_type) { return ((eos_type == "ideal-gas" && (riemann_type == "lmars" || riemann_type == "hllc" || riemann_type == "roe")) || @@ -53,7 +53,7 @@ bool fused_combo_supported(std::string const &eos_type, (eos_type == "shallow-water" && riemann_type == "shallow-roe")); } -FusedEos fused_eos(std::string const &type) { +FusedEos fused_eos(std::string const& type) { if (type == "ideal-gas") return FusedEos::IdealGas; if (type == "ideal-moist") return FusedEos::IdealMoist; if (type == "shallow-water") return FusedEos::ShallowWater; @@ -63,7 +63,7 @@ FusedEos fused_eos(std::string const &type) { type); } -FusedRiemannSolver fused_riemann_solver(std::string const &type) { +FusedRiemannSolver fused_riemann_solver(std::string const& type) { if (type == "lmars") return FusedRiemannSolver::LMARS; if (type == "hllc") return FusedRiemannSolver::HLLC; if (type == "roe") return FusedRiemannSolver::Roe; @@ -74,70 +74,44 @@ FusedRiemannSolver fused_riemann_solver(std::string const &type) { type); } -void ensure_symmetric_group(LayoutImpl const &layout, - std::string const &group_name) { - static std::mutex mutex; - static std::set initialized; - std::lock_guard lock(mutex); - if (initialized.count(group_name)) return; - TORCH_CHECK(layout.comm != nullptr && layout.comm->store.defined(), - "dynamics.fused-recon-riemann cubed-sphere exchange requires an " - "initialized process-group store"); - if (initialized.empty()) { - c10d::symmetric_memory::set_backend("NVSHMEM"); - c10d::symmetric_memory::set_signal_pad_size(std::max( - 1024, 4 * layout.options->process_world_size() * sizeof(uint32_t))); - } - - auto set_group_info_once = [&](std::string const &name) { - if (initialized.count(name)) return; - c10d::symmetric_memory::set_group_info(name, layout.options->process_rank(), - layout.options->process_world_size(), - layout.comm->store); - initialized.insert(name); - }; +constexpr int kFusedCubedSphereExchangeTag = + std::numeric_limits::max() - 113; - // PyTorch's NVSHMEM allocator bootstraps anonymous allocations through this - // default group; rendezvous below still uses the logical snap group. - set_group_info_once("0"); - set_group_info_once(group_name); -} - -struct FusedSymmetricExchangePool { +struct FusedExchangePool { torch::Tensor buffer; - //! non-null when the exchange spans multiple processes (NVSHMEM P2P) - c10::intrusive_ptr symm; - //! single-process fallback: a device int64 tensor holding buffer.data_ptr() - //! so the flux kernel can index buf_ptrs[0] uniformly for same-GPU peers. - torch::Tensor self_ptr; - - void **buffer_ptrs_dev() const { - if (symm) return symm->get_buffer_ptrs_dev(); - return reinterpret_cast(self_ptr.data_ptr()); - } - uint32_t **signal_pads_dev() const { - return symm ? reinterpret_cast(symm->get_signal_pad_ptrs_dev()) - : nullptr; + std::vector peer_buffers; + torch::Tensor buffer_ptrs; + + void** buffer_ptrs_dev() const { + return reinterpret_cast(buffer_ptrs.data_ptr()); } - int rank() const { return symm ? symm->get_rank() : 0; } - int world_size() const { return symm ? symm->get_world_size() : 1; } -}; -uint64_t stable_alloc_id(std::string const &key) { - uint64_t hash = 1469598103934665603ull; - for (unsigned char c : key) { - hash ^= c; - hash *= 1099511628211ull; + void prepare_peer_buffers(std::set const& peers, int rank, + int world_size) { + if (static_cast(peer_buffers.size()) != world_size) { + peer_buffers.resize(world_size); + } + std::vector pointers(world_size, 0); + pointers[rank] = reinterpret_cast(buffer.data_ptr()); + for (int peer : peers) { + TORCH_CHECK(peer >= 0 && peer < world_size, "invalid fused UCX peer ", + peer, " for world size ", world_size); + if (!peer_buffers[peer].defined()) { + peer_buffers[peer] = torch::empty_strided( + buffer.sizes(), buffer.strides(), buffer.options()); + } + pointers[peer] = reinterpret_cast(peer_buffers[peer].data_ptr()); + } + buffer_ptrs = torch::tensor(pointers, torch::dtype(torch::kInt64)) + .to(buffer.device()); } - return hash == 0 ? 1 : hash; -} +}; -FusedSymmetricExchangePool &get_fused_symmetric_exchange_pool( - std::string const &group_name, c10::ScalarType dtype, torch::Device device, - std::vector const &sizes, std::vector const &strides, - bool use_symmetric) { +FusedExchangePool& get_fused_exchange_pool( + std::string const& group_name, c10::ScalarType dtype, torch::Device device, + std::vector const& sizes, std::vector const& strides) { static std::mutex mutex; - static std::unordered_map pools; + static std::unordered_map pools; std::string key = group_name + ":device=" + std::to_string(device.index()) + ":dtype=" + std::to_string(static_cast(dtype)); @@ -148,30 +122,42 @@ FusedSymmetricExchangePool &get_fused_symmetric_exchange_pool( auto it = pools.find(key); if (it != pools.end()) return it->second; - FusedSymmetricExchangePool pool; - if (use_symmetric) { - pool.buffer = c10d::symmetric_memory::empty_strided_p2p( - sizes, strides, dtype, device, std::nullopt, stable_alloc_id(key)); - pool.symm = c10d::symmetric_memory::rendezvous(pool.buffer, group_name); - } else { - // Single process: all panels are co-resident on one GPU, so a plain - // contiguous buffer read by every local block (buf_ptrs[0]) is sufficient; - // no NVSHMEM group or signal-pad handshake is required. - pool.buffer = - torch::empty(sizes, torch::TensorOptions().dtype(dtype).device(device)); - int64_t base = reinterpret_cast(pool.buffer.data_ptr()); - pool.self_ptr = - torch::tensor({base}, torch::dtype(torch::kInt64)).to(device); - } + FusedExchangePool pool; + pool.buffer = torch::empty_strided( + sizes, strides, torch::TensorOptions().dtype(dtype).device(device)); auto [inserted, _] = pools.emplace(key, std::move(pool)); return inserted->second; } +void exchange_fused_buffers_ucx( + FusedExchangePool& pool, std::set const& peers, + std::shared_ptr const& comm) { + TORCH_CHECK(comm != nullptr && comm->initialized() && comm->is_ucx(), + "dynamics.fused-recon-riemann multi-process cubed-sphere " + "exchange requires backend=ucx"); + + std::vector works; + works.reserve(2 * peers.size()); + for (int peer : peers) { + std::vector send_tensors{pool.buffer}; + std::vector recv_tensors{pool.peer_buffers[peer]}; + if (auto work = + comm->send(send_tensors, peer, kFusedCubedSphereExchangeTag)) { + works.push_back(std::move(work)); + } + if (auto work = + comm->recv(recv_tensors, peer, kFusedCubedSphereExchangeTag)) { + works.push_back(std::move(work)); + } + } + for (auto const& work : works) work->wait(); +} + //! \brief Reusable process-local barrier over the blocks_per_process worker -//! threads. The fused cubed-sphere exchange runs pack -> sync -> flux -> -//! release on one shared symmetric buffer; these barriers guarantee every local -//! block has enqueued its pack before the single cross-process visibility sync, -//! and that the sync is enqueued before any block reads peer state. All +//! threads. The fused cubed-sphere exchange runs pack -> UCX exchange -> flux +//! on one shared process-local buffer; these barriers guarantee every local +//! block has enqueued its pack before the leader starts transport and that the +//! transport completes before any block reads peer state. All //! exchange kernels share the device's default stream, so enqueue order fixed //! by these barriers is the execution order, which removes the single-block //! launch ordering deadlock that made the old path require blocks_per_process @@ -200,7 +186,7 @@ class FusedExchangeBarrier { int generation_ = 0; }; -FusedExchangeBarrier &get_fused_exchange_barrier(std::string const &key, +FusedExchangeBarrier& get_fused_exchange_barrier(std::string const& key, int participants) { static std::mutex mutex; static std::unordered_map> @@ -223,8 +209,9 @@ FusedExchangeBarrier &get_fused_exchange_barrier(std::string const &key, //! never reads a freed pointer buffer. struct FusedSeamBatch { //! _flux2/_flux3 are persistent HydroImpl buffers, so raw ptrs are stable. - std::vector flux2_ptr, flux3_ptr; // per local block - std::vector faces; // per local block + std::vector flux2_ptr, flux3_ptr; // per local block + std::vector faces; // per local block + std::vector> remote_peers; // per local block //! side_meta / coords are transient per call; hold the tensors so their //! device pointers stay valid until the async flux-all kernel completes. std::vector side_meta, x2v, x2f, x3v, x3f; @@ -232,15 +219,16 @@ struct FusedSeamBatch { torch::Tensor x2v_dev, x2f_dev, x3v_dev, x3f_dev; }; -FusedSeamBatch &get_fused_seam_batch(std::string const &key, int bpp) { +FusedSeamBatch& get_fused_seam_batch(std::string const& key, int bpp) { static std::mutex mutex; static std::unordered_map batches; std::lock_guard lock(mutex); - auto &b = batches[key]; + auto& b = batches[key]; if (static_cast(b.faces.size()) != bpp) { b.flux2_ptr.assign(bpp, 0); b.flux3_ptr.assign(bpp, 0); b.faces.assign(bpp, 0); + b.remote_peers.assign(bpp, {}); b.side_meta.assign(bpp, torch::Tensor()); b.x2v.assign(bpp, torch::Tensor()); b.x2f.assign(bpp, torch::Tensor()); @@ -250,36 +238,16 @@ FusedSeamBatch &get_fused_seam_batch(std::string const &key, int bpp) { return b; } -//! One-time signal-pad zero after rendezvous. sync_remote_blocks self-restores -//! the pad to zero after each sync+release, so the pad only needs clearing once -//! (removing the per-substage .item() host<->device sync). -void clear_fused_signal_slots_once( - c10::intrusive_ptr const &symm, - std::shared_ptr const &comm, - std::string const &group_name) { - static std::mutex mutex; - static std::set cleared; - { - std::lock_guard lock(mutex); - if (cleared.count(group_name)) return; - cleared.insert(group_name); - } - auto signal_pad = symm->get_signal_pad( - symm->get_rank(), {2 * symm->get_world_size()}, torch::kInt32); - signal_pad.zero_(); - (void)signal_pad.sum().item(); // one-time only - if (comm) comm->barrier(); -} - -torch::Tensor make_side_meta(CubedSphereLayoutImpl const &layout, +torch::Tensor make_side_meta(CubedSphereLayoutImpl const& layout, bool exchange_dim2, bool exchange_dim3, - torch::Device device) { + torch::Device device, + std::vector* remote_peers) { // Per side: [enabled, peer_process, peer_local_block, peer_side, rev]. // neighbor_rank returns a global BLOCK rank; translate it into the owning - // process (to index the symmetric buffer_ptrs array) and the slot within that - // process (to index the shared buffer), so the exchange addresses same-GPU - // and cross-GPU peers uniformly. Must match the CS_META_* layout in the - // kernels. + // process (to index the exchange buffer pointer array) and the slot within + // that process (to index the shared buffer), so the exchange addresses + // same-GPU and cross-GPU peers uniformly. Must match the CS_META_* layout in + // the kernels. constexpr int kStride = 5; std::vector meta(4 * kStride, 0); auto iloc = layout.loc_of(layout.options->rank()); @@ -291,15 +259,19 @@ torch::Tensor make_side_meta(CubedSphereLayoutImpl const &layout, std::tuple{+1, 0, 0}, }; for (int side = 0; side < 4; ++side) { - bool dim_enabled = side <= SIDE_R ? exchange_dim2 : exchange_dim3; - if (!dim_enabled) continue; int nb = layout.neighbor_rank(iloc, offsets[side]); if (nb < 0) continue; auto nb_loc = layout.loc_of(nb); if (std::get<2>(nb_loc) == face) continue; + int peer_process = layout.options->owner_process_rank(nb); + if (peer_process != layout.options->process_rank()) { + remote_peers->push_back(peer_process); + } + bool dim_enabled = side <= SIDE_R ? exchange_dim2 : exchange_dim3; + if (!dim_enabled) continue; auto edge = CS_FACE_EDGES[face][side]; meta[side * kStride + 0] = 1; - meta[side * kStride + 1] = layout.options->owner_process_rank(nb); + meta[side * kStride + 1] = peer_process; meta[side * kStride + 2] = layout.options->local_block_index(nb); meta[side * kStride + 3] = edge.nside; meta[side * kStride + 4] = edge.rev; @@ -310,7 +282,7 @@ torch::Tensor make_side_meta(CubedSphereLayoutImpl const &layout, } // namespace torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, - Variables const &other) { + Variables const& other) { #ifdef NOT_USE_CUDA TORCH_CHECK(false, "dynamics.fused-recon-riemann requires a CUDA-enabled build"); @@ -320,7 +292,7 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, "dynamics.fused-recon-riemann requires CUDA tensors"); TORCH_CHECK(other.count("hydro_w"), "dynamics.fused-recon-riemann requires hydro_w primitives"); - auto const &w = other.at("hydro_w"); + auto const& w = other.at("hydro_w"); TORCH_CHECK(w.is_cuda(), "dynamics.fused-recon-riemann requires CUDA primitive tensors"); TORCH_CHECK(w.is_contiguous() && u.is_contiguous(), @@ -359,22 +331,27 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, playout->options->type()); } else { // Multi-block-per-process is supported: each process may own several panels - // co-resident on one GPU. A single shared per-process symmetric buffer is - // sliced by local block, and cross-panel peers are addressed by (owning - // process, local block). Cross-process exchange still needs the NVSHMEM - // process group; a single-process run (all panels on one GPU) does not. + // co-resident on one GPU. A shared per-process exchange buffer is sliced by + // local block, and cross-panel peers are addressed by (owning process, + // local block). Cross-process exchange uses UCX; a single-process run (all + // panels on one GPU) does not need a process group. TORCH_CHECK( playout->options->process_world_size() == 1 || playout->has_process_group(), "dynamics.fused-recon-riemann cubed-sphere exchange requires an " "initialized process group for multi-process runs"); + TORCH_CHECK( + playout->options->process_world_size() == 1 || + (playout->comm != nullptr && playout->comm->is_ucx()), + "dynamics.fused-recon-riemann multi-process cubed-sphere exchange " + "requires backend=ucx"); } - CubedSphereLayoutImpl const *cs_layout = nullptr; + CubedSphereLayoutImpl const* cs_layout = nullptr; int face = 0; torch::Tensor x2v, x2f, x3v, x3f; if (cubed_sphere_layout) { - cs_layout = dynamic_cast(playout.get()); + cs_layout = dynamic_cast(playout.get()); TORCH_CHECK(cs_layout != nullptr, "expected CubedSphereLayoutImpl for cubed-sphere layout"); face = std::get<2>(cs_layout->loc_of(cs_layout->options->rank())); @@ -405,7 +382,7 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, torch::Tensor inv_mu_ratio_m1, cv_ratio_m1, u0; int nvapor = 0; if (eos == FusedEos::IdealMoist) { - auto ideal_moist = dynamic_cast(peos.get()); + auto ideal_moist = dynamic_cast(peos.get()); TORCH_CHECK(ideal_moist != nullptr, "dynamics.fused-recon-riemann expected IdealMoistImpl"); inv_mu_ratio_m1 = ideal_moist->inv_mu_ratio_m1.to(w.options()); @@ -509,14 +486,13 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, nc1, 1}; - if (multi_process) ensure_symmetric_group(*playout, group_name); - auto &pool = get_fused_symmetric_exchange_pool( - group_name, w.scalar_type(), w.device(), sizes, strides, - multi_process); - auto side_meta = - make_side_meta(*cs_layout, exchange_dim2, exchange_dim3, w.device()); - auto &barrier = get_fused_exchange_barrier(group_name, bpp); - auto &batch = get_fused_seam_batch(group_name, bpp); + auto& pool = get_fused_exchange_pool(group_name, w.scalar_type(), + w.device(), sizes, strides); + std::vector remote_peers; + auto side_meta = make_side_meta(*cs_layout, exchange_dim2, exchange_dim3, + w.device(), &remote_peers); + auto& barrier = get_fused_exchange_barrier(group_name, bpp); + auto& batch = get_fused_seam_batch(group_name, bpp); // register this panel's flux tensors / face / coords / side_meta for the // single process-level seam-flux launch (distinct local-block slots, so @@ -526,6 +502,7 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, batch.flux3_ptr[local_block] = reinterpret_cast(_flux3.data_ptr()); batch.faces[local_block] = face; + batch.remote_peers[local_block] = std::move(remote_peers); batch.side_meta[local_block] = side_meta; batch.x2v[local_block] = x2v; batch.x2f[local_block] = x2f; @@ -545,19 +522,23 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, options->eos()->limiter()}); barrier.wait(); // all panels packed + registered - // ---- Phase B: one leader per process runs the whole seam exchange on - // the GPU: cross-process visibility sync -> ONE flux kernel over ALL - // local panels (overwrites each panel's cross-panel boundary flux) -> - // close the read epoch. No per-panel flux launches and no per-substage - // .item(). ---- + // ---- Phase B: one leader per process exchanges the packed buffer with + // each remote UCX neighbor, then launches ONE flux kernel over ALL local + // panels (overwriting each panel's cross-panel boundary flux). ---- if (is_leader) { + std::set peer_processes; + for (auto const& block_peers : batch.remote_peers) { + peer_processes.insert(block_peers.begin(), block_peers.end()); + } + int process_rank = playout->options->process_rank(); + int process_world_size = playout->options->process_world_size(); + pool.prepare_peer_buffers(peer_processes, process_rank, + process_world_size); if (multi_process) { - clear_fused_signal_slots_once(pool.symm, playout->comm, group_name); - fused_cubed_sphere_sync_cuda(pool.signal_pads_dev(), pool.rank(), - pool.world_size(), w.device()); + exchange_fused_buffers_ucx(pool, peer_processes, playout->comm); } auto i64 = torch::dtype(torch::kInt64); - auto ptr_tensor = [&](std::vector const &ts) { + auto ptr_tensor = [&](std::vector const& ts) { std::vector ptrs(ts.size()); for (size_t p = 0; p < ts.size(); ++p) ptrs[p] = reinterpret_cast(ts[p].data_ptr()); @@ -573,8 +554,8 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, torch::tensor(batch.faces, torch::dtype(torch::kInt32)) .to(w.device()); batch.side_meta_all = torch::stack(batch.side_meta, 0); - auto as_pp = [](torch::Tensor &t) { - return reinterpret_cast(t.data_ptr()); + auto as_pp = [](torch::Tensor& t) { + return reinterpret_cast(t.data_ptr()); }; fused_cubed_sphere_flux_all_cuda( pool.buffer, @@ -589,10 +570,6 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, static_cast(nc1), bpp, w.scalar_type(), w.device(), make_physics_params(recon23_prim, recon23_vel, precon23->pinterp1->options->scale())}); - if (multi_process) { - fused_cubed_sphere_release_cuda(pool.signal_pads_dev(), pool.rank(), - pool.world_size(), w.device()); - } } barrier.wait(); // seam flux enqueued before any panel's divergence } @@ -604,7 +581,7 @@ torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, du.index(interior) = -dt * _div.index(interior); auto temp = peos->compute("W->T", {w}); - for (auto &f : forcings) f.forward(du, w, temp, dt); + for (auto& f : forcings) f.forward(du, w, temp, dt); if (options->grav() && (options->grav()->non_hydrostatic() < 1.)) { du[IVX] += dt * rho_grav * (1. - options->grav()->non_hydrostatic()); diff --git a/src/hydro/hydro_options.cpp b/src/hydro/hydro_options.cpp index 3b8312ab..84b4d310 100644 --- a/src/hydro/hydro_options.cpp +++ b/src/hydro/hydro_options.cpp @@ -40,7 +40,7 @@ bool fused_recon_riemann_supported_by_options(HydroOptions const& op, (eos_type == "shallow-water" && riemann_type == "shallow-roe"); // Cubed-sphere is now supported for any blocks_per_process: each process may // own several panels co-resident on one GPU, exchanged through a shared - // per-process symmetric buffer sliced by local block (see + // per-process packed buffer sliced by local block (see // hydro_forward_fused). return eos_supported && riemann_supported && recon_supported && combo_supported; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 9ff6d6af..a8b10ab5 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -30,17 +30,17 @@ setup_test(test_diffusion_moist) setup_test(test_forcing) setup_parallel_test(test_exchange 2) -if(CUDAToolkit_FOUND) - add_executable(test_fused_cs_sync_smoke.${buildl} - test_fused_cs_sync_smoke.cu) +if(CUDAToolkit_FOUND AND UCX_FOUND) + add_executable(test_fused_cs_ucx_smoke.${buildl} + test_fused_cs_ucx_smoke.cu) - set_target_properties(test_fused_cs_sync_smoke.${buildl} + set_target_properties(test_fused_cs_ucx_smoke.${buildl} PROPERTIES COMPILE_FLAGS ${CMAKE_CXX_FLAGS_${buildu}} CUDA_SEPARABLE_COMPILATION ON) target_include_directories( - test_fused_cs_sync_smoke.${buildl} + test_fused_cs_ucx_smoke.${buildl} PRIVATE ${CMAKE_BINARY_DIR} ${KINTERA_INCLUDE_DIR} ${SNAP_INCLUDE_DIR} @@ -49,24 +49,24 @@ if(CUDAToolkit_FOUND) ${TORCH_INCLUDE_DIR} ${TORCH_API_INCLUDE_DIR}) - target_link_libraries(test_fused_cs_sync_smoke.${buildl} + target_link_libraries(test_fused_cs_ucx_smoke.${buildl} PRIVATE snapy::snap snapy::snap_cu gtest_main) if (UNIX AND NOT APPLE) - target_link_options(test_fused_cs_sync_smoke.${buildl} + target_link_options(test_fused_cs_ucx_smoke.${buildl} PRIVATE -Wl,--no-as-needed) endif() add_test( - NAME test_fused_cs_sync_smoke.${buildl} + NAME test_fused_cs_ucx_smoke.${buildl} COMMAND torchrun --no-python --nproc-per-node=6 - $) - set_tests_properties(test_fused_cs_sync_smoke.${buildl} PROPERTIES - LABELS "cuda;nvshmem;cubed-sphere;fused") + $) + set_tests_properties(test_fused_cs_ucx_smoke.${buildl} PROPERTIES + LABELS "cuda;ucx;cubed-sphere;fused") endif() if(UCX_FOUND) diff --git a/tests/test_fused_cs_sync_smoke.cu b/tests/test_fused_cs_ucx_smoke.cu similarity index 73% rename from tests/test_fused_cs_sync_smoke.cu rename to tests/test_fused_cs_ucx_smoke.cu index 7e11135f..3777815d 100644 --- a/tests/test_fused_cs_sync_smoke.cu +++ b/tests/test_fused_cs_ucx_smoke.cu @@ -1,4 +1,4 @@ -// C/C++ +// C++/CUDA #include #include #include @@ -13,8 +13,6 @@ #include #include #include -#include -#include #include // snap @@ -50,7 +48,7 @@ torch::Device select_cuda_device() { LayoutOptions make_layout_options() { auto opts = LayoutOptionsImpl::create(); opts->type("cubed-sphere"); - opts->backend("gloo"); + opts->backend("ucx"); opts->device("cuda"); opts->px(1); opts->py(1); @@ -135,28 +133,6 @@ bool cuda_6rank_available_or_skip() { return true; } -void initialize_symmetric_memory_group(LayoutImpl const &layout, - std::string const &group_name) { - static std::set initialized; - if (initialized.empty()) { - c10d::symmetric_memory::set_backend("NVSHMEM"); - c10d::symmetric_memory::set_signal_pad_size(std::max( - 1024, layout.options->process_world_size() * sizeof(uint32_t))); - } - - auto set_group_info_once = [&](std::string const &name) { - if (initialized.count(name)) - return; - c10d::symmetric_memory::set_group_info(name, layout.options->process_rank(), - layout.options->process_world_size(), - layout.comm->store); - initialized.insert(name); - }; - - set_group_info_once("0"); - set_group_info_once(group_name); -} - struct SmokeContext { torch::Device device; LayoutOptions options; @@ -178,26 +154,53 @@ SmokeContext make_smoke_context() { return {device, opts, layout, edge_meta}; } -torch::Tensor make_symmetric_buffer(torch::Device device, - std::string const &group_name) { - (void)group_name; +torch::Tensor make_exchange_buffer(torch::Device device) { std::vector sizes = {kSides, kStates, kNvar, kEdgeLen, kNc1}; std::vector strides = {kStates * kNvar * kEdgeLen * kNc1, kNvar * kEdgeLen * kNc1, kEdgeLen * kNc1, kNc1, 1}; - return c10d::symmetric_memory::empty_strided_p2p( - sizes, strides, torch::kFloat64, device, std::nullopt, std::nullopt); + return torch::empty_strided( + sizes, strides, + torch::TensorOptions().dtype(torch::kFloat64).device(device)); } -void clear_signal_slots( - c10::intrusive_ptr const &symm, - std::shared_ptr const &comm) { - auto signal_pad = symm->get_signal_pad( - symm->get_rank(), {2 * symm->get_world_size()}, torch::kInt32); - signal_pad.zero_(); - (void)signal_pad.sum().item(); - if (comm) - comm->barrier(); +struct ExchangedBuffers { + std::vector received; + torch::Tensor pointers; + + void **pointers_dev() const { + return reinterpret_cast(pointers.data_ptr()); + } +}; + +ExchangedBuffers exchange_with_face_neighbors(SmokeContext const &ctx, + torch::Tensor local, int tag) { + int rank = ctx.options->process_rank(); + int world_size = ctx.options->process_world_size(); + auto meta = ctx.edge_meta.cpu(); + std::set peers; + for (int side = 0; side < kSides; ++side) { + peers.insert(meta[side * kMetaStride].item()); + } + + ExchangedBuffers result; + result.received.resize(world_size); + std::vector pointers(world_size, 0); + pointers[rank] = reinterpret_cast(local.data_ptr()); + std::vector works; + for (int peer : peers) { + result.received[peer] = torch::empty_like(local); + pointers[peer] = + reinterpret_cast(result.received[peer].data_ptr()); + std::vector send_tensors{local}; + std::vector recv_tensors{result.received[peer]}; + works.push_back(ctx.layout->comm->send(send_tensors, peer, tag)); + works.push_back(ctx.layout->comm->recv(recv_tensors, peer, tag)); + } + for (auto const &work : works) work->wait(); + result.pointers = + torch::tensor(pointers, torch::dtype(torch::kInt64)).to(ctx.device); + return result; } __device__ double device_payload(int rank, int side, int state, int edge, int i, @@ -222,12 +225,6 @@ __global__ void write_edge_payload_kernel(double *buffer, int rank) { } } -__global__ void sync_previous_kernel_writes(uint32_t **signal_pads, int rank, - int world_size) { - c10d::symmetric_memory::sync_remote_blocks(signal_pads, rank, - world_size); -} - __global__ void verify_remote_edge_kernel(void **buffer_ptrs, int const *meta, int *errors) { int line = blockIdx.x; @@ -318,48 +315,43 @@ __global__ void verify_hydro_remote_constant_kernel(void **buffer_ptrs, } // namespace -TEST(FusedCubedSphereSymmetricMemory, RendezvousCompletes) { +TEST(FusedCubedSphereUCX, ProcessGroupInitializes) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); auto ctx = make_smoke_context(); - std::string group_name = "snapy:test:fused-cs-rendezvous"; - initialize_symmetric_memory_group(*ctx.layout, group_name); - - auto symm_buffer = make_symmetric_buffer(ctx.device, group_name); - auto symm = c10d::symmetric_memory::rendezvous(symm_buffer, group_name); - EXPECT_EQ(symm->get_world_size(), 6); + EXPECT_TRUE(ctx.layout->comm->initialized()); + EXPECT_TRUE(ctx.layout->comm->is_ucx()); + EXPECT_EQ(ctx.options->process_world_size(), 6); if (ctx.layout->comm) ctx.layout->comm->barrier(); } -TEST(FusedCubedSphereSymmetricMemory, PreviousKernelSyncCompletes) { +TEST(FusedCubedSphereUCX, PreviousKernelWritesReachPeers) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); auto ctx = make_smoke_context(); - std::string group_name = "snapy:test:fused-cs-sync"; - initialize_symmetric_memory_group(*ctx.layout, group_name); - - auto symm_buffer = make_symmetric_buffer(ctx.device, group_name); - auto symm = c10d::symmetric_memory::rendezvous(symm_buffer, group_name); - clear_signal_slots(symm, ctx.layout->comm); + auto exchange_buffer = make_exchange_buffer(ctx.device); auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); int blocks = kSides * kEdgeLen * kNc1; write_edge_payload_kernel<<>>( - symm_buffer.data_ptr(), symm->get_rank()); + exchange_buffer.data_ptr(), ctx.options->process_rank()); C10_CUDA_KERNEL_LAUNCH_CHECK(); - sync_previous_kernel_writes<<<1, std::max(32, symm->get_world_size()), 0, - stream>>>( - reinterpret_cast(symm->get_signal_pad_ptrs_dev()), - symm->get_rank(), symm->get_world_size()); + auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 701); + auto errors = + torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); + verify_remote_edge_kernel<<>>( + exchanged.pointers_dev(), ctx.edge_meta.data_ptr(), + errors.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); AT_CUDA_CHECK(cudaStreamSynchronize(stream)); + EXPECT_EQ(errors.cpu().item(), 0); if (ctx.layout->comm) ctx.layout->comm->barrier(); } -TEST(FusedCubedSphereSymmetricMemory, OrientationMetadataMatchesStaged) { +TEST(FusedCubedSphereUCX, OrientationMetadataMatchesStaged) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); auto ctx = make_smoke_context(); @@ -381,7 +373,7 @@ TEST(FusedCubedSphereSymmetricMemory, OrientationMetadataMatchesStaged) { ctx.layout->comm->barrier(); } -TEST(FusedCubedSphereSymmetricMemory, BoundaryStateConventionMatchesStaged) { +TEST(FusedCubedSphereUCX, BoundaryStateConventionMatchesStaged) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); // Staged hydro sends the local right state to lower-side neighbors and the @@ -396,33 +388,23 @@ TEST(FusedCubedSphereSymmetricMemory, BoundaryStateConventionMatchesStaged) { EXPECT_TRUE(fused_exchange_uses_right_interp_for_side(SIDE_T)); } -TEST(FusedCubedSphereSymmetricMemory, RemoteEdgePayloadMatches) { +TEST(FusedCubedSphereUCX, RemoteEdgePayloadMatches) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); auto ctx = make_smoke_context(); - std::string group_name = "snapy:test:fused-cs-remote-read"; - initialize_symmetric_memory_group(*ctx.layout, group_name); - - auto symm_buffer = make_symmetric_buffer(ctx.device, group_name); - auto symm = c10d::symmetric_memory::rendezvous(symm_buffer, group_name); - clear_signal_slots(symm, ctx.layout->comm); + auto exchange_buffer = make_exchange_buffer(ctx.device); auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); int blocks = kSides * kEdgeLen * kNc1; write_edge_payload_kernel<<>>( - symm_buffer.data_ptr(), symm->get_rank()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - sync_previous_kernel_writes<<<1, std::max(32, symm->get_world_size()), 0, - stream>>>( - reinterpret_cast(symm->get_signal_pad_ptrs_dev()), - symm->get_rank(), symm->get_world_size()); + exchange_buffer.data_ptr(), ctx.options->process_rank()); C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 702); auto errors = torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); verify_remote_edge_kernel<<>>( - symm->get_buffer_ptrs_dev(), ctx.edge_meta.data_ptr(), + exchanged.pointers_dev(), ctx.edge_meta.data_ptr(), errors.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); AT_CUDA_CHECK(cudaStreamSynchronize(stream)); @@ -432,35 +414,25 @@ TEST(FusedCubedSphereSymmetricMemory, RemoteEdgePayloadMatches) { ctx.layout->comm->barrier(); } -TEST(FusedCubedSphereSymmetricMemory, RemoteStateSelectionMatchesStaged) { +TEST(FusedCubedSphereUCX, RemoteStateSelectionMatchesStaged) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); auto ctx = make_smoke_context(); - std::string group_name = "snapy:test:fused-cs-remote-state"; - initialize_symmetric_memory_group(*ctx.layout, group_name); - - auto symm_buffer = make_symmetric_buffer(ctx.device, group_name); - auto symm = c10d::symmetric_memory::rendezvous(symm_buffer, group_name); - clear_signal_slots(symm, ctx.layout->comm); + auto exchange_buffer = make_exchange_buffer(ctx.device); auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); int blocks = kSides * kEdgeLen * kNc1; write_edge_payload_kernel<<>>( - symm_buffer.data_ptr(), symm->get_rank()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - sync_previous_kernel_writes<<<1, std::max(32, symm->get_world_size()), 0, - stream>>>( - reinterpret_cast(symm->get_signal_pad_ptrs_dev()), - symm->get_rank(), symm->get_world_size()); + exchange_buffer.data_ptr(), ctx.options->process_rank()); C10_CUDA_KERNEL_LAUNCH_CHECK(); + auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 703); auto errors = torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); auto side_meta = make_hydro_side_meta(*ctx.layout, ctx.device); verify_staged_remote_state_kernel<<>>( - symm->get_buffer_ptrs_dev(), side_meta.data_ptr(), symm->get_rank(), - errors.data_ptr()); + exchanged.pointers_dev(), side_meta.data_ptr(), + ctx.options->process_rank(), errors.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); AT_CUDA_CHECK(cudaStreamSynchronize(stream)); @@ -469,16 +441,11 @@ TEST(FusedCubedSphereSymmetricMemory, RemoteStateSelectionMatchesStaged) { ctx.layout->comm->barrier(); } -TEST(FusedCubedSphereSymmetricMemory, ConstantStateExchangeHasZeroMassFlux) { +TEST(FusedCubedSphereUCX, ConstantStateExchangeHasZeroMassFlux) { if (!cuda_6rank_available_or_skip()) GTEST_SKIP(); auto ctx = make_smoke_context(); - std::string group_name = "snapy:test:fused-cs-constant-state"; - initialize_symmetric_memory_group(*ctx.layout, group_name); - - auto symm_buffer = make_symmetric_buffer(ctx.device, group_name); - auto symm = c10d::symmetric_memory::rendezvous(symm_buffer, group_name); - clear_signal_slots(symm, ctx.layout->comm); + auto exchange_buffer = make_exchange_buffer(ctx.device); auto opts = torch::dtype(torch::kFloat64).device(ctx.device); auto w = torch::zeros({kNvar, kEdgeLen, kEdgeLen, kNc1}, opts); @@ -500,15 +467,13 @@ TEST(FusedCubedSphereSymmetricMemory, ConstantStateExchangeHasZeroMassFlux) { FusedCubedSpherePanelParams panel_params{side_meta, face, local_block, x2v, x2f, x3v, x3f}; fused_cubed_sphere_pack_cuda( - w, symm_buffer, + w, exchange_buffer, FusedCubedSpherePackParams{panel_params, FusedReconScheme::WENO5, FusedReconScheme::WENO5, false, FusedEos::ShallowWater, 0., 0., false}); - fused_cubed_sphere_sync_cuda( - reinterpret_cast(symm->get_signal_pad_ptrs_dev()), - symm->get_rank(), symm->get_world_size(), ctx.device); + auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 704); fused_cubed_sphere_flux_cuda( - w, flux2, flux3, symm_buffer, symm->get_buffer_ptrs_dev(), + w, flux2, flux3, exchange_buffer, exchanged.pointers_dev(), FusedCubedSphereFluxParams{ panel_params, FusedPhysicsParams{FusedReconScheme::WENO5, FusedReconScheme::WENO5, @@ -516,9 +481,6 @@ TEST(FusedCubedSphereSymmetricMemory, ConstantStateExchangeHasZeroMassFlux) { FusedEos::ShallowWater, 1.4, 0., 0., false, torch::Tensor(), torch::Tensor(), torch::Tensor(), 0, 1}}); - fused_cubed_sphere_release_cuda( - reinterpret_cast(symm->get_signal_pad_ptrs_dev()), - symm->get_rank(), symm->get_world_size(), ctx.device); AT_CUDA_CHECK(cudaStreamSynchronize( at::cuda::getCurrentCUDAStream(ctx.device.index()))); @@ -536,13 +498,13 @@ TEST(FusedCubedSphereSymmetricMemory, ConstantStateExchangeHasZeroMassFlux) { EXPECT_LT(state.select(1, IVZ).abs().max().cpu().item(), 1.e-10) << label; }; - expect_constant_state(symm_buffer.select(1, ILT), "left state"); - expect_constant_state(symm_buffer.select(1, IRT), "right state"); + expect_constant_state(exchange_buffer.select(1, ILT), "left state"); + expect_constant_state(exchange_buffer.select(1, IRT), "right state"); auto errors = torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); verify_hydro_remote_constant_kernel<<>>(symm->get_buffer_ptrs_dev(), + stream>>>(exchanged.pointers_dev(), side_meta.data_ptr(), errors.data_ptr()); C10_CUDA_KERNEL_LAUNCH_CHECK(); From 26b35f5b6896cb3916f1d386c18aa1e616083a0d Mon Sep 17 00:00:00 2001 From: dungeon2/cli Date: Thu, 16 Jul 2026 01:12:33 -0400 Subject: [PATCH 2/3] wip --- docs/api/hydro.rst | 12 - docs/user_guide/configuration.rst | 9 - docs/user_guide/distributed.rst | 4 - python/csrc/pyhydro.cpp | 7 +- python/snapy/hydro.pyi | 10 - src/eos/eos_side_quantities.cuh | 59 -- src/hydro/fused_cubed_sphere_dispatch.cu | 444 ---------- src/hydro/fused_dispatch_params.cuh | 121 --- src/hydro/fused_recon_riemann_dispatch.cu | 246 ------ src/hydro/fused_recon_riemann_dispatch.hpp | 126 --- src/hydro/hydro.cpp | 19 - src/hydro/hydro.hpp | 9 +- ...o_forward_staged.cpp => hydro_forward.cpp} | 4 +- src/hydro/hydro_forward_fused.cpp | 595 ------------- src/hydro/hydro_options.cpp | 78 -- src/recon/interp_impl.cuh | 50 -- src/riemann/riemann_dispatch.cu | 108 ++- src/riemann/riemann_dispatch.hpp | 14 +- src/riemann/roe.cpp | 231 +---- src/riemann/roe_dispatch.cpp | 229 +++++ src/riemann/roe_impl.h | 18 +- src/riemann/shallow_roe.cpp | 68 +- src/riemann/shallow_roe_dispatch.cpp | 91 ++ src/snap.h | 20 - tests/CMakeLists.txt | 39 - tests/test_fused_cs_ucx_smoke.cu | 517 ----------- tests/test_hydro_python_api.py | 11 - tests/test_hydrostatic.cpp | 1 - tests/test_riemann.cpp | 805 +----------------- 29 files changed, 491 insertions(+), 3454 deletions(-) delete mode 100644 src/eos/eos_side_quantities.cuh delete mode 100644 src/hydro/fused_cubed_sphere_dispatch.cu delete mode 100644 src/hydro/fused_dispatch_params.cuh delete mode 100644 src/hydro/fused_recon_riemann_dispatch.cu delete mode 100644 src/hydro/fused_recon_riemann_dispatch.hpp rename src/hydro/{hydro_forward_staged.cpp => hydro_forward.cpp} (98%) delete mode 100644 src/hydro/hydro_forward_fused.cpp create mode 100644 src/riemann/roe_dispatch.cpp create mode 100644 src/riemann/shallow_roe_dispatch.cpp delete mode 100644 tests/test_fused_cs_ucx_smoke.cu delete mode 100644 tests/test_hydro_python_api.py diff --git a/docs/api/hydro.rst b/docs/api/hydro.rst index d3ec6585..3a567a2d 100644 --- a/docs/api/hydro.rst +++ b/docs/api/hydro.rst @@ -94,18 +94,6 @@ HydroOptions :return: Disable flux x3 flag :rtype: bool - .. method:: fused_recon_riemann() -> bool - fused_recon_riemann(value: bool) -> HydroOptions - - Get or set the resolved fused reconstruction/Riemann CUDA path status. - For options loaded from YAML, the ``FUSED`` environment variable is the - selection control: ``FUSED=OFF`` disables the path, ``FUSED=ON`` requires - it and errors if unsupported, and ``FUSED=AUTO`` or an unset ``FUSED`` - enables it only when the configuration is supported. - - :return: Resolved fused reconstruction/Riemann status - :rtype: bool - .. method:: eos() -> EquationOfStateOptions eos(value: EquationOfStateOptions) -> HydroOptions diff --git a/docs/user_guide/configuration.rst b/docs/user_guide/configuration.rst index e1f85942..f8b61f7e 100644 --- a/docs/user_guide/configuration.rst +++ b/docs/user_guide/configuration.rst @@ -106,15 +106,6 @@ Hydrodynamic solver configuration: disable_flux_x2: false disable_flux_x3: false -The fused reconstruction/Riemann CUDA path is controlled exclusively by the -``FUSED`` environment variable when loading options from YAML. Set -``FUSED=OFF`` to disable the fused path. Set ``FUSED=ON`` to require it, failing -if the EOS/Riemann/reconstruction combination is unsupported. Set -``FUSED=AUTO`` or leave ``FUSED`` unset to use the fused path when supported and -fall back to the staged path otherwise. Option loading prints the resolved -``FUSED`` mode, whether the fused path is enabled, and whether the configuration -is supported. - Physical Forcing ---------------- diff --git a/docs/user_guide/distributed.rst b/docs/user_guide/distributed.rst index d9d175df..a610b808 100644 --- a/docs/user_guide/distributed.rst +++ b/docs/user_guide/distributed.rst @@ -248,10 +248,6 @@ UCX transports for the cluster: export DEVICE=cuda export UCX_TLS=rc,cuda_copy,cuda_ipc -The fused cubed-sphere CUDA path also uses this UCX backend. It exchanges one -packed edge-state buffer with each neighboring process before launching the -fused seam-flux kernel; no separate symmetric-memory runtime is required. - Load Balancing -------------- diff --git a/python/csrc/pyhydro.cpp b/python/csrc/pyhydro.cpp index 6fead47f..0962ccf7 100644 --- a/python/csrc/pyhydro.cpp +++ b/python/csrc/pyhydro.cpp @@ -13,7 +13,7 @@ namespace py = pybind11; -void bind_hydro(py::module &m) { +void bind_hydro(py::module& m) { auto pyHydroOptions = py::class_(m, "HydroOptions"); @@ -21,7 +21,7 @@ void bind_hydro(py::module &m) { .def_static("from_yaml", &snap::HydroOptionsImpl::from_yaml, py::arg("filename"), py::arg("verbose") = false) .def("__repr__", - [](const snap::HydroOptions &a) { + [](const snap::HydroOptions& a) { std::stringstream ss; a->report(ss); return fmt::format("HydroOptions(\n{})", ss.str()); @@ -30,7 +30,6 @@ void bind_hydro(py::module &m) { .ADD_OPTION(bool, snap::HydroOptionsImpl, disable_flux_x1) .ADD_OPTION(bool, snap::HydroOptionsImpl, disable_flux_x2) .ADD_OPTION(bool, snap::HydroOptionsImpl, disable_flux_x3) - .ADD_OPTION(bool, snap::HydroOptionsImpl, fused_recon_riemann) .ADD_OPTION(snap::ConstGravityOptions, snap::HydroOptionsImpl, grav) .ADD_OPTION(snap::CoriolisOptions, snap::HydroOptionsImpl, coriolis) .ADD_OPTION(snap::DiffusionOptions, snap::HydroOptionsImpl, diffusion) @@ -41,7 +40,7 @@ void bind_hydro(py::module &m) { .ADD_OPTION(snap::ImplicitOptions, snap::HydroOptionsImpl, icorr); ADD_SNAP_MODULE(Hydro, HydroOptions) - .def(py::init(), + .def(py::init(), py::arg("options"), py::arg("block") = nullptr) .def("max_time_step", &snap::HydroImpl::max_time_step); } diff --git a/python/snapy/hydro.pyi b/python/snapy/hydro.pyi index c81d2ee8..7456aa24 100644 --- a/python/snapy/hydro.pyi +++ b/python/snapy/hydro.pyi @@ -83,16 +83,6 @@ class HydroOptions: """Set disable flux x3 flag.""" ... - @overload - def fused_recon_riemann(self) -> bool: - """Get resolved fused reconstruction/Riemann status.""" - ... - - @overload - def fused_recon_riemann(self, value: bool) -> "HydroOptions": - """Set fused reconstruction/Riemann status for manually constructed options.""" - ... - @overload def grav(self) -> ConstGravityOptions: """Get gravity options.""" diff --git a/src/eos/eos_side_quantities.cuh b/src/eos/eos_side_quantities.cuh deleted file mode 100644 index fd066a73..00000000 --- a/src/eos/eos_side_quantities.cuh +++ /dev/null @@ -1,59 +0,0 @@ -#pragma once - -// base -#include - -// snap -#include - -#include "ideal_moist_impl.h" -#include "shallow_water_impl.h" - -namespace snap { - -template -inline DISPATCH_MACRO void eos_side_quantities( - T const* wl, T const* wr, int ny, int nvapor, FusedEos eos, T gammad, - T const* inv_mu_ratio_m1, T const* cv_ratio_m1, T const* u0, T* el, T* er, - T* gl, T* gr, T* cl, T* cr) { - if (eos == FusedEos::IdealGas) { - *el = wl[IPR] / (gammad - 1.); - *er = wr[IPR] / (gammad - 1.); - *gl = gammad; - *gr = gammad; - } else if (eos == FusedEos::ShallowWater) { - *el = 0.; - *er = 0.; - *gl = 0.; - *gr = 0.; - shallow_water_side_quantities(wl[IDN], wr[IDN], cl, cr); - return; - } else { - T yl[32], yr[32]; - T suml = 0., sumr = 0.; - for (int n = 0; n < ny; ++n) { - yl[n] = wl[ICY + n]; - yr[n] = wr[ICY + n]; - suml += yl[n]; - sumr += yr[n]; - } - T fepsl = ideal_moist_feps(yl, ny, nvapor, inv_mu_ratio_m1); - T fepsr = ideal_moist_feps(yr, ny, nvapor, inv_mu_ratio_m1); - T fsigl = ideal_moist_fsig(yl, ny, cv_ratio_m1); - T fsigr = ideal_moist_fsig(yr, ny, cv_ratio_m1); - *el = wl[IPR] * fsigl / fepsl / (gammad - 1.); - *er = wr[IPR] * fsigr / fepsr / (gammad - 1.); - *el += wl[IDN] * (1. - suml) * u0[0]; - *er += wr[IDN] * (1. - sumr) * u0[0]; - for (int n = 0; n < ny; ++n) { - *el += wl[IDN] * yl[n] * u0[1 + n]; - *er += wr[IDN] * yr[n] * u0[1 + n]; - } - *gl = 1. + (gammad - 1.) * fepsl / fsigl; - *gr = 1. + (gammad - 1.) * fepsr / fsigr; - } - *cl = sqrt((*gl) * wl[IPR] / wl[IDN]); - *cr = sqrt((*gr) * wr[IPR] / wr[IDN]); -} - -} // namespace snap diff --git a/src/hydro/fused_cubed_sphere_dispatch.cu b/src/hydro/fused_cubed_sphere_dispatch.cu deleted file mode 100644 index 002aebd1..00000000 --- a/src/hydro/fused_cubed_sphere_dispatch.cu +++ /dev/null @@ -1,444 +0,0 @@ -// torch -#include -#include -#include - -// C/C++ -#include - -// snap -#include - -#include "../coord/gnomonic_equiangle.h" -#include "../eos/eos_side_quantities.cuh" -#include "../layout/cubed_sphere_constants.h" -#include "../recon/interp_impl.cuh" -#include "../riemann/hllc_impl.h" -#include "../riemann/lmars_impl.h" -#include "../riemann/roe_impl.h" -#include "../riemann/shallow_roe_impl.h" -#include "fused_dispatch_params.cuh" -#include "fused_recon_riemann_dispatch.hpp" - -namespace snap { -namespace { - -enum { - CS_SIDE_L = 0, - CS_SIDE_R = 1, - CS_SIDE_B = 2, - CS_SIDE_T = 3, - CS_STATE_LEFT = ILT, - CS_STATE_RIGHT = IRT, - CS_NUM_STATES = 2, - CS_NUM_SIDES = 4, - CS_META_ENABLED = 0, - CS_META_PEER_PROCESS = 1, - CS_META_PEER_LOCAL_BLOCK = 2, - CS_META_PEER_SIDE = 3, - CS_META_REV = 4, - CS_META_STRIDE = 5, -}; - -template -__device__ void cs_coords(int side, int edge_pos, int i, int nc3, int nc2, - int nc1, int nghost, T const *x2v, T const *x2f, - T const *x3v, T const *x3f, T *alpha, T *beta, int *k, - int *j, int *face_pos) { - (void)i; - if (side == CS_SIDE_L || side == CS_SIDE_R) { - *k = edge_pos; - *j = side == CS_SIDE_L ? nghost : nc2 - nghost - 1; - *face_pos = side == CS_SIDE_L ? nghost : nc2 - nghost; - *alpha = x2f[side == CS_SIDE_L ? nghost : nc2 - nghost]; - *beta = x3v[edge_pos]; - } else { - *k = side == CS_SIDE_B ? nghost : nc3 - nghost - 1; - *j = edge_pos; - *face_pos = side == CS_SIDE_B ? nghost : nc3 - nghost; - *alpha = x2v[edge_pos]; - *beta = x3f[side == CS_SIDE_B ? nghost : nc3 - nghost]; - } -} - -template -__global__ void cs_pack_kernel(T const *w, T *buf, - DeviceCubedSpherePackParams params) { - int nvar = params.nvar; - int nc3 = params.nc3; - int nc2 = params.nc2; - int nc1 = params.nc1; - auto panel = params.panel; - int edge_len = max(nc2, nc3); - int line = blockIdx.x; - int i = line % nc1; - int edge_pos = (line / nc1) % edge_len; - int side = line / (nc1 * edge_len); - if (!panel.side_meta[side * CS_META_STRIDE + CS_META_ENABLED]) - return; - if ((side <= CS_SIDE_R && edge_pos >= nc3) || - (side >= CS_SIDE_B && edge_pos >= nc2)) { - return; - } - - int dim = side <= CS_SIDE_R ? 2 : 1; - int axis_size = dim == 2 ? nc2 : nc3; - int stride_dim = dim == 2 ? nc1 : nc1 * nc2; - int stride_var = nc1 * nc2 * nc3; - int pos = threadIdx.x; - int base = side <= CS_SIDE_R ? edge_pos * nc2 * nc1 + i : edge_pos * nc1 + i; - - extern __shared__ unsigned char memory[]; - T *smem = reinterpret_cast(memory); - if (pos < axis_size) { - int in = base + pos * stride_dim; - for (int v = 0; v < nvar; ++v) { - smem[v * axis_size + pos] = w[v * stride_var + in]; - } - } - __syncthreads(); - - int nghost = (params.recon_prim == FusedReconScheme::CP3 || - params.recon_prim == FusedReconScheme::WENO3) - ? 2 - : 3; - int k, j, face_pos; - T alpha, beta; - cs_coords(side, edge_pos, i, nc3, nc2, nc1, nghost, panel.x2v, panel.x2f, - panel.x3v, panel.x3f, &alpha, &beta, &k, &j, &face_pos); - if (pos != face_pos) - return; - - int il = nghost; - int left_start = face_pos - il; - int right_start = face_pos - (il - 1); - T left[64], right[64]; - for (int v = 0; v < nvar; ++v) { - auto scheme = (v == IDN || v >= ICY) ? params.recon_prim : params.recon_vel; - left[v] = interp_shared_fused_impl(smem, v, left_start, axis_size, scheme, - /*right=*/true, params.recon_scale); - right[v] = interp_shared_fused_impl(smem, v, right_start, axis_size, scheme, - /*right=*/false, params.recon_scale); - } - if (params.eos_limiter) { - left[IDN] = max(left[IDN], params.density_floor); - right[IDN] = max(right[IDN], params.density_floor); - if (params.eos != FusedEos::ShallowWater) { - left[IPR] = max(left[IPR], params.pressure_floor); - right[IPR] = max(right[IPR], params.pressure_floor); - for (int v = ICY; v < nvar; ++v) { - left[v] = max(left[v], T(0)); - right[v] = max(right[v], T(0)); - } - } - } - gnomonic_contra_to_sph(left, panel.face, alpha, beta); - gnomonic_contra_to_sph(right, panel.face, alpha, beta); - - int buf_stride_var = edge_len * nc1; - int block_stride = CS_NUM_SIDES * CS_NUM_STATES * nvar * edge_len * nc1; - int block_base = panel.local_block * block_stride; - int left_out = - block_base + - (((side * CS_NUM_STATES + CS_STATE_LEFT) * nvar) * edge_len + edge_pos) * - nc1 + - i; - int right_out = - block_base + - (((side * CS_NUM_STATES + CS_STATE_RIGHT) * nvar) * edge_len + edge_pos) * - nc1 + - i; - for (int v = 0; v < nvar; ++v) { - buf[left_out + v * buf_stride_var] = left[v]; - buf[right_out + v * buf_stride_var] = right[v]; - } -} - -template -__device__ void cs_flux_common(DeviceCubedSphereFluxParams params, int i, - int edge_pos, int side) { - T const *buf = params.buf; - T *flux2 = params.flux2; - T *flux3 = params.flux3; - void **buf_ptrs = params.buf_ptrs; - int nvar = params.nvar; - int nc3 = params.nc3; - int nc2 = params.nc2; - int nc1 = params.nc1; - auto panel = params.panel; - auto physics = params.physics; - int const *side_meta = panel.side_meta; - if (!side_meta[side * CS_META_STRIDE + CS_META_ENABLED]) - return; - if ((side <= CS_SIDE_R && edge_pos >= nc3) || - (side >= CS_SIDE_B && edge_pos >= nc2)) { - return; - } - - int peer_process = side_meta[side * CS_META_STRIDE + CS_META_PEER_PROCESS]; - int peer_local_block = - side_meta[side * CS_META_STRIDE + CS_META_PEER_LOCAL_BLOCK]; - int peer_side = side_meta[side * CS_META_STRIDE + CS_META_PEER_SIDE]; - int rev = side_meta[side * CS_META_STRIDE + CS_META_REV]; - int peer_edge = - rev ? ((side <= CS_SIDE_R ? nc3 : nc2) - 1 - edge_pos) : edge_pos; - auto peer_buf = static_cast(buf_ptrs[peer_process]); - int edge_len = max(nc2, nc3); - int buf_stride_var = edge_len * nc1; - int block_stride = CS_NUM_SIDES * CS_NUM_STATES * nvar * edge_len * nc1; - int local_base = panel.local_block * block_stride; - int peer_base = peer_local_block * block_stride; - int remote_left_off = - peer_base + - (((peer_side * CS_NUM_STATES + CS_STATE_LEFT) * nvar) * edge_len + - peer_edge) * - nc1 + - i; - int remote_right_off = - peer_base + - (((peer_side * CS_NUM_STATES + CS_STATE_RIGHT) * nvar) * edge_len + - peer_edge) * - nc1 + - i; - int local_left_off = - local_base + - (((side * CS_NUM_STATES + CS_STATE_LEFT) * nvar) * edge_len + edge_pos) * - nc1 + - i; - int local_right_off = - local_base + - (((side * CS_NUM_STATES + CS_STATE_RIGHT) * nvar) * edge_len + edge_pos) * - nc1 + - i; - - int dim = side <= CS_SIDE_R ? 2 : 1; - int nghost = (physics.recon_prim == FusedReconScheme::CP3 || - physics.recon_prim == FusedReconScheme::WENO3) - ? 2 - : 3; - int k, j, face_pos; - T alpha, beta; - cs_coords(side, edge_pos, i, nc3, nc2, nc1, nghost, panel.x2v, panel.x2f, - panel.x3v, panel.x3f, &alpha, &beta, &k, &j, &face_pos); - if (threadIdx.x != face_pos) - return; - - T local_left[64], local_right[64], remote_left[64], remote_right[64], wl[64], - wr[64]; - for (int v = 0; v < nvar; ++v) { - local_left[v] = buf[local_left_off + v * buf_stride_var]; - local_right[v] = buf[local_right_off + v * buf_stride_var]; - remote_left[v] = peer_buf[remote_left_off + v * buf_stride_var]; - remote_right[v] = peer_buf[remote_right_off + v * buf_stride_var]; - } - bool lower_side = side == CS_SIDE_L || side == CS_SIDE_B; - bool peer_lower_side = peer_side == CS_SIDE_L || peer_side == CS_SIDE_B; - for (int v = 0; v < nvar; ++v) { - T remote = peer_lower_side ? remote_right[v] : remote_left[v]; - wl[v] = lower_side ? remote : local_left[v]; - wr[v] = lower_side ? local_right[v] : remote; - } - gnomonic_sph_to_contra(wl, panel.face, alpha, beta); - gnomonic_sph_to_contra(wr, panel.face, alpha, beta); - if (physics.eos_limiter) { - wl[IDN] = max(wl[IDN], physics.density_floor); - wr[IDN] = max(wr[IDN], physics.density_floor); - if (physics.eos != FusedEos::ShallowWater) { - wl[IPR] = max(wl[IPR], physics.pressure_floor); - wr[IPR] = max(wr[IPR], physics.pressure_floor); - for (int v = ICY; v < nvar; ++v) { - wl[v] = max(wl[v], T(0)); - wr[v] = max(wr[v], T(0)); - } - } - } - T wl_density = wl[IDN]; - T wr_density = wr[IDN]; - gnomonic_prim2local(wl, dim, alpha, beta); - gnomonic_prim2local(wr, dim, alpha, beta); - wl[IDN] = wl_density; - wr[IDN] = wr_density; - - int flux_k = side <= CS_SIDE_R ? k : face_pos; - int flux_j = side <= CS_SIDE_R ? face_pos : j; - int flux_flat = flux_k * nc2 * nc1 + flux_j * nc1 + i; - int stride_var = nc1 * nc2 * nc3; - T *flux = side <= CS_SIDE_R ? flux2 + flux_flat : flux3 + flux_flat; - if (physics.solver == FusedRiemannSolver::ShallowRoe) { - shallow_roe_impl(flux, wl, wr, dim, physics.shallow_roe_dir_yz, - /*stride_w=*/1, /*stride_f=*/stride_var); - gnomonic_flux2global(flux, dim, alpha, beta, stride_var); - return; - } - - T el = 0., er = 0., gl = 0., gr = 0., cl = 0., cr = 0.; - int ny = nvar - ICY; - eos_side_quantities(wl, wr, ny, physics.nvapor, physics.eos, physics.gammad, - physics.inv_mu_ratio_m1, physics.cv_ratio_m1, physics.u0, - &el, &er, &gl, &gr, &cl, &cr); - - if (physics.solver == FusedRiemannSolver::LMARS) { - lmars_impl(flux, wl, wr, el / wl[IDN], er / wr[IDN], gl, gr, dim, ny, - /*stride_w=*/1, /*stride_f=*/stride_var); - } else if (physics.solver == FusedRiemannSolver::HLLC) { - hllc_impl(flux, wl, wr, el, er, gl, gr, cl, cr, dim, ny, /*stride_w=*/1, - /*stride_f=*/stride_var); - } else { - roe_impl(flux, wl, wr, el, er, gl, gr, cl, cr, dim, ny, physics.eos, - physics.nvapor, physics.gammad, physics.inv_mu_ratio_m1, - physics.cv_ratio_m1, physics.u0, /*stride_w=*/1, - /*stride_f=*/stride_var); - } - gnomonic_flux2global(flux, dim, alpha, beta, stride_var); -} - -template -__global__ void cs_flux_kernel(DeviceCubedSphereFluxParams params) { - (void)params.physics.recon_vel; - int edge_len = max(params.nc2, params.nc3); - int line = blockIdx.x; - int i = line % params.nc1; - int edge_pos = (line / params.nc1) % edge_len; - int side = line / (params.nc1 * edge_len); - cs_flux_common(params, i, edge_pos, side); -} - -template -__global__ void cs_flux_all_kernel(DeviceCubedSphereFluxParams base_params, - void **flux2_ptrs, void **flux3_ptrs, - int const *side_meta_all, int const *faces, - void **x2v_ptrs, void **x2f_ptrs, - void **x3v_ptrs, void **x3f_ptrs) { - (void)base_params.physics.recon_vel; - int edge_len = max(base_params.nc2, base_params.nc3); - int lines_per_panel = CS_NUM_SIDES * edge_len * base_params.nc1; - int panel = blockIdx.x / lines_per_panel; - int line = blockIdx.x % lines_per_panel; - int i = line % base_params.nc1; - int edge_pos = (line / base_params.nc1) % edge_len; - int side = line / (base_params.nc1 * edge_len); - auto params = base_params; - params.flux2 = static_cast(flux2_ptrs[panel]); - params.flux3 = static_cast(flux3_ptrs[panel]); - params.panel.side_meta = - side_meta_all + panel * CS_NUM_SIDES * CS_META_STRIDE; - params.panel.face = faces[panel]; - params.panel.local_block = panel; - params.panel.x2v = static_cast(x2v_ptrs[panel]); - params.panel.x2f = static_cast(x2f_ptrs[panel]); - params.panel.x3v = static_cast(x3v_ptrs[panel]); - params.panel.x3f = static_cast(x3f_ptrs[panel]); - cs_flux_common(params, i, edge_pos, side); -} - -} // namespace - -void fused_cubed_sphere_pack_cuda(torch::Tensor w, - torch::Tensor exchange_buffer, - FusedCubedSpherePackParams const ¶ms) { - at::cuda::CUDAGuard device_guard(w.device()); - int nc3 = w.size(1); - int nc2 = w.size(2); - int nc1 = w.size(3); - int nvar = w.size(0); - int edge_len = std::max(nc2, nc3); - int threads = edge_len; - TORCH_CHECK(threads <= 1024, - "dynamics.fused-recon-riemann cubed-sphere shared-memory " - "exchange requires edge lines to fit in one CUDA block, but got ", - threads); - int blocks = 4 * edge_len * nc1; - auto stream = at::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES(w.scalar_type(), "fused_cubed_sphere_pack", [&] { - size_t shared = static_cast(edge_len) * nvar * sizeof(scalar_t); - DeviceCubedSpherePackParams device_params{ - nvar, - nc3, - nc2, - nc1, - make_device_panel(params.panel), - params.recon_prim, - params.recon_vel, - params.recon_scale, - params.eos, - scalar_t(params.density_floor), - scalar_t(params.pressure_floor), - params.eos_limiter}; - cs_pack_kernel<<>>( - w.data_ptr(), exchange_buffer.data_ptr(), - device_params); - }); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, - torch::Tensor flux3, - torch::Tensor exchange_buffer, - void **exchange_buffer_ptrs_dev, - FusedCubedSphereFluxParams const ¶ms) { - at::cuda::CUDAGuard device_guard(w.device()); - int nc3 = w.size(1); - int nc2 = w.size(2); - int nc1 = w.size(3); - int nvar = w.size(0); - int edge_len = std::max(nc2, nc3); - int threads = edge_len; - TORCH_CHECK(threads <= 1024, - "dynamics.fused-recon-riemann cubed-sphere shared-memory " - "exchange requires edge lines to fit in one CUDA block, but got ", - threads); - int blocks = 4 * edge_len * nc1; - auto stream = at::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES(w.scalar_type(), "fused_cubed_sphere_flux", [&] { - DeviceCubedSphereFluxParams device_params{ - exchange_buffer.data_ptr(), - flux2.defined() ? flux2.data_ptr() : nullptr, - flux3.defined() ? flux3.data_ptr() : nullptr, - exchange_buffer_ptrs_dev, - nvar, - nc3, - nc2, - nc1, - make_device_panel(params.panel), - make_device_physics(params.physics)}; - cs_flux_kernel<<>>(device_params); - }); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -void fused_cubed_sphere_flux_all_cuda( - torch::Tensor exchange_buffer, FusedCubedSphereFluxAllPtrs ptrs, - FusedCubedSphereFluxAllParams const ¶ms) { - at::cuda::CUDAGuard device_guard(params.device); - int edge_len = std::max(params.nc2, params.nc3); - int threads = edge_len; - TORCH_CHECK(threads <= 1024, - "dynamics.fused-recon-riemann cubed-sphere shared-memory " - "exchange requires edge lines to fit in one CUDA block, but got ", - threads); - int blocks = params.bpp * 4 * edge_len * params.nc1; - auto stream = at::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES(params.dtype, "fused_cubed_sphere_flux_all", [&] { - DeviceCubedSphereFluxParams device_params{ - exchange_buffer.data_ptr(), - nullptr, - nullptr, - ptrs.exchange_buffer, - params.nvar, - params.nc3, - params.nc2, - params.nc1, - {nullptr, 0, 0, nullptr, nullptr, nullptr, nullptr}, - make_device_physics(params.physics)}; - cs_flux_all_kernel<<>>( - device_params, ptrs.flux2, ptrs.flux3, - params.side_meta_all.data_ptr(), params.faces.data_ptr(), - ptrs.x2v, ptrs.x2f, ptrs.x3v, ptrs.x3f); - }); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -} // namespace snap diff --git a/src/hydro/fused_dispatch_params.cuh b/src/hydro/fused_dispatch_params.cuh deleted file mode 100644 index 15e6a72e..00000000 --- a/src/hydro/fused_dispatch_params.cuh +++ /dev/null @@ -1,121 +0,0 @@ -#pragma once - -#include "fused_recon_riemann_dispatch.hpp" - -namespace snap { - -template struct DeviceFusedPhysicsParams { - FusedReconScheme recon_prim; - FusedReconScheme recon_vel; - bool recon_scale; - FusedRiemannSolver solver; - FusedEos eos; - T gammad; - T density_floor; - T pressure_floor; - bool eos_limiter; - T const *inv_mu_ratio_m1; - T const *cv_ratio_m1; - T const *u0; - int nvapor; - int shallow_roe_dir_yz; -}; - -template struct DeviceFusedX1RevisionParams { - bool revise_lr; - T const *dx1f; - T *rho_grav; -}; - -template struct DeviceFusedMetricParams { - bool cubed_sphere; - int face; - T const *x2v; - T const *x2f; - T const *x3v; - T const *x3f; -}; - -template struct DeviceFusedReconRiemannParams { - int nvar; - int nc3; - int nc2; - int nc1; - int dim; - T *face_pressure; - DeviceFusedPhysicsParams physics; - DeviceFusedX1RevisionParams x1_revision; - DeviceFusedMetricParams metric; -}; - -template struct DeviceCubedSpherePanelParams { - int const *side_meta; - int face; - int local_block; - T const *x2v; - T const *x2f; - T const *x3v; - T const *x3f; -}; - -template struct DeviceCubedSpherePackParams { - int nvar; - int nc3; - int nc2; - int nc1; - DeviceCubedSpherePanelParams panel; - FusedReconScheme recon_prim; - FusedReconScheme recon_vel; - bool recon_scale; - FusedEos eos; - T density_floor; - T pressure_floor; - bool eos_limiter; -}; - -template struct DeviceCubedSphereFluxParams { - T const *buf; - T *flux2; - T *flux3; - void **buf_ptrs; - int nvar; - int nc3; - int nc2; - int nc1; - DeviceCubedSpherePanelParams panel; - DeviceFusedPhysicsParams physics; -}; - -template -inline DeviceFusedPhysicsParams -make_device_physics(FusedPhysicsParams const &p) { - return {p.recon_prim, - p.recon_vel, - p.recon_scale, - p.solver, - p.eos, - T(p.gammad), - T(p.density_floor), - T(p.pressure_floor), - p.eos_limiter, - p.inv_mu_ratio_m1.defined() ? p.inv_mu_ratio_m1.data_ptr() - : nullptr, - p.cv_ratio_m1.defined() ? p.cv_ratio_m1.data_ptr() : nullptr, - p.u0.defined() ? p.u0.data_ptr() : nullptr, - p.nvapor, - p.shallow_roe_dir_yz}; -} - -template -inline DeviceCubedSpherePanelParams -make_device_panel(FusedCubedSpherePanelParams const &p) { - return {p.side_meta.data_ptr(), - p.face, - p.local_block, - p.x2v.data_ptr(), - p.x2f.data_ptr(), - p.x3v.data_ptr(), - p.x3f.data_ptr()}; -} - -} // namespace snap diff --git a/src/hydro/fused_recon_riemann_dispatch.cu b/src/hydro/fused_recon_riemann_dispatch.cu deleted file mode 100644 index 3f66f481..00000000 --- a/src/hydro/fused_recon_riemann_dispatch.cu +++ /dev/null @@ -1,246 +0,0 @@ -// torch -#include -#include -#include - -// snap -#include - -#include "../coord/gnomonic_equiangle.h" -#include "../eos/eos_side_quantities.cuh" -#include "../recon/interp_impl.cuh" -#include "../riemann/hllc_impl.h" -#include "../riemann/lmars_impl.h" -#include "../riemann/roe_impl.h" -#include "../riemann/shallow_roe_impl.h" -#include "fused_dispatch_params.cuh" -#include "fused_recon_riemann_dispatch.hpp" - -namespace snap { -namespace { - -template -__global__ void fused_kernel(T const *w, T *flux, - DeviceFusedReconRiemannParams params) { - int nvar = params.nvar; - int nc3 = params.nc3; - int nc2 = params.nc2; - int nc1 = params.nc1; - int dim = params.dim; - T *face_pressure = params.face_pressure; - auto physics = params.physics; - auto x1_revision = params.x1_revision; - auto metric = params.metric; - int ncells = nc1 * nc2 * nc3; - int axis_size = dim == 3 ? nc1 : (dim == 2 ? nc2 : nc3); - int stride_dim = dim == 3 ? 1 : (dim == 2 ? nc1 : nc1 * nc2); - int stride_var = ncells; - int pos = threadIdx.x; - int line = blockIdx.x; - int base = 0; - - int i = 0; - int j = 0; - int k = 0; - if (dim == 3) { - j = line % nc2; - k = line / nc2; - base = k * nc2 * nc1 + j * nc1; - } else if (dim == 2) { - i = line % nc1; - k = line / nc1; - base = k * nc2 * nc1 + i; - } else { - i = line % nc1; - j = line / nc1; - base = j * nc1 + i; - } - - extern __shared__ unsigned char memory[]; - T *smem = reinterpret_cast(memory); - T *left_pressure = smem + nvar * axis_size; - T *right_pressure = left_pressure + axis_size; - if (pos < axis_size) { - int flat = base + pos * stride_dim; - for (int v = 0; v < nvar; ++v) { - smem[v * axis_size + pos] = w[v * stride_var + flat]; - } - } - __syncthreads(); - - if (pos >= axis_size) - return; - int flat = base + pos * stride_dim; - int nghost = (physics.recon_prim == FusedReconScheme::CP3 || - physics.recon_prim == FusedReconScheme::WENO3) - ? 2 - : 3; - int il = nghost; - int iu = axis_size - nghost; - int stencil = 2 * nghost - 1; - - bool valid_face = pos >= il && pos <= iu + 1; - - int wl_start = min(max(pos - il, 0), axis_size - stencil); - int wr_start = min(max(pos - (il - 1), 0), axis_size - stencil); - T wl_local[64], wr_local[64]; - for (int v = 0; v < nvar; ++v) { - auto scheme = - (v == IDN || v >= ICY) ? physics.recon_prim : physics.recon_vel; - wl_local[v] = interp_shared_fused_impl(smem, v, wl_start, axis_size, scheme, - /*right=*/true, physics.recon_scale); - wr_local[v] = - interp_shared_fused_impl(smem, v, wr_start, axis_size, scheme, - /*right=*/false, physics.recon_scale); - } - - if (x1_revision.revise_lr && dim == 3 && valid_face) { - if (pos == il) { - wl_local[IPR] = wr_local[IPR]; - wl_local[IDN] = wr_local[IDN]; - } else if (pos == iu) { - wr_local[IPR] = wl_local[IPR]; - wr_local[IDN] = wl_local[IDN]; - } - } - - if (physics.eos_limiter && valid_face) { - wl_local[IDN] = max(wl_local[IDN], physics.density_floor); - wr_local[IDN] = max(wr_local[IDN], physics.density_floor); - if (physics.eos != FusedEos::ShallowWater) { - wl_local[IPR] = max(wl_local[IPR], physics.pressure_floor); - wr_local[IPR] = max(wr_local[IPR], physics.pressure_floor); - for (int v = ICY; v < nvar; ++v) { - wl_local[v] = max(wl_local[v], T(0.)); - wr_local[v] = max(wr_local[v], T(0.)); - } - } - } - - if (x1_revision.revise_lr && dim == 3 && valid_face) { - left_pressure[pos] = wl_local[IPR]; - right_pressure[pos] = wr_local[IPR]; - } - __syncthreads(); - if (x1_revision.revise_lr && dim == 3 && x1_revision.rho_grav != nullptr && - pos >= il && pos < iu) { - x1_revision.rho_grav[flat] = - (left_pressure[pos + 1] - right_pressure[pos]) / x1_revision.dx1f[pos]; - } - - if (!valid_face) { - for (int v = 0; v < nvar; ++v) - flux[v * stride_var + flat] = 0.; - if (face_pressure != nullptr) - face_pressure[flat] = 0.; - return; - } - - T alpha = 0; - T beta = 0; - bool use_cubed_metric = metric.cubed_sphere && dim != 3; - if (use_cubed_metric) { - if (dim == 2) { - alpha = metric.x2f[pos]; - beta = metric.x3v[k]; - } else { - alpha = metric.x2v[j]; - beta = metric.x3f[pos]; - } - T wl_density = wl_local[IDN]; - T wr_density = wr_local[IDN]; - gnomonic_prim2local(wl_local, dim, alpha, beta); - gnomonic_prim2local(wr_local, dim, alpha, beta); - wl_local[IDN] = wl_density; - wr_local[IDN] = wr_density; - } - - if (physics.solver == FusedRiemannSolver::ShallowRoe) { - shallow_roe_impl(flux + flat, wl_local, wr_local, dim, - physics.shallow_roe_dir_yz, /*stride_w=*/1, - /*stride_f=*/stride_var); - } else { - T *face_pressure_out = - face_pressure != nullptr ? face_pressure + flat : nullptr; - T el = 0., er = 0., gl = 0., gr = 0., cl = 0., cr = 0.; - int ny = physics.eos == FusedEos::ShallowWater ? 0 : nvar - ICY; - eos_side_quantities(wl_local, wr_local, ny, physics.nvapor, physics.eos, - physics.gammad, physics.inv_mu_ratio_m1, - physics.cv_ratio_m1, physics.u0, &el, &er, &gl, &gr, - &cl, &cr); - - if (physics.solver == FusedRiemannSolver::LMARS) { - lmars_impl(flux + flat, wl_local, wr_local, el / wl_local[IDN], - er / wr_local[IDN], gl, gr, dim, ny, /*stride_w=*/1, - /*stride_f=*/stride_var, face_pressure_out); - } else if (physics.solver == FusedRiemannSolver::HLLC) { - hllc_impl(flux + flat, wl_local, wr_local, el, er, gl, gr, cl, cr, dim, - ny, /*stride_w=*/1, /*stride_f=*/stride_var, face_pressure_out); - } else { - roe_impl(flux + flat, wl_local, wr_local, el, er, gl, gr, cl, cr, dim, ny, - physics.eos, physics.nvapor, physics.gammad, - physics.inv_mu_ratio_m1, physics.cv_ratio_m1, physics.u0, - /*stride_w=*/1, /*stride_f=*/stride_var, face_pressure_out); - } - } - if (use_cubed_metric) { - gnomonic_flux2global(flux + flat, dim, alpha, beta, stride_var); - } -} - -} // namespace - -void fused_recon_riemann_cuda(torch::Tensor w, torch::Tensor flux, - torch::Tensor face_pressure, - FusedReconRiemannParams const ¶ms) { - at::cuda::CUDAGuard device_guard(w.device()); - int nc3 = w.size(1); - int nc2 = w.size(2); - int nc1 = w.size(3); - int nvar = w.size(0); - int dim = params.dim; - int axis_size = dim == 3 ? nc1 : (dim == 2 ? nc2 : nc3); - TORCH_CHECK(axis_size <= 1024, - "dynamics.fused-recon-riemann shared-memory kernel requires " - "the reconstructed dimension to fit in one CUDA block, but got ", - axis_size); - int threads = axis_size; - int blocks = dim == 3 ? nc2 * nc3 : (dim == 2 ? nc1 * nc3 : nc1 * nc2); - auto stream = at::cuda::getCurrentCUDAStream(); - - AT_DISPATCH_FLOATING_TYPES(w.scalar_type(), "fused_recon_riemann_cuda", [&] { - size_t shared = static_cast(axis_size) * nvar * sizeof(scalar_t); - if (params.x1_revision.revise_lr && dim == 3) { - shared += static_cast(2) * axis_size * sizeof(scalar_t); - } - DeviceFusedReconRiemannParams device_params{ - nvar, - nc3, - nc2, - nc1, - dim, - face_pressure.defined() ? face_pressure.data_ptr() : nullptr, - make_device_physics(params.physics), - {params.x1_revision.revise_lr, - params.x1_revision.dx1f.defined() - ? params.x1_revision.dx1f.data_ptr() - : nullptr, - params.x1_revision.rho_grav.defined() - ? params.x1_revision.rho_grav.data_ptr() - : nullptr}, - {params.metric.cubed_sphere, params.metric.face, - params.metric.x2v.defined() ? params.metric.x2v.data_ptr() - : nullptr, - params.metric.x2f.defined() ? params.metric.x2f.data_ptr() - : nullptr, - params.metric.x3v.defined() ? params.metric.x3v.data_ptr() - : nullptr, - params.metric.x3f.defined() ? params.metric.x3f.data_ptr() - : nullptr}}; - fused_kernel<<>>( - w.data_ptr(), flux.data_ptr(), device_params); - }); - C10_CUDA_KERNEL_LAUNCH_CHECK(); -} - -} // namespace snap diff --git a/src/hydro/fused_recon_riemann_dispatch.hpp b/src/hydro/fused_recon_riemann_dispatch.hpp deleted file mode 100644 index ae80b29f..00000000 --- a/src/hydro/fused_recon_riemann_dispatch.hpp +++ /dev/null @@ -1,126 +0,0 @@ -#pragma once - -// torch -#include - -// snap -#include - -namespace snap { - -struct FusedPhysicsParams { - FusedReconScheme recon_prim; - FusedReconScheme recon_vel; - bool recon_scale; - FusedRiemannSolver solver; - FusedEos eos; - double gammad; - double density_floor; - double pressure_floor; - bool eos_limiter; - torch::Tensor inv_mu_ratio_m1; - torch::Tensor cv_ratio_m1; - torch::Tensor u0; - int nvapor; - int shallow_roe_dir_yz; -}; - -struct FusedX1RevisionParams { - bool revise_lr; - torch::Tensor dx1f; - torch::Tensor rho_grav; -}; - -struct FusedMetricParams { - bool cubed_sphere; - int face; - torch::Tensor x2v; - torch::Tensor x2f; - torch::Tensor x3v; - torch::Tensor x3f; -}; - -struct FusedReconRiemannParams { - int dim; - FusedPhysicsParams physics; - FusedX1RevisionParams x1_revision; - FusedMetricParams metric; -}; - -void fused_recon_riemann_cuda(torch::Tensor w, torch::Tensor flux, - torch::Tensor face_pressure, - FusedReconRiemannParams const& params); - -// Multi-block-per-process cubed-sphere fused exchange. Every local block packs -// edge states into slice `local_block` of a shared per-process buffer. The host -// exchanges that buffer with remote peers, then one kernel computes and -// overwrites every local panel's cross-panel boundary flux. -struct FusedCubedSpherePanelParams { - torch::Tensor side_meta; - int face; - int local_block; - torch::Tensor x2v; - torch::Tensor x2f; - torch::Tensor x3v; - torch::Tensor x3f; -}; - -struct FusedCubedSpherePackParams { - FusedCubedSpherePanelParams panel; - FusedReconScheme recon_prim; - FusedReconScheme recon_vel; - bool recon_scale; - FusedEos eos; - double density_floor; - double pressure_floor; - bool eos_limiter; -}; - -void fused_cubed_sphere_pack_cuda(torch::Tensor w, - torch::Tensor exchange_buffer, - FusedCubedSpherePackParams const& params); - -struct FusedCubedSphereFluxParams { - FusedCubedSpherePanelParams panel; - FusedPhysicsParams physics; -}; - -void fused_cubed_sphere_flux_cuda(torch::Tensor w, torch::Tensor flux2, - torch::Tensor flux3, - torch::Tensor exchange_buffer, - void** exchange_buffer_ptrs_dev, - FusedCubedSphereFluxParams const& params); - -// Process-level seam flux: one launch overwrites the cross-panel boundary flux -// for ALL local panels. flux2_ptrs_dev/flux3_ptrs_dev are device arrays of the -// per-panel flux tensor pointers (indexed by local block), faces is a device -// int array of each panel's face id, and side_meta_all is [bpp, 4, kStride]. -// Coords are the shared equiangular angular grid (same for every panel). -struct FusedCubedSphereFluxAllPtrs { - void** flux2; - void** flux3; - void** exchange_buffer; - void** x2v; - void** x2f; - void** x3v; - void** x3f; -}; - -struct FusedCubedSphereFluxAllParams { - torch::Tensor side_meta_all; - torch::Tensor faces; - int nvar; - int nc3; - int nc2; - int nc1; - int bpp; - c10::ScalarType dtype; - torch::Device device; - FusedPhysicsParams physics; -}; - -void fused_cubed_sphere_flux_all_cuda( - torch::Tensor exchange_buffer, FusedCubedSphereFluxAllPtrs ptrs, - FusedCubedSphereFluxAllParams const& params); - -} // namespace snap diff --git a/src/hydro/hydro.cpp b/src/hydro/hydro.cpp index c12dca7b..6a0a3a2d 100644 --- a/src/hydro/hydro.cpp +++ b/src/hydro/hydro.cpp @@ -12,14 +12,6 @@ #include "hydro.hpp" namespace snap { -namespace { - -bool fused_runtime_supported(torch::Tensor const& u, Variables const& other) { - return u.is_cuda() && !other.count("solid"); -} - -} // namespace - HydroImpl::HydroImpl(const HydroOptions& options_, torch::nn::Module* p) : options(options_) { pmb = dynamic_cast(p); @@ -187,17 +179,6 @@ double HydroImpl::max_time_step(torch::Tensor w, torch::Tensor solid) const { return dt; } -torch::Tensor HydroImpl::forward(double dt, torch::Tensor u, - Variables const& other) { - if (options->fused_recon_riemann()) { - if (!fused_runtime_supported(u, other)) { - return _forward_staged(dt, u, other); - } - return _forward_fused(dt, u, other); - } - return _forward_staged(dt, u, other); -} - torch::Tensor HydroImpl::implicit_mass_correction() const { return picorr ? picorr->mass_correction() : torch::Tensor(); } diff --git a/src/hydro/hydro.hpp b/src/hydro/hydro.hpp index 5b959e17..0b90b904 100644 --- a/src/hydro/hydro.hpp +++ b/src/hydro/hydro.hpp @@ -35,8 +35,7 @@ struct HydroOptionsImpl { os << "* verbose = " << verbose() << "\n" << "* disable_flux_x1 = " << disable_flux_x1() << "\n" << "* disable_flux_x2 = " << disable_flux_x2() << "\n" - << "* disable_flux_x3 = " << disable_flux_x3() << "\n" - << "* fused_recon_riemann = " << fused_recon_riemann() << "\n"; + << "* disable_flux_x3 = " << disable_flux_x3() << "\n"; } //! verbose @@ -45,7 +44,6 @@ struct HydroOptionsImpl { ADD_ARG(bool, disable_flux_x1) = false; ADD_ARG(bool, disable_flux_x2) = false; ADD_ARG(bool, disable_flux_x3) = false; - ADD_ARG(bool, fused_recon_riemann) = false; //! forcing options ADD_ARG(ConstGravityOptions, grav) = nullptr; @@ -145,11 +143,6 @@ class HydroImpl : public torch::nn::Cloneable { //! Register all forcing modules std::vector _register_forcings_module(); - torch::Tensor _forward_staged(double dt, torch::Tensor hydro_u, - Variables const& other); - torch::Tensor _forward_fused(double dt, torch::Tensor hydro_u, - Variables const& other); - torch::Tensor _flux1, _flux2, _flux3, _face_pressure1, _div, _imp; }; diff --git a/src/hydro/hydro_forward_staged.cpp b/src/hydro/hydro_forward.cpp similarity index 98% rename from src/hydro/hydro_forward_staged.cpp rename to src/hydro/hydro_forward.cpp index ff84bd7d..c7d0783d 100644 --- a/src/hydro/hydro_forward_staged.cpp +++ b/src/hydro/hydro_forward.cpp @@ -11,8 +11,8 @@ namespace snap { -torch::Tensor HydroImpl::_forward_staged(double dt, torch::Tensor u, - Variables const& other) { +torch::Tensor HydroImpl::forward(double dt, torch::Tensor u, + Variables const& other) { enum { DIM1 = 3, DIM2 = 2, DIM3 = 1 }; bool has_solid = other.count("solid"); auto start = std::chrono::high_resolution_clock::now(); diff --git a/src/hydro/hydro_forward_fused.cpp b/src/hydro/hydro_forward_fused.cpp deleted file mode 100644 index 70a36be5..00000000 --- a/src/hydro/hydro_forward_fused.cpp +++ /dev/null @@ -1,595 +0,0 @@ -// C/C++ -#include -#include -#include -#include -#include -#include -#include -#include -#include - -// snap -#include - -#include -#include -#include - -#include "../eos/ideal_moist.hpp" -#include "../sedimentation/sedimentation.hpp" -#include "fused_recon_riemann_dispatch.hpp" -#include "hydro.hpp" - -namespace snap { -namespace { - -FusedReconScheme fused_recon_scheme(std::string const& type, - bool velocity_group, bool shock) { - if (type == "weno5") { - return velocity_group && !shock ? FusedReconScheme::CP5 - : FusedReconScheme::WENO5; - } - if (type == "weno3") { - return velocity_group && !shock ? FusedReconScheme::CP3 - : FusedReconScheme::WENO3; - } - if (type == "cp5") return FusedReconScheme::CP5; - if (type == "cp3") return FusedReconScheme::CP3; - TORCH_CHECK(false, - "dynamics.fused-recon-riemann supports cp3, cp5, weno3, and " - "weno5 reconstruction, but got ", - type); -} - -bool fused_combo_supported(std::string const& eos_type, - std::string const& riemann_type) { - return ((eos_type == "ideal-gas" && - (riemann_type == "lmars" || riemann_type == "hllc" || - riemann_type == "roe")) || - (eos_type == "ideal-moist" && - (riemann_type == "lmars" || riemann_type == "hllc" || - riemann_type == "roe")) || - (eos_type == "shallow-water" && riemann_type == "shallow-roe")); -} - -FusedEos fused_eos(std::string const& type) { - if (type == "ideal-gas") return FusedEos::IdealGas; - if (type == "ideal-moist") return FusedEos::IdealMoist; - if (type == "shallow-water") return FusedEos::ShallowWater; - TORCH_CHECK(false, - "dynamics.fused-recon-riemann supports EOS types ideal-gas, " - "ideal-moist, and shallow-water, but got ", - type); -} - -FusedRiemannSolver fused_riemann_solver(std::string const& type) { - if (type == "lmars") return FusedRiemannSolver::LMARS; - if (type == "hllc") return FusedRiemannSolver::HLLC; - if (type == "roe") return FusedRiemannSolver::Roe; - if (type == "shallow-roe") return FusedRiemannSolver::ShallowRoe; - TORCH_CHECK(false, - "dynamics.fused-recon-riemann supports Riemann solvers lmars, " - "hllc, roe, and shallow-roe, but got ", - type); -} - -constexpr int kFusedCubedSphereExchangeTag = - std::numeric_limits::max() - 113; - -struct FusedExchangePool { - torch::Tensor buffer; - std::vector peer_buffers; - torch::Tensor buffer_ptrs; - - void** buffer_ptrs_dev() const { - return reinterpret_cast(buffer_ptrs.data_ptr()); - } - - void prepare_peer_buffers(std::set const& peers, int rank, - int world_size) { - if (static_cast(peer_buffers.size()) != world_size) { - peer_buffers.resize(world_size); - } - std::vector pointers(world_size, 0); - pointers[rank] = reinterpret_cast(buffer.data_ptr()); - for (int peer : peers) { - TORCH_CHECK(peer >= 0 && peer < world_size, "invalid fused UCX peer ", - peer, " for world size ", world_size); - if (!peer_buffers[peer].defined()) { - peer_buffers[peer] = torch::empty_strided( - buffer.sizes(), buffer.strides(), buffer.options()); - } - pointers[peer] = reinterpret_cast(peer_buffers[peer].data_ptr()); - } - buffer_ptrs = torch::tensor(pointers, torch::dtype(torch::kInt64)) - .to(buffer.device()); - } -}; - -FusedExchangePool& get_fused_exchange_pool( - std::string const& group_name, c10::ScalarType dtype, torch::Device device, - std::vector const& sizes, std::vector const& strides) { - static std::mutex mutex; - static std::unordered_map pools; - - std::string key = group_name + ":device=" + std::to_string(device.index()) + - ":dtype=" + std::to_string(static_cast(dtype)); - for (auto size : sizes) key += ":s" + std::to_string(size); - for (auto stride : strides) key += ":t" + std::to_string(stride); - - std::lock_guard lock(mutex); - auto it = pools.find(key); - if (it != pools.end()) return it->second; - - FusedExchangePool pool; - pool.buffer = torch::empty_strided( - sizes, strides, torch::TensorOptions().dtype(dtype).device(device)); - auto [inserted, _] = pools.emplace(key, std::move(pool)); - return inserted->second; -} - -void exchange_fused_buffers_ucx( - FusedExchangePool& pool, std::set const& peers, - std::shared_ptr const& comm) { - TORCH_CHECK(comm != nullptr && comm->initialized() && comm->is_ucx(), - "dynamics.fused-recon-riemann multi-process cubed-sphere " - "exchange requires backend=ucx"); - - std::vector works; - works.reserve(2 * peers.size()); - for (int peer : peers) { - std::vector send_tensors{pool.buffer}; - std::vector recv_tensors{pool.peer_buffers[peer]}; - if (auto work = - comm->send(send_tensors, peer, kFusedCubedSphereExchangeTag)) { - works.push_back(std::move(work)); - } - if (auto work = - comm->recv(recv_tensors, peer, kFusedCubedSphereExchangeTag)) { - works.push_back(std::move(work)); - } - } - for (auto const& work : works) work->wait(); -} - -//! \brief Reusable process-local barrier over the blocks_per_process worker -//! threads. The fused cubed-sphere exchange runs pack -> UCX exchange -> flux -//! on one shared process-local buffer; these barriers guarantee every local -//! block has enqueued its pack before the leader starts transport and that the -//! transport completes before any block reads peer state. All -//! exchange kernels share the device's default stream, so enqueue order fixed -//! by these barriers is the execution order, which removes the single-block -//! launch ordering deadlock that made the old path require blocks_per_process -//! == 1. -class FusedExchangeBarrier { - public: - explicit FusedExchangeBarrier(int participants) - : participants_(participants) {} - void wait() { - std::unique_lock lock(mutex_); - int gen = generation_; - if (++arrived_ == participants_) { - arrived_ = 0; - ++generation_; - cv_.notify_all(); - } else { - cv_.wait(lock, [&] { return generation_ != gen; }); - } - } - - private: - std::mutex mutex_; - std::condition_variable cv_; - int participants_; - int arrived_ = 0; - int generation_ = 0; -}; - -FusedExchangeBarrier& get_fused_exchange_barrier(std::string const& key, - int participants) { - static std::mutex mutex; - static std::unordered_map> - barriers; - std::lock_guard lock(mutex); - auto it = barriers.find(key); - if (it == barriers.end()) { - it = barriers - .emplace(key, std::make_unique(participants)) - .first; - } - return *it->second; -} - -//! \brief Per-process collection of the local panels' flux tensors / face ids / -//! side metadata, so one leader-launched kernel can overwrite every panel's -//! cross-panel boundary flux. Each panel writes its own local-block slot before -//! the barrier; the leader assembles the device pointer arrays after it. The -//! device arrays are kept alive here (persistent) so the async flux-all kernel -//! never reads a freed pointer buffer. -struct FusedSeamBatch { - //! _flux2/_flux3 are persistent HydroImpl buffers, so raw ptrs are stable. - std::vector flux2_ptr, flux3_ptr; // per local block - std::vector faces; // per local block - std::vector> remote_peers; // per local block - //! side_meta / coords are transient per call; hold the tensors so their - //! device pointers stay valid until the async flux-all kernel completes. - std::vector side_meta, x2v, x2f, x3v, x3f; - torch::Tensor flux2_dev, flux3_dev, faces_dev, side_meta_all; - torch::Tensor x2v_dev, x2f_dev, x3v_dev, x3f_dev; -}; - -FusedSeamBatch& get_fused_seam_batch(std::string const& key, int bpp) { - static std::mutex mutex; - static std::unordered_map batches; - std::lock_guard lock(mutex); - auto& b = batches[key]; - if (static_cast(b.faces.size()) != bpp) { - b.flux2_ptr.assign(bpp, 0); - b.flux3_ptr.assign(bpp, 0); - b.faces.assign(bpp, 0); - b.remote_peers.assign(bpp, {}); - b.side_meta.assign(bpp, torch::Tensor()); - b.x2v.assign(bpp, torch::Tensor()); - b.x2f.assign(bpp, torch::Tensor()); - b.x3v.assign(bpp, torch::Tensor()); - b.x3f.assign(bpp, torch::Tensor()); - } - return b; -} - -torch::Tensor make_side_meta(CubedSphereLayoutImpl const& layout, - bool exchange_dim2, bool exchange_dim3, - torch::Device device, - std::vector* remote_peers) { - // Per side: [enabled, peer_process, peer_local_block, peer_side, rev]. - // neighbor_rank returns a global BLOCK rank; translate it into the owning - // process (to index the exchange buffer pointer array) and the slot within - // that process (to index the shared buffer), so the exchange addresses - // same-GPU and cross-GPU peers uniformly. Must match the CS_META_* layout in - // the kernels. - constexpr int kStride = 5; - std::vector meta(4 * kStride, 0); - auto iloc = layout.loc_of(layout.options->rank()); - int face = std::get<2>(iloc); - std::array, 4> offsets = { - std::tuple{0, -1, 0}, - std::tuple{0, +1, 0}, - std::tuple{-1, 0, 0}, - std::tuple{+1, 0, 0}, - }; - for (int side = 0; side < 4; ++side) { - int nb = layout.neighbor_rank(iloc, offsets[side]); - if (nb < 0) continue; - auto nb_loc = layout.loc_of(nb); - if (std::get<2>(nb_loc) == face) continue; - int peer_process = layout.options->owner_process_rank(nb); - if (peer_process != layout.options->process_rank()) { - remote_peers->push_back(peer_process); - } - bool dim_enabled = side <= SIDE_R ? exchange_dim2 : exchange_dim3; - if (!dim_enabled) continue; - auto edge = CS_FACE_EDGES[face][side]; - meta[side * kStride + 0] = 1; - meta[side * kStride + 1] = peer_process; - meta[side * kStride + 2] = layout.options->local_block_index(nb); - meta[side * kStride + 3] = edge.nside; - meta[side * kStride + 4] = edge.rev; - } - return torch::tensor(meta, torch::dtype(torch::kInt32)).to(device); -} - -} // namespace - -torch::Tensor HydroImpl::_forward_fused(double dt, torch::Tensor u, - Variables const& other) { -#ifdef NOT_USE_CUDA - TORCH_CHECK(false, - "dynamics.fused-recon-riemann requires a CUDA-enabled build"); -#endif - - TORCH_CHECK(u.is_cuda(), - "dynamics.fused-recon-riemann requires CUDA tensors"); - TORCH_CHECK(other.count("hydro_w"), - "dynamics.fused-recon-riemann requires hydro_w primitives"); - auto const& w = other.at("hydro_w"); - TORCH_CHECK(w.is_cuda(), - "dynamics.fused-recon-riemann requires CUDA primitive tensors"); - TORCH_CHECK(w.is_contiguous() && u.is_contiguous(), - "dynamics.fused-recon-riemann requires contiguous hydro tensors"); - TORCH_CHECK(_flux1.is_contiguous() && - (!_flux2.defined() || _flux2.is_contiguous()) && - (!_flux3.defined() || _flux3.is_contiguous()), - "dynamics.fused-recon-riemann requires contiguous flux buffers"); - TORCH_CHECK(w.size(0) <= 64, - "dynamics.fused-recon-riemann supports at most 64 hydro " - "variables, but got ", - w.size(0)); - TORCH_CHECK(w.size(0) - ICY <= 32, - "dynamics.fused-recon-riemann supports at most 32 mass " - "fractions, but got ", - w.size(0) - ICY); - TORCH_CHECK(!other.count("solid"), - "dynamics.fused-recon-riemann does not yet support solid " - "internal-boundary state revision"); - - auto eos_type = options->eos()->type(); - auto riemann_type = options->riemann()->type(); - TORCH_CHECK(fused_combo_supported(eos_type, riemann_type), - "dynamics.fused-recon-riemann does not support EOS/Riemann " - "combination ", - eos_type, "/", riemann_type); - - auto playout = pmb->get_layout(); - bool cubed_sphere_layout = playout->options->type() == "cubed-sphere"; - if (!cubed_sphere_layout) { - TORCH_CHECK(pmb->pcoord->options->type() == "cartesian", - "dynamics.fused-recon-riemann currently supports cartesian " - "coordinates or cubed-sphere layouts only, but got coordinate " - "type ", - pmb->pcoord->options->type(), " with layout type ", - playout->options->type()); - } else { - // Multi-block-per-process is supported: each process may own several panels - // co-resident on one GPU. A shared per-process exchange buffer is sliced by - // local block, and cross-panel peers are addressed by (owning process, - // local block). Cross-process exchange uses UCX; a single-process run (all - // panels on one GPU) does not need a process group. - TORCH_CHECK( - playout->options->process_world_size() == 1 || - playout->has_process_group(), - "dynamics.fused-recon-riemann cubed-sphere exchange requires an " - "initialized process group for multi-process runs"); - TORCH_CHECK( - playout->options->process_world_size() == 1 || - (playout->comm != nullptr && playout->comm->is_ucx()), - "dynamics.fused-recon-riemann multi-process cubed-sphere exchange " - "requires backend=ucx"); - } - - CubedSphereLayoutImpl const* cs_layout = nullptr; - int face = 0; - torch::Tensor x2v, x2f, x3v, x3f; - if (cubed_sphere_layout) { - cs_layout = dynamic_cast(playout.get()); - TORCH_CHECK(cs_layout != nullptr, - "expected CubedSphereLayoutImpl for cubed-sphere layout"); - face = std::get<2>(cs_layout->loc_of(cs_layout->options->rank())); - x2v = pmb->pcoord->x2v.to(w.options()); - x2f = pmb->pcoord->x2f.to(w.options()); - x3v = pmb->pcoord->x3v.to(w.options()); - x3f = pmb->pcoord->x3f.to(w.options()); - } - - peos->forward(u, w); - - auto recon1_prim = - fused_recon_scheme(precon1->pinterp1->options->type(), - /*velocity_group=*/false, precon1->options->shock()); - auto recon1_vel = - fused_recon_scheme(precon1->pinterp1->options->type(), - /*velocity_group=*/true, precon1->options->shock()); - auto recon23_prim = - fused_recon_scheme(precon23->pinterp1->options->type(), - /*velocity_group=*/false, precon23->options->shock()); - auto recon23_vel = - fused_recon_scheme(precon23->pinterp1->options->type(), - /*velocity_group=*/true, precon23->options->shock()); - auto solver = fused_riemann_solver(riemann_type); - auto eos = fused_eos(eos_type); - int shallow_roe_dir_yz = options->riemann()->dir() == "yz" ? 1 : 0; - - torch::Tensor inv_mu_ratio_m1, cv_ratio_m1, u0; - int nvapor = 0; - if (eos == FusedEos::IdealMoist) { - auto ideal_moist = dynamic_cast(peos.get()); - TORCH_CHECK(ideal_moist != nullptr, - "dynamics.fused-recon-riemann expected IdealMoistImpl"); - inv_mu_ratio_m1 = ideal_moist->inv_mu_ratio_m1.to(w.options()); - cv_ratio_m1 = ideal_moist->cv_ratio_m1.to(w.options()); - u0 = ideal_moist->u0.to(w.options()); - nvapor = - static_cast(ideal_moist->pthermo->options->vapor_ids().size()) - 1; - } - - torch::Tensor rho_grav = torch::zeros_like(w[IDN]); - bool revise_x1_lr = options->grav() && (options->grav()->grav1() != 0); - torch::Tensor dx1f; - - if (u.size(3) > 1) { - if (revise_x1_lr) { - _revise_x1inner_ghost(w); - _revise_x1outer_ghost(w); - dx1f = pmb->pcoord->dx1f.to(w.options()); - } - } - - auto make_physics_params = [&](FusedReconScheme recon_prim, - FusedReconScheme recon_vel, bool recon_scale) { - return FusedPhysicsParams{recon_prim, - recon_vel, - recon_scale, - solver, - eos, - options->eos()->gammad(), - options->eos()->density_floor(), - options->eos()->pressure_floor(), - options->eos()->limiter(), - inv_mu_ratio_m1, - cv_ratio_m1, - u0, - nvapor, - shallow_roe_dir_yz}; - }; - FusedMetricParams metric_params{ - cubed_sphere_layout, face, x2v, x2f, x3v, x3f}; - auto face_pressure1 = - eos == FusedEos::ShallowWater ? torch::Tensor() : _face_pressure1; - - if (u.size(3) > 1 && !options->disable_flux_x1()) { - fused_recon_riemann_cuda( - w, _flux1, face_pressure1, - FusedReconRiemannParams{ - /*dim=*/3, - make_physics_params(recon1_prim, recon1_vel, - precon1->pinterp1->options->scale()), - FusedX1RevisionParams{revise_x1_lr, dx1f, rho_grav}, - metric_params}); - } - if (u.size(3) > 1 && psed) { - psed->forward(w, _flux1); - } - if (u.size(2) > 1 && !options->disable_flux_x2()) { - fused_recon_riemann_cuda( - w, _flux2, torch::Tensor(), - FusedReconRiemannParams{ - /*dim=*/2, - make_physics_params(recon23_prim, recon23_vel, - precon23->pinterp1->options->scale()), - FusedX1RevisionParams{false, torch::Tensor(), torch::Tensor()}, - metric_params}); - } - if (u.size(1) > 1 && !options->disable_flux_x3()) { - fused_recon_riemann_cuda( - w, _flux3, torch::Tensor(), - FusedReconRiemannParams{ - /*dim=*/1, - make_physics_params(recon23_prim, recon23_vel, - precon23->pinterp1->options->scale()), - FusedX1RevisionParams{false, torch::Tensor(), torch::Tensor()}, - metric_params}); - } - - if (cubed_sphere_layout) { - bool exchange_dim2 = u.size(2) > 1 && !options->disable_flux_x2(); - bool exchange_dim3 = u.size(1) > 1 && !options->disable_flux_x3(); - if (exchange_dim2 || exchange_dim3) { - std::string group_name = "snapy:fused-recon-riemann:cubed-sphere"; - int bpp = std::max(1, playout->options->blocks_per_process()); - int local_block = - playout->options->local_block_index(playout->options->rank()); - bool multi_process = playout->options->process_world_size() > 1; - bool is_leader = local_block == 0; - - int edge_len = std::max(w.size(1), w.size(2)); - constexpr int kSides = 4; - constexpr int kStates = 2; - int64_t nvar = w.size(0), nc1 = w.size(3); - // shared per-process buffer: [blocks_per_process, side, state, var, edge, - // nc1]; every local block owns slice `local_block`. - std::vector sizes = {bpp, kSides, kStates, nvar, edge_len, nc1}; - std::vector strides = { - static_cast(kSides) * kStates * nvar * edge_len * nc1, - static_cast(kStates) * nvar * edge_len * nc1, - nvar * edge_len * nc1, - edge_len * nc1, - nc1, - 1}; - - auto& pool = get_fused_exchange_pool(group_name, w.scalar_type(), - w.device(), sizes, strides); - std::vector remote_peers; - auto side_meta = make_side_meta(*cs_layout, exchange_dim2, exchange_dim3, - w.device(), &remote_peers); - auto& barrier = get_fused_exchange_barrier(group_name, bpp); - auto& batch = get_fused_seam_batch(group_name, bpp); - - // register this panel's flux tensors / face / coords / side_meta for the - // single process-level seam-flux launch (distinct local-block slots, so - // the concurrent workers never touch the same entry) - batch.flux2_ptr[local_block] = - reinterpret_cast(_flux2.data_ptr()); - batch.flux3_ptr[local_block] = - reinterpret_cast(_flux3.data_ptr()); - batch.faces[local_block] = face; - batch.remote_peers[local_block] = std::move(remote_peers); - batch.side_meta[local_block] = side_meta; - batch.x2v[local_block] = x2v; - batch.x2f[local_block] = x2f; - batch.x3v[local_block] = x3v; - batch.x3f[local_block] = x3f; - - // ---- Phase A: every local panel packs its edge states into its slice - // ---- - FusedCubedSpherePanelParams panel_params{ - side_meta, face, local_block, x2v, x2f, x3v, x3f}; - fused_cubed_sphere_pack_cuda( - w, pool.buffer, - FusedCubedSpherePackParams{panel_params, recon23_prim, recon23_vel, - precon23->pinterp1->options->scale(), eos, - options->eos()->density_floor(), - options->eos()->pressure_floor(), - options->eos()->limiter()}); - barrier.wait(); // all panels packed + registered - - // ---- Phase B: one leader per process exchanges the packed buffer with - // each remote UCX neighbor, then launches ONE flux kernel over ALL local - // panels (overwriting each panel's cross-panel boundary flux). ---- - if (is_leader) { - std::set peer_processes; - for (auto const& block_peers : batch.remote_peers) { - peer_processes.insert(block_peers.begin(), block_peers.end()); - } - int process_rank = playout->options->process_rank(); - int process_world_size = playout->options->process_world_size(); - pool.prepare_peer_buffers(peer_processes, process_rank, - process_world_size); - if (multi_process) { - exchange_fused_buffers_ucx(pool, peer_processes, playout->comm); - } - auto i64 = torch::dtype(torch::kInt64); - auto ptr_tensor = [&](std::vector const& ts) { - std::vector ptrs(ts.size()); - for (size_t p = 0; p < ts.size(); ++p) - ptrs[p] = reinterpret_cast(ts[p].data_ptr()); - return torch::tensor(ptrs, i64).to(w.device()); - }; - batch.flux2_dev = torch::tensor(batch.flux2_ptr, i64).to(w.device()); - batch.flux3_dev = torch::tensor(batch.flux3_ptr, i64).to(w.device()); - batch.x2v_dev = ptr_tensor(batch.x2v); - batch.x2f_dev = ptr_tensor(batch.x2f); - batch.x3v_dev = ptr_tensor(batch.x3v); - batch.x3f_dev = ptr_tensor(batch.x3f); - batch.faces_dev = - torch::tensor(batch.faces, torch::dtype(torch::kInt32)) - .to(w.device()); - batch.side_meta_all = torch::stack(batch.side_meta, 0); - auto as_pp = [](torch::Tensor& t) { - return reinterpret_cast(t.data_ptr()); - }; - fused_cubed_sphere_flux_all_cuda( - pool.buffer, - FusedCubedSphereFluxAllPtrs{ - as_pp(batch.flux2_dev), as_pp(batch.flux3_dev), - pool.buffer_ptrs_dev(), as_pp(batch.x2v_dev), - as_pp(batch.x2f_dev), as_pp(batch.x3v_dev), - as_pp(batch.x3f_dev)}, - FusedCubedSphereFluxAllParams{ - batch.side_meta_all, batch.faces_dev, static_cast(nvar), - static_cast(w.size(1)), static_cast(w.size(2)), - static_cast(nc1), bpp, w.scalar_type(), w.device(), - make_physics_params(recon23_prim, recon23_vel, - precon23->pinterp1->options->scale())}); - } - barrier.wait(); // seam flux enqueued before any panel's divergence - } - } - _div.set_(pmb->pcoord->forward(w, _flux1, _flux2, _flux3, _face_pressure1)); - - auto du = torch::zeros_like(_div); - auto interior = pmb->part({0, 0, 0}, PartOptions().exterior(false)); - du.index(interior) = -dt * _div.index(interior); - - auto temp = peos->compute("W->T", {w}); - for (auto& f : forcings) f.forward(du, w, temp, dt); - - if (options->grav() && (options->grav()->non_hydrostatic() < 1.)) { - du[IVX] += dt * rho_grav * (1. - options->grav()->non_hydrostatic()); - du[IPR] += - dt * w[IVX] * rho_grav * (1. - options->grav()->non_hydrostatic()); - } - _apply_implicit_correction(du, w, dt, other); - return du; -} - -} // namespace snap diff --git a/src/hydro/hydro_options.cpp b/src/hydro/hydro_options.cpp index 84b4d310..6cef9e3f 100644 --- a/src/hydro/hydro_options.cpp +++ b/src/hydro/hydro_options.cpp @@ -1,6 +1,4 @@ // C/C++ -#include -#include #include // yaml @@ -13,73 +11,6 @@ #include "hydro.hpp" namespace snap { -namespace { - -bool fused_recon_riemann_supported_by_options(HydroOptions const& op, - YAML::Node const& config) { - (void)config; - auto eos_type = op->eos() ? op->eos()->type() : ""; - auto riemann_type = op->riemann() ? op->riemann()->type() : ""; - auto recon1_type = op->recon1() && op->recon1()->interp() - ? op->recon1()->interp()->type() - : ""; - auto recon23_type = op->recon23() && op->recon23()->interp() - ? op->recon23()->interp()->type() - : ""; - bool eos_supported = eos_type == "ideal-gas" || eos_type == "ideal-moist" || - eos_type == "shallow-water"; - bool riemann_supported = riemann_type == "lmars" || riemann_type == "hllc" || - riemann_type == "shallow-roe"; - bool recon_supported = (recon1_type == "cp3" || recon1_type == "cp5" || - recon1_type == "weno3" || recon1_type == "weno5") && - (recon23_type == "cp3" || recon23_type == "cp5" || - recon23_type == "weno3" || recon23_type == "weno5"); - bool combo_supported = - ((eos_type == "ideal-gas" || eos_type == "ideal-moist") && - (riemann_type == "lmars" || riemann_type == "hllc")) || - (eos_type == "shallow-water" && riemann_type == "shallow-roe"); - // Cubed-sphere is now supported for any blocks_per_process: each process may - // own several panels co-resident on one GPU, exchanged through a shared - // per-process packed buffer sliced by local block (see - // hydro_forward_fused). - return eos_supported && riemann_supported && recon_supported && - combo_supported; -} - -bool fused_env_enabled(bool supported) { - auto const* value = std::getenv("FUSED"); - if (value == nullptr || std::string(value).empty()) return supported; - - std::string flag(value); - for (auto& c : flag) { - c = static_cast(std::tolower(static_cast(c))); - } - - if (flag == "auto") return supported; - if (flag == "off") return false; - if (flag == "on") { - TORCH_CHECK(supported, - "FUSED requests dynamics.fused-recon-riemann, but this " - "configuration is not supported by the fused path"); - return true; - } - - TORCH_CHECK(false, "FUSED must be one of AUTO, ON, or OFF, but got ", value); -} - -std::string fused_env_mode() { - auto const* value = std::getenv("FUSED"); - if (value == nullptr || std::string(value).empty()) return "AUTO"; - - std::string flag(value); - for (auto& c : flag) { - c = static_cast(std::toupper(static_cast(c))); - } - return flag; -} - -} // namespace - HydroOptions HydroOptionsImpl::from_yaml(std::string const& filename, bool verbose) { auto op = HydroOptionsImpl::create(); @@ -115,13 +46,6 @@ HydroOptions HydroOptionsImpl::from_yaml(std::string const& filename, op->disable_flux_x1() = dyn["disable-flux-x1"].as(false); op->disable_flux_x2() = dyn["disable-flux-x2"].as(false); op->disable_flux_x3() = dyn["disable-flux-x3"].as(false); - bool fused_supported = fused_recon_riemann_supported_by_options(op, config); - op->fused_recon_riemann() = fused_env_enabled(fused_supported); - SINFO(HydroOptions) << "FUSED=" << fused_env_mode() - << " fused-recon-riemann=" - << (op->fused_recon_riemann() ? "ON" : "OFF") - << " supported=" << (fused_supported ? "true" : "false") - << "\n"; } // --------------- forcings --------------- // @@ -184,8 +108,6 @@ HydroOptions HydroOptionsImpl::clone() const { op->disable_flux_x1() = disable_flux_x1(); op->disable_flux_x2() = disable_flux_x2(); op->disable_flux_x3() = disable_flux_x3(); - op->fused_recon_riemann() = fused_recon_riemann(); - if (grav()) op->grav() = grav()->clone(); if (coriolis()) op->coriolis() = coriolis()->clone(); if (diffusion()) op->diffusion() = diffusion()->clone(); diff --git a/src/recon/interp_impl.cuh b/src/recon/interp_impl.cuh index f7d2306f..134ce900 100644 --- a/src/recon/interp_impl.cuh +++ b/src/recon/interp_impl.cuh @@ -104,56 +104,6 @@ __device__ T interp_shared_weno5_coeff_impl(T const *line, T const *coeff, (alpha0 + alpha1 + alpha2); } -template -__device__ T interp_shared_fused_impl(T const *line, int v, int start, - int axis_size, FusedReconScheme scheme, - bool right, bool scale) { - if (scheme == FusedReconScheme::CP3) { - constexpr T cm[3] = {1. / 3., 5. / 6., -1. / 6.}; - T c[3]; - for (int k = 0; k < 3; ++k) - c[k] = right ? cm[2 - k] : cm[k]; - return interp_shared_poly_coeff_impl(line, c, v, start, axis_size); - } - if (scheme == FusedReconScheme::CP5) { - constexpr T cm[5] = {-1. / 20., 9. / 20., 47. / 60., -13. / 60., 1. / 30.}; - T c[5]; - for (int k = 0; k < 5; ++k) - c[k] = right ? cm[4 - k] : cm[k]; - return interp_shared_poly_coeff_impl(line, c, v, start, axis_size); - } - if (scheme == FusedReconScheme::WENO3) { - constexpr T cm[4][3] = {{1. / 2., 1. / 2., 0.}, - {0., 3. / 2., -1. / 2.}, - {1., -1., 0.}, - {0., 1., -1.}}; - T c[12]; - for (int r = 0; r < 4; ++r) { - for (int k = 0; k < 3; ++k) { - c[r * 3 + k] = right ? cm[r][2 - k] : cm[r][k]; - } - } - return interp_shared_weno3_coeff_impl(line, c, v, start, axis_size, scale); - } - - constexpr T cm[9][5] = {{-1. / 6., 5. / 6., 1. / 3., 0., 0.}, - {0., 1. / 3., 5. / 6., -1. / 6., 0.}, - {0., 0., 11. / 6., -7. / 6., 1. / 3.}, - {1., -2., 1., 0., 0.}, - {1., -4., 3., 0., 0.}, - {0., 1., -2., 1., 0.}, - {0., -1., 0., 1., 0.}, - {0., 0., 1., -2., 1.}, - {0., 0., 3., -4., 1.}}; - T c[45]; - for (int r = 0; r < 9; ++r) { - for (int k = 0; k < 5; ++k) { - c[r * 5 + k] = right ? cm[r][4 - k] : cm[r][k]; - } - } - return interp_shared_weno5_coeff_impl(line, c, v, start, axis_size, scale); -} - // polynomial template __device__ void interp_poly_impl(T *out, T *inp, T *coeff, int nvar, diff --git a/src/riemann/riemann_dispatch.cu b/src/riemann/riemann_dispatch.cu index be2a0777..a8ca747d 100644 --- a/src/riemann/riemann_dispatch.cu +++ b/src/riemann/riemann_dispatch.cu @@ -5,14 +5,16 @@ #include // snap -#include -#include "lmars_impl.h" #include "hllc_impl.h" +#include "lmars_impl.h" #include "riemann_dispatch.hpp" +#include "roe_impl.h" +#include "shallow_roe_impl.h" +#include namespace snap { -void call_lmars_cuda(at::TensorIterator& iter, int dim) { +void call_lmars_cuda(at::TensorIterator &iter, int dim) { at::cuda::CUDAGuard device_guard(iter.device()); AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "call_lmars_cuda", [&]() { @@ -20,22 +22,21 @@ void call_lmars_cuda(at::TensorIterator& iter, int dim) { auto stride = at::native::ensure_nonempty_stride(iter.output(0), 0); auto ny = nhydro - ICY; - native::gpu_kernel<6>( - iter, [=] GPU_LAMBDA(char* const data[6], unsigned int strides[6]) { - auto out = reinterpret_cast(data[0] + strides[0]); - auto face_pressure = - reinterpret_cast(data[1] + strides[1]); - auto wl = reinterpret_cast(data[2] + strides[2]); - auto wr = reinterpret_cast(data[3] + strides[3]); - auto elr = reinterpret_cast(data[4] + strides[4]); - auto glr = reinterpret_cast(data[5] + strides[5]); - lmars_impl(out, wl, wr, *elr, *(elr + stride), *glr, - *(glr + stride), dim, ny, stride, stride, face_pressure); - }); + native::gpu_kernel<6>(iter, [=] GPU_LAMBDA(char *const data[6], + unsigned int strides[6]) { + auto out = reinterpret_cast(data[0] + strides[0]); + auto face_pressure = reinterpret_cast(data[1] + strides[1]); + auto wl = reinterpret_cast(data[2] + strides[2]); + auto wr = reinterpret_cast(data[3] + strides[3]); + auto elr = reinterpret_cast(data[4] + strides[4]); + auto glr = reinterpret_cast(data[5] + strides[5]); + lmars_impl(out, wl, wr, *elr, *(elr + stride), *glr, *(glr + stride), dim, + ny, stride, stride, face_pressure); + }); }); } -void call_hllc_cuda(at::TensorIterator& iter, int dim) { +void call_hllc_cuda(at::TensorIterator &iter, int dim) { at::cuda::CUDAGuard device_guard(iter.device()); AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "call_hllc_cuda", [&]() { @@ -43,27 +44,72 @@ void call_hllc_cuda(at::TensorIterator& iter, int dim) { auto stride = at::native::ensure_nonempty_stride(iter.output(0), 0); auto ny = nhydro - ICY; - native::gpu_kernel<7>( - iter, [=] GPU_LAMBDA(char* const data[7], unsigned int strides[7]) { - auto out = reinterpret_cast(data[0] + strides[0]); - auto face_pressure = - reinterpret_cast(data[1] + strides[1]); - auto wl = reinterpret_cast(data[2] + strides[2]); - auto wr = reinterpret_cast(data[3] + strides[3]); - auto elr = reinterpret_cast(data[4] + strides[4]); - auto glr = reinterpret_cast(data[5] + strides[5]); - auto clr = reinterpret_cast(data[6] + strides[6]); - hllc_impl(out, wl, wr, *elr, *(elr + stride), *glr, - *(glr + stride), *clr, *(clr + stride), dim, ny, stride, - stride, face_pressure); + native::gpu_kernel<7>(iter, [=] GPU_LAMBDA(char *const data[7], + unsigned int strides[7]) { + auto out = reinterpret_cast(data[0] + strides[0]); + auto face_pressure = reinterpret_cast(data[1] + strides[1]); + auto wl = reinterpret_cast(data[2] + strides[2]); + auto wr = reinterpret_cast(data[3] + strides[3]); + auto elr = reinterpret_cast(data[4] + strides[4]); + auto glr = reinterpret_cast(data[5] + strides[5]); + auto clr = reinterpret_cast(data[6] + strides[6]); + hllc_impl(out, wl, wr, *elr, *(elr + stride), *glr, *(glr + stride), *clr, + *(clr + stride), dim, ny, stride, stride, face_pressure); + }); + }); +} + +void call_roe_cuda(at::TensorIterator &iter, int dim, bool ideal_moist, + int nvapor, double gammad, + torch::Tensor const &inv_mu_ratio_m1, + torch::Tensor const &cv_ratio_m1, torch::Tensor const &u0) { + at::cuda::CUDAGuard device_guard(iter.device()); + + AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "call_roe_cuda", [&]() { + auto nhydro = at::native::ensure_nonempty_size(iter.output(0), 0); + auto stride = at::native::ensure_nonempty_stride(iter.output(0), 0); + auto ny = nhydro - ICY; + auto inv_mu = ideal_moist ? inv_mu_ratio_m1.data_ptr() : nullptr; + auto cv = ideal_moist ? cv_ratio_m1.data_ptr() : nullptr; + auto energy0 = ideal_moist ? u0.data_ptr() : nullptr; + + native::gpu_kernel<7>(iter, [=] GPU_LAMBDA(char *const data[7], + unsigned int strides[7]) { + auto out = reinterpret_cast(data[0] + strides[0]); + auto face_pressure = reinterpret_cast(data[1] + strides[1]); + auto wl = reinterpret_cast(data[2] + strides[2]); + auto wr = reinterpret_cast(data[3] + strides[3]); + auto elr = reinterpret_cast(data[4] + strides[4]); + auto glr = reinterpret_cast(data[5] + strides[5]); + auto clr = reinterpret_cast(data[6] + strides[6]); + roe_impl(out, wl, wr, *elr, *(elr + stride), *glr, *(glr + stride), *clr, + *(clr + stride), dim, ny, ideal_moist, nvapor, scalar_t(gammad), + inv_mu, cv, energy0, stride, stride, face_pressure); + }); + }); +} + +void call_shallow_roe_cuda(at::TensorIterator &iter, int dim, int dir_yz) { + at::cuda::CUDAGuard device_guard(iter.device()); + + AT_DISPATCH_FLOATING_TYPES(iter.common_dtype(), "call_shallow_roe_cuda", [&] { + auto stride = at::native::ensure_nonempty_stride(iter.output(0), 0); + native::gpu_kernel<3>( + iter, [=] GPU_LAMBDA(char *const data[3], unsigned int strides[3]) { + auto out = reinterpret_cast(data[0] + strides[0]); + auto wl = reinterpret_cast(data[1] + strides[1]); + auto wr = reinterpret_cast(data[2] + strides[2]); + shallow_roe_impl(out, wl, wr, dim, dir_yz, stride, stride); }); }); } -} // namespace snap +} // namespace snap namespace at::native { REGISTER_CUDA_DISPATCH(call_lmars, &snap::call_lmars_cuda); REGISTER_CUDA_DISPATCH(call_hllc, &snap::call_hllc_cuda); +REGISTER_CUDA_DISPATCH(call_roe, &snap::call_roe_cuda); +REGISTER_CUDA_DISPATCH(call_shallow_roe, &snap::call_shallow_roe_cuda); -} // namespace at::native +} // namespace at::native diff --git a/src/riemann/riemann_dispatch.hpp b/src/riemann/riemann_dispatch.hpp index 165caa0d..736b92b7 100644 --- a/src/riemann/riemann_dispatch.hpp +++ b/src/riemann/riemann_dispatch.hpp @@ -3,13 +3,23 @@ // torch #include #include +#include namespace at::native { -using rsolver_fn = void (*)(at::TensorIterator &iter, int dim); +using rsolver_fn = void (*)(at::TensorIterator& iter, int dim); DECLARE_DISPATCH(rsolver_fn, call_lmars); DECLARE_DISPATCH(rsolver_fn, call_hllc); -DECLARE_DISPATCH(rsolver_fn, call_roe); + +using roe_fn = void (*)(at::TensorIterator& iter, int dim, bool ideal_moist, + int nvapor, double gammad, + torch::Tensor const& inv_mu_ratio_m1, + torch::Tensor const& cv_ratio_m1, + torch::Tensor const& u0); +DECLARE_DISPATCH(roe_fn, call_roe); + +using shallow_roe_fn = void (*)(at::TensorIterator& iter, int dim, int dir_yz); +DECLARE_DISPATCH(shallow_roe_fn, call_shallow_roe); } // namespace at::native diff --git a/src/riemann/roe.cpp b/src/riemann/roe.cpp index 8fb72bbd..556c65e5 100644 --- a/src/riemann/roe.cpp +++ b/src/riemann/roe.cpp @@ -1,12 +1,12 @@ // snap -#include - #include #include "../eos/ideal_moist.hpp" +#include "riemann_dispatch.hpp" #include "riemann_solver.hpp" namespace snap { + void RoeSolverImpl::reset() { TORCH_CHECK(phydro, "[RoeSolver] parent is nullptr"); } @@ -15,202 +15,47 @@ torch::Tensor RoeSolverImpl::forward(torch::Tensor wl, torch::Tensor wr, int dim, torch::Tensor flx, torch::Tensor face_pressure) { auto peos = phydro->peos; - auto eos_type = peos->options->type(); - int nvar = wl.size(0); - int ny = std::max(nvar - ICY, 0); - - // dim, ivx, ivy, ivz - // 3, IVX, IVY, iVZ - // 2, IVX + 1, IVX + 2, IVX - // 1, IVX + 2, IVX, IVX + 1 - auto ivx = IPR - dim; - auto ivy = IVX + ((ivx - IVX) + 1) % 3; - auto ivz = IVX + ((ivx - IVX) + 2) % 3; - - auto ul = peos->compute("W->U", {wl}); - auto ur = peos->compute("W->U", {wr}); - auto el = ul[IPR]; - auto gammal = peos->compute("W->A", {wl}); - auto cl = peos->compute("WA->L", {wl, gammal}); - - auto er = ur[IPR]; - auto gammar = peos->compute("W->A", {wr}); - auto cr = peos->compute("WA->L", {wr, gammar}); - - auto scalar_view = wl.sizes().vec(); - scalar_view[0] = 1; - - auto sqrtdl = torch::sqrt(wl[IDN]); - auto sqrtdr = torch::sqrt(wr[IDN]); - auto isdlpdr = 1.0 / (sqrtdl + sqrtdr); - auto rhobar = sqrtdl * sqrtdr; - - //--- Step 2. Compute Roe-averaged data from left- and right-states - auto wroe = torch::zeros_like(wl); - - wroe[IDN] = sqrtdl * sqrtdr; - wroe.narrow(0, IVX, 3) = - (sqrtdl * wl.narrow(0, IVX, 3) + sqrtdr * wr.narrow(0, IVX, 3)) * isdlpdr; - - // Following Roe(1981), the enthalpy H=(E+P)/d is averaged for adiabatic - // flows, rather than E or P directly. sqrtdl*hl = sqrtdl*(el+pl)/dl = - // (el+pl)/sqrtdl - auto hl = (el + wl[IPR]) / wl[IDN]; - auto hr = (er + wr[IPR]) / wr[IDN]; - wroe[IPR] = (hl * sqrtdl + hr * sqrtdr) * isdlpdr; - if (face_pressure.defined()) { - face_pressure.copy_(wroe[IPR]); - } - - //--- Step 3. Compute L/R fluxes - auto fl = torch::zeros_like(wl); - auto fr = torch::zeros_like(wr); - - fl[IDN] = ul[IDN] * wl[ivx]; - fr[IDN] = ur[IDN] * wr[ivx]; - if (ny > 0) { - fl.narrow(0, ICY, ny) = ul.narrow(0, ICY, ny) * wl[ivx].view(scalar_view); - fr.narrow(0, ICY, ny) = ur.narrow(0, ICY, ny) * wr[ivx].view(scalar_view); - } - - fl.narrow(0, IVX, 3) = wl[IDN] * wl[ivx] * wl.narrow(0, IVX, 3); - fr.narrow(0, IVX, 3) = wr[IDN] * wr[ivx] * wr.narrow(0, IVX, 3); - - fl[ivx] += wl[IPR]; - fr[ivx] += wr[IPR]; - - fl[IPR] = (el + wl[IPR]) * wl[ivx]; - fr[IPR] = (er + wr[IPR]) * wr[ivx]; - - //--- Step 4. Compute Roe fluxes. - auto du = ur - ul; - - auto out = 0.5 * (fl + fr); - - auto vsq = wroe.narrow(0, IVX, 3).square().sum(0); - auto gamma_roe = 0.5 * (gammal + gammar); - auto offset = torch::zeros_like(wroe[IPR]); - auto qbar_dry = (ul[IDN] / sqrtdl + ur[IDN] / sqrtdr) * isdlpdr; - auto qbar = torch::Tensor(); - auto alpha_species = torch::Tensor(); - - if (ny > 0) { - auto sqrtdl_v = sqrtdl.view(scalar_view); - auto sqrtdr_v = sqrtdr.view(scalar_view); - auto isdlpdr_v = isdlpdr.view(scalar_view); - qbar = - (ul.narrow(0, ICY, ny) / sqrtdl_v + ur.narrow(0, ICY, ny) / sqrtdr_v) * - isdlpdr_v; - } - - if (eos_type == "ideal-moist") { + auto ideal_moist = peos->options->type() == "ideal-moist"; + int nvapor = 0; + torch::Tensor inv_mu_ratio_m1, cv_ratio_m1, u0; + if (ideal_moist) { auto moist = dynamic_cast(peos.get()); TORCH_CHECK(moist != nullptr, "[RoeSolver] ideal-moist EOS cast failed"); - int nvapor = moist->pthermo->options->vapor_ids().size() - 1; - if (ny > 0) { - auto feps = torch::ones_like(wroe[IPR]); - if (nvapor > 0) { - feps += (qbar.narrow(0, 0, nvapor) * - moist->inv_mu_ratio_m1.narrow(0, 0, nvapor) - .to(qbar) - .view({nvapor, 1, 1, 1})) - .sum(0); - } - if (ny > nvapor) { - feps -= qbar.narrow(0, nvapor, ny - nvapor).sum(0); - } - auto fsig = - torch::ones_like(wroe[IPR]) + - (qbar * moist->cv_ratio_m1.to(qbar).view({ny, 1, 1, 1})).sum(0); - gamma_roe = 1.0 + (moist->options->gammad() - 1.0) * feps / fsig; - offset = qbar_dry * moist->u0[0].to(wl); - offset += (qbar * moist->u0.narrow(0, 1, ny).to(qbar).view({ny, 1, 1, 1})) - .sum(0); - } - } - - auto q = wroe[IPR] - 0.5 * vsq - offset; - auto cs_sq = torch::clamp_min((gamma_roe - 1.0) * q, 1.0e-10); - auto cs = torch::sqrt(cs_sq); - auto lam_m = wroe[ivx] - cs; - auto lam_0 = wroe[ivx]; - auto lam_p = wroe[ivx] + cs; - auto spd_m = torch::abs(lam_m); - auto spd_0 = torch::abs(lam_0); - auto spd_p = torch::abs(lam_p); - - auto duv = wr[ivx] - wl[ivx]; - auto dv = wr[ivy] - wl[ivy]; - auto dw = wr[ivz] - wl[ivz]; - auto dp = wr[IPR] - wl[IPR]; - - auto alpha_m = -0.5 * rhobar / cs * duv + 0.5 * dp / cs_sq; - auto alpha_p = +0.5 * rhobar / cs * duv + 0.5 * dp / cs_sq; - auto alpha_v = rhobar * dv; - auto alpha_w = rhobar * dw; - - auto alpha_dry = ur[IDN] - ul[IDN] - dp / cs_sq * qbar_dry; - auto alpha0 = alpha_dry.clone(); - auto llf_flag = - torch::logical_or(ul[IDN] + alpha_m * qbar_dry < 0.0, - ul[IDN] + alpha_m * qbar_dry + alpha_dry < 0.0); - - out[IDN] -= 0.5 * (spd_m * alpha_m * qbar_dry + spd_0 * alpha_dry + - spd_p * alpha_p * qbar_dry); - - if (ny > 0) { - auto sq_view = cs_sq.view(scalar_view); - auto spd_m_v = spd_m.view(scalar_view); - auto spd_0_v = spd_0.view(scalar_view); - auto spd_p_v = spd_p.view(scalar_view); - auto alpha_m_v = alpha_m.view(scalar_view); - auto alpha_p_v = alpha_p.view(scalar_view); - auto dp_v = dp.view(scalar_view); - alpha_species = - ur.narrow(0, ICY, ny) - ul.narrow(0, ICY, ny) - dp_v * qbar / sq_view; - alpha0 += alpha_species.sum(0); - - out.narrow(0, ICY, ny) -= - 0.5 * (spd_m_v * alpha_m_v * qbar + spd_0_v * alpha_species + - spd_p_v * alpha_p_v * qbar); - - auto species_first = ul.narrow(0, ICY, ny) + alpha_m_v * qbar; - auto species_second = species_first + alpha_species; - llf_flag = torch::logical_or( - llf_flag, torch::logical_or(torch::any(species_first < 0.0, 0), - torch::any(species_second < 0.0, 0))); + nvapor = moist->pthermo->options->vapor_ids().size() - 1; + inv_mu_ratio_m1 = moist->inv_mu_ratio_m1.to(wl); + cv_ratio_m1 = moist->cv_ratio_m1.to(wl); + u0 = moist->u0.to(wl); } - out[ivx] -= - 0.5 * (spd_m * alpha_m * (wroe[ivx] - cs) + spd_0 * (wroe[ivx] * alpha0) + - spd_p * alpha_p * (wroe[ivx] + cs)); - out[ivy] -= 0.5 * (spd_m * alpha_m * wroe[ivy] + - spd_0 * (wroe[ivy] * alpha0 + alpha_v) + - spd_p * alpha_p * wroe[ivy]); - out[ivz] -= 0.5 * (spd_m * alpha_m * wroe[ivz] + - spd_0 * (wroe[ivz] * alpha0 + alpha_w) + - spd_p * alpha_p * wroe[ivz]); - out[IPR] -= 0.5 * (spd_m * alpha_m * (wroe[IPR] - wroe[ivx] * cs) + - spd_0 * (rhobar * (hr - hl - wroe[ivx] * duv) + - alpha0 * wroe[IPR] - dp) + - spd_p * alpha_p * (wroe[IPR] + wroe[ivx] * cs)); - - //--- Step 5. Overwrite with upwind flux if flow is supersonic - auto evi = (lam_m > 0).to(out.scalar_type()).view(scalar_view); - out = evi * fl + (1 - evi) * out; - - evi = (lam_p < 0).to(out.scalar_type()).view(scalar_view); - out = evi * fr + (1 - evi) * out; - - //--- Step 6. Overwrite with LLF flux if any of intermediate states are - // negative - auto cmax = - 0.5 * torch::max(torch::abs(wl[ivx]) + cl, torch::abs(wr[ivx]) + cr); - auto llf_mask = llf_flag.to(out.scalar_type()).view(scalar_view); - out = llf_mask * (0.5 * (fl + fr) - cmax.view(scalar_view) * du) + - (1 - llf_mask) * out; + auto eil = peos->compute("W->I", {wl}); + auto eir = peos->compute("W->I", {wr}); + auto gammal = peos->compute("W->A", {wl}); + auto gammar = peos->compute("W->A", {wr}); + auto cl = peos->compute("WA->L", {wl, gammal}); + auto cr = peos->compute("WA->L", {wr, gammar}); - flx.copy_(out); + auto face_pressure_out = + face_pressure.defined() ? face_pressure : torch::empty_like(wl[IDN]); + auto elr = torch::stack({eil, eir}); + auto glr = torch::stack({gammal, gammar}); + auto clr = torch::stack({cl, cr}); + auto iter = at::TensorIteratorConfig() + .resize_outputs(false) + .check_all_same_dtype(true) + .declare_static_shape(flx.sizes(), /*squash_dims=*/0) + .add_output(flx) + .add_owned_output(face_pressure_out.unsqueeze(0)) + .add_input(wl) + .add_input(wr) + .add_input(elr) + .add_input(glr) + .add_input(clr) + .build(); + + at::native::call_roe(flx.device().type(), iter, dim, ideal_moist, nvapor, + peos->options->gammad(), inv_mu_ratio_m1, cv_ratio_m1, + u0); return flx; } + } // namespace snap diff --git a/src/riemann/roe_dispatch.cpp b/src/riemann/roe_dispatch.cpp new file mode 100644 index 00000000..9457d99d --- /dev/null +++ b/src/riemann/roe_dispatch.cpp @@ -0,0 +1,229 @@ +// torch +#include +#include +#include +#include + +// snap +#include + +#include "riemann_dispatch.hpp" +#include "roe_impl.h" + +namespace snap { + +void call_roe_cpu(at::TensorIterator& iter, int dim, bool ideal_moist, + int nvapor, double gammad, + torch::Tensor const& inv_mu_ratio_m1, + torch::Tensor const& cv_ratio_m1, torch::Tensor const& u0) { + int grain_size = iter.numel() / at::get_num_threads(); + + AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "call_roe_cpu", [&] { + auto nhydro = at::native::ensure_nonempty_size(iter.output(0), 0); + auto stride = at::native::ensure_nonempty_stride(iter.output(0), 0); + auto ny = nhydro - ICY; + auto inv_mu = ideal_moist ? inv_mu_ratio_m1.data_ptr() : nullptr; + auto cv = ideal_moist ? cv_ratio_m1.data_ptr() : nullptr; + auto energy0 = ideal_moist ? u0.data_ptr() : nullptr; + + iter.for_each( + [&](char** data, const int64_t* strides, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + auto out = reinterpret_cast(data[0] + i * strides[0]); + auto face_pressure = + reinterpret_cast(data[1] + i * strides[1]); + auto wl = reinterpret_cast(data[2] + i * strides[2]); + auto wr = reinterpret_cast(data[3] + i * strides[3]); + auto elr = reinterpret_cast(data[4] + i * strides[4]); + auto glr = reinterpret_cast(data[5] + i * strides[5]); + auto clr = reinterpret_cast(data[6] + i * strides[6]); + roe_impl(out, wl, wr, *elr, *(elr + stride), *glr, *(glr + stride), + *clr, *(clr + stride), dim, ny, ideal_moist, nvapor, + scalar_t(gammad), inv_mu, cv, energy0, stride, stride, + face_pressure); + } + }, + grain_size); + }); +} + +void call_roe_mps(at::TensorIterator& iter, int dim, bool ideal_moist, + int nvapor, double gammad, + torch::Tensor const& inv_mu_ratio_m1, + torch::Tensor const& cv_ratio_m1, torch::Tensor const& u0) { + auto flx = iter.output(0); + auto face_pressure = iter.output(1).squeeze(0); + auto wl = iter.input(0); + auto wr = iter.input(1); + auto elr = iter.input(2); + auto glr = iter.input(3); + auto clr = iter.input(4); + + int ny = wl.size(0) - ICY; + auto ivx = IPR - dim; + auto ivy = IVX + ((ivx - IVX) + 1) % 3; + auto ivz = IVX + ((ivx - IVX) + 2) % 3; + auto scalar_view = wl.sizes().vec(); + scalar_view[0] = 1; + + auto sqrtdl = torch::sqrt(wl[IDN]); + auto sqrtdr = torch::sqrt(wr[IDN]); + auto isdlpdr = 1.0 / (sqrtdl + sqrtdr); + auto rhobar = sqrtdl * sqrtdr; + + auto wroe = torch::zeros_like(wl); + wroe[IDN] = rhobar; + wroe.narrow(0, IVX, 3) = + (sqrtdl * wl.narrow(0, IVX, 3) + sqrtdr * wr.narrow(0, IVX, 3)) * isdlpdr; + + auto el = elr[ILT] + 0.5 * wl[IDN] * wl.narrow(0, IVX, 3).square().sum(0); + auto er = elr[IRT] + 0.5 * wr[IDN] * wr.narrow(0, IVX, 3).square().sum(0); + auto hl = (el + wl[IPR]) / wl[IDN]; + auto hr = (er + wr[IPR]) / wr[IDN]; + wroe[IPR] = (hl * sqrtdl + hr * sqrtdr) * isdlpdr; + + auto ul = torch::zeros_like(wl); + auto ur = torch::zeros_like(wr); + auto dryl = torch::ones_like(wl[IDN]); + auto dryr = torch::ones_like(wr[IDN]); + if (ny > 0) { + dryl -= wl.narrow(0, ICY, ny).sum(0); + dryr -= wr.narrow(0, ICY, ny).sum(0); + ul.narrow(0, ICY, ny) = wl[IDN] * wl.narrow(0, ICY, ny); + ur.narrow(0, ICY, ny) = wr[IDN] * wr.narrow(0, ICY, ny); + } + ul[IDN] = wl[IDN] * dryl; + ur[IDN] = wr[IDN] * dryr; + ul.narrow(0, IVX, 3) = wl[IDN] * wl.narrow(0, IVX, 3); + ur.narrow(0, IVX, 3) = wr[IDN] * wr.narrow(0, IVX, 3); + ul[IPR] = el; + ur[IPR] = er; + + auto fl = torch::zeros_like(wl); + auto fr = torch::zeros_like(wr); + fl[IDN] = ul[IDN] * wl[ivx]; + fr[IDN] = ur[IDN] * wr[ivx]; + if (ny > 0) { + fl.narrow(0, ICY, ny) = ul.narrow(0, ICY, ny) * wl[ivx].view(scalar_view); + fr.narrow(0, ICY, ny) = ur.narrow(0, ICY, ny) * wr[ivx].view(scalar_view); + } + fl.narrow(0, IVX, 3) = wl[IDN] * wl[ivx] * wl.narrow(0, IVX, 3); + fr.narrow(0, IVX, 3) = wr[IDN] * wr[ivx] * wr.narrow(0, IVX, 3); + fl[ivx] += wl[IPR]; + fr[ivx] += wr[IPR]; + fl[IPR] = (el + wl[IPR]) * wl[ivx]; + fr[IPR] = (er + wr[IPR]) * wr[ivx]; + + auto du = ur - ul; + auto out = 0.5 * (fl + fr); + auto vsq = wroe.narrow(0, IVX, 3).square().sum(0); + auto gamma_roe = 0.5 * (glr[ILT] + glr[IRT]); + auto offset = torch::zeros_like(wroe[IPR]); + auto qbar_dry = (ul[IDN] / sqrtdl + ur[IDN] / sqrtdr) * isdlpdr; + auto qbar = torch::Tensor(); + auto alpha_species = torch::Tensor(); + if (ny > 0) { + qbar = (ul.narrow(0, ICY, ny) / sqrtdl.view(scalar_view) + + ur.narrow(0, ICY, ny) / sqrtdr.view(scalar_view)) * + isdlpdr.view(scalar_view); + } + + if (ideal_moist) { + auto feps = torch::ones_like(wroe[IPR]); + if (nvapor > 0) { + feps += (qbar.narrow(0, 0, nvapor) * + inv_mu_ratio_m1.narrow(0, 0, nvapor).view({nvapor, 1, 1, 1})) + .sum(0); + } + if (ny > nvapor) { + feps -= qbar.narrow(0, nvapor, ny - nvapor).sum(0); + } + auto fsig = torch::ones_like(wroe[IPR]) + + (qbar * cv_ratio_m1.view({ny, 1, 1, 1})).sum(0); + gamma_roe = 1.0 + (gammad - 1.0) * feps / fsig; + offset = qbar_dry * u0[0]; + if (ny > 0) { + offset += (qbar * u0.narrow(0, 1, ny).view({ny, 1, 1, 1})).sum(0); + } + } + + auto q = wroe[IPR] - 0.5 * vsq - offset; + auto cs_sq = torch::clamp_min((gamma_roe - 1.0) * q, 1.0e-10); + auto cs = torch::sqrt(cs_sq); + face_pressure.copy_(0.5 * + (wl[IPR] + wr[IPR] + rhobar * cs * (wl[ivx] - wr[ivx]))); + + auto lam_m = wroe[ivx] - cs; + auto lam_0 = wroe[ivx]; + auto lam_p = wroe[ivx] + cs; + auto spd_m = torch::abs(lam_m); + auto spd_0 = torch::abs(lam_0); + auto spd_p = torch::abs(lam_p); + auto duv = wr[ivx] - wl[ivx]; + auto dv = wr[ivy] - wl[ivy]; + auto dw = wr[ivz] - wl[ivz]; + auto dp = wr[IPR] - wl[IPR]; + auto alpha_m = -0.5 * rhobar / cs * duv + 0.5 * dp / cs_sq; + auto alpha_p = 0.5 * rhobar / cs * duv + 0.5 * dp / cs_sq; + auto alpha_v = rhobar * dv; + auto alpha_w = rhobar * dw; + auto alpha_dry = ur[IDN] - ul[IDN] - dp / cs_sq * qbar_dry; + auto alpha0 = alpha_dry.clone(); + auto llf_flag = + torch::logical_or(ul[IDN] + alpha_m * qbar_dry < 0.0, + ul[IDN] + alpha_m * qbar_dry + alpha_dry < 0.0); + + out[IDN] -= 0.5 * (spd_m * alpha_m * qbar_dry + spd_0 * alpha_dry + + spd_p * alpha_p * qbar_dry); + if (ny > 0) { + alpha_species = ur.narrow(0, ICY, ny) - ul.narrow(0, ICY, ny) - + dp.view(scalar_view) * qbar / cs_sq.view(scalar_view); + alpha0 += alpha_species.sum(0); + out.narrow(0, ICY, ny) -= + 0.5 * (spd_m.view(scalar_view) * alpha_m.view(scalar_view) * qbar + + spd_0.view(scalar_view) * alpha_species + + spd_p.view(scalar_view) * alpha_p.view(scalar_view) * qbar); + auto species_first = + ul.narrow(0, ICY, ny) + alpha_m.view(scalar_view) * qbar; + auto species_second = species_first + alpha_species; + llf_flag = torch::logical_or( + llf_flag, torch::logical_or(torch::any(species_first < 0.0, 0), + torch::any(species_second < 0.0, 0))); + } + + out[ivx] -= + 0.5 * (spd_m * alpha_m * (wroe[ivx] - cs) + spd_0 * wroe[ivx] * alpha0 + + spd_p * alpha_p * (wroe[ivx] + cs)); + out[ivy] -= 0.5 * (spd_m * alpha_m * wroe[ivy] + + spd_0 * (wroe[ivy] * alpha0 + alpha_v) + + spd_p * alpha_p * wroe[ivy]); + out[ivz] -= 0.5 * (spd_m * alpha_m * wroe[ivz] + + spd_0 * (wroe[ivz] * alpha0 + alpha_w) + + spd_p * alpha_p * wroe[ivz]); + out[IPR] -= 0.5 * (spd_m * alpha_m * (wroe[IPR] - wroe[ivx] * cs) + + spd_0 * (rhobar * (hr - hl - wroe[ivx] * duv) + + alpha0 * wroe[IPR] - dp) + + spd_p * alpha_p * (wroe[IPR] + wroe[ivx] * cs)); + + auto left_upwind = (lam_m > 0).to(out.scalar_type()).view(scalar_view); + out = left_upwind * fl + (1 - left_upwind) * out; + auto right_upwind = (lam_p < 0).to(out.scalar_type()).view(scalar_view); + out = right_upwind * fr + (1 - right_upwind) * out; + + auto cmax = 0.5 * torch::max(torch::abs(wl[ivx]) + clr[ILT], + torch::abs(wr[ivx]) + clr[IRT]); + auto llf_mask = llf_flag.to(out.scalar_type()).view(scalar_view); + out = llf_mask * (0.5 * (fl + fr) - cmax.view(scalar_view) * du) + + (1 - llf_mask) * out; + flx.copy_(out); +} + +} // namespace snap + +namespace at::native { + +DEFINE_DISPATCH(call_roe); +REGISTER_ALL_CPU_DISPATCH(call_roe, &snap::call_roe_cpu); +REGISTER_MPS_DISPATCH(call_roe, &snap::call_roe_mps); + +} // namespace at::native diff --git a/src/riemann/roe_impl.h b/src/riemann/roe_impl.h index 5395db35..d20be5ed 100644 --- a/src/riemann/roe_impl.h +++ b/src/riemann/roe_impl.h @@ -15,10 +15,15 @@ namespace snap { +template +inline DISPATCH_MACRO T roe_max(T a, T b) { + return a > b ? a : b; +} + template void DISPATCH_MACRO roe_impl(T* flx, T const* wl, T const* wr, T el, T er, T gammal, T gammar, T cl, T cr, int dim, int ny, - FusedEos eos, int nvapor, T gammad, + bool ideal_moist, int nvapor, T gammad, T const* inv_mu_ratio_m1, T const* cv_ratio_m1, T const* u0, int stride_w, int stride_f, T* face_pressure = nullptr) { @@ -42,7 +47,6 @@ void DISPATCH_MACRO roe_impl(T* flx, T const* wl, T const* wr, T el, T er, T hl = (etl + WL(IPR)) / WL(IDN); T hr = (etr + WR(IPR)) / WR(IDN); T h = (hl * sqrtdl + hr * sqrtdr) * isdlpdr; - if (face_pressure != nullptr) *face_pressure = h; T dryl = T(1); T dryr = T(1); @@ -79,7 +83,7 @@ void DISPATCH_MACRO roe_impl(T* flx, T const* wl, T const* wr, T el, T er, T gamma_roe = T(0.5) * (gammal + gammar); T offset = T(0); - if (eos == FusedEos::IdealMoist) { + if (ideal_moist) { T feps = ideal_moist_feps(qbar + 1, ny, nvapor, inv_mu_ratio_m1); T fsig = ideal_moist_fsig(qbar + 1, ny, cv_ratio_m1); gamma_roe = T(1) + (gammad - T(1)) * feps / fsig; @@ -110,8 +114,12 @@ void DISPATCH_MACRO roe_impl(T* flx, T const* wl, T const* wr, T el, T er, T vsq = v1 * v1 + v2 * v2 + v3 * v3; T gm1_roe = gamma_roe - T(1); T q = h - T(0.5) * vsq - offset; - T cs_sq = max(gm1_roe * max(q, TINY_NUMBER), TINY_NUMBER); + T cs_sq = roe_max(gm1_roe * roe_max(q, TINY_NUMBER), TINY_NUMBER); T cs = sqrt(cs_sq); + if (face_pressure != nullptr) { + *face_pressure = + T(0.5) * (WL(IPR) + WR(IPR) + rhobar * cs * (WL(ivx) - WR(ivx))); + } T lam_m = v1 - cs; T lam_0 = v1; @@ -187,7 +195,7 @@ void DISPATCH_MACRO roe_impl(T* flx, T const* wl, T const* wr, T el, T er, } if (llf_flag) { - T a = T(0.5) * max(abs(WL(ivx)) + cl, abs(WR(ivx)) + cr); + T a = T(0.5) * roe_max(abs(WL(ivx)) + cl, abs(WR(ivx)) + cr); FLX(IDN) = T(0.5) * (fl0 + fr0) - a * (rhor0 - rhol0); FLX(ivx) = T(0.5) * (fl1 + fr1) - a * du1; FLX(ivy) = T(0.5) * (fl2 + fr2) - a * du2; diff --git a/src/riemann/shallow_roe.cpp b/src/riemann/shallow_roe.cpp index 4b9469c3..c5a83013 100644 --- a/src/riemann/shallow_roe.cpp +++ b/src/riemann/shallow_roe.cpp @@ -6,6 +6,7 @@ #include // snap +#include "riemann_dispatch.hpp" #include "riemann_solver.hpp" namespace snap { @@ -27,19 +28,9 @@ torch::Tensor ShallowRoeSolverImpl::forward(torch::Tensor wl, torch::Tensor wr, torch::Tensor face_pressure) { TORCH_CHECK(!face_pressure.defined(), "Face-pressure output is not implemented for ShallowRoeSolver"); - auto peos = phydro->peos; auto pcoord = phydro->pmb->pcoord; - int ivx, ivy, ivz; - if (options->dir() == "xy") { - ivx = dim == 3 ? 1 : 2; - ivy = dim == 3 ? 2 : 1; - ivz = 3; - } else if (options->dir() == "yz") { - ivx = dim == 2 ? 2 : 3; - ivy = dim == 2 ? 3 : 2; - ivz = 1; - } else { + if (options->dir() != "xy" && options->dir() != "yz") { TORCH_CHECK(false, "ShallowRoeSolver takes options->dir() = 'xy' or 'yz'" " but got options->dir() = ", @@ -63,50 +54,17 @@ torch::Tensor ShallowRoeSolverImpl::forward(torch::Tensor wl, torch::Tensor wr, TORCH_CHECK(false, "Invalid dimension: ", dim); } - auto sqrtdl = torch::sqrt(wl[0]); - auto sqrtdr = torch::sqrt(wr[0]); - auto isdlpdr = 1.0 / (sqrtdl + sqrtdr); - - auto ubar = (wl[ivx] * sqrtdl + wr[ivx] * sqrtdr) * isdlpdr; - auto vbar = (wl[ivy] * sqrtdl + wr[ivy] * sqrtdr) * isdlpdr; - auto cbar = torch::sqrt(0.5 * (wl[0] + wr[0])); - - auto del = wr - wl; - auto hbar = torch::sqrt(wl[0] * wr[0]); - - auto a1 = 0.5 * (cbar * del[0] - hbar * del[ivx]) / cbar; - auto a2 = hbar * del[ivy]; - auto a3 = 0.5 * (cbar * del[0] + hbar * del[ivx]) / cbar; - - auto wave0 = torch::zeros_like(del); - auto wave1 = torch::zeros_like(del); - auto wave2 = torch::zeros_like(del); - - wave0[0] = a1; - wave0[ivx] = a1 * (ubar - cbar); - wave0[ivy] = a1 * vbar; - - wave1[0] = 0.; - wave1[ivx] = 0.; - wave1[ivy] = a2; - - wave2[0] = a3; - wave2[ivx] = a3 * (ubar + cbar); - wave2[ivy] = a3 * vbar; - - auto speed = torch::zeros_like(del); - - speed[0] = torch::abs(ubar - cbar); - speed[1] = torch::abs(ubar); - speed[2] = torch::abs(ubar + cbar); - - flx[0] = 0.5 * (wl[0] * wl[ivx] + wr[0] * wr[ivx]); - flx[ivx] = 0.5 * (wl[0] * wl[ivx].square() + 0.5 * wl[0].square() + - wr[0] * wr[ivx].square() + 0.5 * wr[0].square()); - flx[ivy] = 0.5 * (wl[0] * wl[ivx] * wl[ivy] + wr[0] * wr[ivx] * wr[ivy]); - flx[ivz] = 0.; - - flx -= 0.5 * (speed[0] * wave0 + speed[1] * wave1 + speed[2] * wave2); + auto iter = at::TensorIteratorConfig() + .resize_outputs(false) + .check_all_same_dtype(true) + .declare_static_shape(flx.sizes(), /*squash_dims=*/0) + .add_output(flx) + .add_input(wl) + .add_input(wr) + .build(); + + at::native::call_shallow_roe(flx.device().type(), iter, dim, + options->dir() == "yz"); switch (dim) { case 1: diff --git a/src/riemann/shallow_roe_dispatch.cpp b/src/riemann/shallow_roe_dispatch.cpp new file mode 100644 index 00000000..e05247a6 --- /dev/null +++ b/src/riemann/shallow_roe_dispatch.cpp @@ -0,0 +1,91 @@ +// torch +#include +#include +#include +#include + +// snap +#include + +#include "riemann_dispatch.hpp" +#include "shallow_roe_impl.h" + +namespace snap { + +void call_shallow_roe_cpu(at::TensorIterator& iter, int dim, int dir_yz) { + int grain_size = iter.numel() / at::get_num_threads(); + + AT_DISPATCH_FLOATING_TYPES(iter.dtype(), "call_shallow_roe_cpu", [&] { + auto stride = at::native::ensure_nonempty_stride(iter.output(0), 0); + iter.for_each( + [&](char** data, const int64_t* strides, int64_t n) { + for (int64_t i = 0; i < n; ++i) { + auto out = reinterpret_cast(data[0] + i * strides[0]); + auto wl = reinterpret_cast(data[1] + i * strides[1]); + auto wr = reinterpret_cast(data[2] + i * strides[2]); + shallow_roe_impl(out, wl, wr, dim, dir_yz, stride, stride); + } + }, + grain_size); + }); +} + +void call_shallow_roe_mps(at::TensorIterator& iter, int dim, int dir_yz) { + auto flx = iter.output(0); + auto wl = iter.input(0); + auto wr = iter.input(1); + + int ivx, ivy, ivz; + if (dir_yz) { + ivx = dim == 2 ? IVY : IVZ; + ivy = dim == 2 ? IVZ : IVY; + ivz = IVX; + } else { + ivx = dim == 3 ? IVX : IVY; + ivy = dim == 3 ? IVY : IVX; + ivz = IVZ; + } + + auto sqrtdl = torch::sqrt(wl[IDN]); + auto sqrtdr = torch::sqrt(wr[IDN]); + auto isdlpdr = 1.0 / (sqrtdl + sqrtdr); + auto ubar = (wl[ivx] * sqrtdl + wr[ivx] * sqrtdr) * isdlpdr; + auto vbar = (wl[ivy] * sqrtdl + wr[ivy] * sqrtdr) * isdlpdr; + auto cbar = torch::sqrt(0.5 * (wl[IDN] + wr[IDN])); + auto del = wr - wl; + auto hbar = torch::sqrt(wl[IDN] * wr[IDN]); + + auto a1 = 0.5 * (cbar * del[IDN] - hbar * del[ivx]) / cbar; + auto a2 = hbar * del[ivy]; + auto a3 = 0.5 * (cbar * del[IDN] + hbar * del[ivx]) / cbar; + auto wave0 = torch::zeros_like(del); + auto wave1 = torch::zeros_like(del); + auto wave2 = torch::zeros_like(del); + wave0[IDN] = a1; + wave0[ivx] = a1 * (ubar - cbar); + wave0[ivy] = a1 * vbar; + wave1[ivy] = a2; + wave2[IDN] = a3; + wave2[ivx] = a3 * (ubar + cbar); + wave2[ivy] = a3 * vbar; + + auto speed0 = torch::abs(ubar - cbar); + auto speed1 = torch::abs(ubar); + auto speed2 = torch::abs(ubar + cbar); + flx[IDN] = 0.5 * (wl[IDN] * wl[ivx] + wr[IDN] * wr[ivx]); + flx[ivx] = 0.5 * (wl[IDN] * wl[ivx].square() + 0.5 * wl[IDN].square() + + wr[IDN] * wr[ivx].square() + 0.5 * wr[IDN].square()); + flx[ivy] = 0.5 * (wl[IDN] * wl[ivx] * wl[ivy] + wr[IDN] * wr[ivx] * wr[ivy]); + flx[ivz] = 0.0; + flx -= 0.5 * (speed0 * wave0 + speed1 * wave1 + speed2 * wave2); +} + +} // namespace snap + +namespace at::native { + +DEFINE_DISPATCH(call_shallow_roe); +REGISTER_ALL_CPU_DISPATCH(call_shallow_roe, &snap::call_shallow_roe_cpu); +REGISTER_MPS_DISPATCH(call_shallow_roe, &snap::call_shallow_roe_mps); + +} // namespace at::native diff --git a/src/snap.h b/src/snap.h index 8a82f9c1..a867dcf0 100644 --- a/src/snap.h +++ b/src/snap.h @@ -70,24 +70,4 @@ enum { VEL3 = 2 //!< Third velocity component }; -enum class FusedReconScheme : int { - CP3 = 0, - CP5 = 1, - WENO3 = 2, - WENO5 = 3, -}; - -enum class FusedRiemannSolver : int { - LMARS = 0, - HLLC = 1, - Roe = 2, - ShallowRoe = 3, -}; - -enum class FusedEos : int { - IdealGas = 0, - IdealMoist = 1, - ShallowWater = 2, -}; - } // namespace snap diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a8b10ab5..2922d941 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -30,45 +30,6 @@ setup_test(test_diffusion_moist) setup_test(test_forcing) setup_parallel_test(test_exchange 2) -if(CUDAToolkit_FOUND AND UCX_FOUND) - add_executable(test_fused_cs_ucx_smoke.${buildl} - test_fused_cs_ucx_smoke.cu) - - set_target_properties(test_fused_cs_ucx_smoke.${buildl} - PROPERTIES - COMPILE_FLAGS ${CMAKE_CXX_FLAGS_${buildu}} - CUDA_SEPARABLE_COMPILATION ON) - - target_include_directories( - test_fused_cs_ucx_smoke.${buildl} - PRIVATE ${CMAKE_BINARY_DIR} - ${KINTERA_INCLUDE_DIR} - ${SNAP_INCLUDE_DIR} - ${HARP_INCLUDE_DIR} - ${NETCDF_INCLUDES} - ${TORCH_INCLUDE_DIR} - ${TORCH_API_INCLUDE_DIR}) - - target_link_libraries(test_fused_cs_ucx_smoke.${buildl} - PRIVATE snapy::snap - snapy::snap_cu - gtest_main) - - if (UNIX AND NOT APPLE) - target_link_options(test_fused_cs_ucx_smoke.${buildl} - PRIVATE -Wl,--no-as-needed) - endif() - - add_test( - NAME test_fused_cs_ucx_smoke.${buildl} - COMMAND torchrun - --no-python - --nproc-per-node=6 - $) - set_tests_properties(test_fused_cs_ucx_smoke.${buildl} PROPERTIES - LABELS "cuda;ucx;cubed-sphere;fused") -endif() - if(UCX_FOUND) add_test( NAME test_exchange_ucx diff --git a/tests/test_fused_cs_ucx_smoke.cu b/tests/test_fused_cs_ucx_smoke.cu deleted file mode 100644 index 3777815d..00000000 --- a/tests/test_fused_cs_ucx_smoke.cu +++ /dev/null @@ -1,517 +0,0 @@ -// C++/CUDA -#include -#include -#include -#include -#include -#include - -// external -#include - -// torch -#include -#include -#include -#include - -// snap -#include "../src/hydro/fused_recon_riemann_dispatch.hpp" -#include - -using namespace snap; - -namespace { - -constexpr int kSides = 4; -constexpr int kNvar = 4; -constexpr int kEdgeLen = 8; -constexpr int kNc1 = 3; -constexpr int kStates = 2; -constexpr int kMetaStride = 5; -// hydro side_meta: [enabled, peer_process, peer_local_block, peer_side, rev] -constexpr int kHydroMetaStride = 5; - -int env_int(char const *name, int fallback) { - char const *value = std::getenv(name); - if (value == nullptr || value[0] == '\0') - return fallback; - return std::stoi(value); -} - -torch::Device select_cuda_device() { - int local_rank = env_int("LOCAL_RANK", 0); - c10::cuda::set_device(local_rank); - return torch::Device(torch::kCUDA, local_rank); -} - -LayoutOptions make_layout_options() { - auto opts = LayoutOptionsImpl::create(); - opts->type("cubed-sphere"); - opts->backend("ucx"); - opts->device("cuda"); - opts->px(1); - opts->py(1); - opts->pz(1); - opts->blocks_per_process(1); - opts->verbose(false); - return opts; -} - -torch::Tensor make_edge_meta(CubedSphereLayoutImpl const &layout, - torch::Device device) { - auto iloc = layout.loc_of(layout.options->rank()); - int face = std::get<2>(iloc); - std::array, kSides> offsets = { - std::tuple{0, -1, 0}, - std::tuple{0, +1, 0}, - std::tuple{-1, 0, 0}, - std::tuple{+1, 0, 0}, - }; - - std::vector meta(kSides * kMetaStride); - for (int side = 0; side < kSides; ++side) { - auto edge = CS_FACE_EDGES[face][side]; - EXPECT_EQ(layout.neighbor_rank(iloc, offsets[side]), edge.nface); - meta[side * kMetaStride + 0] = edge.nface; - meta[side * kMetaStride + 1] = edge.nside; - meta[side * kMetaStride + 2] = edge.rev; - meta[side * kMetaStride + 3] = (side % 2) == (edge.nside % 2); - meta[side * kMetaStride + 4] = (side - 1.5) * (edge.nside - 1.5) < 0; - } - return torch::tensor(meta, torch::dtype(torch::kInt32)).to(device); -} - -torch::Tensor make_hydro_side_meta(CubedSphereLayoutImpl const &layout, - torch::Device device) { - auto iloc = layout.loc_of(layout.options->rank()); - int face = std::get<2>(iloc); - std::array, kSides> offsets = { - std::tuple{0, -1, 0}, - std::tuple{0, +1, 0}, - std::tuple{-1, 0, 0}, - std::tuple{+1, 0, 0}, - }; - - std::vector meta(kSides * kHydroMetaStride); - for (int side = 0; side < kSides; ++side) { - int nb = layout.neighbor_rank(iloc, offsets[side]); - EXPECT_GE(nb, 0); - auto nb_loc = layout.loc_of(nb); - EXPECT_NE(std::get<2>(nb_loc), face); - auto edge = CS_FACE_EDGES[face][side]; - EXPECT_EQ(nb, edge.nface); - meta[side * kHydroMetaStride + 0] = 1; - meta[side * kHydroMetaStride + 1] = layout.options->owner_process_rank(nb); - meta[side * kHydroMetaStride + 2] = layout.options->local_block_index(nb); - meta[side * kHydroMetaStride + 3] = edge.nside; - meta[side * kHydroMetaStride + 4] = edge.rev; - } - return torch::tensor(meta, torch::dtype(torch::kInt32)).to(device); -} - -bool fused_exchange_uses_right_state_start_for_side(int side) { - return side == SIDE_L || side == SIDE_B; -} - -bool fused_exchange_uses_right_interp_for_side(int side) { - return !fused_exchange_uses_right_state_start_for_side(side); -} - -bool cuda_6rank_available_or_skip() { - if (!torch::cuda::is_available()) { - return false; - } - int world_size = env_int("WORLD_SIZE", 1); - if (world_size != 6) { - return false; - } - int device_count = c10::cuda::device_count(); - if (device_count < world_size) { - return false; - } - return true; -} - -struct SmokeContext { - torch::Device device; - LayoutOptions options; - std::shared_ptr layout; - torch::Tensor edge_meta; -}; - -SmokeContext make_smoke_context() { - auto device = select_cuda_device(); - auto opts = make_layout_options(); - auto layout = std::make_shared(opts); - static std::shared_ptr shared_comm = - ProcessGroupContext::create(opts); - layout->comm = shared_comm; - EXPECT_TRUE(layout->has_process_group()); - EXPECT_NE(layout->comm, nullptr); - EXPECT_TRUE(layout->comm->store.defined()); - auto edge_meta = make_edge_meta(*layout, device); - return {device, opts, layout, edge_meta}; -} - -torch::Tensor make_exchange_buffer(torch::Device device) { - std::vector sizes = {kSides, kStates, kNvar, kEdgeLen, kNc1}; - std::vector strides = {kStates * kNvar * kEdgeLen * kNc1, - kNvar * kEdgeLen * kNc1, kEdgeLen * kNc1, - kNc1, 1}; - return torch::empty_strided( - sizes, strides, - torch::TensorOptions().dtype(torch::kFloat64).device(device)); -} - -struct ExchangedBuffers { - std::vector received; - torch::Tensor pointers; - - void **pointers_dev() const { - return reinterpret_cast(pointers.data_ptr()); - } -}; - -ExchangedBuffers exchange_with_face_neighbors(SmokeContext const &ctx, - torch::Tensor local, int tag) { - int rank = ctx.options->process_rank(); - int world_size = ctx.options->process_world_size(); - auto meta = ctx.edge_meta.cpu(); - std::set peers; - for (int side = 0; side < kSides; ++side) { - peers.insert(meta[side * kMetaStride].item()); - } - - ExchangedBuffers result; - result.received.resize(world_size); - std::vector pointers(world_size, 0); - pointers[rank] = reinterpret_cast(local.data_ptr()); - std::vector works; - for (int peer : peers) { - result.received[peer] = torch::empty_like(local); - pointers[peer] = - reinterpret_cast(result.received[peer].data_ptr()); - std::vector send_tensors{local}; - std::vector recv_tensors{result.received[peer]}; - works.push_back(ctx.layout->comm->send(send_tensors, peer, tag)); - works.push_back(ctx.layout->comm->recv(recv_tensors, peer, tag)); - } - for (auto const &work : works) work->wait(); - result.pointers = - torch::tensor(pointers, torch::dtype(torch::kInt64)).to(ctx.device); - return result; -} - -__device__ double device_payload(int rank, int side, int state, int edge, int i, - int v) { - return static_cast(100000 * rank + 10000 * side + 1000 * state + - 100 * edge + 10 * i + v); -} - -__global__ void write_edge_payload_kernel(double *buffer, int rank) { - int line = blockIdx.x; - int i = line % kNc1; - int edge = (line / kNc1) % kEdgeLen; - int side = line / (kNc1 * kEdgeLen); - int stride_var = kEdgeLen * kNc1; - for (int state = 0; state < kStates; ++state) { - int base = - (((side * kStates + state) * kNvar) * kEdgeLen + edge) * kNc1 + i; - for (int v = 0; v < kNvar; ++v) { - buffer[base + v * stride_var] = - device_payload(rank, side, state, edge, i, v); - } - } -} - -__global__ void verify_remote_edge_kernel(void **buffer_ptrs, int const *meta, - int *errors) { - int line = blockIdx.x; - int i = line % kNc1; - int edge = (line / kNc1) % kEdgeLen; - int side = line / (kNc1 * kEdgeLen); - int peer_rank = meta[side * kMetaStride + 0]; - int peer_side = meta[side * kMetaStride + 1]; - int rev = meta[side * kMetaStride + 2]; - int peer_edge = rev ? (kEdgeLen - 1 - edge) : edge; - int stride_var = kEdgeLen * kNc1; - auto peer_buffer = static_cast(buffer_ptrs[peer_rank]); - for (int state = 0; state < kStates; ++state) { - int remote_base = - (((peer_side * kStates + state) * kNvar) * kEdgeLen + peer_edge) * - kNc1 + - i; - for (int v = 0; v < kNvar; ++v) { - double actual = peer_buffer[remote_base + v * stride_var]; - double expected = - device_payload(peer_rank, peer_side, state, peer_edge, i, v); - if (actual != expected) - atomicAdd(errors, 1); - } - } -} - -__global__ void verify_staged_remote_state_kernel(void **buffer_ptrs, - int const *meta, int rank, - int *errors) { - int line = blockIdx.x; - int i = line % kNc1; - int edge = (line / kNc1) % kEdgeLen; - int side = line / (kNc1 * kEdgeLen); - int peer_rank = meta[side * kHydroMetaStride + 1]; - int peer_side = meta[side * kHydroMetaStride + 3]; - int rev = meta[side * kHydroMetaStride + 4]; - int peer_edge = rev ? (kEdgeLen - 1 - edge) : edge; - bool lower_side = side == SIDE_L || side == SIDE_B; - bool peer_lower_side = peer_side == SIDE_L || peer_side == SIDE_B; - int remote_state = peer_lower_side ? IRT : ILT; - int local_state = lower_side ? IRT : ILT; - int stride_var = kEdgeLen * kNc1; - int remote_base = - (((peer_side * kStates + remote_state) * kNvar) * kEdgeLen + peer_edge) * - kNc1 + - i; - int local_base = - (((side * kStates + local_state) * kNvar) * kEdgeLen + edge) * kNc1 + i; - auto peer_buffer = static_cast(buffer_ptrs[peer_rank]); - auto local_buffer = static_cast(buffer_ptrs[rank]); - for (int v = 0; v < kNvar; ++v) { - double remote_actual = peer_buffer[remote_base + v * stride_var]; - double remote_expected = - device_payload(peer_rank, peer_side, remote_state, peer_edge, i, v); - double local_actual = local_buffer[local_base + v * stride_var]; - double local_expected = device_payload(rank, side, local_state, edge, i, v); - if (remote_actual != remote_expected || local_actual != local_expected) { - atomicAdd(errors, 1); - } - } -} - -__global__ void verify_hydro_remote_constant_kernel(void **buffer_ptrs, - int const *meta, - int *errors) { - int line = blockIdx.x; - int i = line % kNc1; - int edge = (line / kNc1) % kEdgeLen; - int side = line / (kNc1 * kEdgeLen); - int peer_rank = meta[side * kHydroMetaStride + 1]; - int peer_side = meta[side * kHydroMetaStride + 3]; - int rev = meta[side * kHydroMetaStride + 4]; - int peer_edge = rev ? (kEdgeLen - 1 - edge) : edge; - int stride_var = kEdgeLen * kNc1; - int remote_base = - (((peer_side * kStates) * kNvar) * kEdgeLen + peer_edge) * kNc1 + i; - auto peer_buffer = static_cast(buffer_ptrs[peer_rank]); - double density = peer_buffer[remote_base + IDN * stride_var]; - if (fabs(density - 500.) > 1.e-10) - atomicAdd(errors, 1); - for (int v = IVX; v <= IVZ; ++v) { - double velocity = peer_buffer[remote_base + v * stride_var]; - if (fabs(velocity) > 1.e-10) - atomicAdd(errors, 1); - } -} - -} // namespace - -TEST(FusedCubedSphereUCX, ProcessGroupInitializes) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - auto ctx = make_smoke_context(); - EXPECT_TRUE(ctx.layout->comm->initialized()); - EXPECT_TRUE(ctx.layout->comm->is_ucx()); - EXPECT_EQ(ctx.options->process_world_size(), 6); - if (ctx.layout->comm) - ctx.layout->comm->barrier(); -} - -TEST(FusedCubedSphereUCX, PreviousKernelWritesReachPeers) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - auto ctx = make_smoke_context(); - auto exchange_buffer = make_exchange_buffer(ctx.device); - - auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); - int blocks = kSides * kEdgeLen * kNc1; - write_edge_payload_kernel<<>>( - exchange_buffer.data_ptr(), ctx.options->process_rank()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - - auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 701); - auto errors = - torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); - verify_remote_edge_kernel<<>>( - exchanged.pointers_dev(), ctx.edge_meta.data_ptr(), - errors.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - AT_CUDA_CHECK(cudaStreamSynchronize(stream)); - EXPECT_EQ(errors.cpu().item(), 0); - if (ctx.layout->comm) - ctx.layout->comm->barrier(); -} - -TEST(FusedCubedSphereUCX, OrientationMetadataMatchesStaged) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - auto ctx = make_smoke_context(); - - auto meta = ctx.edge_meta.cpu(); - auto iloc = ctx.layout->loc_of(ctx.layout->options->rank()); - int face = std::get<2>(iloc); - for (int side = 0; side < kSides; ++side) { - auto edge = CS_FACE_EDGES[face][side]; - EXPECT_EQ(meta[side * kMetaStride + 0].item(), edge.nface); - EXPECT_EQ(meta[side * kMetaStride + 1].item(), edge.nside); - EXPECT_EQ(meta[side * kMetaStride + 2].item(), edge.rev); - EXPECT_EQ(meta[side * kMetaStride + 3].item(), - (side % 2) == (edge.nside % 2)); - EXPECT_EQ(meta[side * kMetaStride + 4].item(), - (side - 1.5) * (edge.nside - 1.5) < 0); - } - if (ctx.layout->comm) - ctx.layout->comm->barrier(); -} - -TEST(FusedCubedSphereUCX, BoundaryStateConventionMatchesStaged) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - // Staged hydro sends the local right state to lower-side neighbors and the - // local left state to upper-side neighbors. - EXPECT_TRUE(fused_exchange_uses_right_state_start_for_side(SIDE_L)); - EXPECT_FALSE(fused_exchange_uses_right_interp_for_side(SIDE_L)); - EXPECT_FALSE(fused_exchange_uses_right_state_start_for_side(SIDE_R)); - EXPECT_TRUE(fused_exchange_uses_right_interp_for_side(SIDE_R)); - EXPECT_TRUE(fused_exchange_uses_right_state_start_for_side(SIDE_B)); - EXPECT_FALSE(fused_exchange_uses_right_interp_for_side(SIDE_B)); - EXPECT_FALSE(fused_exchange_uses_right_state_start_for_side(SIDE_T)); - EXPECT_TRUE(fused_exchange_uses_right_interp_for_side(SIDE_T)); -} - -TEST(FusedCubedSphereUCX, RemoteEdgePayloadMatches) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - auto ctx = make_smoke_context(); - auto exchange_buffer = make_exchange_buffer(ctx.device); - - auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); - int blocks = kSides * kEdgeLen * kNc1; - write_edge_payload_kernel<<>>( - exchange_buffer.data_ptr(), ctx.options->process_rank()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 702); - - auto errors = - torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); - verify_remote_edge_kernel<<>>( - exchanged.pointers_dev(), ctx.edge_meta.data_ptr(), - errors.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - AT_CUDA_CHECK(cudaStreamSynchronize(stream)); - - EXPECT_EQ(errors.cpu().item(), 0); - if (ctx.layout->comm) - ctx.layout->comm->barrier(); -} - -TEST(FusedCubedSphereUCX, RemoteStateSelectionMatchesStaged) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - auto ctx = make_smoke_context(); - auto exchange_buffer = make_exchange_buffer(ctx.device); - - auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); - int blocks = kSides * kEdgeLen * kNc1; - write_edge_payload_kernel<<>>( - exchange_buffer.data_ptr(), ctx.options->process_rank()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 703); - - auto errors = - torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); - auto side_meta = make_hydro_side_meta(*ctx.layout, ctx.device); - verify_staged_remote_state_kernel<<>>( - exchanged.pointers_dev(), side_meta.data_ptr(), - ctx.options->process_rank(), errors.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - AT_CUDA_CHECK(cudaStreamSynchronize(stream)); - - EXPECT_EQ(errors.cpu().item(), 0); - if (ctx.layout->comm) - ctx.layout->comm->barrier(); -} - -TEST(FusedCubedSphereUCX, ConstantStateExchangeHasZeroMassFlux) { - if (!cuda_6rank_available_or_skip()) - GTEST_SKIP(); - auto ctx = make_smoke_context(); - auto exchange_buffer = make_exchange_buffer(ctx.device); - - auto opts = torch::dtype(torch::kFloat64).device(ctx.device); - auto w = torch::zeros({kNvar, kEdgeLen, kEdgeLen, kNc1}, opts); - w[IDN].fill_(500.); - auto flux2 = torch::zeros_like(w); - auto flux3 = torch::zeros_like(w); - auto side_meta = make_hydro_side_meta(*ctx.layout, ctx.device); - auto x2v = torch::linspace(-0.7, 0.7, kEdgeLen, opts); - auto x3v = torch::linspace(-0.7, 0.7, kEdgeLen, opts); - auto x2f = torch::linspace(-0.8, 0.8, kEdgeLen + 1, opts); - auto x3f = torch::linspace(-0.8, 0.8, kEdgeLen + 1, opts); - - auto iloc = ctx.layout->loc_of(ctx.layout->options->rank()); - int face = std::get<2>(iloc); - // One panel per process (blocks_per_process == 1), so local_block == 0 and - // the flat [side, state, var, edge, nc1] buffer is exactly this block's - // slice. - int local_block = 0; - FusedCubedSpherePanelParams panel_params{side_meta, face, local_block, x2v, - x2f, x3v, x3f}; - fused_cubed_sphere_pack_cuda( - w, exchange_buffer, - FusedCubedSpherePackParams{panel_params, FusedReconScheme::WENO5, - FusedReconScheme::WENO5, false, - FusedEos::ShallowWater, 0., 0., false}); - auto exchanged = exchange_with_face_neighbors(ctx, exchange_buffer, 704); - fused_cubed_sphere_flux_cuda( - w, flux2, flux3, exchange_buffer, exchanged.pointers_dev(), - FusedCubedSphereFluxParams{ - panel_params, - FusedPhysicsParams{FusedReconScheme::WENO5, FusedReconScheme::WENO5, - false, FusedRiemannSolver::ShallowRoe, - FusedEos::ShallowWater, 1.4, 0., 0., false, - torch::Tensor(), torch::Tensor(), torch::Tensor(), - 0, 1}}); - AT_CUDA_CHECK(cudaStreamSynchronize( - at::cuda::getCurrentCUDAStream(ctx.device.index()))); - - auto expect_constant_state = [](torch::Tensor state, char const *label) { - EXPECT_LT(std::abs(state.select(1, IDN).min().cpu().item() - 500.), - 1.e-10) - << label; - EXPECT_LT(std::abs(state.select(1, IDN).max().cpu().item() - 500.), - 1.e-10) - << label; - EXPECT_LT(state.select(1, IVX).abs().max().cpu().item(), 1.e-10) - << label; - EXPECT_LT(state.select(1, IVY).abs().max().cpu().item(), 1.e-10) - << label; - EXPECT_LT(state.select(1, IVZ).abs().max().cpu().item(), 1.e-10) - << label; - }; - expect_constant_state(exchange_buffer.select(1, ILT), "left state"); - expect_constant_state(exchange_buffer.select(1, IRT), "right state"); - auto errors = - torch::zeros({1}, torch::dtype(torch::kInt32).device(ctx.device)); - auto stream = at::cuda::getCurrentCUDAStream(ctx.device.index()); - verify_hydro_remote_constant_kernel<<>>(exchanged.pointers_dev(), - side_meta.data_ptr(), - errors.data_ptr()); - C10_CUDA_KERNEL_LAUNCH_CHECK(); - AT_CUDA_CHECK(cudaStreamSynchronize(stream)); - EXPECT_EQ(errors.cpu().item(), 0); - EXPECT_LT(flux2[IDN].abs().max().cpu().item(), 1.e-10); - EXPECT_LT(flux3[IDN].abs().max().cpu().item(), 1.e-10); - if (ctx.layout->comm) - ctx.layout->comm->barrier(); -} diff --git a/tests/test_hydro_python_api.py b/tests/test_hydro_python_api.py deleted file mode 100644 index 966a61a6..00000000 --- a/tests/test_hydro_python_api.py +++ /dev/null @@ -1,11 +0,0 @@ -import snapy - - -def test_hydro_options_exposes_fused_recon_riemann(): - options = snapy.HydroOptions() - - assert options.fused_recon_riemann() is False - assert options.fused_recon_riemann(True) is options - assert options.fused_recon_riemann() is True - assert options.fused_recon_riemann(False) is options - assert options.fused_recon_riemann() is False diff --git a/tests/test_hydrostatic.cpp b/tests/test_hydrostatic.cpp index ef449aec..c4e71100 100644 --- a/tests/test_hydrostatic.cpp +++ b/tests/test_hydrostatic.cpp @@ -126,7 +126,6 @@ void initialize_hydrostatic_column(MeshBlock const& block, Variables& vars) { } // namespace TEST(HydrostaticAtmosphere, cubed_sphere_remains_at_rest) { - setenv("FUSED", "OFF", 1); torch::set_num_threads(1); torch::set_num_interop_threads(1); diff --git a/tests/test_riemann.cpp b/tests/test_riemann.cpp index bb54c4e8..6c3bb030 100644 --- a/tests/test_riemann.cpp +++ b/tests/test_riemann.cpp @@ -409,267 +409,6 @@ boundary-condition: using namespace snap; -namespace { - -int env_int(char const* name, int fallback) { - char const* value = std::getenv(name); - if (value == nullptr || value[0] == '\0') return fallback; - return std::stoi(value); -} - -struct ScopedEnv { - explicit ScopedEnv(char const* name_) : name(name_) { - auto* value = std::getenv(name); - if (value) { - had_value = true; - old_value = value; - } - } - ~ScopedEnv() { - if (had_value) { - setenv(name, old_value.c_str(), 1); - } else { - unsetenv(name); - } - } - - char const* name; - bool had_value = false; - std::string old_value; -}; - -void expect_fused_reconstruct_options_match_staged(torch::Device device, - torch::Dtype dtype, - bool scale, bool shock) { - auto replacement = std::string("scale: ") + (scale ? "true" : "false") + - ", shock: " + (shock ? "true" : "false"); - auto config = - std::regex_replace(std::string(small_ideal_gas_config), - std::regex("scale: false, shock: false"), replacement); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream file(fname); - file << config; - file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(fname); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fname); - staged_opts->hydro()->fused_recon_riemann() = false; - fused_opts->hydro()->fused_recon_riemann() = true; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(1200 + 10 * static_cast(scale) + - static_cast(shock)); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].mul_(0.2).add_(1.0); - w[IPR].mul_(2.e4).add_(8.e4); - w.narrow(0, IVX, 3).sub_(0.5).mul_(20.0); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - if (dtype == torch::kFloat64) { - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-6, 1.e-6)); - } - auto pcoord = staged_block->pcoord; - EXPECT_TRUE(torch::allclose( - fused_block->phydro->flux1().slice(DIM1, pcoord->il(), pcoord->iu() + 2), - staged_block->phydro->flux1().slice(DIM1, pcoord->il(), pcoord->iu() + 2), - 1.e-6, 1.e-6)); - EXPECT_TRUE(torch::allclose( - fused_block->phydro->flux2().slice(DIM2, pcoord->jl(), pcoord->ju() + 2), - staged_block->phydro->flux2().slice(DIM2, pcoord->jl(), pcoord->ju() + 2), - 1.e-6, 1.e-6)); - EXPECT_TRUE(torch::allclose( - fused_block->phydro->flux3().slice(DIM3, pcoord->kl(), pcoord->ku() + 2), - staged_block->phydro->flux3().slice(DIM3, pcoord->kl(), pcoord->ku() + 2), - 1.e-6, 1.e-6)); - - std::remove(fname); -} - -} // namespace - -TEST(HydroOptions, auto_detects_supported_fused_recon_riemann) { - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << small_ideal_gas_config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_TRUE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, fused_env_disables_supported_fused_recon_riemann) { - ScopedEnv env("FUSED"); - setenv("FUSED", "OFF", 1); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << small_ideal_gas_config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_FALSE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, fused_env_auto_preserves_auto_detection) { - ScopedEnv env("FUSED"); - setenv("FUSED", "AUTO", 1); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << small_ideal_gas_config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_TRUE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, fused_env_auto_falls_back_for_unsupported_configuration) { - ScopedEnv env("FUSED"); - setenv("FUSED", "AUTO", 1); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << block_config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_FALSE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, fused_env_rejects_invalid_value) { - ScopedEnv env("FUSED"); - setenv("FUSED", "maybe", 1); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << small_ideal_gas_config; - outfile.close(); - - EXPECT_THROW((void)MeshBlockOptionsImpl::from_yaml(fname), c10::Error); - - std::remove(fname); -} - -TEST(HydroOptions, fused_env_on_rejects_unsupported_configuration) { - ScopedEnv env("FUSED"); - setenv("FUSED", "ON", 1); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << block_config; - outfile.close(); - - EXPECT_THROW((void)MeshBlockOptionsImpl::from_yaml(fname), c10::Error); - - std::remove(fname); -} - -TEST(HydroOptions, auto_detects_fused_recon_riemann_for_cartesian_multiblock) { - auto config = std::string( - "distribute:\n" - " layout: slab\n" - " blocks_per_process: 3\n\n") + - small_ideal_gas_config; - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_TRUE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, auto_detects_fused_recon_riemann_for_cubed_sphere_bpp1) { - auto config = std::string( - "distribute:\n" - " layout: cubed-sphere\n" - " blocks_per_process: 1\n\n") + - small_ideal_gas_config; - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_TRUE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, - auto_detects_fused_recon_riemann_for_cubed_sphere_multiblock) { - auto config = std::string( - "distribute:\n" - " layout: cubed-sphere\n" - " blocks_per_process: 6\n\n") + - small_ideal_gas_config; - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_TRUE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - -TEST(HydroOptions, leaves_unsupported_fused_recon_riemann_off) { - auto config = std::regex_replace( - block_config, std::regex("riemann-solver:\\n type: lmars"), - "fused-recon-riemann: true\n\n riemann-solver:\n type: lmars"); - - char fname[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fname); - std::ofstream outfile(fname); - outfile << config; - outfile.close(); - - auto op_block = MeshBlockOptionsImpl::from_yaml(fname); - EXPECT_FALSE(op_block->hydro()->fused_recon_riemann()); - - std::remove(fname); -} - TEST(RiemannSolver, hllc_writes_face_pressure_output) { auto config = std::regex_replace(std::string(small_ideal_gas_config), std::regex("type: lmars"), "type: hllc"); @@ -810,9 +549,17 @@ TEST(RiemannSolver, roe_writes_face_pressure_output) { auto ur = block->phydro->peos->compute("W->U", {wr}); auto sqrtdl = torch::sqrt(wl[IDN]); auto sqrtdr = torch::sqrt(wr[IDN]); + auto isdlpdr = 1.0 / (sqrtdl + sqrtdr); + auto hbar = + ((ul[IPR] + wl[IPR]) / sqrtdl + (ur[IPR] + wr[IPR]) / sqrtdr) * isdlpdr; + auto vbar = + (sqrtdl * wl.narrow(0, IVX, 3) + sqrtdr * wr.narrow(0, IVX, 3)) * isdlpdr; + auto gamma_l = block->phydro->peos->compute("W->A", {wl}); + auto gamma_r = block->phydro->peos->compute("W->A", {wr}); + auto cs = torch::sqrt((0.5 * (gamma_l + gamma_r) - 1.0) * + (hbar - 0.5 * vbar.square().sum(0))); auto expected = - ((ul[IPR] + wl[IPR]) / sqrtdl + (ur[IPR] + wr[IPR]) / sqrtdr) / - (sqrtdl + sqrtdr); + 0.5 * (wl[IPR] + wr[IPR] + sqrtdl * sqrtdr * cs * (wl[IVX] - wr[IVX])); auto flx = torch::zeros_like(wl); auto face_pressure = torch::zeros({nc3, nc2, nc1}, torch::kFloat64); @@ -858,14 +605,9 @@ TEST(RiemannSolver, roe_writes_face_pressure_output_ideal_moist) { wr[ICY].mul_(1.e-3).add_(1.5e-2); wl[ICY + 1].mul_(1.e-3).add_(4.e-3); wr[ICY + 1].mul_(1.e-3).add_(4.e-3); + wr.copy_(wl); - auto ul = block->phydro->peos->compute("W->U", {wl}); - auto ur = block->phydro->peos->compute("W->U", {wr}); - auto sqrtdl = torch::sqrt(wl[IDN]); - auto sqrtdr = torch::sqrt(wr[IDN]); - auto expected = - ((ul[IPR] + wl[IPR]) / sqrtdl + (ur[IPR] + wr[IPR]) / sqrtdr) / - (sqrtdl + sqrtdr); + auto expected = wl[IPR].clone(); auto flx = torch::zeros_like(wl); auto face_pressure = torch::zeros({nc3, nc2, nc1}, torch::kFloat64); @@ -875,529 +617,6 @@ TEST(RiemannSolver, roe_writes_face_pressure_output_ideal_moist) { EXPECT_GT(torch::abs(flx).sum().item(), 0.0); } -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_ideal_gas_roe) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - auto staged_config = - std::regex_replace(std::string(small_ideal_gas_config), - std::regex("type: lmars"), "type: roe"); - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << staged_config; - staged_file.close(); - - auto fused_config = std::regex_replace( - staged_config, std::regex("riemann-solver:\\n type: roe"), - "fused-recon-riemann: true\n\n riemann-solver:\n type: roe"); - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << fused_config; - fused_file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(205); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].add_(1.0); - w[IPR].add_(1.0); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-7, 1.e-7)); - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_ideal_gas_lmars) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_ideal_gas_config; - staged_file.close(); - - auto fused_config = std::regex_replace( - std::string(small_ideal_gas_config), - std::regex("riemann-solver:\\n type: lmars"), - "fused-recon-riemann: true\n\n riemann-solver:\n type: lmars"); - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << fused_config; - fused_file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(202); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].add_(1.0); - w[IPR].add_(1.0); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-8, 1.e-8)); - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_honors_weno_scale_and_shock) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - for (auto const& options : {std::pair{true, false}, std::pair{false, true}, - std::pair{true, true}}) { - SCOPED_TRACE(testing::Message() - << "scale=" << options.first << " shock=" << options.second); - expect_fused_reconstruct_options_match_staged(device, dtype, options.first, - options.second); - } -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_ideal_gas_gravity) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_ideal_gas_gravity_config; - staged_file.close(); - - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << small_ideal_gas_gravity_config; - fused_file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(203); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].add_(1.0); - w[IPR].add_(1.0); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - int outer_face = staged_block->pcoord->iu() + 1; - EXPECT_TRUE(torch::allclose( - fused_block->phydro->flux1().select(DIM1, outer_face), - staged_block->phydro->flux1().select(DIM1, outer_face), 1.e-6, 1.e-6)); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-8, 1.e-8)); - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_cubed_sphere_hllc) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - if (env_int("WORLD_SIZE", 1) != 6) { - GTEST_SKIP() - << "cubed-sphere fused reconstruction/Riemann test requires 6 ranks"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_ideal_gas_cubed_sphere_hllc_config; - staged_file.close(); - - auto fused_config = std::regex_replace( - std::string(small_ideal_gas_cubed_sphere_hllc_config), - std::regex("riemann-solver:\\n type: hllc"), - "fused-recon-riemann: true\n\n riemann-solver:\n type: hllc"); - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << fused_config; - fused_file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(902); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].add_(1.0); - w[IPR].add_(1.0); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-7, 1.e-7)); - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_ideal_gas_implicit) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_ideal_gas_implicit_config; - staged_file.close(); - - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << small_ideal_gas_implicit_config; - fused_file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(204); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].add_(1.0); - w[IPR].add_(1.0); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-8, 1.e-8)); - EXPECT_TRUE(torch::allclose(fused_block->phydro->implicit_mass_correction(), - staged_block->phydro->implicit_mass_correction(), - 1.e-8, 1.e-8)); - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, - fused_recon_riemann_matches_staged_ideal_moist_sedimentation) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_ideal_moist_sedimentation_config; - staged_file.close(); - - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << small_ideal_moist_sedimentation_config; - fused_file.close(); - - kintera::init_species_from_yaml(staged_name); - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - kintera::init_species_from_yaml(fused_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - staged_opts->hydro()->disable_flux_x1() = true; - staged_opts->hydro()->disable_flux_x2() = true; - staged_opts->hydro()->disable_flux_x3() = true; - staged_opts->hydro()->grav()->grav1(-9.8); - fused_opts->hydro()->disable_flux_x1() = true; - fused_opts->hydro()->disable_flux_x2() = true; - fused_opts->hydro()->disable_flux_x3() = true; - fused_opts->hydro()->grav()->grav1(-9.8); - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(205); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].mul_(0.1).add_(1.0); - w[IPR].mul_(1000.0).add_(1.e5); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - if (nvar > ICY) { - w.narrow(0, ICY, nvar - ICY).mul_(1.e-3); - } - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - double tol = dtype == torch::kFloat32 ? 1.e-4 : 1.e-8; - EXPECT_TRUE(torch::allclose(fused_du, staged_du, tol, tol)); - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_ideal_moist_lmars_cloud) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_ideal_moist_lmars_cloud_config; - staged_file.close(); - - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << small_ideal_moist_lmars_cloud_config; - fused_file.close(); - - kintera::init_species_from_yaml(staged_name); - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - kintera::init_species_from_yaml(fused_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(206); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].mul_(0.1).add_(1.0); - w[IPR].mul_(1000.0).add_(1.e5); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - w[ICY].mul_(1.e-3).add_(1.5e-2); - w[ICY + 1].mul_(1.e-3).add_(4.e-3); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - double tol = dtype == torch::kFloat32 ? 1.e-4 : 1.e-8; - auto max_diff = (fused_du - staged_du).abs().max().item(); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, tol, tol)) - << "max diff = " << max_diff; - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_ideal_moist_roe_cloud) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - auto staged_config = - std::regex_replace(std::string(small_ideal_moist_lmars_cloud_config), - std::regex("type: lmars"), "type: roe"); - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << staged_config; - staged_file.close(); - - auto fused_config = std::regex_replace( - staged_config, std::regex("riemann-solver:\\n type: roe"), - "fused-recon-riemann: true\n\n riemann-solver:\n type: roe"); - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << fused_config; - fused_file.close(); - - kintera::init_species_from_yaml(staged_name); - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - kintera::init_species_from_yaml(fused_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(207); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].mul_(0.1).add_(1.0); - w[IPR].mul_(1000.0).add_(1.e5); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - w[ICY].mul_(1.e-3).add_(1.5e-2); - w[ICY + 1].mul_(1.e-3).add_(4.e-3); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - double tol = dtype == torch::kFloat32 ? 1.e-4 : 1.e-8; - auto max_diff = (fused_du - staged_du).abs().max().item(); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, tol, tol)) - << "max diff = " << max_diff; - - std::remove(staged_name); - std::remove(fused_name); -} - -TEST_P(DeviceTest, fused_recon_riemann_matches_staged_shallow_water_roe) { - if (device.type() != torch::kCUDA) { - GTEST_SKIP() << "fused reconstruction/Riemann path is CUDA-only"; - } - - char staged_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(staged_name); - std::ofstream staged_file(staged_name); - staged_file << small_shallow_water_config; - staged_file.close(); - - char fused_name[80] = "/tmp/tempfile.XXXXXX"; - mkstemp(fused_name); - std::ofstream fused_file(fused_name); - fused_file << small_shallow_water_config; - fused_file.close(); - - auto staged_opts = MeshBlockOptionsImpl::from_yaml(staged_name); - auto fused_opts = MeshBlockOptionsImpl::from_yaml(fused_name); - staged_opts->hydro()->fused_recon_riemann() = false; - auto staged_block = MeshBlock(staged_opts); - auto fused_block = MeshBlock(fused_opts); - staged_block->to(device, dtype); - fused_block->to(device, dtype); - - int nc1 = staged_block->pcoord->options->nc1(); - int nc2 = staged_block->pcoord->options->nc2(); - int nc3 = staged_block->pcoord->options->nc3(); - int nvar = staged_block->phydro->peos->nvar(); - - torch::manual_seed(204); - auto w = - torch::rand({nvar, nc3, nc2, nc1}, torch::device(device).dtype(dtype)); - w[IDN].add_(1.0); - w.narrow(0, IVX, 3).sub_(0.5).mul_(0.1); - - auto u = staged_block->phydro->peos->compute("W->U", {w}); - Variables staged_vars, fused_vars; - staged_vars["hydro_w"] = torch::empty_like(w); - fused_vars["hydro_w"] = torch::empty_like(w); - - auto staged_du = staged_block->phydro->forward(1.e-4, u.clone(), staged_vars); - auto fused_du = fused_block->phydro->forward(1.e-4, u.clone(), fused_vars); - EXPECT_TRUE(torch::allclose(fused_du, staged_du, 1.e-8, 1.e-8)); - - std::remove(staged_name); - std::remove(fused_name); -} - TEST_P(DeviceTest, test_lmars) { // create a temporary file char fname[80] = "/tmp/tempfile.XXXXXX"; From 7856dc61dd328b6eb84af027a6a69243e7da7c3e Mon Sep 17 00:00:00 2001 From: dungeon2/cli Date: Thu, 16 Jul 2026 02:40:24 -0400 Subject: [PATCH 3/3] wip --- tests/CMakeLists.txt | 9 --------- 1 file changed, 9 deletions(-) diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2922d941..495843e0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -80,15 +80,6 @@ set_tests_properties(test_mesh_exchange_python PROPERTIES LABELS "python;exchange" ) -add_test( - NAME test_hydro_python_api - COMMAND ${Python3_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/test_hydro_python_api.py -) -set_tests_properties(test_hydro_python_api PROPERTIES - WORKING_DIRECTORY ${CMAKE_BINARY_DIR}/tests - LABELS "python;api;hydro" -) - add_test( NAME test_cubed_sphere_vertical_velocity_exchange_python COMMAND ${Python3_EXECUTABLE}