-
Notifications
You must be signed in to change notification settings - Fork 191
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
General-purpose batched tridiagonal solver #533
Conversation
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. |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Looks pretty good, and nice code. I'm impressed.
ϕ[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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
These symbols don't render on github; all I see is cₖ₋₁
, for example.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
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.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Why do we need the [:]
?
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
ϕ_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.