-
Notifications
You must be signed in to change notification settings - Fork 9
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Merge branch 'main' into gd/allow-ci-failure
- Loading branch information
Showing
11 changed files
with
329 additions
and
18 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,46 @@ | ||
| name: Benchmarks | ||
| permissions: | ||
| contents: write # contents permission to update benchmark contents in gh-pages branch | ||
| statuses: read | ||
| deployments: write # deployments permission to deploy GitHub pages website | ||
| pull-requests: write | ||
|
|
||
| on: | ||
| pull_request: | ||
|
|
||
| push: | ||
| branches: | ||
| - main | ||
|
|
||
| jobs: | ||
| benchmark: | ||
| if: ${{ !contains(github.event.head_commit.message, '[skip benchmarks]') }} | ||
| runs-on: ubuntu-latest | ||
| steps: | ||
| - uses: actions/checkout@v4 | ||
| - name: Download Buildkite Artifacts | ||
| id: download | ||
| uses: EnricoMi/download-buildkite-artifact-action@v1 | ||
| with: | ||
| buildkite_token: ${{ secrets.BUILDKITE_TOKEN }} | ||
| output_path: artifacts | ||
|
|
||
| - name: Locate Benchmarks Artifact | ||
| id: locate | ||
| if: ${{ steps.download.outputs.download-state == 'success' }} | ||
| run: echo "path=$(find artifacts -type f -name combinedbenchmarks.json 2>/dev/null)" >> $GITHUB_OUTPUT | ||
|
|
||
| - name: Upload Benchmark Results | ||
| if: ${{ steps.locate.outputs.path != '' }} | ||
| uses: benchmark-action/github-action-benchmark@v1 | ||
| with: | ||
| name: Reactant.jl Benchmarks | ||
| tool: "julia" | ||
| output-file-path: ${{ steps.locate.outputs.path }} | ||
| benchmark-data-dir-path: "benchmarks" | ||
| github-token: ${{ secrets.GITHUB_TOKEN }} | ||
| comment-always: true | ||
| summary-always: true | ||
| alert-threshold: "150%" | ||
| fail-on-alert: false | ||
| auto-push: ${{ github.event_name != 'pull_request' }} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -273,3 +273,5 @@ deps/ReactantExtra/MODULE.bazel.lock | |
| external | ||
|
|
||
| archive/ | ||
|
|
||
| benchmark/results/* | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,36 @@ | ||
| using BenchmarkTools | ||
|
|
||
| const BACKENDS = ["CPU", "CUDA"] | ||
|
|
||
| const CPU_Results = joinpath(dirname(@__FILE__), "results", "CPUbenchmarks.json") | ||
| @assert(ispath(CPU_Results)) | ||
|
|
||
| const RESULTS = BenchmarkTools.load(CPU_Results)[1] | ||
| @assert RESULTS isa BenchmarkTools.BenchmarkGroup | ||
|
|
||
| for backend in BACKENDS[2:end] | ||
| @info "Aggregating results for $(backend)" | ||
| filename = string(backend, "benchmarks.json") | ||
| filepath = joinpath(dirname(@__FILE__), "results", filename) | ||
| if !ispath(filepath) | ||
| @warn "No file found at path: $(filepath)" | ||
| else | ||
| backend_results = BenchmarkTools.load(filepath)[1] | ||
| if backend_results isa BenchmarkTools.BenchmarkGroup | ||
| # <benchmark name>/<forward or reverse>/<backend>/<reactant or package> | ||
| for benchmark in keys(RESULTS) | ||
| for pass in keys(RESULTS[benchmark]) | ||
| for pkg in keys(backend_results[benchmark][pass][backend]) | ||
| RESULTS[benchmark][pass][backend][pkg] = backend_results[benchmark][pass][backend][pkg] | ||
| end | ||
| end | ||
| end | ||
| else | ||
| @warn "Unexpected file format for file at path: $(filepath)" | ||
| end | ||
| end | ||
| end | ||
|
|
||
| BenchmarkTools.save( | ||
| joinpath(dirname(@__FILE__), "results", "combinedbenchmarks.json"), RESULTS | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,57 @@ | ||
| # Accelerator Support for testing non-Reactant performance | ||
| using LuxCUDA | ||
|
|
||
| using BenchmarkTools: BenchmarkTools, BenchmarkGroup, @btime, @benchmarkable | ||
| using CpuId: CpuId | ||
| using InteractiveUtils: versioninfo | ||
| using LinearAlgebra: BLAS | ||
| using Reactant: Reactant | ||
| using Statistics: median | ||
|
|
||
| # To run benchmarks on a specific GPU backend, add AMDGPU / CUDA / Metal / oneAPI | ||
| # to benchmarks/Project.toml and change BENCHMARK_GROUP to the backend name | ||
| const BENCHMARK_GROUP = get(ENV, "BENCHMARK_GROUP", "CPU") | ||
| @info "Running benchmarks for $BENCHMARK_GROUP" | ||
|
|
||
| BenchmarkTools.DEFAULT_PARAMETERS.seconds = 20 | ||
|
|
||
| if BENCHMARK_GROUP == "CPU" | ||
| if Sys.isapple() && (Sys.ARCH == :aarch64 || Sys.ARCH == :arm64) | ||
| @info "Running benchmarks on Apple with ARM CPUs. Using `AppleAccelerate.jl`." | ||
| using AppleAccelerate: AppleAccelerate | ||
| end | ||
|
|
||
| if Sys.ARCH == :x86_64 && occursin("intel", lowercase(CpuId.cpubrand())) | ||
| @info "Running benchmarks on Intel CPUs. Loading `MKL.jl`." | ||
| using MKL: MKL | ||
| end | ||
| end | ||
|
|
||
| const BENCHMARK_CPU_THREADS = Threads.nthreads() | ||
| BLAS.set_num_threads(BENCHMARK_CPU_THREADS) | ||
|
|
||
| @info sprint(versioninfo) | ||
| @info "BLAS threads: $(BLAS.get_num_threads())" | ||
|
|
||
| const SUITE = BenchmarkGroup() | ||
|
|
||
| if BENCHMARK_GROUP == "CUDA" | ||
| Reactant.set_default_backend("gpu") | ||
| @info "Running CUDA benchmarks" maxlog = 1 | ||
| CUDA.versioninfo() | ||
| else | ||
| @info "Running CPU benchmarks with $(BENCHMARK_CPU_THREADS) thread(s)" maxlog = 1 | ||
| end | ||
|
|
||
| # Main benchmark files | ||
| include("setup.jl") | ||
| setup_benchmarks!(SUITE, BENCHMARK_GROUP) | ||
|
|
||
| results = BenchmarkTools.run(SUITE; verbose=true) | ||
|
|
||
| filepath = joinpath(dirname(@__FILE__), "results") | ||
| mkpath(filepath) | ||
| filename = string(BENCHMARK_GROUP, "benchmarks.json") | ||
| BenchmarkTools.save(joinpath(filepath, filename), median(results)) | ||
|
|
||
| @info "Saved results to $(joinpath(filepath, filename))" |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,90 @@ | ||
| using Boltz: Vision | ||
| using Lux: Lux | ||
| using MLDataDevices: AbstractDevice, CPUDevice, CUDADevice | ||
| using Random: Random | ||
| using Reactant: Reactant, @compile | ||
|
|
||
| using Enzyme: Enzyme | ||
| using Zygote: Zygote | ||
|
|
||
| # Helper Functions | ||
| @inline synchronize(::CPUDevice) = nothing | ||
| @inline synchronize(::CUDADevice) = CUDA.synchronize() | ||
|
|
||
| @inline reclaim(::CPUDevice) = GC.gc() | ||
| @inline reclaim(::CUDADevice) = CUDA.reclaim() | ||
|
|
||
| @inline sumabs2(model, x, p, st) = sum(abs2, first(Lux.apply(model, x, p, st))) | ||
| @inline sumabs2(model, x) = sum(abs2, model(x)) | ||
|
|
||
| function benchmark_group_to_backend(benchmark_group::String) | ||
| benchmark_group == "CPU" && return CPUDevice() | ||
| benchmark_group == "CUDA" && return CUDADevice() | ||
| return error("Unknown backend: $(benchmark_group)") | ||
| end | ||
|
|
||
| function general_lux_setup(model, x_dims) | ||
| rng = Random.default_rng() # don't use any other rng | ||
| ps, st = Lux.setup(rng, model) | ||
| x_dims === nothing && return ps, st | ||
| x = randn(rng, Float32, x_dims) | ||
| return x, ps, st | ||
| end | ||
|
|
||
| function setup_benchmarks!(suite::BenchmarkGroup, backend::String) | ||
| dev = benchmark_group_to_backend(backend) | ||
|
|
||
| setup_vit_benchmark!(suite, backend, dev) | ||
|
|
||
| return nothing | ||
| end | ||
|
|
||
| # Lux Benchmarks | ||
| function setup_vit_benchmark!(suite::BenchmarkGroup, backend, dev::AbstractDevice) | ||
| for mode in (:tiny, :small, :base), bsize in (4, 16, 32) | ||
| benchmark_name = "ViT $(mode) (256 x 256 x 3 x $(bsize))" | ||
|
|
||
| setup_lux_forward_pass_benchmark!( | ||
| suite, benchmark_name, backend, Vision.ViT(mode), (256, 256, 3, bsize), dev | ||
| ) | ||
| end | ||
| end | ||
|
|
||
| function setup_lux_forward_pass_benchmark!( | ||
| suite::BenchmarkGroup, | ||
| benchmark_name::String, | ||
| backend::String, | ||
| model, | ||
| x_dims, | ||
| dev::AbstractDevice, | ||
| ) | ||
| suite[benchmark_name]["forward"][backend]["Lux"] = @benchmarkable begin | ||
| Lux.apply($model, x, ps, st_test) | ||
| synchronize($dev) | ||
| end setup = begin | ||
| GC.gc() | ||
| reclaim($dev) | ||
| x, ps, st = $dev(general_lux_setup($model, $x_dims)) | ||
| st_test = Lux.testmode(st) | ||
| GC.gc() | ||
| reclaim($dev) | ||
| end | ||
|
|
||
| suite[benchmark_name]["forward"][backend]["Reactant"] = @benchmarkable begin | ||
| y, _ = apply_compiled($model, x_ra, ps_ra, st_test_ra) | ||
| Reactant.synchronize(y) | ||
| end setup = begin | ||
| GC.gc() | ||
| reclaim($dev) | ||
| x, ps, st = general_lux_setup($model, $x_dims) | ||
| st_test = Lux.testmode(st) | ||
| x_ra = Reactant.to_rarray(x) | ||
| ps_ra = Reactant.to_rarray(ps) | ||
| st_test_ra = Reactant.to_rarray(st_test) | ||
| apply_compiled = @compile Lux.apply($model, x_ra, ps_ra, st_test_ra) | ||
| GC.gc() | ||
| reclaim($dev) | ||
| end | ||
|
|
||
| return nothing | ||
| end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.