Skip to content

Commit d11d9ee

Browse files
maleadtmaleadt
authored
Ensure global variables have a safe name. (#507)
1 parent c92ec6d commit d11d9ee

File tree

3 files changed

+120
-114
lines changed

3 files changed

+120
-114
lines changed

‎src/GPUCompiler.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ using Core: MethodInstance, CodeInstance, CodeInfo
1717
const use_newpm = LLVM.has_newpm()
1818

1919
include("utils.jl")
20+
include("mangling.jl")
2021

2122
# compiler interface and implementations
2223
include("interface.jl")

‎src/irgen.jl‎

Lines changed: 7 additions & 114 deletions
Original file line numberDiff line numberDiff line change
@@ -44,15 +44,13 @@ function irgen(@nospecialize(job::CompilerJob))
4444
Base.depwarn("GPUCompiler.process_module! is deprecated; implement GPUCompiler.finish_module! instead", :process_module)
4545
end
4646

47-
# sanitize function names
48-
# FIXME: Julia should do this, but sometimes fails (see maleadt/LLVM.jl#201)
49-
for f in functions(mod)
50-
isdeclaration(f) && continue
51-
llvmfn = LLVM.name(f)
52-
llvmfn′ = safe_name(llvmfn)
53-
if llvmfn != llvmfn′
54-
@assert !haskey(functions(mod), llvmfn′) "Cannot rename $llvmfn to $llvmfn′, already exists"
55-
LLVM.name!(f, llvmfn′)
47+
# sanitize global values (Julia doesn't when using the external codegen policy)
48+
for val in [collect(globals(mod)); collect(functions(mod))]
49+
isdeclaration(val) && continue
50+
old_name = LLVM.name(val)
51+
new_name = safe_name(old_name)
52+
if old_name != new_name
53+
LLVM.name!(val, new_name)
5654
end
5755
end
5856

@@ -123,111 +121,6 @@ function irgen(@nospecialize(job::CompilerJob))
123121
end
124122

125123

126-
## name mangling
127-
128-
# we generate function names that look like C++ functions, because many NVIDIA tools
129-
# support them, e.g., grouping different instantiations of the same kernel together.
130-
131-
function mangle_param(t, substitutions=String[])
132-
t == Nothing && return "v"
133-
134-
if isa(t, DataType) && t <: Ptr
135-
tn = mangle_param(eltype(t), substitutions)
136-
"P$tn"
137-
elseif isa(t, DataType)
138-
tn = safe_name(t)
139-
140-
# handle substitutions
141-
sub = findfirst(isequal(tn), substitutions)
142-
if sub === nothing
143-
str = "$(length(tn))$tn"
144-
push!(substitutions, tn)
145-
elseif sub == 1
146-
str = "S_"
147-
else
148-
str = "S$(sub-2)_"
149-
end
150-
151-
# encode typevars as template parameters
152-
if !isempty(t.parameters)
153-
str *= "I"
154-
for t in t.parameters
155-
str *= mangle_param(t, substitutions)
156-
end
157-
str *= "E"
158-
end
159-
160-
str
161-
elseif isa(t, Union)
162-
tn = "Union"
163-
164-
# handle substitutions
165-
sub = findfirst(isequal(tn), substitutions)
166-
if sub === nothing
167-
str = "$(length(tn))$tn"
168-
push!(substitutions, tn)
169-
elseif sub == 1
170-
str = "S_"
171-
else
172-
str = "S$(sub-2)_"
173-
end
174-
175-
# encode union types as template parameters
176-
if !isempty(Base.uniontypes(t))
177-
str *= "I"
178-
for t in Base.uniontypes(t)
179-
str *= mangle_param(t, substitutions)
180-
end
181-
str *= "E"
182-
end
183-
184-
str
185-
elseif isa(t, Integer)
186-
t > 0 ? "Li$(t)E" : "Lin$(abs(t))E"
187-
else
188-
tn = safe_name(t) # TODO: actually does support digits...
189-
if startswith(tn, r"\d")
190-
# C++ classes cannot start with a digit, so mangling doesn't support it
191-
tn = "_$(tn)"
192-
end
193-
"$(length(tn))$tn"
194-
end
195-
end
196-
197-
function mangle_sig(sig)
198-
ft, tt... = sig.parameters
199-
200-
# mangle the function name
201-
fn = safe_name(ft)
202-
str = "_Z$(length(fn))$fn"
203-
204-
# mangle each parameter
205-
substitutions = String[]
206-
for t in tt
207-
str *= mangle_param(t, substitutions)
208-
end
209-
210-
return str
211-
end
212-
213-
# make names safe for ptxas
214-
safe_name(fn::String) = replace(fn, r"[^A-Za-z0-9]"=>"_")
215-
safe_name(t::DataType) = safe_name(String(nameof(t)))
216-
function safe_name(t::Type{<:Function})
217-
# like Base.nameof, but for function types
218-
mt = t.name.mt
219-
fn = if mt === Symbol.name.mt
220-
# uses shared method table, so name is not unique to this function type
221-
nameof(t)
222-
else
223-
mt.name
224-
end
225-
safe_name(string(fn))
226-
end
227-
safe_name(::Type{Union{}}) = "Bottom"
228-
safe_name(x) = safe_name(repr(x))
229-
230-
231124
## exception handling
232125

233126
# this pass lowers `jl_throw` and friends to GPU-compatible exceptions.

‎src/mangling.jl‎

Lines changed: 112 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,112 @@
1+
# name mangling
2+
3+
# safe name generation
4+
5+
# LLVM doesn't like names with special characters, so we need to sanitize them.
6+
# note that we are stricter than LLVM, because of `ptxas`.
7+
8+
safe_name(fn::String) = replace(fn, r"[^A-Za-z0-9]"=>"_")
9+
10+
safe_name(t::DataType) = safe_name(String(nameof(t)))
11+
function safe_name(t::Type{<:Function})
12+
# like Base.nameof, but for function types
13+
mt = t.name.mt
14+
fn = if mt === Symbol.name.mt
15+
# uses shared method table, so name is not unique to this function type
16+
nameof(t)
17+
else
18+
mt.name
19+
end
20+
safe_name(string(fn))
21+
end
22+
safe_name(::Type{Union{}}) = "Bottom"
23+
24+
safe_name(x) = safe_name(repr(x))
25+
26+
27+
# C++ mangling
28+
29+
# we generate function names that look like C++ functions, because many tools, like NVIDIA's
30+
# profilers, support them (grouping different instantiations of the same kernel together).
31+
32+
function mangle_param(t, substitutions=String[])
33+
t == Nothing && return "v"
34+
35+
if isa(t, DataType) && t <: Ptr
36+
tn = mangle_param(eltype(t), substitutions)
37+
"P$tn"
38+
elseif isa(t, DataType)
39+
tn = safe_name(t)
40+
41+
# handle substitutions
42+
sub = findfirst(isequal(tn), substitutions)
43+
if sub === nothing
44+
str = "$(length(tn))$tn"
45+
push!(substitutions, tn)
46+
elseif sub == 1
47+
str = "S_"
48+
else
49+
str = "S$(sub-2)_"
50+
end
51+
52+
# encode typevars as template parameters
53+
if !isempty(t.parameters)
54+
str *= "I"
55+
for t in t.parameters
56+
str *= mangle_param(t, substitutions)
57+
end
58+
str *= "E"
59+
end
60+
61+
str
62+
elseif isa(t, Union)
63+
tn = "Union"
64+
65+
# handle substitutions
66+
sub = findfirst(isequal(tn), substitutions)
67+
if sub === nothing
68+
str = "$(length(tn))$tn"
69+
push!(substitutions, tn)
70+
elseif sub == 1
71+
str = "S_"
72+
else
73+
str = "S$(sub-2)_"
74+
end
75+
76+
# encode union types as template parameters
77+
if !isempty(Base.uniontypes(t))
78+
str *= "I"
79+
for t in Base.uniontypes(t)
80+
str *= mangle_param(t, substitutions)
81+
end
82+
str *= "E"
83+
end
84+
85+
str
86+
elseif isa(t, Integer)
87+
t > 0 ? "Li$(t)E" : "Lin$(abs(t))E"
88+
else
89+
tn = safe_name(t) # TODO: actually does support digits...
90+
if startswith(tn, r"\d")
91+
# C++ classes cannot start with a digit, so mangling doesn't support it
92+
tn = "_$(tn)"
93+
end
94+
"$(length(tn))$tn"
95+
end
96+
end
97+
98+
function mangle_sig(sig)
99+
ft, tt... = sig.parameters
100+
101+
# mangle the function name
102+
fn = safe_name(ft)
103+
str = "_Z$(length(fn))$fn"
104+
105+
# mangle each parameter
106+
substitutions = String[]
107+
for t in tt
108+
str *= mangle_param(t, substitutions)
109+
end
110+
111+
return str
112+
end

0 commit comments

Comments
 (0)