Simplify name mangling, and add support for Unions.

Author: maleadt

src/irgen.jl

@@ -61,9 +61,8 @@ function irgen(@nospecialize(job::CompilerJob); ctx::JuliaContextType)
     # rename and process the entry point
     if job.config.name !== nothing
         LLVM.name!(entry, safe_name(string("julia_", job.config.name)))
-    end
-    if job.config.kernel
-        LLVM.name!(entry, mangle_call(entry, job.source.specTypes))
+    elseif job.config.kernel
+        LLVM.name!(entry, mangle_sig(job.source.specTypes))
     end
     entry = process_entry!(job, mod, entry)
     if job.config.entry_abi === :specfunc
@@ -110,13 +109,13 @@ end
 # we generate function names that look like C++ functions, because many NVIDIA tools
 # support them, e.g., grouping different instantiations of the same kernel together.
 
-function mangle_param(t, substitutions)
+function mangle_param(t, substitutions=String[])
     t == Nothing && return "v"
 
     if isa(t, DataType) && t <: Ptr
         tn = mangle_param(eltype(t), substitutions)
         "P$tn"
-    elseif isa(t, DataType) || isa(t, Core.Function)
+    elseif isa(t, DataType)
         tn = safe_name(t)
 
         # handle substitutions
@@ -139,6 +138,30 @@ function mangle_param(t, substitutions)
             str *= "E"
         end
 
+        str
+    elseif isa(t, Union)
+        tn = "Union"
+
+        # handle substitutions
+        sub = findfirst(isequal(tn), substitutions)
+        if sub === nothing
+            str = "$(length(tn))$tn"
+            push!(substitutions, tn)
+        elseif sub == 1
+            str = "S_"
+        else
+            str = "S$(sub-2)_"
+        end
+
+        # encode union types as template parameters
+        if !isempty(Base.uniontypes(t))
+            str *= "I"
+            for t in Base.uniontypes(t)
+                str *= mangle_param(t, substitutions)
+            end
+            str *= "E"
+        end
+
         str
     elseif isa(t, Integer)
         t > 0 ? "Li$(t)E" : "Lin$(abs(t))E"
@@ -152,23 +175,36 @@ function mangle_param(t, substitutions)
     end
 end
 
-# TODO: tt is sug, drop f
-function mangle_call(f, tt)
-    fn = safe_name(f)
+function mangle_sig(sig)
+    ft, tt... = sig.parameters
+
+    # mangle the function name
+    fn = safe_name(ft)
     str = "_Z$(length(fn))$fn"
 
+    # mangle each parameter
     substitutions = String[]
-    for t in tt.parameters
+    for t in tt
         str *= mangle_param(t, substitutions)
     end
 
     return str
 end
 
 # make names safe for ptxas
-safe_name(fn::String) = replace(fn, r"[^A-Za-z0-9_]"=>"_")
-safe_name(f::Union{Core.Function,DataType}) = safe_name(String(nameof(f)))
-safe_name(f::LLVM.Function) = safe_name(LLVM.name(f))
+safe_name(fn::String) = replace(fn, r"[^A-Za-z0-9]"=>"_")
+safe_name(t::DataType) = safe_name(String(nameof(t)))
+function safe_name(t::Type{<:Function})
+    # like Base.nameof, but for function types
+    mt = t.name.mt
+    fn = if mt === Symbol.name.mt
+        # uses shared method table, so name is not unique to this function type
+        nameof(t)
+    else
+        mt.name
+    end
+    safe_name(string(fn))
+end
 safe_name(x) = safe_name(repr(x))
 
 

test/gcn.jl

@@ -84,9 +84,9 @@ end
     end
 
     asm = sprint(io->gcn_code_native(io, entry, Tuple{Int64}; dump_module=true, kernel=true))
-    @test occursin(r"\.amdhsa_kernel _Z\d*julia_entry", asm)
-    @test !occursin(r"\.amdhsa_kernel _Z\d*julia_nonentry", asm)
-    @test occursin(r"\.type.*julia_nonentry_\d*,@function", asm)
+    @test occursin(r"\.amdhsa_kernel \w*entry", asm)
+    @test !occursin(r"\.amdhsa_kernel \w*nonentry", asm)
+    @test occursin(r"\.type.*\w*nonentry\w*,@function", asm)
 end
 
 @testset "child function reuse" begin

