Add support for using FileCheck in tests.

Author: maleadt

test/Project.toml

@@ -2,8 +2,10 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
+LLVM_jll = "86de99a1-58d6-5da7-8064-bd56ce2e322c"
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"

test/helpers/bpf.jl

@@ -25,6 +25,15 @@ function code_native(io::IO, @nospecialize(func), @nospecialize(types); kwargs..
     GPUCompiler.code_native(io, job; kwargs...)
 end
 
+# aliases without ::IO argument
+for method in (:code_llvm, :code_native)
+    method = Symbol("$(method)")
+    @eval begin
+        $method(@nospecialize(func), @nospecialize(types); kwargs...) =
+            $method(stdout, func, types; kwargs...)
+    end
+end
+
 # simulates codegen for a kernel function: validates by default
 function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
     job, kwargs = create_job(func, types; kwargs...)

test/helpers/gcn.jl

@@ -35,6 +35,15 @@ function code_native(io::IO, @nospecialize(func), @nospecialize(types); kwargs..
     GPUCompiler.code_native(io, job; kwargs...)
 end
 
+# aliases without ::IO argument
+for method in (:code_warntype, :code_llvm, :code_native)
+    method = Symbol("$(method)")
+    @eval begin
+        $method(@nospecialize(func), @nospecialize(types); kwargs...) =
+            $method(stdout, func, types; kwargs...)
+    end
+end
+
 # simulates codegen for a kernel function: validates by default
 function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
     job, kwargs = create_job(func, types; kernel=true, kwargs...)

test/helpers/metal.jl

@@ -35,6 +35,15 @@ function code_native(io::IO, @nospecialize(func), @nospecialize(types); kwargs..
     GPUCompiler.code_native(io, job; kwargs...)
 end
 
+# aliases without ::IO argument
+for method in (:code_warntype, :code_llvm, :code_native)
+    method = Symbol("$(method)")
+    @eval begin
+        $method(@nospecialize(func), @nospecialize(types); kwargs...) =
+            $method(stdout, func, types; kwargs...)
+    end
+end
+
 # simulates codegen for a kernel function: validates by default
 function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
     job, kwargs = create_job(func, types; kernel=true, kwargs...)

test/helpers/ptx.jl

@@ -67,6 +67,15 @@ function code_native(io::IO, @nospecialize(func), @nospecialize(types); kwargs..
     GPUCompiler.code_native(io, job; kwargs...)
 end
 
+# aliases without ::IO argument
+for method in (:code_warntype, :code_llvm, :code_native)
+    method = Symbol("$(method)")
+    @eval begin
+        $method(@nospecialize(func), @nospecialize(types); kwargs...) =
+            $method(stdout, func, types; kwargs...)
+    end
+end
+
 # simulates codegen for a kernel function: validates by default
 function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
     job, kwargs = create_job(func, types; kernel=true, kwargs...)

test/helpers/spirv.jl

@@ -38,6 +38,15 @@ function code_native(io::IO, @nospecialize(func), @nospecialize(types); kwargs..
     GPUCompiler.code_native(io, job; kwargs...)
 end
 
+# aliases without ::IO argument
+for method in (:code_warntype, :code_llvm, :code_native)
+    method = Symbol("$(method)")
+    @eval begin
+        $method(@nospecialize(func), @nospecialize(types); kwargs...) =
+            $method(stdout, func, types; kwargs...)
+    end
+end
+
 # simulates codegen for a kernel function: validates by default
 function code_execution(@nospecialize(func), @nospecialize(types); kwargs...)
     job, kwargs = create_job(func, types; kernel=true, kwargs...)

test/helpers/test.jl

@@ -35,3 +35,123 @@ end
                      ret i64 %value"""
     return :(Base.llvmcall($llvmcall_str, T, Tuple{T}, i))
 end
+
+# filecheck utils
+
+module FileCheck
+    import LLVM_jll
+    import IOCapture
+    using GPUCompiler, LLVM
+    using Test
+
+    export filecheck, @filecheck, @check_str
+
+    global filecheck_path::String
+    function __init__()
+        # TODO: Windows
+        global filecheck_path = joinpath(LLVM_jll.artifact_dir, "tools", "FileCheck")
+    end
+
+    function filecheck_exe(; adjust_PATH::Bool=true, adjust_LIBPATH::Bool=true)
+        env = Base.invokelatest(
+            LLVM_jll.JLLWrappers.adjust_ENV!,
+            copy(ENV),
+            LLVM_jll.PATH[],
+            LLVM_jll.LIBPATH[],
+            adjust_PATH,
+            adjust_LIBPATH
+        )
+
+        return Cmd(Cmd([filecheck_path]); env)
+    end
+
+    const julia_typed_pointers = JuliaContext() do ctx
+        supports_typed_pointers(ctx)
+    end
+
+    function filecheck(f, input)
+        # FileCheck assumes that the input is available as a file
+        mktemp() do path, input_io
+            write(input_io, input)
+            close(input_io)
+
+            # capture the output of `f` and write it into a temporary buffer
+            result = IOCapture.capture(rethrow=Union{}) do
+                f(input)
+            end
+            output_io = IOBuffer()
+            write(output_io, result.output)
+            println(output_io)
+
+            if result.error
+                # if the function errored, also render the exception and backtrace
+                showerror(output_io, result.value, result.backtrace)
+            elseif result.value !== nothing
+                # also show the returned value; some APIs don't print
+                write(output_io, string(result.value))
+            end
+
+            # determine some useful prefixes for FileCheck
+            prefixes = ["CHECK",
+                        "JULIA$(VERSION.major)_$(VERSION.minor)",
+                        "LLVM$(Base.libllvm_version.major)"]
+            ## whether we use typed pointers or opaque pointers
+            if julia_typed_pointers
+                push!(prefixes, "TYPED")
+            else
+                push!(prefixes, "OPAQUE")
+            end
+            ## whether we pass pointers as integers or as actual pointers
+            if VERSION >= v"1.12.0-DEV.225"
+                push!(prefixes, "PTR_ABI")
+            else
+                push!(prefixes, "INTPTR_ABI")
+            end
+
+            # now pass the collected output to FileCheck
+            seekstart(output_io)
+            filecheck_io = Pipe()
+            cmd = ```$(filecheck_exe())
+                     --color
+                     --allow-unused-prefixes
+                     --check-prefixes $(join(prefixes, ','))
+                     $path```
+            proc = run(pipeline(ignorestatus(cmd); stdin=output_io, stdout=filecheck_io, stderr=filecheck_io); wait=false)
+            close(filecheck_io.in)
+
+            # collect the output of FileCheck
+            reader = Threads.@spawn String(read(filecheck_io))
+            Base.wait(proc)
+            log = strip(fetch(reader))
+
+            # error out if FileCheck did not succeed.
+            # otherwise, return true so that `@test @filecheck` works as expected.
+            if !success(proc)
+                error(log)
+            end
+            return true
+        end
+    end
+
+    # collect checks used in the @filecheck block by piggybacking on macro expansion
+    const checks = String[]
+    macro check_str(str)
+        push!(checks, str)
+        nothing
+    end
+
+    macro filecheck(ex)
+        ex = Base.macroexpand(__module__, ex)
+        if isempty(checks)
+            error("No checks provided within the @filecheck macro block")
+        end
+        check_str = join(checks, "\n")
+        empty!(checks)
+
+        esc(quote
+            filecheck($check_str) do _
+                $ex
+            end
+        end)
+    end
+end

test/setup.jl

@@ -9,6 +9,7 @@ for file in readdir(joinpath(@__DIR__, "helpers"))
         include(joinpath(@__DIR__, "helpers", file))
     end
 end
+using .FileCheck
 
 
 ## entry point

test/utils.jl

@@ -92,3 +92,26 @@ end
         @test occursin(ansi_color, highlighted) skip = !can_highlight
     end
 end
+
+# Test FileCheck
+@testset "FileCheck" begin
+    @test @filecheck begin
+        check"CHECK: works"
+        println("works")
+    end
+
+    @test_throws "expected string not found in input" @filecheck begin
+        check"CHECK: works"
+        println("doesn't work")
+    end
+
+    @test @filecheck begin
+        check"CHECK: errors"
+        error("errors")
+    end
+
+    @test_throws "expected string not found in input" @filecheck begin
+        check"CHECK: works"
+        error("errors")
+    end
+end