General-purpose batched tridiagonal solver #533
Conversation
ali-ramadhan
added
feature 🌟
Codecov Report
@@ Coverage Diff @@
## master #533 +/- ##
==========================================
+ Coverage 71.09% 71.47% +0.37%
==========================================
Files 69 70 +1
Lines 1958 1991 +33
==========================================
+ Hits 1392 1423 +31
- Misses 566 568 +2
Continue to review full report at Codecov.
|
|
@glwagner @suyashbire1 This PR should be ready for review. Let me know what you guys think. Right now it's a pretty standalone solver but I'd like to convert the example in PR #306 to a pressure solver that uses Future improvements for solving batched tridiagonal systems on the GPU:
|
|
GitLab CI seems to have crapped out but GPU tests pass on Cyclops: https://gist.github.com/ali-ramadhan/ec039bd71c21daff23c140b8aec180eb |
There's some people that are heavily using GPUs on |
|
@maleadt Oh that's awesome, thanks for the update! So we just update to candidates = [(dev=dev,
cap=capability(dev),
mem=CuContext(ctx->CUDAdrv.available_memory(), dev))
for dev in devices()]
thorough = parse(Bool, get(ENV, "CI_THOROUGH", "false"))
if thorough
sort!(candidates, by=x->(x.cap, x.mem))
else
sort!(candidates, by=x->x.mem)
end
pick = last(candidates)
device!(pick.dev) |
Like that. By default, CuArrays follows the default CUDA order, I think it only makes sense for CI to pick a device with as much memory as possible. |
|
Ah thanks for the tip @maleadt! I'll make a PR to improve the GPU CI. |
| ϕ[i, j, 1] = f₁ / β | ||
|
|
||
| @unroll for k = 2:Nz | ||
| cₖ₋₁ = get_coefficient(c, i, j, k-1, grid, p) |
There was a problem hiding this comment.
These symbols don't render on github; all I see is cₖ₋₁, for example.
There was a problem hiding this comment.
Ah hmmm it's just subscript k-1. Looks fine on my browser so must be font-dependent.
Could change to superscript as we seem to have had fewer issues with superscript in the past.
It's pretty low-level/private code so I'll leave it with subscripts for now but if becomes an issue outside of GitHub once more then we can change it.
test/test_solvers.jl
Outdated
|
|
||
| solve_batched_tridiagonal_system!(ϕ, arch, btsolver) | ||
|
|
||
| return ϕ[:] ≈ ϕ_correct |
There was a problem hiding this comment.
ϕ_correct is a 1D array because backslash returns a vector while ϕ is a 1×1×Nz array so I use ϕ[:] to flatten it to a 1D vector for ≈ to work.
This PR adds a general-purpose batched tridiagonal solver which we can use for vertically stretched pressure solves, implicit vertical diffusion, sub-stepping of atmospheric acoustic waves, etc.
The coefficients and right-hand-side can be specified as a 1D array (shared by all the tridiagonal systems), a 3D array (different for each tridiagonal system), or a function (coefficient is calculated on the fly).
Should probably add proper docstrings before merging.
The vertically stretched pressure solver (PR #306) can be built on top of this batched tridiagonal solver.