-
Notifications
You must be signed in to change notification settings - Fork 63
Expand file tree
/
Copy pathPrecompile.jl
More file actions
69 lines (62 loc) · 2.26 KB
/
Precompile.jl
File metadata and controls
69 lines (62 loc) · 2.26 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
using PrecompileTools: @setup_workload, @compile_workload
# Precompilation on 1.10 hits an apparent bug: https://github.com/JuliaLang/julia/issues/56947
function precompilation_supported()
return VERSION >= v"1.10.8"
end
function clear_oc_cache() end
if Reactant_jll.is_available()
@setup_workload begin
initialize_dialect()
if XLA.REACTANT_XLA_RUNTIME == "PJRT"
client = Accelerators.CPU.make_pjrt_client()
elseif XLA.REACTANT_XLA_RUNTIM
628B
E == "IFRT"
client = Accelerators.CPU.make_ifrt_client()
else
error("Unsupported runtime: $(XLA.REACTANT_XLA_RUNTIME)")
end
@compile_workload begin
@static if precompilation_supported()
x = ConcreteRNumber(2.0; client)
@static if VERSION >= v"1.11"
compile(sin, (x,); client, optimize=:all)
else
try
compile(sin, (x,); client, optimize=:all)
catch e
if !(e isa ReactantPrecompilationException)
rethrow()
end
end
end
if x isa ConcreteIFRTNumber
XLA.free_buffer(x.data.buffer)
x.data.buffer.buffer = C_NULL
else
for dat in x.data
XLA.free_buffer(dat.buffer)
dat.buffer.buffer = C_NULL
end
end
y = ConcreteRArray([2.0]; client)
try
compile(Base.sum, (y,); client, optimize=:all)
catch e
if !(e isa ReactantPrecompilationException)
rethrow()
end
end
if y isa ConcreteIFRTArray
XLA.free_buffer(y.data.buffer)
y.data.buffer.buffer = C_NULL
else
for dat in y.data
XLA.free_buffer(dat.buffer)
dat.buffer.buffer = C_NULL
end
end
end
end
XLA.free_client(client)
deinitialize_dialect()
end
end