Add disk cache infrastructure for Julia 1.11 #557
Merged
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
vchuravy
commented
Apr 16, 2024
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #557 +/- ##
==========================================
- Coverage 74.75% 71.80% -2.95%
==========================================
Files 24 24
Lines 3414 3469 +55
==========================================
- Hits 2552 2491 -61
- Misses 862 978 +116 ☔ View full report in Codecov by Sentry. |
|
Okay the big caveat is that this will only work with code that has been precompiled. So we will need to make sure our inference caches are hot. |
|
Using this small code module GemmDenseCUDA
using PrecompileTools
import CUDA
import CUDA: i32
BLOCK_SIZE = 32
function gemm!(A, B, C)
row = (CUDA.blockIdx().x - Int32(1)) * CUDA.blockDim().x + CUDA.threadIdx().x
col = (CUDA.blockIdx().y - Int32(1)) * CUDA.blockDim().y + CUDA.threadIdx().y
sum = zero(eltype(C))
if row <= size(A, 1) && col <= size(B, 2)
for i in 1:size(A, 2)
@inbounds sum += A[row, i] * B[i, col]
end
@inbounds C[row, col] = sum
end
return
end
@setup_workload let
if CUDA.functional()
A16 = CUDA.CuArray{Float16,2}(undef, 0, 0)
A32 = CUDA.CuArray{Float32,2}(undef, 0, 0)
A64 = CUDA.CuArray{Float64,2}(undef, 0, 0)
@compile_workload begin
CUDA.@cuda launch=false gemm!(A16, A16, A16)
CUDA.@cuda launch=false gemm!(A32, A32, A32)
CUDA.@cuda launch=false gemm!(A64, A64, A32)
end
end
end
end #moduleTesting in REPLFirst compilationSecond compilationWithout disk-cache (just inference caching)First compilationSecond compilationWould be interesting to see the impact of JuliaGPU/CUDA.jl#2325 With disk-cacheFirst compilationSecond Compilation
|
|
|
So precompiling with disk_cache = true, leads to this interesting hang. |
|
sigh |
|
The issue is that our The issue is that during image creation we sometimes broaden our horizon and then try to compile based on methods |
This was referenced Apr 19, 2024
WaterLily TGV example
ClimaOcean OMIPNote that I tried to make the model as "small" as possible,
|
maleadt
reviewed
Jun 24, 2024
maleadt
reviewed
Jun 28, 2024
|
Yay! |
|
Ah we might still need to add an interface for backends to opt out. E.g. Enzyme |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Replaces #351
Using JuliaLang/julia#52233 for ensuring inference correctness.
Note the CodeInstances should be precompiled, which we don't yet have a nice-ish infrastructure for.
Might need JuliaLang/julia#53943 so that we can encode the dependencies of the CodeInstance,
instead of just
GPUCompiler.TODO:
precompilefunctionCompilerConfigto be robust against compiler params changes