test/metal.jl

@@ -24,22 +24,22 @@ end
     kernel(x) = return
 
     ir = sprint(io->metal_code_llvm(io, kernel, Tuple{Tuple{Int}}))
-    @test occursin(r"@.*julia.*kernel.*\(({ i64 }|\[1 x i64\])\*", ir)
+    @test occursin(r"@\w*kernel\w*\(({ i64 }|\[1 x i64\])\*", ir)
 
     # for kernels, every pointer argument needs to take an address space
     ir = sprint(io->metal_code_llvm(io, kernel, Tuple{Tuple{Int}}; kernel=true))
-    @test occursin(r"@.*julia.*kernel.*\(({ i64 }|\[1 x i64\]) addrspace\(1\)\*", ir)
+    @test occursin(r"@\w*kernel\w*\(({ i64 }|\[1 x i64\]) addrspace\(1\)\*", ir)
 end
 
 @testset "byref primitives" begin
     kernel(x) = return
 
     ir = sprint(io->metal_code_llvm(io, kernel, Tuple{Int}))
-    @test occursin(r"@.*julia.*kernel.*\(i64 ", ir)
+    @test occursin(r"@\w*kernel\w*\(i64 ", ir)
 
     # for kernels, every pointer argument needs to take an address space
     ir = sprint(io->metal_code_llvm(io, kernel, Tuple{Int}; kernel=true))
-    @test occursin(r"@.*julia.*kernel.*\(i64 addrspace\(1\)\*", ir)
+    @test occursin(r"@\w*kernel\w*\(i64 addrspace\(1\)\*", ir)
 end
 
 @testset "module metadata" begin
@@ -71,11 +71,11 @@ end
     end
 
     ir = sprint(io->metal_code_llvm(io, kernel, Tuple{Core.LLVMPtr{Int,1}}))
-    @test occursin(r"@.*julia.*kernel.*\(.* addrspace\(1\)\* %0\)", ir)
+    @test occursin(r"@\w*kernel\w*\(.* addrspace\(1\)\* %0\)", ir)
     @test occursin(r"call i32 @julia.air.thread_position_in_threadgroup.i32", ir)
 
     ir = sprint(io->metal_code_llvm(io, kernel, Tuple{Core.LLVMPtr{Int,1}}; kernel=true))
-    @test occursin(r"@.*julia.*kernel.*\(.* addrspace\(1\)\* %0, i32 %thread_position_in_threadgroup\)", ir)
+    @test occursin(r"@\w*kernel\w*\(.* addrspace\(1\)\* %0, i32 %thread_position_in_threadgroup\)", ir)
     @test !occursin(r"call i32 @julia.air.thread_position_in_threadgroup.i32", ir)
 end
 
@@ -94,8 +94,8 @@ Sys.isapple() && @testset "asm" begin
                                     dump_module=true, kernel=true))
     @test occursin("[header]", asm)
     @test occursin("[program]", asm)
-    @test occursin(r"name: .*julia.*kernel.*", asm)
-    @test occursin(r"define void @.*julia.*kernel.*", asm)
+    @test occursin(r"name: \w*kernel\w*", asm)
+    @test occursin(r"define void @\w*kernel\w*", asm)
 end
 
 end

test/ptx.jl

@@ -25,10 +25,10 @@ end
     end
 
     ir = sprint(io->ptx_code_llvm(io, kernel, Tuple{Aggregate}))
-    @test occursin(r"@.*julia_kernel.+\(({ i64 }|\[1 x i64\])\* ", ir)
+    @test occursin(r"@\w*kernel\w*\(({ i64 }|\[1 x i64\])\* ", ir)
 
     ir = sprint(io->ptx_code_llvm(io, kernel, Tuple{Aggregate}; kernel=true))
