GPU cleanup improvements

.github/workflows/CI.yml

@@ -23,22 +23,34 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.10'
           - '1.9'
-          - 'nightly'
+          - '1'
+          - 'pre'
         os:
           - ubuntu-latest
         arch:
           - x64
     steps:
       - uses: actions/checkout@v4
+      - uses: actions/setup-python@v5 
+        with:
+          python-version: '3.13'
+      - run: # For syntax highlighting of PTX and GCN
+          pip install pygments
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
       - uses: julia-actions/cache@v2
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
+      - uses: actions/upload-artifact@v4
+        if: always()
+        with:
+          name: "tests-output-jl-${{ matrix.version }}-${{ matrix.os }}-${{ matrix.arch }}"
+          path: |
+            test/tests_output.txt
+            test/references_mismatch/
       - uses: julia-actions/julia-processcoverage@v1
       - uses: codecov/codecov-action@v4
         env:

.gitignore

@@ -1,6 +1,8 @@
 *.jl.*.cov
 *.jl.cov
 *.jl.mem
-/Manifest.toml
-/docs/Manifest.toml
+/Manifest*.toml
+/docs/Manifest*.toml
 /docs/build/
+/test/references_mismatch/
+/test/tests_output.txt

Project.toml

@@ -11,6 +11,8 @@ MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 StringDistances = "88034a9c-02f8-509d-84a9-84ec65e18404"
 WidthLimitedIO = "b8c1c048-cf81-46c6-9da0-18c1d99e41f2"
+YAML = "ddb6d928-2868-570f-bddf-ab3f9cf99eb6"
+demumble_jll = "1e29f10c-031c-5a83-9565-69cddfc27673"
 
 [weakdeps]
 AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
@@ -21,13 +23,13 @@ oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b"
 
 [extensions]
 CodeDiffsAMDGPU = "AMDGPU"
-CodeDiffsCUDA = "CUDA"
+CodeDiffsCUDA = ["CUDA", "KernelAbstractions"]
 CodeDiffsKernelAbstractions = "KernelAbstractions"
 CodeDiffsMetal = "Metal"
 CodeDiffsOneAPI = "oneAPI"
 
 [compat]
-AMDGPU = "0.9"
+AMDGPU = "1"
 Aqua = "0.7"
 CUDA = "4, 5"
 CodeTracking = "1"
@@ -41,12 +43,13 @@ ReferenceTests = "0.10"
 StringDistances = "0.11"
 Test = "1"
 WidthLimitedIO = "1"
+YAML = "0.4"
+demumble_jll = "1"
 julia = "1.9"
 oneAPI = "1.5"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -56,4 +59,4 @@ Revise = "295af30f-e4ad-537b-8983-00126c2a3abe"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "CUDA", "InteractiveUtils", "KernelAbstractions", "OhMyREPL", "ReferenceTests", "Revise", "Test"]
+test = ["Aqua", "InteractiveUtils", "KernelAbstractions", "OhMyREPL", "ReferenceTests", "Revise", "Test"]

docs/make.jl

@@ -16,7 +16,13 @@ makedocs(;
     ),
     pages=[
         "Home" => "index.md",
-        "Extensions" => "extensions.md"
+        "GPU Extensions" => "extensions.md",
+        "API" => [
+            "Main API" => "code_disp_and_diff.md",
+            "Cleanup" => "cleanup.md",
+            "Code stats" => "stats.md",
+            "Defining a new extension" => "new_extension.md",
+        ],
     ],
     doctest = can_doctest
 )

docs/src/cleanup.md

@@ -0,0 +1,22 @@
+```@meta
+CurrentModule = CodeDiffs.Cleanup
+```
+
+# Cleanup
+
+```@docs
+cleanup_code
+```
+
+## LLVM and mangling
+
+```@docs
+replace_llvm_module_name
+function_unique_gen_name_regex
+global_var_unique_gen_name_regex
+demangle
+demangle_all
+mangled_base_name
+clean_function_name
+cleanup_inline_llvmcall_modules
+```

docs/src/code_disp_and_diff.md

@@ -0,0 +1,36 @@
+```@meta
+CurrentModule = CodeDiffs
+```
+
+# Comparison entry points
+
+```@docs
+@code_diff
+code_diff
+code_for_diff
+CodeDiff
+```
+
+# Code fetching
+
+```@docs
+@code_for
+code_native
+code_llvm
+code_typed
+code_ast
+get_code
+```
+
+# Highlighting
+
+```@docs
+code_highlighter
+```
+
+# Diff display
+
+```@docs
+optimize_line_changes!
+side_by_side_diff
+```

docs/src/extensions.md

