Skip to content

Commit 57a528e

Browse files
AhdhnAhdhn
authored
Merge pull request #77 from owensgroup/tet
Adding Tet support for sparse matrices, solvers, and autodiff
2 parents e468c34 + 548413e commit 57a528e

33 files changed

Lines changed: 1333 additions & 1473 deletions

‎apps/CMakeLists.txt‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@ add_subdirectory(SurfaceTracking)
1515
add_subdirectory(SCP)
1616
add_subdirectory(ARAP)
1717
add_subdirectory(Heat)
18+
add_subdirectory(TetHeat)
1819
add_subdirectory(Param)
1920
add_subdirectory(ManiOpt)
2021
add_subdirectory(MassSpring)

‎apps/TetHeat/CMakeLists.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
set(SOURCE_LIST
2+
tet_heat.cu
3+
)
4+
5+
rxmesh_add_app(TetHeat
6+
SOURCES ${SOURCE_LIST}
7+
LIBS RXMesh CLI11::CLI11
8+
)

‎apps/TetHeat/tet_heat.cu‎

Lines changed: 337 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,337 @@
1+
#include <CLI/CLI.hpp>
2+
3+
#include <algorithm>
4+
#include <cmath>
5+
#include <limits>
6+
#include <string>
7+
#include <type_traits>
8+
9+
#include "rxmesh/matrix/cg_solver.h"
10+
#ifdef USE_CUDSS
11+
#include "rxmesh/matrix/cudss_cholesky_solver.h"
12+
#endif
13+
#include "rxmesh/matrix/pcg_solver.h"
14+
#include "rxmesh/matrix/sparse_matrix.h"
15+
#include "rxmesh/reduce_handle.h"
16+
#include "rxmesh/rxmesh_static.h"
17+
#include "rxmesh/util/timer.h"
18+
19+
using namespace rxmesh;
20+
21+
struct SolveStats
22+
{
23+
float pre_solve_ms = 0.0f;
24+
float time_stepping_ms = 0.0f;
25+
uint64_t total_iterations = 0;
26+
double final_solver_residual = 0.0;
27+
bool converged = true;
28+
};
29+
30+
template <typename T>
31+
void assemble_heat_system(const RXMeshStatic& rx,
32+
const VertexAttribute<T>& coordinates,
33+
const T dt,
34+
SparseMatrix<T>& system,
35+
VertexAttribute<T>& mass)
36+
{
37+
system.reset(T(0), DEVICE);
38+
mass.reset(T(0), DEVICE);
39+
40+
rx.for_each<Op::TV, 256>([=] __device__(const TetHandle&,
41+
const VertexIterator& tv) mutable {
42+
const vec3<T> x[4] = {coordinates.template to_glm<3>(tv[0]),
43+
coordinates.template to_glm<3>(tv[1]),
44+
coordinates.template to_glm<3>(tv[2]),
45+
coordinates.template to_glm<3>(tv[3])};
46+
47+
T max_edge = T(0);
48+
for (int i = 0; i < 4; ++i) {
49+
for (int j = i + 1; j < 4; ++j) {
50+
max_edge = std::max(max_edge, glm::length(x[j] - x[i]));
51+
}
52+
}
53+
54+
const vec3<T> e1 = x[1] - x[0];
55+
const vec3<T> e2 = x[2] - x[0];
56+
const vec3<T> e3 = x[3] - x[0];
57+
const T det = glm::dot(e1, glm::cross(e2, e3));
58+
59+
// const T det_tolerance = T(32) * std::numeric_limits<T>::epsilon() *
60+
// max_edge * max_edge * max_edge;
61+
// if (!isfinite(det) || !isfinite(max_edge) ||
62+
// std::abs(det) <= det_tolerance) {
63+
// return;
64+
// }
65+
66+
const T volume = std::abs(det) / T(6);
67+
const vec3<T> g1 = glm::cross(e2, e3) / det;
68+
const vec3<T> g2 = glm::cross(e3, e1) / det;
69+
const vec3<T> g3 = glm::cross(e1, e2) / det;
70+
const vec3<T> gradient[4] = {-(g1 + g2 + g3), g1, g2, g3};
71+
72+
for (int i = 0; i < 4; ++i) {
73+
for (int j = 0; j < 4; ++j) {
74+
::atomicAdd(&system(tv[i], tv[j]),
75+
dt * (volume * glm::dot(gradient[i], gradient[j])));
76+
}
77+
}
78+
79+
const T vertex_mass = volume / T(4);
80+
for (int i = 0; i < 4; ++i) {
81+
::atomicAdd(&mass(tv[i]), vertex_mass);
82+
::atomicAdd(&system(tv[i], tv[i]), vertex_mass);
83+
}
84+
});
85+
}
86+
87+
template <typename T>
88+
void build_rhs(const RXMeshStatic& rx,
89+
const VertexAttribute<T>& mass,
90+
const DenseMatrix<T>& temperature,
91+
DenseMatrix<T>& rhs)
92+
{
93+
rx.for_each_vertex(DEVICE, [=] __device__(const VertexHandle vh) mutable {
94+
rhs(vh) = mass(vh) * temperature(vh);
95+
});
96+
}
97+
98+
99+
template <typename T, typename SolverT>
100+
SolveStats solve_iterative(const RXMeshStatic& rx,
101+
const VertexAttribute<T>& mass,
102+
DenseMatrix<T>& rhs,
103+
DenseMatrix<T>& temperature,
104+
const int steps,
105+
const int max_iterations,
106+
SolverT& solver)
107+
{
108+
SolveStats stats;
109+
GPUTimer timer;
110+
timer.start();
111+
for (int step = 0; step < steps; ++step) {
112+
build_rhs(rx, mass, temperature, rhs);
113+
solver.pre_solve(rhs, temperature);
114+
solver.solve(rhs, temperature);
115+
stats.total_iterations += solver.iter_taken();
116+
stats.final_solver_residual = solver.final_residual();
117+
stats.converged = stats.converged &&
118+
solver.iter_taken() < max_iterations &&
119+
std::isfinite(stats.final_solver_residual);
120+
}
121+
timer.stop();
122+
stats.time_stepping_ms = timer.elapsed_millis();
123+
return stats;
124+
}
125+
126+
int main(int argc, char** argv)
127+
{
128+
using T = rx_coord_t;
129+
130+
CLI::App app{"Implicit heat diffusion on a tet mesh"};
131+
132+
std::string mesh_path = STRINGIFY(INPUT_DIR) "car.msh";
133+
std::string solver_name = "pcg";
134+
uint32_t device_id = 0;
135+
uint32_t source_vid = 0;
136+
T t_factor = T(1);
137+
int steps = 10;
138+
int cg_max_iter = 10000;
139+
140+
141+
app.add_option("-i,--input", mesh_path, "Input tetrahedral MSH file")
142+
->default_val(mesh_path);
143+
app.add_option("-d,--device_id", device_id, "GPU device ID")
144+
->default_val(device_id);
145+
app.add_option("-s,--source", source_vid, "Initial hot vertex")
146+
->default_val(source_vid);
147+
app.add_option("-t,--t-factor",
148+
t_factor,
149+
"Time-step multiplier on mean edge length squared")
150+
->default_val(t_factor);
151+
app.add_option("-n,--steps", steps, "Number of implicit time steps")
152+
->default_val(steps);
153+
app.add_option("-l,--solver", solver_name, "Solver: cg, pcg, or cudss")
154+
->default_val(solver_name);
155+
app.add_option(
156+
"-c,--cg_max_iter", cg_max_iter, "Maximum iterations for CG and PCG")
157+
->default_val(cg_max_iter);
158+
159+
try {
160+
app.parse(argc, argv);
161+
} catch (const CLI::ParseError& e) {
162+
return app.exit(e);
163+
}
164+
165+
if (solver_name != "cg" && solver_name != "pcg" && solver_name != "cudss") {
166+
RXMESH_ERROR("Unsupported solver '{}'. Use cg, pcg, or cudss",
167+
solver_name);
168+
return EXIT_FAILURE;
169+
}
170+
#ifndef USE_CUDSS
171+
if (solver_name == "cudss") {
172+
RXMESH_ERROR("The cudss solver requires RX_USE_CUDSS=ON");
173+
return EXIT_FAILURE;
174+
}
175+
#endif
176+
177+
rx_init(device_id);
178+
179+
RXMESH_INFO("input = {}", mesh_path);
180+
RXMESH_INFO("device_id = {}", device_id);
181+
RXMESH_INFO("source = {}", source_vid);
182+
RXMESH_INFO("t_factor = {}", t_factor);
183+
RXMESH_INFO("steps = {}", steps);
184+
RXMESH_INFO("solver = {}", solver_name);
185+
RXMESH_INFO("cg_max_iter = {}", cg_max_iter);
186+
187+
RXMeshStatic rx(mesh_path);
188+
189+
const uint32_t num_vertices = rx.get_num_vertices();
190+
const uint32_t num_edges = rx.get_num_edges();
191+
const uint32_t num_tets = rx.get_num_tets();
192+
193+
// allocate attributes and sparse matrices
194+
auto coordinates = *rx.get_input_vertex_coordinates();
195+
auto edge_length = *rx.add_edge_attribute<T>("edge_length", 1);
196+
auto mass = *rx.add_vertex_attribute<T>("mass", 1);
197+
DenseMatrix<T> rhs(rx, num_vertices, 1, DEVICE);
198+
DenseMatrix<T> temperature(rx, num_vertices, 1, LOCATION_ALL);
199+
SparseMatrix<T> system(rx, Op::VV);
200+
201+
if (rx.get_num_components() > 1) {
202+
RXMESH_WARN("Input mesh has {} components", rx.get_num_components());
203+
}
204+
205+
// pick the handle
206+
VertexHandle source_handle;
207+
rx.for_each_vertex(HOST, [&](const VertexHandle vh) {
208+
if (rx.map_to_global(vh) == source_vid) {
209+
source_handle = vh;
210+
}
211+
});
212+
temperature.reset(T(0), LOCATION_ALL);
213+
temperature(source_handle) = T(100);
214+
temperature.move(HOST, DEVICE);
215+
216+
// calc mean edge len
217+
rx.for_each<Op::EV, 256>(
218+
[=] __device__(const EdgeHandle& eh, const VertexIterator& ev) mutable {
219+
const vec3<T> a = coordinates.template to_glm<3>(ev[0]);
220+
const vec3<T> b = coordinates.template to_glm<3>(ev[1]);
221+
edge_length(eh) = glm::length(a - b);
222+
});
223+
EdgeReduceHandle<T> edge_reducer(edge_length);
224+
const T edge_length_sum =
225+
edge_reducer.reduce(edge_length, cub::Sum(), T(0));
226+
const T mean_edge_length = edge_length_sum / static_cast<T>(num_edges);
227+
228+
const T dt = t_factor * mean_edge_length * mean_edge_length;
229+
230+
RXMESH_INFO("#vertices = {}, #edges = {}, #tets = {}",
231+
num_vertices,
232+
num_edges,
233+
num_tets);
234+
RXMESH_INFO("Mean edge length = {}, dt = {}", mean_edge_length, dt);
235+
236+
237+
// assemble the system
238+
GPUTimer assembly_timer;
239+
assembly_timer.start();
240+
assemble_heat_system(rx, coordinates, dt, system, mass);
241+
assembly_timer.stop();
242+
const float assembly_ms = assembly_timer.elapsed_millis();
243+
244+
// solve
245+
const T iterative_abs_tolerance = std::numeric_limits<T>::min();
246+
const T iterative_rel_tolerance =
247+
std::is_same_v<T, double> ? T(1e-24) : T(1e-12);
248+
249+
SolveStats solve_stats;
250+
if (solver_name == "pcg") {
251+
PCGSolver<T> solver(system,
252+
1,
253+
cg_max_iter,
254+
iterative_abs_tolerance,
255+
iterative_rel_tolerance);
256+
solve_stats = solve_iterative(
257+
rx, mass, rhs, temperature, steps, cg_max_iter, solver);
258+
} else if (solver_name == "cg") {
259+
CGSolver<T> solver(system,
260+
1,
261+
cg_max_iter,
262+
iterative_abs_tolerance,
263+
iterative_rel_tolerance);
264+
solve_stats = solve_iterative(
265+
rx, mass, rhs, temperature, steps, cg_max_iter, solver);
266+
#ifdef USE_CUDSS
267+
} else if (solver_name == "cudss") {
268+
cuDSSCholeskySolver<SparseMatrix<T>> solver(&system);
269+
build_rhs(rx, mass, temperature, rhs);
270+
271+
GPUTimer pre_solve_timer;
272+
pre_solve_timer.start();
273+
solver.pre_solve(rx, rhs, temperature);
274+
pre_solve_timer.stop();
275+
solve_stats.pre_solve_ms = pre_solve_timer.elapsed_millis();
276+
277+
GPUTimer time_stepping_timer;
278+
time_stepping_timer.start();
279+
for (int step = 0; step < steps; ++step) {
280+
if (step != 0) {
281+
build_rhs(rx, mass, temperature, rhs);
282+
}
283+
solver.solve(rhs, temperature);
284+
}
285+
time_stepping_timer.stop();
286+
solve_stats.time_stepping_ms = time_stepping_timer.elapsed_millis();
287+
#endif
288+
}
289+
290+
if (!solve_stats.converged) {
291+
RXMESH_WARN("An iterative heat solve did not converge");
292+
}
293+
294+
// calc residual
295+
DenseMatrix<T> residual(rx, num_vertices, 1, DEVICE);
296+
system.multiply(temperature, residual);
297+
residual.axpy(rhs, T(-1));
298+
const double rhs_norm = static_cast<double>(rhs.norm2());
299+
const double residual_norm = static_cast<double>(residual.norm2());
300+
const double relative_residual = residual_norm / rhs_norm;
301+
302+
RXMESH_INFO("Assembly took {} ms", assembly_ms);
303+
if (solver_name == "cudss") {
304+
RXMESH_INFO("cuDSS analysis and factorization took {} ms",
305+
solve_stats.pre_solve_ms);
306+
} else {
307+
RXMESH_INFO("Solver took {} iterations and final solver residual = {}",
308+
solve_stats.total_iterations,
309+
solve_stats.final_solver_residual);
310+
}
311+
RXMESH_INFO("{} heat steps took {} ms ({} ms/step)",
312+
steps,
313+
solve_stats.time_stepping_ms,
314+
solve_stats.time_stepping_ms / static_cast<float>(steps));
315+
316+
317+
RXMESH_INFO("Final ordinary relative residual = {}", relative_residual);
318+
319+
320+
#if USE_POLYSCOPE
321+
temperature.move(DEVICE, HOST);
322+
auto final_temperature = *rx.add_vertex_attribute<T>("temperature", 1);
323+
final_temperature.from_matrix(&temperature);
324+
rx.get_polyscope_volume_mesh()
325+
->addVertexScalarQuantity("temperature", final_temperature)
326+
->setEnabled(true);
327+
polyscope::show();
328+
329+
#endif
330+
331+
residual.release();
332+
rhs.release();
333+
temperature.release();
334+
system.release();
335+
336+
return 0;
337+
}

