Fix and test mangling (#632)

Author: maleadt

src/mangling.jl

@@ -29,24 +29,37 @@ safe_name(x) = safe_name(repr(x))
 # we generate function names that look like C++ functions, because many tools, like NVIDIA's
 # profilers, support them (grouping different instantiations of the same kernel together).
 
-function mangle_param(t, substitutions=String[])
+function mangle_param(t, substitutions=Any[])
     t == Nothing && return "v"
 
+    function find_substitution(x)
+        sub = findfirst(isequal(x), substitutions)
+        if sub === nothing
+            nothing
+        elseif sub == 1
+            "S_"
+        else
+            seq_id = uppercase(string(sub-2; base=36))
+            "S$(seq_id)_"
+        end
+    end
+
     if isa(t, DataType) && t <: Ptr
         tn = mangle_param(eltype(t), substitutions)
         "P$tn"
     elseif isa(t, DataType)
-        tn = safe_name(t)
+        # check if we already know this type
+        str = find_substitution(t)
+        if str !== nothing
+            return str
+        end
 
-        # handle substitutions
-        sub = findfirst(isequal(tn), substitutions)
-        if sub === nothing
+        # check if we already know this base type
+        str = find_substitution(t.name.wrapper)
+        if str === nothing
+            tn = safe_name(t)
             str = "$(length(tn))$tn"
-            push!(substitutions, tn)
-        elseif sub == 1
-            str = "S_"
-        else
-            str = "S$(sub-2)_"
+            push!(substitutions, t.name.wrapper)
         end
 
         # encode typevars as template parameters
@@ -56,21 +69,24 @@ function mangle_param(t, substitutions=String[])
                 str *= mangle_param(t, substitutions)
             end
             str *= "E"
+
+            push!(substitutions, t)
         end
 
         str
     elseif isa(t, Union)
-        tn = "Union"
+        # check if we already know this union type
+        str = find_substitution(t)
+        if str !== nothing
+            return str
+        end
 
-        # handle substitutions
-        sub = findfirst(isequal(tn), substitutions)
-        if sub === nothing
+        # check if we already know the Union name
+        str = find_substitution(Union)
+        if str === nothing
+            tn = "Union"
             str = "$(length(tn))$tn"
             push!(substitutions, tn)
-        elseif sub == 1
-            str = "S_"
-        else
-            str = "S$(sub-2)_"
         end
 
         # encode union types as template parameters
@@ -80,9 +96,13 @@ function mangle_param(t, substitutions=String[])
                 str *= mangle_param(t, substitutions)
             end
             str *= "E"
+
+            push!(substitutions, t)
         end
 
         str
+    elseif isa(t, UnionAll)
+        mangle_param(t.body, substitutions)
     elseif isa(t, Union{Bool, Cchar, Cuchar, Cshort, Cushort, Cint, Cuint, Clong, Culong, Clonglong, Culonglong, Int128, UInt128})
         ts = t isa Bool       ? 'b' : # bool
              t isa Cchar      ? 'a' : # signed char
@@ -99,10 +119,20 @@ function mangle_param(t, substitutions=String[])
              t isa UInt128    ? 'o' : # unsigned __int128
              error("Invalid type")
         tn = string(abs(t), base=10)
+        # for legibility, encode Julia-native integers as C-native integers, if possible
+        if t isa Int && typemin(Cint) <= t <= typemax(Cint)
+            ts = 'i'
+        end
         if t < 0
             tn = 'n'*tn
         end
         "L$(ts)$(tn)E"
+    elseif t isa Float32
+        bits = string(reinterpret(UInt32, t); base=16)
+        "Lf$(bits)E"
+    elseif t isa Float64
+        bits = string(reinterpret(UInt64, t); base=16)
+        "Ld$(bits)E"
     else
         tn = safe_name(t)   # TODO: actually does support digits...
         if startswith(tn, r"\d")
@@ -121,7 +151,7 @@ function mangle_sig(sig)
     str = "_Z$(length(fn))$fn"
 
     # mangle each parameter
-    substitutions = String[]
+    substitutions = []
     for t in tt
         str *= mangle_param(t, substitutions)
     end

test/Project.toml

@@ -10,6 +10,7 @@ SPIRV_LLVM_Translator_unified_jll = "85f0d8ed-5b39-5caa-b1ae-7472de402361"
 SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+demumble_jll = "1e29f10c-031c-5a83-9565-69cddfc27673"
 
 [compat]
 Aqua = "0.8"

test/util_tests.jl

@@ -9,14 +9,43 @@
     @test groups[3] == [:(d=4)]
 end
 
-@testset "mangle" begin
-    struct XX{T} end
-    # values checked with c++filt / cu++filt
-    @test GPUCompiler.mangle_sig(Tuple{typeof(sin), XX{false}})       == "_Z3sin2XXILb0EE"    # "sin(XX<false>)"
-    @test GPUCompiler.mangle_sig(Tuple{typeof(sin), XX{true}})        == "_Z3sin2XXILb1EE"    # "sin(XX<true>)"
-    @test GPUCompiler.mangle_sig(Tuple{typeof(sin), XX{Cshort(10)}})  == "_Z3sin2XXILs10EE"   # "sin(XX<(short)10>)"
-    @test GPUCompiler.mangle_sig(Tuple{typeof(sin), XX{Cshort(0)}})   == "_Z3sin2XXILs0EE"    # "sin(XX<(short)l>)"
-    @test GPUCompiler.mangle_sig(Tuple{typeof(sin), XX{Cshort(-10)}}) == "_Z3sin2XXILsn10EE"  # "sin(XX<(short)-10>)"
+@testset "mangling" begin
+    using demumble_jll
+
+    function mangle(f, argtyps...)
+        mangled = GPUCompiler.mangle_sig(Tuple{typeof(f), argtyps...})
+        chomp(read(`$(demumble_jll.demumble()) $mangled`, String))
+    end
+
+    # basic stuff
+    @test mangle(identity) == "identity"
+    @test mangle(identity, Nothing) == "identity()"
+
+    # primitive types
+    @test mangle(identity, Int32) == "identity(Int32)"
+    @test mangle(identity, Int64) == "identity(Int64)"
+
+    # literals
+    @test mangle(identity, Val{1}) == "identity(Val<1>)"
+    @test mangle(identity, Val{-1}) == "identity(Val<-1>)"
+    @test mangle(identity, Val{Cshort(1)}) == "identity(Val<(short)1>)"
+    @test mangle(identity, Val{1.0}) == "identity(Val<0x1p+0>)"
+    @test mangle(identity, Val{1f0}) == "identity(Val<0x1p+0f>)"
+
+    # unions
+    @test mangle(identity, Union{Int32, Int64}) == "identity(Union<Int32, Int64>)"
+
+    # union alls
+    @test mangle(identity, Array) == "identity(Array<T, N>)"
+
+    # many substitutions
+    @test mangle(identity, Val{1}, Val{2}, Val{3}, Val{4}, Val{5}, Val{6}, Val{7}, Val{8},
+                           Val{9}, Val{10}, Val{11}, Val{12}, Val{13}, Val{14}, Val{15},
+                           Val{16}, Val{16}) ==
+          "identity(Val<1>, Val<2>, Val<3>, Val<4>, Val<5>, Val<6>, Val<7>, Val<8>, Val<9>, Val<10>, Val<11>, Val<12>, Val<13>, Val<14>, Val<15>, Val<16>, Val<16>)"
+
+    # problematic examples
+    @test mangle(identity, String, Matrix{Float32}, Broadcast.Broadcasted{Broadcast.ArrayStyle{Matrix{Float32}}, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}, typeof(Base.literal_pow), Tuple{Base.RefValue{typeof(sin)}, Broadcast.Extruded{Matrix{Float32}, Tuple{Bool, Bool}, Tuple{Int64, Int64}}}}) == "identity(String, Array<Float32, 2>, Broadcasted<ArrayStyle<Array<Float32, 2>>, Tuple<OneTo<Int64>, OneTo<Int64>>, literal_pow, Tuple<RefValue<sin>, Extruded<Array<Float32, 2>, Tuple<Bool, Bool>, Tuple<Int64, Int64>>>>)"
 end
 
 @testset "safe loggers" begin