@@ -2,7 +2,9 @@
 CurrentModule = CodeDiffs
 ```
 
-# KernelAbstractions.jl
+# GPU Extensions
+
+## KernelAbstractions.jl
 
 [`@code_diff`](@ref) will automatically detect calls to [KernelAbstractions.jl](https://github.com/JuliaGPU/KernelAbstractions.jl)
 and get the code for the actual underlying kernel function (whatever the backend is).
@@ -12,7 +14,7 @@ or when calling the kernel (`gpu_kernel(a, b, c; ndrange=1000)`).
 
 There is no support for AST comparison with KA.jl kernels.
 
-# GPU kernels
+## GPU kernels
 
 [`@code_diff`](@ref) supports functions compiled in a GPU context with any of the GPU packages:
 
@@ -32,36 +34,34 @@ CUDA has one additional layer of assembly code, SASS, available with `:sass`.
 !!! info
 
     Unlike with the `@device_code_*` macros, no kernel code is executed by [`@code_diff`](@ref).
-    This also means that kernels launched indirectly by the function will be ignored.
+    The `@device_code_*` macros work by capturing kernel launches, while `@code_diff` or `@code_for`
+    work with the kernel function directly: this means that kernels launched indirectly by the
+    function call will be ignored.
 
 !!! info
 
     Note that behind the scenes, `GPUCompiler.jl` only cares about the most recent methods.
     Hence the `world` keyword is unsupported for all GPU backends, as we cannot compile back in time.
 
-# Defining a new extension
-
-Defining a new `code_type` involves four functions:
-- `CodeDiffs.get_code_dispatch(::Val{code_type}, f, types; kwargs...)` (**not** `get_code`!)
-  should return a printable object (usually a `String`) representing the code for `f(types)`.
-  `kwargs` are the options passed to `@code_diff`.
-- `CodeDiffs.cleanup_code(::Val{:code_type}, obj)` does some cleanup on the code object to
-  make it more `diff`-able
-- `CodeDiffs.code_highlighter(::Val{code_type})` returns a `f(io, obj)` to print the `obj`
-  to as text in `io`. This is done twice: once without highlighting (`get(io, :color, false) == false`),
-  and another with highlighting.
-- `CodeDiffs.argconvert(::Val{code_type}, arg)` converts `arg` as needed (by default `arg` is unchanged)
-
-Defining a new pre-processing step for functions and its arguments (like for KernelAbstractions.jl kernels)
-involves two functions:
-- `CodeDiffs.extract_extra_options(f, kwargs)` returns some additional `kwargs` which are passed to `get_code`
-- `CodeDiffs.get_code(code_type, f, types; kwargs...)` allows to change `f` depending on its type.
-  To avoid method ambiguities, do not put type constraints on `code_type`.
-
-Defining a new object type which can be put as an argument to `@code_diff` or `@code_for`
-invoves at one function: `CodeDiffs.code_for_diff(obj::YourType; kwargs...)`.
-It must return two `String`s, one without and the other without highlighting.
-When calling `@code_for obj`, [`code_for_diff(obj)`](@ref) will be called only if `obj` is
-not a call expression or a quoted `Expr`.
-`kwargs` are the options passed to `@code_for` or the options passed to `@code_diff` for
-the side of `obj`.
+### GPU kernel statistics
+
+With the `:cuda_stats` code type, you can get an overview of your CUDA kernel through statistics
+inferred from its PTX and SASS code.
+
+Other supported types are `:cuda_stats`, `:ptx_stats` and `:sass_stats` for CUDA kernels, and
+`:gcn_stats` for AMDGPU kernels.
+See [CodeDiffs.Stats.extract_stats](@ref) for more about them.
+
+Example usage:
+```julia
+@code_for :cuda_stats some_kernel(a, b, c)
+```
+Output:
+```@eval
+using Markdown
+using CodeDiffs
+ptx_source  = readchomp("../../test/samples/extern_func_with_no_params.ptx")
+sass_source = readchomp("../../test/samples/extern_func_with_no_params.sass")
+cuda_stats = CodeDiffs.Stats.extract_stats(Val(:cuda_stats), (ptx_source, sass_source))
+Markdown.MD(Markdown.julia, Markdown.Code("", sprint(Base.show, MIME"text/plain"(), cuda_stats)))
+```

docs/src/index.md

@@ -12,86 +12,36 @@ Supports:
  - LLVM IR (output of `@code_llvm`, highlighted by `InteractiveUtils.print_llvm`)
  - Typed Julia IR (output of `@code_typed`, highlighted through the `Base.show` method of `Core.CodeInfo`)
  - Julia AST (an `Expr`), highlighting is done with OhMyREPL.jl's Julia syntax highlighting in Markdown code blocks
+ - GPU typed Julia IR / LLVM IR / native assembly (see [GPU Extensions](@ref))
 
-The [`@code_diff`](@ref) macro is the main entry point. If possible, the code type will be
-detected automatically, otherwise add e.g. `type=:llvm` for LLVM IR comparison:
+`CodeDiffs.jl` exports two macros:
+ - [`@code_for`](@ref) will display the code for a function call. The output is
+   cleaned and highlighted to maximize clarity.
+ - [`@code_diff`](@ref) will compare the code of two function calls.
 
-```jldoctest f1_vs_f2; setup=:(using CodeDiffs)
-julia> f1(a) = a + 1
-f1 (generic function with 1 method)
+Both support all code types.
+If possible, the code type will be detected automatically, otherwise add e.g.
+`type=:llvm` for LLVM IR comparison:
 