‎cmake/RXMeshConfig.cmake.in‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@ include(CMakeFindDependencyMacro)
1010
include("${CMAKE_CURRENT_LIST_DIR}/RXMeshApp.cmake")
1111

1212
find_dependency(CUDAToolkit REQUIRED)
13+
if(@RX_USE_CUDSS@)
14+
find_dependency(cudss REQUIRED)
15+
endif()
1316
find_package(OpenMP QUIET)
1417

1518
set(_RXMESH_INCLUDE_DIR "${PACKAGE_PREFIX_DIR}/@RXMESH_INSTALL_INCLUDEDIR@")
@@ -32,6 +35,7 @@ endfunction()
3235

3336
_rxmesh_import_static_target(RXMesh::GKlib GKlib)
3437
_rxmesh_import_static_target(RXMesh::metis metis)
38+
_rxmesh_import_static_target(RXMesh::mshio mshio)
3539
_rxmesh_import_static_target(RXMesh::polyscope polyscope)
3640
_rxmesh_import_static_target(RXMesh::imgui imgui)
3741
_rxmesh_import_static_target(RXMesh::glad glad)
@@ -98,10 +102,22 @@ set(_rxmesh_link_libraries
98102
CUDA::cusparse
99103
CUDA::cusolver)
100104

105+
if(@RX_USE_CUDSS@)
106+
if(WIN32)
107+
list(APPEND _rxmesh_link_libraries cudss)
108+
else()
109+
list(APPEND _rxmesh_link_libraries cudss_static)
110+
endif()
111+
endif()
112+
101113
if(TARGET RXMesh::metis)
102114
list(APPEND _rxmesh_link_libraries RXMesh::metis)
103115
endif()
104116

117+
if(TARGET RXMesh::mshio)
118+
list(APPEND _rxmesh_link_libraries RXMesh::mshio)
119+
endif()
120+
105121
if(@RX_USE_POLYSCOPE@ AND TARGET RXMesh::polyscope)
106122
list(APPEND _rxmesh_link_libraries RXMesh::polyscope)
107123
endif()

0 commit comments

Comments
 (0)