-    @test occursin(r"@.*julia_kernel.+\(.*({ i64 }|\[1 x i64\]) ", ir)
+    @test occursin(r"@\w*kernel\w*\(.*({ i64 }|\[1 x i64\]) ", ir)
 end
 
 @testset "property_annotations" begin
@@ -83,10 +83,10 @@ end
     kernel() = return
 
     ir = sprint(io->ptx_code_llvm(io, kernel, Tuple{}))
-    @test any(occursin(r"@.*kernel.+\(\)", a) for a in split(ir, "\n"))
+    @test occursin(r"@\w*kernel\w*\(\)", ir)
 
     ir = sprint(io->ptx_code_llvm(io, kernel, Tuple{}; kernel=true))
-    @test any(occursin(r"@.*kernel.+\(\[1 x i64\] %state\)", a) for a in split(ir, "\n"))
+    @test occursin(r"@\w*kernel\w*\(\[1 x i64\] %state\)", ir)
 
     # state should only passed to device functions that use it
 
@@ -106,16 +106,16 @@ end
                                   kernel=true, dump_module=true))
 
     # kernel should take state argument before all else
-    @test any(occursin(r"@.*kernel.+\(\[1 x i64\] %state", a) for a in split(ir, "\n"))
+    @test occursin(r"@\w*kernel\w*\(\[1 x i64\] %state", ir)
 
     # child1 doesn't use the state
-    @test any(occursin(r"@.*child1.+\(i64", a) for a in split(ir, "\n"))
+    @test occursin(r"@\w*child1\w*\(i64", ir)
 
     # child2 does
-    @test any(occursin(r"@.*child2.+\(\[1 x i64\] %state", a) for a in split(ir, "\n"))
+    @test occursin(r"@\w*child2\w*\(\[1 x i64\] %state", ir)
 
     # can't have the unlowered intrinsic
-    @test !any(occursin("julia.gpu.state_getter", a) for a in split(ir, "\n"))
+    @test !occursin("julia.gpu.state_getter", ir)
 end
 end
 
@@ -146,9 +146,9 @@ end
     end
 
     asm = sprint(io->ptx_code_native(io, entry, Tuple{Int64}; kernel=true))
-    @test occursin(r"\.visible \.entry .*julia_entry", asm)
-    @test !occursin(r"\.visible \.func .*julia_nonentry", asm)
-    @test occursin(r"\.func .*julia_nonentry", asm)
+    @test occursin(r"\.visible \.entry \w*entry", asm)
+    @test !occursin(r"\.visible \.func \w*nonentry", asm)
+    @test occursin(r"\.func \w*nonentry", asm)
 
 @testset "property_annotations" begin
     asm = sprint(io->ptx_code_native(io, entry, Tuple{Int64}; kernel=true))
@@ -185,15 +185,15 @@ end
     end
 
     asm = sprint(io->ptx_code_native(io, parent1, Tuple{Int}))
-    @test occursin(r".func julia_.*child_", asm)
+    @test occursin(r".func \w*child_", asm)
 
     function parent2(i)
         child(i+1)
         return
     end
 
     asm = sprint(io->ptx_code_native(io, parent2, Tuple{Int}))
-    @test occursin(r".func julia_.*child_", asm)
+    @test occursin(r".func \w*child_", asm)
 end
 
 @testset "child function reuse bis" begin

test/spirv.jl

@@ -24,10 +24,10 @@ end
     kernel(x) = return
 
     ir = sprint(io->spirv_code_llvm(io, kernel, Tuple{Tuple{Int}}))
-    @test occursin(r"@.*julia_.+_kernel.+\(({ i64 }|\[1 x i64\])\*", ir)
+    @test occursin(r"@\w*kernel\w*\(({ i64 }|\[1 x i64\])\*", ir)
 
     ir = sprint(io->spirv_code_llvm(io, kernel, Tuple{Tuple{Int}}; kernel=true))
-    @test occursin(r"@.*julia_.+_kernel.+\(.*{ ({ i64 }|\[1 x i64\]) }\*.+byval", ir)
+    @test occursin(r"@\w*kernel\w*\(.*{ ({ i64 }|\[1 x i64\]) }\*.+byval", ir)
 end
 
 @testset "byval bug" begin