-julia> @code_diff type=:llvm debuginfo=:none color=false f1(Int64(1)) f1(Int8(1))
-define i64 @f1(i64 signext %0) #0 {   ⟪╋⟫define i64 @f1(i8 signext %0) #0 {
-top:                                   ┃ top:
-                                       ┣⟫  %1 = sext i8 %0 to i64
-  %1 = add i64 %0, 1                  ⟪╋⟫  %2 = add nsw i64 %1, 1
-  ret i64 %1                          ⟪╋⟫  ret i64 %2
-}                                      ┃ }
-
-julia> f2(a) = a - 1
-f2 (generic function with 1 method)
-
-julia> @code_diff type=:llvm debuginfo=:none color=false f1(1) f2(1)
-define i64 @f1(i64 signext %0) #0 {   ⟪╋⟫define i64 @f2(i64 signext %0) #0 {
-top:                                   ┃ top:
-  %1 = add i64 %0, 1                  ⟪╋⟫  %1 = add i64 %0, -1
-  ret i64 %1                           ┃   ret i64 %1
-}                                      ┃ }
+```@repl 1
+using CodeDiffs  # hide
+f1(a) = a + 1
+@code_diff type=:llvm debuginfo=:none color=false f1(Int64(1)) f1(Int8(1))
+f2(a) = a - 1
+@code_diff type=:llvm debuginfo=:none color=false f1(1) f2(1)
 ```
 
 Setting the environment variable `"CODE_DIFFS_LINE_NUMBERS"` to `true` will display line
 numbers on each side:
 
-```jldoctest f1_vs_f2
-julia> ENV["CODE_DIFFS_LINE_NUMBERS"] = true
-true
-
-julia> @code_diff type=:llvm debuginfo=:none color=false f1(1) f2(1)
-1 define i64 @f1(i64 signext %0) #0 { ⟪╋⟫define i64 @f2(i64 signext %0) #0 { 1
-2 top:                                 ┃ top:                                2
-3   %1 = add i64 %0, 1                ⟪╋⟫  %1 = add i64 %0, -1               3
-4   ret i64 %1                         ┃   ret i64 %1                        4
-5 }                                    ┃ }                                   5
+```@repl 1
+ENV["CODE_DIFFS_LINE_NUMBERS"] = true
+@code_diff type=:llvm debuginfo=:none color=false f1(1) f2(1)
 ```
 
-# Comparison entry points
+# Main API
 
-```@docs
-@code_diff
+```@docs; canonical=false
 @code_for
-code_diff
-code_for_diff
-CodeDiff
-```
-
-# Code fetching
-
-```@docs
-code_native
-code_llvm
-code_typed
-code_ast
-get_code
-```
-
-# Highlighting
-
-```@docs
-code_highlighter
-```
-
-# Cleanup
-
-```@docs
-cleanup_code
-replace_llvm_module_name
-function_unique_gen_name_regex
-global_var_unique_gen_name_regex
-```
-
-# Diff display
-
-```@docs
-optimize_line_changes!
-side_by_side_diff
+@code_diff
 ```

docs/src/new_extension.md

@@ -0,0 +1,30 @@
+```@meta
+CurrentModule = CodeDiffs
+```
+
+# Defining a new extension
+
+Defining a new `code_type` involves four functions:
+- `CodeDiffs.get_code_dispatch(::Val{code_type}, f, types; kwargs...)` (**not** `get_code`!)
+  should return a printable object (usually a `String`) representing the code for `f(types)`.
+  `kwargs` are the options passed to `@code_diff`.
+- `CodeDiffs.Cleanup.cleanup_code(::Val{:code_type}, code, dbinfo, cleanup_opts)` does some cleanup
+  on the code object to make it more `diff`-able.
+- `CodeDiffs.code_highlighter(::Val{code_type})` returns a `f(io, obj)` to print the `obj`
+  to as text in `io`. This is done twice: once without highlighting (`get(io, :color, false) == false`),
+  and another with highlighting.
+- `CodeDiffs.argconvert(::Val{code_type}, arg)` converts `arg` as needed (by default `arg` is unchanged)
+
+Defining a new pre-processing step for functions and its arguments (like for KernelAbstractions.jl kernels)
+involves two functions:
+- `CodeDiffs.extract_extra_options(f, kwargs)` returns some additional `kwargs` which are passed to `get_code`
+- `CodeDiffs.get_code(code_type, f, types; kwargs...)` allows to change `f` depending on its type.
+  To avoid method ambiguities, do not put type constraints on `code_type`.
+
+Defining a new object type which can be put as an argument to `@code_diff` or `@code_for`
+invoves at one function: `CodeDiffs.code_for_diff(obj::YourType; kwargs...)`.
+It must return two `String`s, one without and the other without highlighting.
+When calling `@code_for obj`, [`code_for_diff(obj)`](@ref) will be called only if `obj` is
+not a call expression or a quoted `Expr`.
+`kwargs` are the options passed to `@code_for` or the options passed to `@code_diff` for
+the side of `obj`.

docs/src/stats.md

@@ -0,0 +1,9 @@
+```@meta
+CurrentModule = CodeDiffs.Stats
+```
+
+# Code statistics
+
+```@docs
+extract_stats
+```