More cleanup fixes for PTX and GCN, more tests

Author: Keluaa

src/cleanup/gcn.jl

@@ -1,6 +1,6 @@
 
 # The GCN output of LLVM's AMDGPU backend has a comment with the mangled function name
-const GCN_FUNC_NAME_COMMENT_REGEX = r"; -- Begin function "m * MANGLED_NAME_REGEX * r"$"m
+const GCN_FUNC_NAME_COMMENT_REGEX = r".globl\s+"m * MANGLED_NAME_REGEX * r"\b"m
 
 
 function cleanup_code(::Val{:gcn}, c, dbinfo, cleanup_opts)
@@ -48,12 +48,13 @@ function cleanup_code(::Val{:gcn}, c, dbinfo, cleanup_opts)
         push!(extra_patterns, align_instruction_operand(align_operands))
     end
 
-    return replace(c,
+    c = replace(c,
         llvm_module_name_patterns()...,
         # This "begin" comment is redundant
         r"\s*; -- Begin function.+$"m => "",
         extra_patterns...,
     )
+    return rstrip(c)  # remove trailing newlines
 end
 
 

src/cleanup/ptx.jl

@@ -25,7 +25,10 @@ function cleanup_code(::Val{:ptx}, c, dbinfo, cleanup_opts)
     if !get(cleanup_opts, :keep_block_comments, false)
         # Block comments are placed after labels, or at the beginning of a line
         # They look like "some_label: // %blabla" or "// %blabla".
-        push!(extra_patterns, r"^([a-zA-Z0-9_$%]+:)?(?:\s*)\/\/ %.+"m => s"\1")
+        push!(extra_patterns,
+            r"^([a-zA-Z0-9_$%]+:)(?:\s*)\/\/ %.+"m => s"\1",
+            r"^\/\/ %.+\R"m => "",  # when the whole line is a comment like such, remove the newline as well
+        )
     end
     if (align_preds = get(cleanup_opts, :align_preds, 8); align_preds > 0)
         # Align guard predicates such that the beginning of the instruction is at the
@@ -170,8 +173,17 @@ function indent_ptx_function_calls()
         indent = " "^8
         call_attrs = strip(m[1])  # remove any '\r' on Windows
         func_name = strip(m[2])  # and again
-        params = replace(m[3], "\n" => "\n" * indent)
-        return "call" * m[1] * m[2] * ", (" * params * "\n" * indent * ");"
+        if m[3] == "\n"
+            # No parameters
+            params = ""
+        else
+            params = replace(m[3], "\n" => "\n" * indent) * "\n" * indent
+        end
+        # For some reason, the `replace` in `cleanup_code(::Val{:ptx})` doesn't replace
+        # the function name here, so we have to do it manually.
+        func_name = replace(func_name, llvm_module_name_patterns()...)
+        # Rebuild the function call in a cleaner way
+        return "call" * call_attrs * " " * func_name * ", (" * params * ");"
     end
 
     return ptx_call_regex => indent_call_params

test/cleanup.jl

@@ -5,7 +5,7 @@ import CodeDiffs.Cleanup as CDC
 is_removed(needle, present_there, not_there) =
         occursin(needle, present_there) && !occursin(needle, not_there)
 
-has_trailling_spaces(str) = occursin(r"\h+(\n|$)"m, str)
+has_trailing_spaces(str) = occursin(r"\h+(\n|$)"m, str)
 
 
 @testset "Demangling" begin
@@ -200,10 +200,10 @@ end
         @test is_removed(r"\t@%", ptx_sample, cleaned_ptx)
 
         # Reformat function calls
-        @test is_removed(r"call.+?,\R", ptx_sample, cleaned_ptx)
+        @test is_removed(r"call\S+\s+\R", ptx_sample, cleaned_ptx)
 
         # Others
-        @test !has_trailling_spaces(cleaned_ptx)
+        @test !has_trailing_spaces(cleaned_ptx)
         @test count(r"\R{2,}", cleaned_ptx) == 0  # no empty lines
 
         # Make sure we didn't remove any instruction by mistake
@@ -220,15 +220,17 @@ end
         println(TEST_IO, cleaned_ptx)
 
         # Proper cleanup of each parameter name
-        @test count(func_name * r"\d+_param_\d+\b", ptx_sample) == 3
-        @test count(func_name * r"_param_\d+\b", cleaned_ptx) == 3
+        @test count(func_name * r"_\d+_param_\d+\b", ptx_sample) == 6
+        @test count(r"\bparam_\d+\b", cleaned_ptx) == 6
+
+        # Others
+        @test !has_trailing_spaces(cleaned_ptx)
+        @test count(r"\R{2,}", cleaned_ptx) == 0  # no empty lines
+        @test !endswith(cleaned_ptx, r"\R")  # no trailing newlines
 
         # Make sure we didn't remove any instruction by mistake
         @test count(';', ptx_sample) == count(';', cleaned_ptx)
     end
-
-
-    # TODO: calls with 0 parameters are not reformatted correctly
 end
 
 
@@ -257,9 +259,14 @@ end
         # Code-gen comments are all removed
         @test is_removed(r"; %L\d+", gcn_sample, cleaned_gcn)
         @test is_removed(r"; %bb\.\d+:", gcn_sample, cleaned_gcn)
-        @test is_removed("divergent unreachable", gcn_sample, cleaned_gcn)
+        @test is_removed("; divergent unreachable", gcn_sample, cleaned_gcn)
         @test is_removed("; -- Begin function", gcn_sample, cleaned_gcn)
         @test is_removed("; -- End function", gcn_sample, cleaned_gcn)
+
+        # Others
+        @test !has_trailing_spaces(cleaned_gcn)
+        @test count(r"\R{2,}", cleaned_gcn) == 0  # no empty lines
+        @test !endswith(cleaned_gcn, r"\R")  # no trailing newlines
     end
 end