373 lines
12 KiB
Plaintext
373 lines
12 KiB
Plaintext
|
//===----------------------------------------------------------------------===//
|
|||
|
|
//
|
|||
|
|
// Part of CUDASTF in CUDA C++ Core Libraries,
|
|||
|
|
// under the Apache License v2.0 with LLVM Exceptions.
|
|||
|
|
// See https://llvm.org/LICENSE.txt for license information.
|
|||
|
|
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
|
|||
|
|
// SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES.
|
|||
|
|
//
|
|||
|
|
//===----------------------------------------------------------------------===//
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* @file
|
|||
|
|
* @brief Sparse conjugate gradient algorithm
|
|||
|
|
*/
|
|||
|
|
|
|||
|
|
#include <cuda/experimental/stf.cuh>
|
|||
|
|
|
|||
|
|
#include <chrono>
|
|||
|
|
#include <iostream>
|
|||
|
|
#include <string>
|
|||
|
|
#include <type_traits>
|
|||
|
|
#include <vector>
|
|||
|
|
|
|||
|
|
#include "dot.cuh"
|
|||
|
|
#include "newton_solver.cuh"
|
|||
|
|
|
|||
|
|
using namespace cuda::experimental::stf;
|
|||
|
|
|
|||
|
|
#if !_CCCL_CTK_BELOW(12, 4)
|
|||
|
|
|
|||
|
|
void build_full_csr_structure(size_t* row_offsets, size_t* col_indices, size_t N)
|
|||
|
|
{
|
|||
|
|
size_t nnz = 0;
|
|||
|
|
row_offsets[0] = 0;
|
|||
|
|
|
|||
|
|
for (size_t row = 0; row < N; row++)
|
|||
|
|
{
|
|||
|
|
if (row == 0 || row == N - 1)
|
|||
|
|
{
|
|||
|
|
// Boundary rows: only diagonal entry (identity for BC: u[i] = prescribed_value)
|
|||
|
|
col_indices[nnz++] = row;
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
// Interior rows: tridiagonal structure (left, center, right)
|
|||
|
|
col_indices[nnz++] = row - 1; // left neighbor
|
|||
|
|
col_indices[nnz++] = row; // center (diagonal)
|
|||
|
|
col_indices[nnz++] = row + 1; // right neighbor
|
|||
|
|
}
|
|||
|
|
row_offsets[row + 1] = nnz;
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
template <typename ctx_t>
|
|||
|
|
void assemble_jacobian_full(
|
|||
|
|
ctx_t& ctx, vector_t<double> U, vector_t<double> values, size_t N, double h, double dt, double nu)
|
|||
|
|
{
|
|||
|
|
ctx.parallel_for(box(N), U.read(), values.write()).set_symbol("assemble_jacobian_full")
|
|||
|
|
->*[N, h, dt, nu] __device__(size_t row, auto dU, auto dvalues) {
|
|||
|
|
if (row == 0)
|
|||
|
|
{
|
|||
|
|
// Left boundary: u[0] = 0 (homogeneous Dirichlet)
|
|||
|
|
// Jacobian row: [1, 0, 0, ..., 0]
|
|||
|
|
size_t val_idx = 0; // First entry in CSR values array
|
|||
|
|
dvalues[val_idx] = 1.0;
|
|||
|
|
}
|
|||
|
|
else if (row == N - 1)
|
|||
|
|
{
|
|||
|
|
// Right boundary: u[N-1] = 0 (homogeneous Dirichlet)
|
|||
|
|
// Jacobian row: [0, ..., 0, 1]
|
|||
|
|
size_t val_idx = 1 + 3 * (N - 2); // Last entry in CSR values array
|
|||
|
|
dvalues[val_idx] = 1.0;
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
// Interior point: Burger's equation discretization
|
|||
|
|
double u_i = dU[row];
|
|||
|
|
double u_ip1 = dU[row + 1];
|
|||
|
|
double u_im1 = dU[row - 1];
|
|||
|
|
|
|||
|
|
// Jacobian entries: ∂F_i/∂u_{i-1}, ∂F_i/∂u_i, ∂F_i/∂u_{i+1}
|
|||
|
|
double left = -u_i / (2 * h) - nu / (h * h);
|
|||
|
|
double center = 1.0 / dt + (u_ip1 - u_im1) / (2 * h) + 2.0 * nu / (h * h);
|
|||
|
|
double right = u_i / (2 * h) - nu / (h * h);
|
|||
|
|
|
|||
|
|
// CSR indexing for interior row i: starts at 1 + 3*(i-1)
|
|||
|
|
size_t val_idx = 1 + 3 * (row - 1);
|
|||
|
|
dvalues[val_idx] = left; // ∂F_i/∂u_{i-1}
|
|||
|
|
dvalues[val_idx + 1] = center; // ∂F_i/∂u_i
|
|||
|
|
dvalues[val_idx + 2] = right; // ∂F_i/∂u_{i+1}
|
|||
|
|
}
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// residual: length N (full system including boundaries)
|
|||
|
|
template <typename ctx_t, typename T>
|
|||
|
|
void compute_residual_full(
|
|||
|
|
ctx_t& ctx, vector_t<T> U, vector_t<T> U_prev, vector_t<T> residual, size_t N, double h, double dt, double nu)
|
|||
|
|
{
|
|||
|
|
ctx.parallel_for(box(N), residual.write(), U.read(), U_prev.read()).set_symbol("compute_residual_full")
|
|||
|
|
->*[N, h, dt, nu] __device__(size_t i, auto dresidual, auto dU, auto dU_prev) {
|
|||
|
|
if (i == 0)
|
|||
|
|
{
|
|||
|
|
// Left boundary condition: u[0] = 0
|
|||
|
|
dresidual(i) = dU(i) - 0.0;
|
|||
|
|
}
|
|||
|
|
else if (i == N - 1)
|
|||
|
|
{
|
|||
|
|
// Right boundary condition: u[N-1] = 0
|
|||
|
|
dresidual(i) = dU(i) - 0.0;
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
// Interior point: Burger's equation F_i = ∂u/∂t + u*∂u/∂x - nu*∂²u/∂x²
|
|||
|
|
double u_i = dU(i);
|
|||
|
|
double u_ip1 = dU(i + 1);
|
|||
|
|
double u_im1 = dU(i - 1);
|
|||
|
|
|
|||
|
|
double term_time = (u_i - dU_prev(i)) / dt; // ∂u/∂t
|
|||
|
|
double term_conv = u_i * (u_ip1 - u_im1) / (2 * h); // u * ∂u/∂x (nonlinear convection)
|
|||
|
|
double term_diff = -nu * (u_im1 - 2 * u_i + u_ip1) / (h * h); // -nu * ∂²u/∂x²
|
|||
|
|
|
|||
|
|
dresidual(i) = term_time + term_conv + term_diff;
|
|||
|
|
}
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Callback function objects for Burger's equation
|
|||
|
|
struct BurgerResidualCallback
|
|||
|
|
{
|
|||
|
|
size_t N;
|
|||
|
|
double h, dt, nu;
|
|||
|
|
|
|||
|
|
template <typename ctx_t>
|
|||
|
|
void
|
|||
|
|
operator()(ctx_t& ctx, const vector_t<double>& x, const vector_t<double>& x_prev, vector_t<double>& residual) const
|
|||
|
|
{
|
|||
|
|
compute_residual_full(ctx, x, x_prev, residual, N, h, dt, nu);
|
|||
|
|
}
|
|||
|
|
};
|
|||
|
|
|
|||
|
|
struct BurgerJacobianCallback
|
|||
|
|
{
|
|||
|
|
size_t N;
|
|||
|
|
double h, dt, nu;
|
|||
|
|
|
|||
|
|
template <typename ctx_t>
|
|||
|
|
void operator()(ctx_t& ctx, const vector_t<double>& x, vector_t<double>& jacobian_values) const
|
|||
|
|
{
|
|||
|
|
assemble_jacobian_full(ctx, x, jacobian_values, N, h, dt, nu);
|
|||
|
|
}
|
|||
|
|
};
|
|||
|
|
|
|||
|
|
// Initialize the solution output file (call once at simulation start)
|
|||
|
|
void initialize_solution_file(const char* filename, size_t N, double h)
|
|||
|
|
{
|
|||
|
|
FILE* fp = fopen(filename, "w");
|
|||
|
|
if (fp)
|
|||
|
|
{
|
|||
|
|
fprintf(fp, "# Burger equation solution - block format\n");
|
|||
|
|
fprintf(fp, "# Each timestep is a separate block, separated by blank lines\n");
|
|||
|
|
fprintf(fp, "# Format: x_coordinate u(x,t)\n");
|
|||
|
|
fprintf(fp, "# Grid points: %zu, h=%.6e\n", N, h);
|
|||
|
|
fprintf(fp,
|
|||
|
|
"# Use in gnuplot: plot for [i=0:*] 'solution.dat' index i with lines title sprintf('step %%d', i*10)\n");
|
|||
|
|
fprintf(fp, "#\n");
|
|||
|
|
fclose(fp);
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
printf("Error: Could not create %s for writing\n", filename);
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Function to append timestep block to solution file (simple and reliable)
|
|||
|
|
template <typename ctx_t>
|
|||
|
|
void dump_solution(
|
|||
|
|
ctx_t& ctx, vector_t<double>& U, size_t timestep, size_t N, double h, double dt, const char* filename = "solution.dat")
|
|||
|
|
{
|
|||
|
|
ctx.host_launch(U.read()).set_symbol("dump solution")->*[timestep, h, N, dt, filename](auto hU) {
|
|||
|
|
FILE* fp = fopen(filename, "a"); // Simple append - no read/modify/write
|
|||
|
|
if (fp)
|
|||
|
|
{
|
|||
|
|
fprintf(fp, "# Timestep %zu, t=%.6e\n", timestep, timestep * dt);
|
|||
|
|
|
|||
|
|
for (size_t i = 0; i < N; i++)
|
|||
|
|
{
|
|||
|
|
double x = i * h;
|
|||
|
|
fprintf(fp, "%.10e %.10e\n", x, hU(i));
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
fprintf(fp, "\n"); // Blank line to separate datasets
|
|||
|
|
fclose(fp);
|
|||
|
|
|
|||
|
|
printf("Appended timestep %zu (t=%.4e) to %s\n", timestep, timestep * dt, filename);
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
printf("Error: Could not open %s for appending\n", filename);
|
|||
|
|
}
|
|||
|
|
};
|
|||
|
|
}
|
|||
|
|
#endif
|
|||
|
|
|
|||
|
|
int main([[maybe_unused]] int argc, [[maybe_unused]] char** argv)
|
|||
|
|
{
|
|||
|
|
#if _CCCL_CTK_BELOW(12, 4)
|
|||
|
|
fprintf(stderr, "Waiving test: conditional nodes are only available since CUDA 12.4.\n");
|
|||
|
|
return 0;
|
|||
|
|
#else
|
|||
|
|
// Usage: ./burger [N] [nsteps] [nu]
|
|||
|
|
// N = Grid points (default: 100000)
|
|||
|
|
// nsteps = Time steps (default: 10000)
|
|||
|
|
// nu = Viscosity (default: 0.05, try 0.001 for shocks)
|
|||
|
|
|
|||
|
|
stackable_ctx ctx;
|
|||
|
|
|
|||
|
|
size_t N = 2560;
|
|||
|
|
if (argc > 1)
|
|||
|
|
{
|
|||
|
|
N = atoi(argv[1]);
|
|||
|
|
fprintf(stderr, "N = %zu\n", N);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
size_t nsteps = 200;
|
|||
|
|
if (argc > 2)
|
|||
|
|
{
|
|||
|
|
nsteps = atol(argv[2]);
|
|||
|
|
fprintf(stderr, "nsteps = %ld\n", nsteps);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Set reasonable parameters - implicit method allows larger time steps
|
|||
|
|
double nu = 0.05; // Default viscosity
|
|||
|
|
if (argc > 3)
|
|||
|
|
{
|
|||
|
|
nu = atof(argv[3]);
|
|||
|
|
fprintf(stderr, "nu = %e\n", nu);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
ssize_t output_freq = -1;
|
|||
|
|
if (argc > 4)
|
|||
|
|
{
|
|||
|
|
output_freq = atoi(argv[4]);
|
|||
|
|
fprintf(stderr, "output_freq %ld\n", output_freq);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// use_while = 0 => no while; 1 => while in CG, 2 => while in Newton and CG
|
|||
|
|
int use_while = 2;
|
|||
|
|
if (argc > 5)
|
|||
|
|
{
|
|||
|
|
use_while = atoi(argv[5]);
|
|||
|
|
fprintf(stderr, "use_while = %d\n", use_while);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
double h = 1.0 / (N - 1);
|
|||
|
|
|
|||
|
|
double dt_diffusion = 0.5 * h * h / nu; // Diffusion-limited time step
|
|||
|
|
double dt_fixed = 0.001; // Fixed reasonable time step
|
|||
|
|
double dt = std::max(dt_diffusion, dt_fixed); // Use larger of the two
|
|||
|
|
|
|||
|
|
// For very fine grids, cap the time step to prevent tiny steps
|
|||
|
|
if (N > 10000)
|
|||
|
|
{
|
|||
|
|
dt = std::min(dt, 0.01); // Cap at 0.01 for large grids
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
double total_time = nsteps * dt;
|
|||
|
|
|
|||
|
|
fprintf(stderr, "=== Simulation Parameters ===\n");
|
|||
|
|
fprintf(stderr, "Grid: N=%zu, h=%e\n", N, h);
|
|||
|
|
fprintf(stderr, "Time: dt=%e, nsteps=%zu, total_time=%e\n", dt, nsteps, total_time);
|
|||
|
|
fprintf(stderr, "Physics: nu=%e (viscosity)\n", nu);
|
|||
|
|
fprintf(stderr, "Diffusion number: nu*dt/h^2 = %e\n", nu * dt / (h * h));
|
|||
|
|
fprintf(stderr, "=============================\n");
|
|||
|
|
|
|||
|
|
// Full N×N system: boundary rows have 1 entry each, interior rows have 3 entries each
|
|||
|
|
// Total: 2*1 + (N-2)*3 = 3*N - 4 non-zeros
|
|||
|
|
size_t nz = 3 * N - 4;
|
|||
|
|
|
|||
|
|
size_t* row_offsets;
|
|||
|
|
size_t* col_indices;
|
|||
|
|
cuda_safe_call(cudaHostAlloc(&row_offsets, (N + 1) * sizeof(size_t), cudaHostAllocMapped));
|
|||
|
|
cuda_safe_call(cudaHostAlloc(&col_indices, nz * sizeof(size_t), cudaHostAllocMapped));
|
|||
|
|
|
|||
|
|
build_full_csr_structure(row_offsets, col_indices, N);
|
|||
|
|
|
|||
|
|
auto csr_row_offsets = ctx.logical_data(make_slice(row_offsets, N + 1)).set_symbol("csr_row");
|
|||
|
|
auto csr_col_ind = ctx.logical_data(make_slice(col_indices, nz)).set_symbol("csr_col");
|
|||
|
|
auto csr_values = ctx.logical_data(shape_of<slice<double>>(nz)).set_symbol("csr_val");
|
|||
|
|
|
|||
|
|
auto U = ctx.logical_data(shape_of<slice<double>>(N)).set_symbol("U");
|
|||
|
|
|
|||
|
|
// This will prevent erroneous modifications and may allow access from concurrent graphs
|
|||
|
|
csr_row_offsets.set_read_only();
|
|||
|
|
csr_col_ind.set_read_only();
|
|||
|
|
|
|||
|
|
// Initial condition
|
|||
|
|
ctx.parallel_for(U.shape(), U.write()).set_symbol("init conditions")->*[h, N] __device__(size_t i, auto dU) {
|
|||
|
|
double x = i * h;
|
|||
|
|
if (i == 0 || i == N - 1)
|
|||
|
|
{
|
|||
|
|
dU(i) = 0.0; // Homogeneous Dirichlet boundary conditions
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
dU(i) = sin(M_PI * x);
|
|||
|
|
}
|
|||
|
|
};
|
|||
|
|
|
|||
|
|
// Initialize solution output file
|
|||
|
|
initialize_solution_file("solution.dat", N, h);
|
|||
|
|
|
|||
|
|
auto start = std::chrono::high_resolution_clock::now();
|
|||
|
|
cuda_safe_call(cudaStreamSynchronize(ctx.fence()));
|
|||
|
|
|
|||
|
|
// Parameters are now set above with auto-scaling
|
|||
|
|
size_t substeps = (output_freq > 0) ? output_freq : nsteps;
|
|||
|
|
size_t outer_iterations = nsteps / substeps;
|
|||
|
|
|
|||
|
|
if (use_while == 2)
|
|||
|
|
{
|
|||
|
|
for (size_t outer = 0; outer < outer_iterations; outer++)
|
|||
|
|
{
|
|||
|
|
auto g = ctx.graph_scope();
|
|||
|
|
|
|||
|
|
// Repeat substeps inner iterations using STF repeat block
|
|||
|
|
{
|
|||
|
|
auto repeat_guard = ctx.repeat_graph_scope(substeps);
|
|||
|
|
|
|||
|
|
// Create callback function objects for Burger's equation
|
|||
|
|
BurgerResidualCallback residual_callback{N, h, dt, nu};
|
|||
|
|
BurgerJacobianCallback jacobian_callback{N, h, dt, nu};
|
|||
|
|
|
|||
|
|
// Solve the nonlinear system using generic Newton solver
|
|||
|
|
newton_solver(ctx, U, csr_values, csr_row_offsets, csr_col_ind, residual_callback, jacobian_callback);
|
|||
|
|
} // repeat_guard automatically manages the loop condition
|
|||
|
|
|
|||
|
|
// Dump solution after each substep block
|
|||
|
|
size_t current_timestep = (outer + 1) * substeps;
|
|||
|
|
dump_solution(ctx, U, current_timestep, N, h, dt);
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
else
|
|||
|
|
{
|
|||
|
|
for (size_t outer = 0; outer < outer_iterations; outer++)
|
|||
|
|
{
|
|||
|
|
// Repeat substeps inner iterations using STF repeat block
|
|||
|
|
for (size_t substep = 0; substep < substeps; substep++)
|
|||
|
|
{
|
|||
|
|
// Create callback function objects for Burger's equation
|
|||
|
|
BurgerResidualCallback residual_callback{N, h, dt, nu};
|
|||
|
|
BurgerJacobianCallback jacobian_callback{N, h, dt, nu};
|
|||
|
|
|
|||
|
|
// Solve the nonlinear system using generic Newton solver
|
|||
|
|
newton_solver_no_while(
|
|||
|
|
ctx, U, csr_values, csr_row_offsets, csr_col_ind, residual_callback, jacobian_callback, use_while == 1);
|
|||
|
|
} // repeat_guard automatically manages the loop condition
|
|||
|
|
|
|||
|
|
// Dump solution after each substep block
|
|||
|
|
size_t current_timestep = (outer + 1) * substeps;
|
|||
|
|
dump_solution(ctx, U, current_timestep, N, h, dt);
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
cuda_safe_call(cudaStreamSynchronize(ctx.fence()));
|
|||
|
|
auto end = std::chrono::high_resolution_clock::now();
|
|||
|
|
auto duration = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
|
|||
|
|
std::cout << "Duration: " << duration << " milliseconds" << '\n';
|
|||
|
|
|
|||
|
|
ctx.finalize();
|
|||
|
|
#endif
|
|||
|
|
}
|