Use Pygments' highlighting for PTX and GCN (#659)

Co-authored-by: Tim Besard <[email protected]>

Author: Keluaa

res/pygments/ptx.py

@@ -12,6 +12,8 @@ GCNCompilerTarget(dev_isa; features="") = GCNCompilerTarget(dev_isa, features)
 
 llvm_triple(::GCNCompilerTarget) = "amdgcn-amd-amdhsa"
 
+source_code(target::GCNCompilerTarget) = "gcn"
+
 function llvm_machine(target::GCNCompilerTarget)
     @static if :AMDGPU ∉ LLVM.backends()
         return nothing

src/gcn.jl

@@ -16,23 +16,92 @@ function pygmentize()
     return _pygmentize[]
 end
 
+const _pygmentize_version = Ref{Union{VersionNumber, Nothing}}()
+function pygmentize_version()
+    isassigned(_pygmentize_version) && return _pygmentize_version[]
+
+    pygmentize_cmd = pygmentize()
+    if isnothing(pygmentize_cmd)
+        return _pygmentize_version[] = nothing
+    end
+
+    cmd = `$pygmentize_cmd -V`
+    @static if Sys.iswindows()
+        # Avoid encoding issues with pipes on Windows by using cmd.exe to capture stdout for us
+        cmd = `cmd.exe /C $cmd`
+        cmd = addenv(cmd, "PYTHONUTF8" => 1)
+    end
+    version_str = readchomp(cmd)
+
+    pos = findfirst("Pygments version ", version_str)
+    if !isnothing(pos)
+        version_start = last(pos) + 1
+        version_end = findnext(',', version_str, version_start) - 1
+        version = tryparse(VersionNumber, version_str[version_start:version_end])
+    else
+        version = nothing
+    end
+
+    if isnothing(version)
+        @warn "Could not parse Pygments version:\n$version_str"
+    end
+
+    return _pygmentize_version[] = version
+end
+
+function pygmentize_support(lexer)
+    highlighter_ver = pygmentize_version()
+    if isnothing(highlighter_ver)
+        @warn "Syntax highlighting of $lexer code relies on Pygments.\n\
+               Use `pip install pygments` to install the lastest version" maxlog = 1
+        return false
+    elseif lexer == "ptx"
+        if highlighter_ver < v"2.16"
+            @warn "Pygments supports PTX highlighting starting from version 2.16\n\
+                   Detected version $highlighter_ver\n\
+                   Please update with `pip install pygments -U`" maxlog = 1
+            return false
+        end
+        return true
+    elseif lexer == "gcn"
+        if highlighter_ver < v"2.8"
+            @warn "Pygments supports GCN highlighting starting from version 2.8\n\
+                   Detected version $highlighter_ver\n\
+                   Please update with `pip install pygments -U`" maxlog = 1
+            return false
+        end
+        return true
+    else
+        return false
+    end
+end
+
 function highlight(io::IO, code, lexer)
-    highlighter = pygmentize()
     have_color = get(io, :color, false)
     if !have_color
         print(io, code)
     elseif lexer == "llvm"
         InteractiveUtils.print_llvm(io, code)
-    elseif highlighter !== nothing
-        custom_lexer = joinpath(dirname(@__DIR__), "res", "pygments", "$lexer.py")
-        if isfile(custom_lexer)
-            lexer = `$custom_lexer -x`
+    elseif pygmentize_support(lexer)
+        lexer = lexer == "gcn" ? "amdgpu" : lexer
+        pygments_args = String[pygmentize(), "-f", "terminal", "-P", "bg=dark", "-l", lexer]
+        @static if Sys.iswindows()
+            # Avoid encoding issues with pipes on Windows by using a temporary file
+            mktemp() do tmp_path, tmp_io
+                println(tmp_io, code)
+                close(tmp_io)
+                push!(pygments_args, "-o", tmp_path, tmp_path)
+                cmd = Cmd(pygments_args)
+                wait(open(cmd))  # stdout and stderr go to devnull
+                print(io, read(tmp_path, String))
+            end
+        else
+            cmd = Cmd(pygments_args)
+            pipe = open(cmd, "r+")
+            print(pipe, code)
+            close(pipe.in)
+            print(io, read(pipe, String))
         end
-
-        pipe = open(`$highlighter -f terminal -P bg=dark -l $lexer`, "r+")
-        print(pipe, code)
-        close(pipe.in)
-        print(io, read(pipe, String))
     else
         print(io, code)
     end

src/reflection.jl

@@ -45,3 +45,30 @@ end
     # 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 "highlighting" begin
+    ansi_color = "\x1B[3"  # beginning of any foreground color change
+
+    @testset "PTX" begin
+        sample = """
+            max.s64         %rd24, %rd18, 0;
+            add.s64         %rd7, %rd23, -1;
+            setp.lt.u64     %p2, %rd7, %rd24;
+            @%p2 bra        \$L__BB0_3;
+        """
+        can_highlight = GPUCompiler.pygmentize_support("ptx")
+        highlighted = sprint(GPUCompiler.highlight, sample, "ptx"; context = (:color => true))
+        @test occursin(ansi_color, highlighted) skip = !can_highlight
+    end
+
+    @testset "GCN" begin
+        sample = """
+            v_add_u32     v3, vcc, s0, v0
+            v_mov_b32     v4, s1
+            v_addc_u32    v4, vcc, v4, 0, vcc
+        """
+        can_highlight = GPUCompiler.pygmentize_support("gcn")
+        highlighted = sprint(GPUCompiler.highlight, sample, "gcn"; context = (:color => true))
+        @test occursin(ansi_color, highlighted) skip = !can_highlight
+    end
+end