diff --git a/.buildkite/jll.pipeline.yml b/.buildkite/jll.pipeline.yml index 27656870d..7d591fb46 100644 --- a/.buildkite/jll.pipeline.yml +++ b/.buildkite/jll.pipeline.yml @@ -23,7 +23,7 @@ steps: fi julia --project -e 'using Pkg; Pkg.resolve(); Pkg.instantiate()' - JuliaCI/julia-test#v1: - test_args: "--quickfail --jobs=8 --verbose" + test_args: "--jobs=8 --verbose" - JuliaCI/julia-coverage#v1: dirs: - src diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 8fab17e69..7fc8e4fef 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -10,6 +10,6 @@ steps: command: ".buildkite/upload_gpu_ci.sh" agents: queue: "cuda" - if: build.message !~ /\[skip tests\]/ + if: build.message !~ /\[skip ci\]/ if_changed: "{src/**,scripts/**,deps/build.jl,Project.toml,lib/CNPreferences/src/**,lib/cunumeric_jl_wrapper/**}" timeout_in_minutes: 5 diff --git a/.buildkite/run_developer_ci.sh b/.buildkite/run_developer_ci.sh index e4004be1f..4d130393c 100755 --- a/.buildkite/run_developer_ci.sh +++ b/.buildkite/run_developer_ci.sh @@ -86,5 +86,5 @@ cp LocalPreferences.toml test/LocalPreferences.toml julia --color=yes --project=. -e ' using Pkg - Pkg.test("cuNumeric"; test_args = ["--quickfail", "--jobs=8", "--verbose"]) + Pkg.test("cuNumeric"; test_args = ["--jobs=8", "--verbose"]) ' diff --git a/.buildkite/upload_gpu_ci.sh b/.buildkite/upload_gpu_ci.sh index 3e8546d1e..dc1fd3ec7 100755 --- a/.buildkite/upload_gpu_ci.sh +++ b/.buildkite/upload_gpu_ci.sh @@ -15,13 +15,24 @@ message="${BUILDKITE_MESSAGE:-}" run_jll=true run_developer=true -# Keep both suites for main and PRs into main. For non-main PRs, select the -# suite whose wrapper matches the code under test. +if [[ "$message" =~ \[skip[[:space:]]ci\] ]]; then + echo "Skipping all GPU CI because the build message requests it." + exit 0 +fi + +if [[ "$message" =~ \[skip[[:space:]]jll\] ]]; then + echo "Skipping JLL GPU CI because the build message contains [skip jll]." + run_jll=false +fi +if [[ "$message" =~ \[skip[[:space:]]dev\] ]]; then + echo "Skipping developer GPU CI because the build message contains [skip dev]." + run_developer=false +fi + +# Keep both suites for main and PRs into main. For non-main PRs and post-merge +# develop builds, select the suite whose wrapper matches the code under test. if [[ "$branch" != "main" && "$base_branch" != "main" ]]; then - if [[ "$message" =~ \[skip[[:space:]]jll\] ]]; then - echo "Skipping JLL GPU CI because the build message contains [skip jll]." - run_jll=false - elif [[ "$pull_request" != "false" && -n "$base_branch" ]]; then + if [[ ("$pull_request" != "false" && -n "$base_branch") || "$branch" == "develop" ]]; then base_ref="refs/remotes/origin/$WRAPPER_BASE_BRANCH" # The published wrapper JLL tracks main, so compare against main even # when the pull request targets develop. diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 72d381fbb..a4eb983fc 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -1,4 +1,4 @@ -name: CI CPU +name: CI on: workflow_dispatch: @@ -18,7 +18,10 @@ on: - 'scripts/**' - 'deps/build.jl' - 'Project.toml' + - 'lib/cunumeric_jl_wrapper/**' - 'lib/CNPreferences/src/**' + - '.github/workflows/ci.yml' + - '.github/workflows/developer.yml' tags: - 'v*' branches: @@ -29,13 +32,19 @@ on: - 'scripts/**' - 'deps/build.jl' - 'Project.toml' + - 'lib/cunumeric_jl_wrapper/**' - 'lib/CNPreferences/src/**' + - '.github/workflows/ci.yml' + - '.github/workflows/developer.yml' jobs: - pkg_resolve: + resolve: + name: Package resolution + if: ${{ !contains(toJSON(github.event), '[skip ci]') }} uses: ./.github/workflows/pkg_resolve.yml - check_changes: - name: Check for wrapper changes + wrapper_changes: + name: Wrapper change detection + if: ${{ !contains(toJSON(github.event), '[skip ci]') }} runs-on: ubuntu-latest outputs: wrapper_changed: ${{ steps.wrapper-changes.outputs.changed }} @@ -60,10 +69,22 @@ jobs: fi fi - test: - name: Julia ${{ matrix.julia }} - ${{ matrix.os }} - needs: [pkg_resolve, check_changes] - if: ${{ github.base_ref == 'main' || needs.check_changes.outputs.wrapper_changed != 'true' }} + developer_tests: + name: Developer wrapper tests + needs: [resolve, wrapper_changes] + if: ${{ !contains(toJSON(github.event), '[skip ci]') && !contains(toJSON(github.event), '[skip dev]') && (github.event_name != 'pull_request' || github.base_ref == 'main' || needs.wrapper_changes.outputs.wrapper_changed == 'true') }} + permissions: + contents: read + packages: write + attestations: write + id-token: write + actions: write + uses: ./.github/workflows/developer.yml + + jll_tests: + name: JLL wrapper tests - Julia ${{ matrix.julia }} - ${{ matrix.os }} + needs: [resolve, wrapper_changes] + if: ${{ !contains(toJSON(github.event), '[skip ci]') && !contains(toJSON(github.event), '[skip jll]') && (github.base_ref == 'main' || needs.wrapper_changes.outputs.wrapper_changed != 'true') }} runs-on: ${{ matrix.os }} strategy: fail-fast: false @@ -125,4 +146,4 @@ jobs: LEGATE_AUTO_CONFIG: "0" LEGATE_SKIP_RUNTIME: "true" LEGATE_CONFIG: "--cpus 1 --utility 1 --sysmem 500" - run: julia --project -e 'using Pkg; Pkg.test(test_args=["--quickfail", "--jobs=2", "--verbose"])' + run: julia --project -e 'using Pkg; Pkg.test(test_args=["--jobs=2", "--verbose"])' diff --git a/.github/workflows/container.yml b/.github/workflows/container.yml index 07a68c927..323a1e539 100644 --- a/.github/workflows/container.yml +++ b/.github/workflows/container.yml @@ -12,13 +12,13 @@ on: required: false default: false workflow_run: - workflows: ['CI CPU'] + workflows: ['CI'] types: [completed] branches: - main jobs: push_to_registry: - if: ${{ github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success' }} + if: ${{ !contains(toJSON(github.event), '[skip ci]') && (github.event_name == 'workflow_dispatch' || github.event.workflow_run.conclusion == 'success') }} name: Container for ${{ matrix.platform }} - Julia ${{ matrix.julia }} - CUDA ${{ matrix.cuda }} permissions: contents: read diff --git a/.github/workflows/developer.yml b/.github/workflows/developer.yml index 93cd090ce..a1f9c16ce 100644 --- a/.github/workflows/developer.yml +++ b/.github/workflows/developer.yml @@ -1,55 +1,12 @@ -# Develeper CI test. This will build the workflow using Jlls and building wrappers from SRC -name: Develeper CI test +# Developer wrapper tests build the wrappers from source instead of using JLLs. +name: Developer Wrapper Tests on: - workflow_dispatch: - inputs: - tag: - description: 'Tag to build instead' - required: false - default: '' - mark_as_latest: - description: 'Mark as latest' - type: boolean - required: false - default: false - push: - paths: - - 'src/**' - - 'scripts/**' - - 'deps/build.jl' - - 'Project.toml' - - 'lib/cunumeric_jl_wrapper/src/**' - - 'lib/cunumeric_jl_wrapper/include/**' - - 'lib/CNPreferences/src/**' - - '.github/workflows/developer.yml' - tags: - - 'v*' - branches: - - main - pull_request: - paths: - - 'src/**' - - 'scripts/**' - - 'deps/build.jl' - - 'Project.toml' - - 'lib/cunumeric_jl_wrapper/src/**' - - 'lib/cunumeric_jl_wrapper/include/**' - - 'lib/CNPreferences/src/**' - - '.github/workflows/developer.yml' -jobs: - pkg_resolve: - uses: ./.github/workflows/pkg_resolve.yml + workflow_call: - docs: - name: Developer CI test - Julia ${{ matrix.julia }} - needs: pkg_resolve - permissions: - contents: read - packages: write - attestations: write - id-token: write - actions: write +jobs: + test: + name: Julia ${{ matrix.julia }} strategy: fail-fast: false matrix: @@ -157,4 +114,4 @@ jobs: cp LocalPreferences.toml test/LocalPreferences.toml - julia --color=yes --project=. -e 'using Pkg; Pkg.test("cuNumeric"; test_args=["--quickfail", "--jobs=2", "--verbose"])' + julia --color=yes --project=. -e 'using Pkg; Pkg.test("cuNumeric"; test_args=["--jobs=2", "--verbose"])' diff --git a/.github/workflows/docs-tags.yml b/.github/workflows/docs-tags.yml index ead34e481..5e982c79d 100644 --- a/.github/workflows/docs-tags.yml +++ b/.github/workflows/docs-tags.yml @@ -9,6 +9,7 @@ on: jobs: docs: name: Documentation + if: ${{ !contains(toJSON(github.event), '[skip ci]') }} permissions: actions: write contents: write diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 3317cea6a..84da51106 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -16,6 +16,7 @@ on: jobs: docs: name: Documentation + if: ${{ !contains(toJSON(github.event), '[skip ci]') }} permissions: actions: write contents: write diff --git a/.github/workflows/pkg_resolve.yml b/.github/workflows/pkg_resolve.yml index 47c5302d1..6eb18b0bf 100644 --- a/.github/workflows/pkg_resolve.yml +++ b/.github/workflows/pkg_resolve.yml @@ -1,11 +1,11 @@ -name: Pkg Resolve +name: Package Resolution on: workflow_call: jobs: resolve: - name: Pkg.resolve + name: Resolve dependencies runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 diff --git a/.github/workflows/version_check.yml b/.github/workflows/version_check.yml index 346bc7d6b..a738ada5d 100644 --- a/.github/workflows/version_check.yml +++ b/.github/workflows/version_check.yml @@ -8,6 +8,7 @@ on: jobs: version-check: name: Version Check + if: ${{ !contains(toJSON(github.event), '[skip ci]') }} runs-on: ubuntu-latest steps: - uses: actions/checkout@v4 diff --git a/.gitignore b/.gitignore index d431bc365..cf02abbe8 100644 --- a/.gitignore +++ b/.gitignore @@ -16,6 +16,8 @@ logging logging/* debug debug/* +!benchmark/debug/ +!benchmark/debug/grayscott_accelerate.jl # example outputs (examples/data and the docs copy of gray-scott.gif are tracked) examples/*.h5 @@ -31,6 +33,7 @@ benchmark/results/* benchmark/plots** compile_wrapper.sh +__plot_results.jl *.tar.gz # generated by CMake diff --git a/README.md b/README.md index 39dd0896b..cad89042a 100644 --- a/README.md +++ b/README.md @@ -1,128 +1,89 @@

cuNumeric.jl - cuNumeric.jl -

-

- cuNumeric.jl - cuNumeric.jl + cuNumeric.jl

-[![Documentation dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://julialegate.github.io/cuNumeric.jl/dev/) [![codecov](https://codecov.io/github/julialegate/cuNumeric.jl/branch/main/graph/badge.svg)](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT) -[![Documentation dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://julialegate.github.io/cuNumeric.jl/dev/) [![codecov](https://codecov.io/github/julialegate/cuNumeric.jl/branch/main/graph/badge.svg)](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT) +[![Documentation dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://julialegate.github.io/cuNumeric.jl/dev) [![codecov](https://codecov.io/github/julialegate/cuNumeric.jl/branch/main/graph/badge.svg)](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT) -cuNumeric.jl wraps and extends the [cuPyNumeric](https://github.com/nv-legate/cupynumeric) library from NVIDIA to bring distributed array computing on GPUs and CPUs to Julia. The central type is `NDArray`, which behaves like Julia's `Array` or the `CuArray` from [CUDA.jl](https://github.com/juliagpu/cuda.jl), but executes across multiple GPUs/CPUs. We implement array-level operations on `NDArray` which can be composed into larger programs without the need for explicit MPI calls or writing CUDA kernels. cuNumeric.jl wraps and extends the [cuPyNumeric](https://github.com/nv-legate/cupynumeric) library from NVIDIA to bring distributed array computing on GPUs and CPUs to Julia. The central type is `NDArray`, which behaves like Julia's `Array` or the `CuArray` from [CUDA.jl](https://github.com/juliagpu/cuda.jl), but executes across multiple GPUs/CPUs. We implement array-level operations on `NDArray` which can be composed into larger programs without the need for explicit MPI calls or writing CUDA kernels. -cuNumeric.jl requires x86 Linux, an NVIDIA GPU, and Julia >= 1.10. If ARM support is of interest open an issue. cuNumeric.jl requires x86 Linux, an NVIDIA GPU, and Julia >= 1.10. If ARM support is of interest open an issue. ### Quick Start -cuNumeric.jl can be installed with the Julia package manager. Activate your preferred environment and then from the Julia REPL run: - - cuNumeric.jl can be installed with the Julia package manager. Activate your preferred environment and then from the Julia REPL run: ```julia using Pkg Pkg.add(url = "https://github.com/JuliaLegate/cuNumeric.jl", rev = "main") -using Pkg -Pkg.add(url = "https://github.com/JuliaLegate/cuNumeric.jl", rev = "main") ``` -The first time might take awhile as it has to install multiple large dependencies such as the CUDA SDK (if you have an NVIDIA GPU). To use a local build of cupynumeric.so, see [Build Modes](./install.md). - -The first time might take awhile as it has to install multiple large dependencies such as the CUDA SDK (if you have an NVIDIA GPU). To use a local build of cupynumeric.so, see [Build Modes](./install.md). +The first installation can take a while because it includes several large dependencies, such as the CUDA SDK. To use a local cupynumeric build, see [Build Modes](https://julialegate.github.io/cuNumeric.jl/dev/install). ```julia using cuNumeric -using cuNumeric cuNumeric.versioninfo() ``` > [!WARNING] > Starting more than one instance of cuNumeric.jl can lead to a hard-crash. The default hardware configuration reserves all available resources. -For more details, see [Hardware](./configuration/hardware.md). +For more details, see [Hardware](https://julialegate.github.io/cuNumeric.jl/dev/configuration/hardware). ### How `NDArray`s work The semantics of `NDArray` closely mirror Julia's `Array`, and in most cases it is a drop-in replacement. You can use the same constructors (i.e., `zeros`, `ones`, `rand`), broadcasting, slicing, and linear algebra. Under the hood a few details differ from Base, and knowing them can help you write fast code. -**Data may live across many devices.** An `NDArray` is a logical array whose physical buffers can be partitioned over GPUs and CPUs by the Legate runtime. You write ordinary array code and Legate decides where the data lives and how/when it is communicated between devices. As a result, elementwise indexing (i.e. `arr[1]`) is slow (and is prevented by default). Scalar indexing like this forces synchronization and blocks other tasks from executing. Functions like `println` result in data being copied to the host and can also be slow. +**Data may live across many devices.** An `NDArray` is a logical array whose physical buffers can be partitioned over GPUs and CPUs by the Legate runtime. You write ordinary array code and Legate decides where the data lives and how/when it is communicated between devices. As a result, elementwise indexing (i.e. `arr[1]`) is slow (and is prevented by default). Scalar indexing like this forces synchronization and blocks other tasks from executing. **Slices are views.** Indexing an `NDArray` with ranges returns a view onto the same store, not a copy. That differs from Base Julia, where `A[1:n]` allocates a new `Array`. Mutations through an `NDArray` slice are visible through other aliases of the same data. -**Reductions return arrays, not Julia scalars.** Reductions such as `sum(A)` produce a **0D or 1D** `NDArray` (axis reductions produce a lower-rank `NDArray`), rather than a bare `Float64` / `Float32`. That keeps the Legate task graph asynchronous instead of forcing synchronization to communite with the Julia runtime. When you need a plain Julia number, call `unwrap` or `only`: +**Reductions return arrays, not Julia scalars.** Reductions such as `sum(A)` produce a **0D or 1D** `NDArray` (axis reductions produce a lower-rank `NDArray`), rather than a bare `Float64` / `Float32`. That keeps the Legate task graph asynchronous instead of forcing synchronization to communicate with the Julia runtime. When you need a plain Julia number, call `unwrap`: ```julia s = sum(A) # NDArray{T,0} x = unwrap(s) # T, e.g. Float32 -x2 = only(s) -``` -s = sum(A) # NDArray{T,0} -x = unwrap(s) # T, e.g. Float32 ``` -**The Legate runtime builds a DAG asynchronously.** Calling `cuNumeric.zeros` or `A .+ B` records work into Legate's task graph rather than blocking until every GPU kernel finishes. Results are materialized when you need them (for example `println`, `unwrap`, or converting with `Array(A)`). Hiding latency enables performant code. **The Legate runtime builds a DAG asynchronously.** Calling `cuNumeric.zeros` or `A .+ B` records work into Legate's task graph rather than blocking until every GPU kernel finishes. Results are materialized when you need them (for example `println`, `unwrap`, or converting with `Array(A)`). Hiding latency enables performant code. -For API details see [Initialization](./api_initialization.md) and [NDArray Reference](./api.md). For anti-patterns that kill performance, see [Patterns to Avoid](./perf/patterns_to_avoid.md). -For API details see [Initialization](./api_initialization.md) and [NDArray Reference](./api.md). For anti-patterns that kill performance, see [Patterns to Avoid](./perf/patterns_to_avoid.md). +For API details see [Initialization](https://julialegate.github.io/cuNumeric.jl/dev/api_initialization) and [NDArray Reference](https://julialegate.github.io/cuNumeric.jl/dev/api). For common performance pitfalls, see [Patterns to Avoid](https://julialegate.github.io/cuNumeric.jl/dev/perf/patterns_to_avoid). -### Kernel Fusion ### Kernel Fusion -Nested broadcast expressions fuse into a single kernel by default when on GPU. Prefer `@.` for multi-op elementwise code so every operator is dotted and the expression stays completely fused. Even just forgetting the `.` on unary negation (i.e., `y .= -a .+ b .* c`) will result in unfused code. Use the following pattern instead. Nested broadcast expressions fuse into a single kernel by default when on GPU. Prefer `@.` for multi-op elementwise code so every operator is dotted and the expression stays completely fused. Even just forgetting the `.` on unary negation (i.e., `y .= -a .+ b .* c`) will result in unfused code. Use the following pattern instead. -```julia -y .= @. -a + b * c ```julia y .= @. -a + b * c ``` -See [Kernel Fusion](./perf/kernel_fusion.md) and [Debugging](./debugging.md) for controls and pretty printers. - -See [Kernel Fusion](./perf/kernel_fusion.md) and [Debugging](./debugging.md) for controls and pretty printers. - -### Helping the Garbage Collector - -Many calls such as array slicing and un-fused broadcasts allocate a new `NDArray`. The Legate runtime keeps track of all references to the underlying data and will not free the memory until Julia's GC frees the `NDArray` handles. Because Julia's GC runs on memory pressure and an `NDArray` only stores a pointer (i.e., Julia's GC does not know the true size), many dead buffers accumulate and can cause out-of-memory errors. -Many calls such as array slicing and un-fused broadcasts allocate a new `NDArray`. The Legate runtime keeps track of all references to the underlying data and will not free the memory until Julia's GC frees the `NDArray` handles. Because Julia's GC runs on memory pressure and an `NDArray` only stores a pointer (i.e., Julia's GC does not know the true size), many dead buffers accumulate and can cause out-of-memory errors. - -`@analyze_lifetimes` performs a **static last-use analysis** at macro-expansion time and inserts eager calls to immediately free unused `NDArrays`. These buffers can then be reused by legate later for same-sized allocations. -`@analyze_lifetimes` performs a **static last-use analysis** at macro-expansion time and inserts eager calls to immediately free unused `NDArrays`. These buffers can then be reused by legate later for same-sized allocations. +See [Kernel Fusion](https://julialegate.github.io/cuNumeric.jl/dev/perf/kernel_fusion) and [Debugging](https://julialegate.github.io/cuNumeric.jl/dev/debugging) for controls and diagnostics. -```julia -@analyze_lifetimes begin - result = @. A[1:end, :] + B[1:end, :] - C .= @. result * 2.0f0 - result = @. A[1:end, :] + B[1:end, :] - C .= @. result * 2.0f0 -end -``` +### The `@accelerate` macro -### Performance at a glance +`@accelerate` fuses eligible GPU broadcasts within and across statements, then releases materialized temporary `NDArray`s after their last use on CPU or GPU. See [The `@accelerate` Macro](https://julialegate.github.io/cuNumeric.jl/dev/perf/reduce_allocations) for usage guidance. -A representative benchmark figure will go here (add something like `docs/src/images/benchmarks-overview.png` when ready). +### Benchmarks -Numbers, plots, and how to reproduce them live under [Benchmark Results](./benchmarks/results.md) and [How to Benchmark](./benchmarks/howto.md). +Results and reproduction instructions live under [Benchmark Results](https://julialegate.github.io/cuNumeric.jl/dev/benchmarks/results) and [How to Benchmark](https://julialegate.github.io/cuNumeric.jl/dev/benchmarks/howto). ### Try an example ```julia using cuNumeric -integrand = (x) -> @. exp(-x^2) +integrand(x) = @. exp(-x^2) + +@accelerate function monte_carlo(N, x_max) + Ω = 2 * x_max + raw_samples = cuNumeric.rand(N) + samples = @. Ω * raw_samples - x_max + return (Ω / N) * sum(integrand(samples)) +end N = 1_000_000 x_max = 10.0f0 -Ω = 2 * x_max - -samples = Ω .* cuNumeric.rand(N) -samples = samples .- x_max -estimate = (Ω / N) .* sum(integrand(samples)) +estimate = monte_carlo(N, x_max) println("Monte-Carlo Estimate: $(estimate)") ``` @@ -131,11 +92,3 @@ More worked examples (initialization, Gray-Scott, …) are in the documentation ### Known Limitations - There is no support for `Float16` or `ComplexF16` -- Arrays with 4 or more dimensions might have worse performance -- Maximum array dimension is 6 - - -### Known Deviations from Base Julia -- Reductions return 0D stores instead of scalars -- Slices return views -- `inv` does not throw `SingularException` for singular matrices diff --git a/benchmark/README.md b/benchmark/README.md index 205cd08ca..eb3b346eb 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -5,7 +5,7 @@ Benchmarks are declared in `benchmarks.toml`. `run.jl` parses it. ## Running ```bash -julia --project run.jl # runs whatever benchmarks.toml configures +julia --project=. run.jl # runs whatever benchmarks.toml configures ``` `run.jl` runs each (benchmark, backend) pair in its own process via @@ -35,7 +35,7 @@ n_warmup = 5 n_iter = 1000 n_trial = 5 -[[gemm]] # name registered in src/benchmarks.jl +[[gemm]] # name registered under src/benchmarks/ T = "Float32" # element type gpus = 1 cpus = 2 @@ -59,6 +59,10 @@ two axes: default `true`); it only affects cuNumeric, so comparison backends run once, not per variant. +Benchmark names are defined by the registered benchmark implementations. A +benchmark may expose baseline, optimized, backend-specific, or other variants; +the harness treats each name uniformly and records each result independently. + Each zipped field must be one of: - a scalar or single-element list (`cpus = 2` or `[2]`) -> broadcast to every config @@ -82,3 +86,14 @@ M = [150, 300, 600] # When `T = ["Float32", "Float64"]` and a length-2 `N`/`M` sweep you get all **4** combinations, not a paired `Float32 -> N[1], Float64 -> N[2]`. To pin a type to a specific size, use separate `[[name]]` blocks. + +## Plotting + +```bash +julia --project=benchmark benchmark/plot_results.jl +``` + +The plotter reads the result files in the selected results directory and writes +one weak-scaling figure per benchmark, plus aggregate fusion and no-fusion +figures when those result groups are present. Outputs are grouped under a +subdirectory named for the shared benchmark prefix. diff --git a/benchmark/__plot_results.jl b/benchmark/__plot_results.jl deleted file mode 100644 index fd1ed646a..000000000 --- a/benchmark/__plot_results.jl +++ /dev/null @@ -1,301 +0,0 @@ -#!/usr/bin/env julia -# Weak-scaling plots (1/2/4/8 GPUs) for the benchmark result CSVs. -# One figure per benchmark, three panels: throughput, time/step, parallel efficiency. -# -# CSV schema (see src/core.jl save_result): -# implementation,gpus,N,M,trial,time_ms,throughput,correctness -# `throughput` is the benchmark's `total_flops` divided by elapsed time. For -# Gray-Scott that unit is Gpoint-updates/s; other benchmarks report GFLOP/s. -# -# save_result appends and does NOT encode the code-path variant, so a cuNumeric CSV -# holds alternating runs: baseline, @accelerate, baseline, @accelerate, ... -# (a run boundary = the GPU count resetting downward). -# cuPyNumeric / CUDA.jl have no accelerated path -> a single block. -# -# Encoding: color = implementation; line style = code path -# solid = baseline, dashed = @accelerate. - -using Plots -using Statistics - -gr() - -function parse_args(args) - results_dir = "results" - out_dir = nothing - single_cunumeric_run = :baseline - hide_baseline = false - output_suffix = "" - - for arg in args - if startswith(arg, "--single-cu=") - value = Symbol(lowercase(last(split(arg, "="; limit=2)))) - value in (:baseline, :accelerated) || - error("--single-cu must be baseline or accelerated") - single_cunumeric_run = value - elseif arg == "--hide-baseline" - hide_baseline = true - elseif startswith(arg, "--out=") - out_dir = last(split(arg, "="; limit=2)) - elseif startswith(arg, "--suffix=") - output_suffix = last(split(arg, "="; limit=2)) - else - results_dir = arg - end - end - isempty(output_suffix) && hide_baseline && (output_suffix = "_no_baseline") - - results_dir = isabspath(results_dir) ? results_dir : joinpath(@__DIR__, results_dir) - if out_dir === nothing - out_dir = if basename(normpath(results_dir)) == "results" - joinpath(@__DIR__, "plots") - else - joinpath(@__DIR__, "plots", basename(normpath(results_dir))) - end - else - out_dir = isabspath(out_dir) ? out_dir : joinpath(@__DIR__, out_dir) - end - return (; results_dir, out_dir, single_cunumeric_run, hide_baseline, output_suffix) -end - -const CONFIG = parse_args(ARGS) -const RESULTS_DIR = CONFIG.results_dir -const OUT_DIR = CONFIG.out_dir -const SINGLE_CUNUMERIC_RUN = CONFIG.single_cunumeric_run -const HIDE_BASELINE = CONFIG.hide_baseline -const OUTPUT_SUFFIX = CONFIG.output_suffix - -# filekey, family label, color, marker, can_contain_accelerated_blocks -const FAMILIES = [ - ("cunumeric", "cuNumeric.jl (fused)", "#2a78d6", :circle, true), - ("cunumeric_nofusion", "cuNumeric.jl (unfused)", "#4a3aa7", :diamond, true), - ("cupynumeric", "cuPyNumeric", "#eb6834", :rect, false), - ("CUDA.jl", "CUDA.jl", "#008300", :utriangle, false), -] - -const INK = "#0b0b0b" -const MUTED = "#898781" -const GRIDCOL = "#e1e0d9" -const IDEALCOL = "#c3c2b7" - -struct Row - gpus::Int - time_ms::Float64 - thr::Float64 -end - -# Parse a CSV into runs, split wherever the GPU count resets to a smaller value. -function load_runs(path) - rows = Row[] - for line in eachline(path) - isempty(strip(line)) && continue - f = split(line, ',') - push!(rows, Row(parse(Int, f[2]), parse(Float64, f[6]), parse(Float64, f[7]))) - end - isempty(rows) && return Vector{Row}[] - runs = [Row[]] - for (i, r) in enumerate(rows) - i > 1 && r.gpus < rows[i - 1].gpus && push!(runs, Row[]) - push!(runs[end], r) - end - return runs -end - -# Aggregate trials per GPU count -> sorted vector of (gpus, t, tsd, h, hsd). -function aggregate(rows) - by = Dict{Int,Vector{Row}}() - for r in rows - push!(get!(by, r.gpus, Row[]), r) - end - sd(x) = length(x) > 1 ? std(x) : 0.0 - return [ - (gpus=g, t=mean(getfield.(by[g], :time_ms)), tsd=sd(getfield.(by[g], :time_ms)), - h=mean(getfield.(by[g], :thr)), hsd=sd(getfield.(by[g], :thr))) - for g in sort(collect(keys(by))) - ] -end - -# Build the series (color+marker+linestyle+agg) present for one benchmark. -function series_for(bench) - series = [] # NamedTuple(label,color,marker,ls,agg) - for (key, fam, color, marker, splits) in FAMILIES - path = joinpath(RESULTS_DIR, "$(bench)_$(key).csv") - isfile(path) || continue - runs = load_runs(path) - isempty(runs) && continue - if splits && bench == "grayscott" - if length(runs) == 1 - label = - SINGLE_CUNUMERIC_RUN === :accelerated ? "$fam · accelerated" : - "$fam · baseline" - ls = SINGLE_CUNUMERIC_RUN === :accelerated ? :dash : :solid - push!( - series, (label=label, color=color, marker=marker, ls=ls, - agg=aggregate(runs[1])) - ) - else - # Repeated harness runs append alternating baseline/accelerated blocks. - baseline_rows = reduce(vcat, runs[1:2:end]) - accelerated_rows = reduce(vcat, runs[2:2:end]) - push!( - series, - (label="$fam · baseline", color=color, marker=marker, - ls=:solid, agg=aggregate(baseline_rows)), - ) - push!(series, - (label="$fam · accelerated", color=color, marker=marker, - ls=:dash, agg=aggregate(accelerated_rows))) - end - else - push!( - series, - (label=fam, color=color, marker=marker, ls=:solid, - agg=aggregate(reduce(vcat, runs))), - ) - end - end - return series -end - -function throughput_label(bench) - return bench == "grayscott" ? - "Throughput (Gpoint-updates/s)" : "Throughput (GFLOP/s)" -end - -function addline!(p, s, y; kw...) - return plot!(p, getfield.(s.agg, :gpus), y; color=s.color, - lw=2.2, ls=s.ls, marker=s.marker, ms=6, msc=s.color, markerstrokewidth=0.8, - label=s.label, kw...) -end - -# One legend key: a short line sample (+ optional marker) with a text label. -function swatch!(p, x, y, color, ls, marker, label) - plot!(p, [x, x + 0.032], [y, y]; color=color, lw=2.6, ls=ls, label="") - marker !== nothing && scatter!(p, [x + 0.016], [y]; color=color, marker=marker, - ms=6, msc=color, markerstrokewidth=0.8, label="") - return annotate!(p, x + 0.045, y, text(label, 9, INK, :left)) -end - -# Grouped legend: color/marker = implementation, line style = code path. -function build_legend(series) - pl = plot(; framestyle=:none, legend=false, xlims=(0, 1), ylims=(0, 1), - grid=false, ticks=false) - # implementations present, in FAMILIES order, matched by color - present = [ - (fam, color, marker) for (key, fam, color, marker, _) in FAMILIES - if any(s.color == color for s in series) - ] - annotate!(pl, 0.015, 0.74, text("Implementation", 10, INK, :left)) - xs = range(0.18, 0.80; length=max(length(present), 1)) - for ((fam, color, marker), x) in zip(present, xs) - swatch!(pl, x, 0.74, color, :solid, marker, fam) - end - annotate!(pl, 0.015, 0.26, text("Line style", 10, INK, :left)) - has_baseline = any(endswith(s.label, "· baseline") for s in series) - has_accelerated = any(endswith(s.label, "· accelerated") for s in series) - if has_baseline && has_accelerated - swatch!(pl, 0.18, 0.26, MUTED, :solid, nothing, "baseline") - swatch!(pl, 0.40, 0.26, MUTED, :dash, nothing, "@accelerate") - swatch!(pl, 0.70, 0.26, IDEALCOL, :dashdot, nothing, "efficiency = 1") - elseif has_accelerated - swatch!(pl, 0.18, 0.26, MUTED, :dash, nothing, "@accelerate") - swatch!(pl, 0.52, 0.26, IDEALCOL, :dashdot, nothing, "efficiency = 1") - elseif has_baseline - swatch!(pl, 0.18, 0.26, MUTED, :solid, nothing, "baseline") - swatch!(pl, 0.48, 0.26, IDEALCOL, :dashdot, nothing, "efficiency = 1") - else - swatch!(pl, 0.18, 0.26, IDEALCOL, :dashdot, nothing, "efficiency = 1") - end - return pl -end - -function positive_ylim(vals; pad=0.12) - isempty(vals) && return (0, 1) - hi = maximum(vals) - hi > 0 || return (0, 1) - return (0, hi * (1 + pad)) -end - -function main() - mkpath(OUT_DIR) - files = filter(f -> endswith(f, ".csv"), readdir(RESULTS_DIR)) - benches = unique( - String[ - m.captures[1] for f in files for (key, _, _, _, _) in FAMILIES - for m in (match(Regex("^(.*)_" * replace(key, "." => "\\.") * "\\.csv\$"), f),) - if m !== nothing - ], - ) - - for bench in benches - series = series_for(bench) - if HIDE_BASELINE - series = filter(s -> !endswith(s.label, "· baseline"), series) - end - isempty(series) && continue - - common = (xscale=:log2, xticks=([1, 2, 4, 8], ["1", "2", "4", "8"]), xlabel="GPUs", - framestyle=:box, grid=true, gridcolor=GRIDCOL, gridalpha=1.0, - foreground_color_text=INK, tickfontcolor=MUTED, legend=false, - xlims=(0.85, 9.4)) - - # Panel 1: throughput (higher better) - throughput = [x.h for s in series for x in s.agg] - p1 = plot(; ylabel=throughput_label(bench), title="Throughput", - ylims=positive_ylim(throughput), common...) - for s in series - addline!(p1, s, getfield.(s.agg, :h); yerror=getfield.(s.agg, :hsd)) - end - - # Panel 2: time per step (lower better; ideal = flat) - p2 = plot(; ylabel="Time / step (ms)", title="Time per step", common...) - for s in series - addline!(p2, s, getfield.(s.agg, :t); yerror=getfield.(s.agg, :tsd)) - end - - # Panel 3: parallel efficiency = thr(p)/(p*thr(1)); ideal = 1.0 - efficiencies = Float64[] - for s in series - i1 = findfirst(x -> x.gpus == 1, s.agg) - i1 === nothing && continue - base = s.agg[i1].h - append!(efficiencies, [x.h/(x.gpus*base) for x in s.agg]) - end - p3 = plot(; ylabel="Parallel efficiency", title="Weak-scaling efficiency", - ylims=positive_ylim(vcat(efficiencies, [1.0])), common...) - hline!(p3, [1.0]; color=IDEALCOL, ls=:dashdot, lw=1.4, label="") - for s in series - i1 = findfirst(x -> x.gpus == 1, s.agg) - i1 === nothing && continue - base = s.agg[i1].h - addline!(p3, s, [x.h/(x.gpus*base) for x in s.agg]) - end - - # grouped legend panel: color/marker = implementation, style = code path - pl = build_legend(series) - has_baseline = any(endswith(s.label, "· baseline") for s in series) - has_accelerated = any(endswith(s.label, "· accelerated") for s in series) - style_title = if has_baseline && has_accelerated - "solid = baseline · dashed = @accelerate" - elseif has_accelerated - "dashed = @accelerate" - elseif has_baseline - "solid = baseline" - else - "implementation comparison" - end - - fig = plot(p1, p2, p3, pl; layout=@layout([grid(1, 3); leg{0.16h}]), - size=(1400, 600), dpi=200, - plot_title=titlecase(bench) * " — weak scaling ($style_title)", - plot_titlefontsize=12, left_margin=6Plots.mm, - bottom_margin=6Plots.mm, top_margin=4Plots.mm, - background_color="#fcfcfb") - - out = joinpath(OUT_DIR, "$(bench)_weak_scaling$(OUTPUT_SUFFIX).png") - savefig(fig, out) - println("wrote $out") - end -end - -main() diff --git a/benchmark/benchmarks.toml b/benchmark/benchmarks.toml index 1d8b3194f..747989673 100644 --- a/benchmark/benchmarks.toml +++ b/benchmark/benchmarks.toml @@ -36,16 +36,40 @@ T = "Float32" gpus = [1, 2, 4, 8] cpus = 16 fusion = [true, false] -N = [2000, 2832, 4000, 5656] -M = [2000, 2832, 4000, 5656] +N = [24000, 33944, 48000, 67888] +M = [24000, 33944, 48000, 67888] -[[grayscott_lifetimes]] +[[grayscott_function_accelerated]] T = "Float32" gpus = [1, 2, 4, 8] cpus = 16 -fusion = false -N = [2000, 2832, 4000, 5656] -M = [2000, 2832, 4000, 5656] +fusion = [true, false] +N = [24000, 33944, 48000, 67888] +M = [24000, 33944, 48000, 67888] + +[[grayscott_begin_accelerated]] +T = "Float32" +gpus = [1, 2, 4, 8] +cpus = 16 +fusion = [true, false] +N = [24000, 33944, 48000, 67888] +M = [24000, 33944, 48000, 67888] + +[[grayscott_let_accelerated]] +T = "Float32" +gpus = [1, 2, 4, 8] +cpus = 16 +fusion = [true, false] +N = [24000, 33944, 48000, 67888] +M = [24000, 33944, 48000, 67888] + +[[grayscott_expression_accelerated]] +T = "Float32" +gpus = [1, 2, 4, 8] +cpus = 16 +fusion = [true, false] +N = [24000, 33944, 48000, 67888] +M = [24000, 33944, 48000, 67888] ################################# # DMD # @@ -72,11 +96,11 @@ fusion = [true, false] N = [50000, 100000, 200000, 400000] M = 512 -[[dmd_lifetimes]] +[[dmd_accelerated]] T = "Float32" gpus = [1, 2, 4, 8] cpus = 16 -fusion = false +fusion = [true, false] N = [50000, 100000, 200000, 400000] M = 512 diff --git a/benchmark/debug/grayscott_accelerate.jl b/benchmark/debug/grayscott_accelerate.jl new file mode 100644 index 000000000..53a276bb0 --- /dev/null +++ b/benchmark/debug/grayscott_accelerate.jl @@ -0,0 +1,66 @@ +#!/usr/bin/env julia + +# Print @accelerate lifetime rewrites and the fused kernels launched by each +# Gray–Scott macro form. Run with a small grid, for example: +# julia --project=benchmark benchmark/debug/grayscott_accelerate.jl 64 + +using cuNumeric + +const BENCHMARK_SRC = joinpath(@__DIR__, "..", "src") +include(joinpath(BENCHMARK_SRC, "core.jl")) +include(joinpath(BENCHMARK_SRC, "benchmarks", "grayscott.jl")) +include(joinpath(BENCHMARK_SRC, "benchmarks", "grayscott_accelerate_forms.jl")) + +const N = length(ARGS) >= 1 ? parse(Int, ARGS[1]) : 64 +N >= 4 || error("N must be at least 4") + +const FORMS = ( + (:function, "function", GrayScottFunctionAccelerated), + (:begin, "begin", GrayScottBeginAccelerated), + (:let, "let", GrayScottLetAccelerated), + (:expression, "expression", GrayScottExpressionAccelerated), +) + +function lifetime_expansion(kind) + body = deepcopy(GRAYSCOTT_STEP_BODY) + input = if kind === :function + Expr(:function, Expr(:call, :debug_step, :u, :v, :u_new, :v_new, :args), body) + elseif kind === :let + Expr(:let, body) + else + body + end + return cuNumeric._accelerate_expand(input, @__MODULE__) +end + +function print_lifetimes(kind, label) + println("\n", "="^80, "\n", uppercase(label), " — lifetime analysis\n", "="^80) + if kind === :expression + # Expression form accelerates each assignment independently. + for statement in GRAYSCOTT_STEP_BODY.args + statement isa LineNumberNode && continue + statement isa Expr && statement.head === :(=) || continue + lhs, rhs = statement.args + println("\nRHS: ", lhs) + expansion = cuNumeric._accelerate_expand(rhs, @__MODULE__) + cuNumeric.print_lifetime_analysis(expansion) + end + else + cuNumeric.print_lifetime_analysis(lifetime_expansion(kind)) + end +end + +function run_form(label, type) + println("\n", "="^80, "\n", uppercase(label), " — runtime kernels\n", "="^80) + b = build_benchmark(type, Float32, N, N) + state = only(initialize(b; deterministic=true)) + return run!(b, state) +end + +println("Gray–Scott @accelerate debug; grid=$(N)x$(N)") +cuNumeric.BCAST_FUSION_DEBUG[] = true +println("BCAST_FUSION_DEBUG = ", cuNumeric.BCAST_FUSION_DEBUG[]) +for (kind, label, type) in FORMS + print_lifetimes(kind, label) + run_form(label, type) +end diff --git a/benchmark/plot_results.jl b/benchmark/plot_results.jl new file mode 100644 index 000000000..b53ac3618 --- /dev/null +++ b/benchmark/plot_results.jl @@ -0,0 +1,284 @@ +#!/usr/bin/env julia +# Generate weak-scaling plots from benchmark CSVs. +# Each figure shows throughput, time per step, and parallel efficiency. + +using Plots +using Statistics + +gr() + +function parse_args(args) + results_dir = "results" + out_dir = nothing + output_suffix = "" + + for arg in args + if startswith(arg, "--out=") + out_dir = last(split(arg, "="; limit=2)) + elseif startswith(arg, "--suffix=") + output_suffix = last(split(arg, "="; limit=2)) + else + results_dir = arg + end + end + results_dir = isabspath(results_dir) ? results_dir : joinpath(@__DIR__, results_dir) + if out_dir === nothing + out_dir = if basename(normpath(results_dir)) == "results" + joinpath(@__DIR__, "plots") + else + joinpath(@__DIR__, "plots", basename(normpath(results_dir))) + end + else + out_dir = isabspath(out_dir) ? out_dir : joinpath(@__DIR__, out_dir) + end + return (; results_dir, out_dir, output_suffix) +end + +const CONFIG = parse_args(ARGS) +const RESULTS_DIR = CONFIG.results_dir +const OUT_DIR = CONFIG.out_dir +const OUTPUT_SUFFIX = CONFIG.output_suffix + +# file key, display label, color, marker +const FAMILIES = [ + ("cunumeric", "cuNumeric.jl (fused)", "#2a78d6", :circle), + ("cunumeric_nofusion", "cuNumeric.jl (unfused)", "#4a3aa7", :diamond), + ("cupynumeric", "cuPyNumeric", "#eb6834", :rect), + ("CUDA.jl", "CUDA.jl", "#008300", :utriangle), +] + +const INK = "#0b0b0b" +const MUTED = "#898781" +const GRIDCOL = "#e1e0d9" +const IDEALCOL = "#c3c2b7" +const GROUP_COLORS = ["#2a78d6", "#eb6834", "#008300", "#8b3fb0", "#c47f00", "#159a9c"] +const GROUP_MARKERS = [:circle, :diamond, :utriangle, :rect, :star5, :hexagon] + +struct Row + gpus::Int + time_ms::Float64 + thr::Float64 +end + +# Parse a CSV into runs, split wherever the GPU count resets to a smaller value. +function load_runs(path) + rows = Row[] + for line in eachline(path) + isempty(strip(line)) && continue + f = split(line, ',') + push!(rows, Row(parse(Int, f[2]), parse(Float64, f[6]), parse(Float64, f[7]))) + end + isempty(rows) && return Vector{Row}[] + runs = [Row[]] + for (i, r) in enumerate(rows) + i > 1 && r.gpus < rows[i - 1].gpus && push!(runs, Row[]) + push!(runs[end], r) + end + return runs +end + +# Aggregate trials per GPU count -> sorted vector of (gpus, t, tsd, h, hsd). +function aggregate(rows) + by = Dict{Int,Vector{Row}}() + for r in rows + push!(get!(by, r.gpus, Row[]), r) + end + sd(x) = length(x) > 1 ? std(x) : 0.0 + return [ + (gpus=g, t=mean(getfield.(by[g], :time_ms)), tsd=sd(getfield.(by[g], :time_ms)), + h=mean(getfield.(by[g], :thr)), hsd=sd(getfield.(by[g], :thr))) + for g in sort(collect(keys(by))) + ] +end + +_all_rows(runs) = reduce(vcat, runs; init=Row[]) + +function make_series(label, color, marker, ls, runs) + isempty(runs) && return nothing + return (label=label, color=color, marker=marker, ls=ls, + agg=aggregate(_all_rows(runs))) +end + +# Build the available implementation series for one benchmark. +function load_series(bench, specs; filename=(b, k) -> "$(b)_$(k).csv") + series = [] + for (key, label, color, marker) in specs + path = joinpath(RESULTS_DIR, filename(bench, key)) + isfile(path) || continue + runs = load_runs(path) + isempty(runs) && continue + + push!(series, make_series(label, color, marker, :solid, runs)) + end + return filter(!isnothing, series) +end + +series_for(bench) = load_series(bench, FAMILIES) + +function common_prefix(names) + words = split.(names, '_') + n = minimum(length, words) + i = 0 + while i < n && all(w -> w[i + 1] == words[1][i + 1], words) + i += 1 + end + return i == 0 ? String[] : words[1][1:i] +end + +function family_stem(name, prefix) + words = split(name, '_') + length(words) > length(prefix) || return name + return join(words[(length(prefix) + 1):end], "_") +end + +function grouped_series(benches, key) + series = [] + prefix = common_prefix(benches) + for (i, bench) in enumerate(benches) + path = joinpath(RESULTS_DIR, "$(bench)_$(key).csv") + isfile(path) || continue + runs = load_runs(path) + isempty(runs) && continue + color = GROUP_COLORS[mod1(i, length(GROUP_COLORS))] + marker = GROUP_MARKERS[mod1(i, length(GROUP_MARKERS))] + words = split(bench, '_') + length(words) > length(prefix) && (words = words[(length(prefix) + 1):end]) + label = titlecase(join(words, ' ')) + push!(series, make_series(label, color, marker, :solid, runs)) + end + return filter(!isnothing, series) +end + +function throughput_label() + return "Throughput" +end + +function addline!(p, s, y; kw...) + return plot!(p, getfield.(s.agg, :gpus), y; color=s.color, + lw=2.2, ls=s.ls, marker=s.marker, ms=6, msc=s.color, markerstrokewidth=0.8, + label=s.label, kw...) +end + +# One legend key: a short line sample (+ optional marker) with a text label. +function swatch!(p, x, y, color, ls, marker, label) + plot!(p, [x, x + 0.032], [y, y]; color=color, lw=2.6, ls=ls, label="") + marker !== nothing && scatter!(p, [x + 0.016], [y]; color=color, marker=marker, + ms=6, msc=color, markerstrokewidth=0.8, label="") + return annotate!(p, x + 0.045, y, text(label, 9, INK, :left)) +end + +# Build a legend from the series currently being plotted. +function build_legend(series) + pl = plot(; framestyle=:none, legend=false, xlims=(0, 1), ylims=(0, 1), + grid=false, ticks=false) + present = [(s.label, s.color, s.marker) for s in series] + annotate!(pl, 0.015, 0.74, text("Series", 10, INK, :left)) + xs = range(0.18, 0.80; length=max(length(present), 1)) + for ((fam, color, marker), x) in zip(present, xs) + swatch!(pl, x, 0.74, color, :solid, marker, fam) + end + swatch!(pl, 0.18, 0.26, IDEALCOL, :dashdot, nothing, "efficiency = 1") + return pl +end + +function positive_ylim(vals; pad=0.12) + isempty(vals) && return (0, 1) + hi = maximum(vals) + hi > 0 || return (0, 1) + return (0, hi * (1 + pad)) +end + +function weak_scaling_figure(bench, series, legend_panel; plot_title) + common = (xscale=:log2, xticks=([1, 2, 4, 8], ["1", "2", "4", "8"]), xlabel="GPUs", + framestyle=:box, grid=true, gridcolor=GRIDCOL, gridalpha=1.0, + foreground_color_text=INK, tickfontcolor=MUTED, legend=false, + xlims=(0.85, 9.4)) + + # Panel 1: throughput (higher better) + throughput = [x.h for s in series for x in s.agg] + p1 = plot(; ylabel=throughput_label(), title="Throughput", + ylims=positive_ylim(throughput), common...) + for s in series + addline!(p1, s, getfield.(s.agg, :h); yerror=getfield.(s.agg, :hsd)) + end + + # Panel 2: time per step (lower better; ideal = flat) + p2 = plot(; ylabel="Time / step (ms)", title="Time per step", common...) + for s in series + addline!(p2, s, getfield.(s.agg, :t); yerror=getfield.(s.agg, :tsd)) + end + + # Panel 3: parallel efficiency = thr(p)/(p*thr(1)); ideal = 1.0 + efficiencies = Float64[] + for s in series + i1 = findfirst(x -> x.gpus == 1, s.agg) + i1 === nothing && continue + base = s.agg[i1].h + append!(efficiencies, [x.h/(x.gpus*base) for x in s.agg]) + end + p3 = plot(; ylabel="Parallel efficiency", title="Weak-scaling efficiency", + ylims=positive_ylim(vcat(efficiencies, [1.0])), common...) + hline!(p3, [1.0]; color=IDEALCOL, ls=:dashdot, lw=1.4, label="") + for s in series + i1 = findfirst(x -> x.gpus == 1, s.agg) + i1 === nothing && continue + base = s.agg[i1].h + addline!(p3, s, [x.h/(x.gpus*base) for x in s.agg]) + end + + return plot( + p1, p2, p3, legend_panel; + layout=@layout([grid(1, 3); leg{0.16h}]), + size=(1400, 600), dpi=200, plot_title, + plot_titlefontsize=12, left_margin=6Plots.mm, + bottom_margin=6Plots.mm, top_margin=4Plots.mm, + background_color="#fcfcfb", + ) +end + +function main() + mkpath(OUT_DIR) + files = filter(f -> endswith(f, ".csv"), readdir(RESULTS_DIR)) + benches = unique( + String[ + m.captures[1] for f in files for (key, _, _, _) in FAMILIES + for m in (match(Regex("^(.*)_" * replace(key, "." => "\\.") * "\\.csv\$"), f),) + if m !== nothing + ], + ) + family = common_prefix(benches) + benchmark_out = joinpath(OUT_DIR, isempty(family) ? "benchmarks" : join(family, "_")) + mkpath(benchmark_out) + + for bench in benches + series = series_for(bench) + isempty(series) && continue + + fig = weak_scaling_figure( + bench, series, build_legend(series); + plot_title=titlecase(bench) * " — weak scaling", + ) + + stem = family_stem(bench, family) + out = joinpath(benchmark_out, "$(stem)_weak_scaling$(OUTPUT_SUFFIX).png") + savefig(fig, out) + println("wrote $out") + end + + # Add aggregate views that compare all benchmark variants for each mode. + for (key, title, stem) in (("cunumeric", "Fusion enabled", "fusion"), + ("cunumeric_nofusion", "Fusion disabled", "no_fusion")) + series = grouped_series(benches, key) + isempty(series) && continue + fig = weak_scaling_figure( + stem, series, build_legend(series); + plot_title=title * " — weak scaling", + ) + out = joinpath(benchmark_out, "$(stem)_weak_scaling$(OUTPUT_SUFFIX).png") + savefig(fig, out) + println("wrote $out") + end + return nothing +end + +main() diff --git a/benchmark/run.jl b/benchmark/run.jl index cb748ccde..36c4c2661 100644 --- a/benchmark/run.jl +++ b/benchmark/run.jl @@ -23,8 +23,8 @@ const POSARGS = filter(a -> a ∉ VERBOSE_FLAGS, ARGS) banner(msg) = println("\n", "="^128, "\n", msg, "\n", "="^128) -# `_lifetimes` is a cuNumeric-only code-path variant (@analyze_lifetimes) -cunumeric_only(name) = endswith(name, "_lifetimes") +# `_accelerated` is a cuNumeric-only code-path variant (`@accelerate`). +cunumeric_only(name) = endswith(name, "_accelerated") const LAST_FUSION_TOGGLE = Ref{Union{Nothing,Bool}}(nothing) diff --git a/benchmark/src/benchmarks/dmd.jl b/benchmark/src/benchmarks/dmd.jl index 8099f34c5..7f5b41fea 100644 --- a/benchmark/src/benchmarks/dmd.jl +++ b/benchmark/src/benchmarks/dmd.jl @@ -13,13 +13,13 @@ Base.@kwdef struct DMDBaseline{T} <: AbstractDMD{T} M::Int end -Base.@kwdef struct DMDLifetimes{T} <: AbstractDMD{T} +Base.@kwdef struct DMDAccelerated{T} <: AbstractDMD{T} N::Int M::Int end name(::DMDBaseline) = "dmd_baseline" -name(::DMDLifetimes) = "dmd_lifetimes" +name(::DMDAccelerated) = "dmd_accelerated" dims(b::AbstractDMD) = (b.N, b.M) data(b::AbstractDMD{T}) where {T} = "DMD with T=$(T), N=$(b.N), M=$(b.M)" allowed_types(::Type{<:AbstractDMD}) = cuNumeric.SUPPORTED_FLOAT_TYPES @@ -65,8 +65,8 @@ _dmd_row(v) = v isa NDArray ? cuNumeric.reshape(v, (1, length(v))) : reshape(v, # svd / eigen return factorizations whose stores the lifetime rewriter cannot # see, so those stay outside the macro. The GEMM lift is wrapped. # -# Do not form Diagonal(1 ./ S) inside @analyze_lifetimes: the rewriter treats -# `1 ./ S` as a last-used temp and destroy!s it, while Diagonal still holds +# Do not form Diagonal(1 ./ S) inside @accelerate: the rewriter treats +# `1 ./ S` as a last-used temporary and frees it while Diagonal still holds # that same vector. Scale columns with a broadcast instead (same math). function _dmd_factors(X, r) n = size(X, 2) @@ -84,7 +84,11 @@ let body = quote (B, Ã) end @eval _dmd_project(::DMDBaseline, X, X2, U, Vt, S) = $body - @eval _dmd_project(::DMDLifetimes, X, X2, U, Vt, S) = @analyze_lifetimes $body + @eval @accelerate function _dmd_project( + ::DMDAccelerated, X, X2, U, Vt, S + ) + $body + end end function _dmd_compute!(b::AbstractDMD, X, r) @@ -108,9 +112,9 @@ function check_benchmark_correctness( Xh = rand(T, b.N, b.M) X = NDArray(Xh) r = _dmd_rank(b) - # Values, not lifetimes: compare against the baseline body on host and device. + # Values, not lifetimes: compare the selected device path against the host baseline. ref = DMDBaseline{T}(; N=b.N, M=b.M) - λ, _ = _dmd_compute!(ref, X, r) + λ, _ = _dmd_compute!(b, X, r) λh, _ = _dmd_compute!(ref, Xh, r) mag = sort(abs.(Array(λ)); rev=true) @@ -119,4 +123,4 @@ function check_benchmark_correctness( end register_benchmark("dmd_baseline", DMDBaseline) -register_benchmark("dmd_lifetimes", DMDLifetimes) +register_benchmark("dmd_accelerated", DMDAccelerated) diff --git a/benchmark/src/benchmarks/grayscott.jl b/benchmark/src/benchmarks/grayscott.jl index 3ba6e6398..2fee6146e 100644 --- a/benchmark/src/benchmarks/grayscott.jl +++ b/benchmark/src/benchmarks/grayscott.jl @@ -18,7 +18,7 @@ Base.@kwdef struct GrayScottBaseline{T} <: AbstractGrayScott{T} M::Int end -Base.@kwdef struct GrayScottLifetimes{T} <: AbstractGrayScott{T} +Base.@kwdef struct GrayScottAccelerated{T} <: AbstractGrayScott{T} N::Int M::Int end @@ -92,66 +92,69 @@ function check_benchmark_correctness( return (u_ok && v_ok) ? "pass" : "fail" end -# VARIANT DESCRIPTION -# baseline: as written -# lifetimes: step wrapped in @analyze_lifetimes -let body = quote - # currently we don't have NDArray^x working yet. every operator is dotted - # so each rhs fuses into a single broadcast kernel rather than shattering - # into bare +/-/* binary tasks. - F_u = ( - ( - .-u[2:(end - 1), 2:(end - 1)] .* - (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) - ) .+ args.f .* (1.0f0 .- u[2:(end - 1), 2:(end - 1)]) - ) - F_v = ( - ( - u[2:(end - 1), 2:(end - 1)] .* - (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) - ) .- (args.f + args.k) .* v[2:(end - 1), 2:(end - 1)] - ) - # 2-D Laplacian via slicing, excluding boundaries - u_lap = ( - ( - u[3:end, 2:(end - 1)] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ - u[1:(end - 2), 2:(end - 1)] - ) ./ args.dx^2 .+ - ( - u[2:(end - 1), 3:end] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ - u[2:(end - 1), 1:(end - 2)] - ) ./ args.dx^2 - ) - v_lap = ( - ( - v[3:end, 2:(end - 1)] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+ - v[1:(end - 2), 2:(end - 1)] - ) ./ args.dx^2 .+ - ( - v[2:(end - 1), 3:end] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+ - v[2:(end - 1), 1:(end - 2)] - ) ./ args.dx^2 - ) - - # Forward-Euler step for all interior points - u_new[2:(end - 1), 2:(end - 1)] = - ((args.c_u .* u_lap) .+ F_u) .* args.dt .+ u[2:(end - 1), 2:(end - 1)] - v_new[2:(end - 1), 2:(end - 1)] = - ((args.c_v .* v_lap) .+ F_v) .* args.dt .+ v[2:(end - 1), 2:(end - 1)] - - # Periodic boundary conditions - u_new[:, 1] = u[:, end - 1] - u_new[:, end] = u[:, 2] - u_new[1, :] = u[end - 1, :] - u_new[end, :] = u[2, :] - v_new[:, 1] = v[:, end - 1] - v_new[:, end] = v[:, 2] - v_new[1, :] = v[end - 1, :] - v_new[end, :] = v[2, :] - end +# Shared syntax tree keeps every Gray-Scott variant on the exact same workload. +const GRAYSCOTT_STEP_BODY = quote + # currently we don't have NDArray^x working yet. every operator is dotted + # so each rhs fuses into a single broadcast kernel rather than shattering + # into bare +/-/* binary tasks. + F_u = ( + ( + .-u[2:(end - 1), 2:(end - 1)] .* + (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) + ) .+ args.f .* (1.0f0 .- u[2:(end - 1), 2:(end - 1)]) + ) + F_v = ( + ( + u[2:(end - 1), 2:(end - 1)] .* + (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) + ) .- (args.f + args.k) .* v[2:(end - 1), 2:(end - 1)] + ) + # 2-D Laplacian via slicing, excluding boundaries + u_lap = ( + ( + u[3:end, 2:(end - 1)] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ + u[1:(end - 2), 2:(end - 1)] + ) ./ args.dx^2 .+ + ( + u[2:(end - 1), 3:end] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ + u[2:(end - 1), 1:(end - 2)] + ) ./ args.dx^2 + ) + v_lap = ( + ( + v[3:end, 2:(end - 1)] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+ + v[1:(end - 2), 2:(end - 1)] + ) ./ args.dx^2 .+ + ( + v[2:(end - 1), 3:end] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+ + v[2:(end - 1), 1:(end - 2)] + ) ./ args.dx^2 + ) + + # Forward-Euler step for all interior points + u_new[2:(end - 1), 2:(end - 1)] = + ((args.c_u .* u_lap) .+ F_u) .* args.dt .+ u[2:(end - 1), 2:(end - 1)] + v_new[2:(end - 1), 2:(end - 1)] = + ((args.c_v .* v_lap) .+ F_v) .* args.dt .+ v[2:(end - 1), 2:(end - 1)] + + # Periodic boundary conditions + u_new[:, 1] = u[:, end - 1] + u_new[:, end] = u[:, 2] + u_new[1, :] = u[end - 1, :] + u_new[end, :] = u[2, :] + v_new[:, 1] = v[:, end - 1] + v_new[:, end] = v[:, 2] + v_new[1, :] = v[end - 1, :] + v_new[end, :] = v[2, :] +end + +# Original baseline and recommended function-form benchmark. +let body = deepcopy(GRAYSCOTT_STEP_BODY) @eval _gs_step!(b::GrayScottBaseline, u, v, u_new, v_new, args::GSParams) = $body - @eval _gs_step!(b::GrayScottLifetimes, u, v, u_new, v_new, args::GSParams) = - @analyze_lifetimes $body + definition = _define_accelerated_definition( + :(_gs_step!(b::GrayScottAccelerated, u, v, u_new, v_new, args::GSParams)), body + ) + @eval $definition end function run!(b::AbstractGrayScott, st::GrayScottState) @@ -163,4 +166,4 @@ function run!(b::AbstractGrayScott, st::GrayScottState) end register_benchmark("grayscott_baseline", GrayScottBaseline) -register_benchmark("grayscott_lifetimes", GrayScottLifetimes) +register_benchmark("grayscott_accelerated", GrayScottAccelerated) diff --git a/benchmark/src/benchmarks/grayscott_accelerate_forms.jl b/benchmark/src/benchmarks/grayscott_accelerate_forms.jl new file mode 100644 index 000000000..187dc299f --- /dev/null +++ b/benchmark/src/benchmarks/grayscott_accelerate_forms.jl @@ -0,0 +1,95 @@ +# Compare the four scope contracts of `@accelerate` on one shared Gray-Scott step. +# Each type has a distinct result name so benchmark runs produce separate CSVs. + +abstract type AbstractGrayScottAccelerateForm{T} <: AbstractGrayScott{T} end + +Base.@kwdef struct GrayScottFunctionAccelerated{T} <: + AbstractGrayScottAccelerateForm{T} + N::Int + M::Int +end + +Base.@kwdef struct GrayScottBeginAccelerated{T} <: AbstractGrayScottAccelerateForm{T} + N::Int + M::Int +end + +Base.@kwdef struct GrayScottLetAccelerated{T} <: AbstractGrayScottAccelerateForm{T} + N::Int + M::Int +end + +Base.@kwdef struct GrayScottExpressionAccelerated{T} <: + AbstractGrayScottAccelerateForm{T} + N::Int + M::Int +end + +name(::GrayScottFunctionAccelerated) = "grayscott_function_accelerated" +name(::GrayScottBeginAccelerated) = "grayscott_begin_accelerated" +name(::GrayScottLetAccelerated) = "grayscott_let_accelerated" +name(::GrayScottExpressionAccelerated) = "grayscott_expression_accelerated" + +# Function form is the reusable default: arguments and the return value survive, +# while non-returned locals may fuse across statements or die after their last use. +let body = deepcopy(GRAYSCOTT_STEP_BODY) + definition = _define_accelerated_definition( + :(_gs_step!(b::GrayScottFunctionAccelerated, u, v, u_new, v_new, args::GSParams)), + body, + :function, + ) + @eval $definition +end + +# `begin` adds no scope. Every named local remains visible, so it measures the +# multi-output/materialized path rather than eliminating named intermediates. +let body = deepcopy(GRAYSCOTT_STEP_BODY) + definition = _define_accelerated_definition( + :(_gs_step!(b::GrayScottBeginAccelerated, u, v, u_new, v_new, args::GSParams)), + body, + :begin, + ) + @eval $definition +end + +# `let` is a hard one-off scope. Only its result escapes, allowing aggressive +# inter-statement fusion and last-use cleanup for all other local temporaries. +let body = deepcopy(GRAYSCOTT_STEP_BODY) + definition = _define_accelerated_definition( + :(_gs_step!(b::GrayScottLetAccelerated, u, v, u_new, v_new, args::GSParams)), + body, + :let, + ) + @eval $definition +end + +# Expression form has no multi-statement scope. Accelerating each RHS preserves +# fusion inside that expression but deliberately materializes statement results, +# isolating intra-expression fusion from the inter-statement rewrite cases above. +function accelerate_grayscott_rhs(body::Expr) + statements = Any[] + for statement in body.args + if statement isa LineNumberNode + push!(statements, statement) + elseif statement isa Expr && statement.head === :(=) + lhs, rhs = statement.args + push!(statements, :($lhs = @accelerate $rhs)) + else + error("Gray-Scott expression benchmark expected assignments; got $(repr(statement))") + end + end + return Expr(:block, statements...) +end + +let body = accelerate_grayscott_rhs(deepcopy(GRAYSCOTT_STEP_BODY)) + @eval function _gs_step!( + b::GrayScottExpressionAccelerated, u, v, u_new, v_new, args::GSParams + ) + $body + end +end + +register_benchmark("grayscott_function_accelerated", GrayScottFunctionAccelerated) +register_benchmark("grayscott_begin_accelerated", GrayScottBeginAccelerated) +register_benchmark("grayscott_let_accelerated", GrayScottLetAccelerated) +register_benchmark("grayscott_expression_accelerated", GrayScottExpressionAccelerated) diff --git a/benchmark/src/core.jl b/benchmark/src/core.jl index 526ce4eb1..7a0aacbc4 100644 --- a/benchmark/src/core.jl +++ b/benchmark/src/core.jl @@ -39,6 +39,16 @@ function total_flops end function initialize end function run! end +# Internal adapter for benchmark generators that share a quoted step body. +function _define_accelerated_definition(signature, body, form=:function) + if form === :function + return cuNumeric._accelerate_expand(Expr(:function, signature, body), @__MODULE__) + end + scoped = form === :begin ? Expr(:block, body.args...) : Expr(:let, body) + call = Expr(:macrocall, Symbol("@accelerate"), LineNumberNode(0), scoped) + return Expr(:function, signature, Expr(:block, Base.macroexpand(@__MODULE__, call))) +end + # Maps a benchmarks.toml table name to its benchmark type. Each benchmark file # registers itself via `register_benchmark`. const BENCHMARKS = Dict{String,Type}() diff --git a/docs/make.jl b/docs/make.jl index 628a272bc..5ce5328be 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -42,7 +42,7 @@ makedocs(; ], "Performance Tips" => [ "Kernel Fusion" => "perf/kernel_fusion.md", - "Reduce Allocations" => "perf/reduce_allocations.md", + "The @accelerate Macro" => "perf/reduce_allocations.md", "Patterns to Avoid" => "perf/patterns_to_avoid.md", ], "Configuration" => [ diff --git a/docs/src/api.md b/docs/src/api.md index a187a79e1..d8197626d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -1,9 +1,9 @@ # NDArray Reference -Indexing, reshaping, reductions, comparisons, memory helpers, lifetime macros, and related utilities. For constructors (`zeros`, `ones`, `rand`, …) see [Initialization](./api_initialization.md). For RNG engines and `default_rng`, see [Random](./api_random.md). +Indexing, reshaping, reductions, comparisons, memory helpers, acceleration macros, and related utilities. For constructors (`zeros`, `ones`, `rand`, …) see [Initialization](./api_initialization.md). For RNG engines and `default_rng`, see [Random](./api_random.md). ```@autodocs Modules = [cuNumeric] -Pages = ["ndarray/ndarray.jl", "ndarray/linalg.jl", "ndarray/batched_linalg.jl", "cuNumeric.jl", "warnings.jl", "util.jl", "memory.jl", "scoping/scoping.jl"] +Pages = ["ndarray/ndarray.jl", "ndarray/linalg.jl", "ndarray/batched_linalg.jl", "cuNumeric.jl", "warnings.jl", "util.jl", "memory.jl", "scoping/scoping.jl", "scoping/accelerate.jl"] Filter = t -> !(t isa Function && nameof(t) in (:zeros, :ones, :fill, :trues, :falses, :eye, :rand, :rand!, :randn, :randn!, :randexp, :randexp!, :default_rng, :random, :random!)) ``` diff --git a/docs/src/api_cuda.md b/docs/src/api_cuda.md index 6e912cd8a..eedb9806c 100644 --- a/docs/src/api_cuda.md +++ b/docs/src/api_cuda.md @@ -58,7 +58,7 @@ allowscalar() do end ``` -See `examples/custom_cuda.jl` for a more complete example with multiple kernels. +See `examples/custom_cuda.jl` for a runnable two-kernel example. ## API Reference diff --git a/docs/src/api_unary.md b/docs/src/api_unary.md index 3ea0442b3..7582f4ce1 100644 --- a/docs/src/api_unary.md +++ b/docs/src/api_unary.md @@ -1,7 +1,7 @@ # Unary Operations >[!NOTE] -> Prefer `@.` for multi-op elementwise expressions so every operator is dotted (especially unary negation). This ensures broadcast operations are fused. See [Kernel Fusion](./perf/kernel_fusion.md). +> Prefer `@.` for multi-op elementwise expressions so every operator is dotted (especially unary negation). This makes eligible CUDA broadcasts fusion-friendly. See [Kernel Fusion](./perf/kernel_fusion.md). The following unary operations are supported and can be broadcast over `NDArray`: diff --git a/docs/src/benchmarks/howto.md b/docs/src/benchmarks/howto.md index 851f3bdd3..efce38239 100644 --- a/docs/src/benchmarks/howto.md +++ b/docs/src/benchmarks/howto.md @@ -47,7 +47,10 @@ n_correctness_iter = 5 - `cupynumeric` / `cuda`: optional comparison backends - `check_correctness`: one CPU-reference check per config (not per timed iter), recorded in the CSV -Each `[[name]]` block is a registered benchmark (`gemm`, `montecarlo`, `dmd_baseline`, `dmd_lifetimes`, `grayscott_baseline`, `grayscott_lifetimes`, …). Names must match what `src/benchmarks/*.jl` registers. +Each `[[name]]` block is a registered benchmark (`gemm`, `montecarlo`, +`dmd_baseline`, `dmd_accelerated`, `grayscott_baseline`, +`grayscott_function_accelerated`, …). Names must match what +`src/benchmarks/*.jl` registers. DMD's `N` is the number of spatial degrees of freedom (rows of the snapshot matrix), not a grid side length. The SVD is of the tall-skinny `N × (M-1)` matrix `X1`. Thin SVD plus the rank-`r` lift is `Θ(N)` when `M` and `r` are fixed, so weak scaling is `N ∝ P` (same idea as Monte Carlo, not GEMM's `N ∝ P^{1/3}`). The flop count is in `src/benchmarks/dmd.jl`. @@ -78,7 +81,12 @@ M = [150, 300, 600] That is 2 types × 3 sweep points = **6 runs**. -`fusion` toggles cuNumeric broadcast fusion (`true`/`false` or `"on"`/`"off"`, default `true`). Comparison backends ignore fusion and run once (on the fused pass), not per variant. Names ending in `_lifetimes` are cuNumeric-only code-path variants. +`fusion` toggles cuNumeric broadcast fusion (`true`/`false` or `"on"`/`"off"`, +default `true`). Comparison backends ignore fusion and run once (on the fused +pass), not per variant. Entries ending in `_accelerated` are cuNumeric-only. +The Gray-Scott function, `begin`, `let`, and expression entries compare the +four `@accelerate` scope contracts on the same step; `dmd_accelerated` applies +the recommended function form to the DMD projection. Gotcha: when `T = ["Float32", "Float64"]` and a length-2 `N`/`M` sweep you get all **4** combinations, not a paired `Float32 -> N[1]`. To pin a type to a size, use separate `[[name]]` blocks. diff --git a/docs/src/benchmarks/results.md b/docs/src/benchmarks/results.md index fd62aedcf..08cc9e3c2 100644 --- a/docs/src/benchmarks/results.md +++ b/docs/src/benchmarks/results.md @@ -1,6 +1,6 @@ # Benchmark Results -For JuliaCon2025 we benchmarks cuNumeric.jl on 8 A100 GPUs (single-node) and compared it to the Python library cuPyNumeric and other relevant benchmarks depending on the problem. All results shown are weak scaling. We hope to have multi-node benchmarks soon! +These historical JuliaCon 2025 results compare cuNumeric.jl with cuPyNumeric and problem-specific alternatives on one node with eight A100 GPUs. All plots show weak scaling. See [How to Benchmark](./howto.md) for the current harness and its baseline, `@accelerate`, fused, and unfused variants. ## SGEMM @@ -24,7 +24,7 @@ mul!(C, A, B) ## Monte-Carlo Integration -Monte-Carlo integration is embaressingly parallel and should scale perfectly. We do not know the exact number of operations in `exp` so the GFLOPs is off by a constant factor. +Monte-Carlo integration is embarrassingly parallel. Because the exact operation count of `exp` is implementation-dependent, the plotted operation rate is scaled by an approximate constant. Code Outline: ```julia diff --git a/docs/src/configuration/hardware.md b/docs/src/configuration/hardware.md index 0102ee183..4ddacc4c2 100644 --- a/docs/src/configuration/hardware.md +++ b/docs/src/configuration/hardware.md @@ -1,6 +1,6 @@ # Hardware Configuration -There is no programmatic way to set the hardware configuration used by CuPyNumeric (as of 26.01). By default, the hardware configuration is set automatically by Legate. This configuration can be manipulated through the following environment variables: +Legate chooses the hardware configuration automatically by default. Set these environment variables before starting Julia to override it: - `LEGATE_SHOW_CONFIG` : When set to 1, the Legate config is printed to stdout - `LEGATE_AUTO_CONFIG`: When set to 1, Legate will automatically choose the hardware configuration diff --git a/docs/src/debugging.md b/docs/src/debugging.md index 6b25941d5..3e868c907 100644 --- a/docs/src/debugging.md +++ b/docs/src/debugging.md @@ -6,7 +6,7 @@ Debug the layer that matches the problem: |---|---| | Which operations did Legate submit, and when did they run? | [Legate logs and profiles](#trace-legate-runtime-work) | | How were broadcasts fused? | [`BCAST_FUSION_DEBUG`](#inspect-fused-broadcasts-with-bcast_fusion_debug) | -| Where does `@analyze_lifetimes` free temporaries? | [`@show_lifetimes`](#inspect-lifetime-rewrites-with-show_lifetimes) | +| How does `@accelerate` rewrite code and free temporaries? | [`@show_lifetimes`](#inspect-lifetime-rewrites-with-show_lifetimes) | ## Trace Legate runtime work @@ -78,7 +78,7 @@ cuNumeric already supplies names for individual operations when task-scope namin When broadcast fusion is on, set `cuNumeric.BCAST_FUSION_DEBUG[] = true` to print inter-statement rewrites and each fused kernel's expression tree, arguments, and launch geometry. Inter-statement rewrites are reported when -`@analyze_lifetimes` expands, so enable the flag before defining or evaluating +`@accelerate` expands, so enable the flag before defining or evaluating the expression you want to inspect. Kernel details are reported at runtime. ```julia @@ -91,15 +91,16 @@ A = cuNumeric.ones(Float32, N, N) B = cuNumeric.ones(Float32, N, N) C = cuNumeric.zeros(Float32, N, N) -@analyze_lifetimes begin +@accelerate function combine!(C, A, B) product = A[2:end-1, 2:end-1] .* B[2:end-1, 2:end-1] C[2:end-1, 2:end-1] = product .+ 2.0f0 + return C end cuNumeric.BCAST_FUSION_DEBUG[] = false ``` -For example, a single-use producer inside `@analyze_lifetimes` is reported as: +For example, a single-use producer inside `@accelerate` is reported as: ```text ======================================== inter-broadcast fusion rewrite @@ -147,38 +148,23 @@ See [Kernel Fusion](./perf/kernel_fusion.md) for `@.` / fusion usage, and [Inter ## Inspect lifetime rewrites with `@show_lifetimes` -`@analyze_lifetimes` rewrites a block so temps are freed after their last use. `@show_lifetimes` prints the re-written code (without execution). It is pure AST work, so it works even without a GPU. +`@accelerate` rewrites straight-line code so eligible broadcasts combine and non-returned temporaries are freed after their final use. `@show_lifetimes` prints the exact expansion without executing it, so it works without a GPU. ```julia using cuNumeric -@show_lifetimes begin +@show_lifetimes function update!(C, A, B) result = A[1:end, :] .+ B[1:end, :] C .= result .* 2.0 + return C end ``` -Example output when broadcast fusion is enabled (fusion-aware analysis): - -```text -@analyze_lifetimes expansion (fusion-aware analysis) ------------------------------------------------------------- - 1 tmp1 = A[1:end, :] - 2 tmp2 = B[1:end, :] - 3 tmp3 = tmp1 .+ tmp2 - ✗ free tmp1 - ✗ free tmp2 - 4 result = tmp3 - 5 res3 = (C .= result .* 2.0) - ✗ free tmp3 - 6 res3 ------------------------------------------------------------- -``` - How to read it: -- Numbered lines are the rewritten statements. +- The header identifies the exact function, `let`, block, or expression form expanded. +- Numbered lines are rewritten statements. - Red `✗ free tmpN` lines are the inserted `maybe_insert_delete` calls. -- With fusion enabled, dotted intermediates stay as broadcast expressions instead of being treated as many separate allocations. With fusion disabled, the header says `plain analysis` and more call sites are hoisted. +- With fusion enabled, dotted intermediates stay as broadcast expressions instead of being treated as separate allocations. With fusion disabled, the header says `plain` and more call sites are hoisted. Use this when a hot loop still looks allocation-heavy, or when you want to confirm that a value is freed before it escapes the block. diff --git a/docs/src/developer_mode.md b/docs/src/developer_mode.md index 5cac1631d..913977942 100644 --- a/docs/src/developer_mode.md +++ b/docs/src/developer_mode.md @@ -69,4 +69,4 @@ Restart Julia. You do not need `Pkg.build` for pure JLL mode (the build script e - [Build Modes](./install.md): JLL, developer, and conda providers - [CNPreferences](./api_preferences.md): preference defaults and function reference - [Debugging](./debugging.md): fusion and lifetime printers while developing -- [Internals](./internals.md): how fusion and `@analyze_lifetimes` work +- [Internals](./internals.md): how fusion and `@accelerate` work diff --git a/docs/src/examples/grayscott.md b/docs/src/examples/grayscott.md index 59ef00463..7c230143b 100644 --- a/docs/src/examples/grayscott.md +++ b/docs/src/examples/grayscott.md @@ -1,59 +1,27 @@ # Gray-Scott Reaction Diffusion -```julia -# found in examples/gray-scott.jl -using cuNumeric -using Plots - -# Flattened u snapshots land here for examples/dmd.jl to analyze. -const SNAPSHOT_FILE = "gray-scott.h5" - -struct Params{T} - dx::T - dt::T - c_u::T - c_v::T - f::T - k::T - - function Params(dx=1.0f0, c_u=1.0f0, c_v=0.3f0, f=0.03f0, k=0.06f0) - new{Float32}(dx, dx/5, c_u, c_v, f, k) - end -end - -function bc!(u_new, v_new, u, v) - u_new[:,1] = u[:,end-1] - u_new[:,end] = u[:,2] - u_new[1,:] = u[end-1,:] - u_new[end,:] = u[2,:] - v_new[:,1] = v[:,end-1] - v_new[:,end] = v[:,2] - v_new[1,:] = v[end-1,:] - v_new[end,:] = v[2,:] -end - -function step!(u, v, u_new, v_new, args::Params) - @analyze_lifetimes begin - # Prefer @. so every op is dotted and the tree can fuse - F_u = @. -u[2:end-1, 2:end-1] * (v[2:end-1, 2:end-1]^2) + - args.f * (1.0f0 - u[2:end-1, 2:end-1]) - F_v = @. u[2:end-1, 2:end-1] * (v[2:end-1, 2:end-1]^2) - - (args.f + args.k) * v[2:end-1, 2:end-1] - - u_lap = @. ( - (u[3:end, 2:end-1] - 2 * u[2:end-1, 2:end-1] + u[1:end-2, 2:end-1]) / args.dx^2 + - (u[2:end-1, 3:end] - 2 * u[2:end-1, 2:end-1] + u[2:end-1, 1:end-2]) / args.dx^2 - ) - v_lap = @. ( - (v[3:end, 2:end-1] - 2 * v[2:end-1, 2:end-1] + v[1:end-2, 2:end-1]) / args.dx^2 + - (v[2:end-1, 3:end] - 2 * v[2:end-1, 2:end-1] + v[2:end-1, 1:end-2]) / args.dx^2 - ) - - u_new[2:end-1, 2:end-1] = @. (args.c_u * u_lap + F_u) * args.dt + u[2:end-1, 2:end-1] - v_new[2:end-1, 2:end-1] = @. (args.c_v * v_lap + F_v) * args.dt + v[2:end-1, 2:end-1] - end +The runnable example in `examples/gray-scott.jl` evolves two chemical fields with periodic boundaries. Its update is a straight-line function, so the recommended function form of `@accelerate` can release reaction and Laplacian temporaries after their final use and fuse eligible CUDA broadcasts. +```julia +@accelerate function step!(u, v, u_new, v_new, args::Params) + F_u = @. -u[2:end-1, 2:end-1] * v[2:end-1, 2:end-1]^2 + + args.f * (1.0f0 - u[2:end-1, 2:end-1]) + F_v = @. u[2:end-1, 2:end-1] * v[2:end-1, 2:end-1]^2 - + (args.f + args.k) * v[2:end-1, 2:end-1] + + u_lap = @. ( + (u[3:end, 2:end-1] - 2u[2:end-1, 2:end-1] + u[1:end-2, 2:end-1]) / args.dx^2 + + (u[2:end-1, 3:end] - 2u[2:end-1, 2:end-1] + u[2:end-1, 1:end-2]) / args.dx^2 + ) + v_lap = @. ( + (v[3:end, 2:end-1] - 2v[2:end-1, 2:end-1] + v[1:end-2, 2:end-1]) / args.dx^2 + + (v[2:end-1, 3:end] - 2v[2:end-1, 2:end-1] + v[2:end-1, 1:end-2]) / args.dx^2 + ) + + u_new[2:end-1, 2:end-1] = @. (args.c_u * u_lap + F_u) * args.dt + u[2:end-1, 2:end-1] + v_new[2:end-1, 2:end-1] = @. (args.c_v * v_lap + F_v) * args.dt + v[2:end-1, 2:end-1] bc!(u_new, v_new, u, v) + return nothing end function gray_scott() @@ -103,11 +71,9 @@ function gray_scott() cuNumeric.Legate.runtime_sync() return u, v - -end - -u, v = gray_scott() + end ``` + ![Simulation Output](../gray-scott.gif) The snapshots written to `gray-scott.h5` are the input to diff --git a/docs/src/examples/montecarlo.md b/docs/src/examples/montecarlo.md index 5a66572db..eb45188a0 100644 --- a/docs/src/examples/montecarlo.md +++ b/docs/src/examples/montecarlo.md @@ -1,40 +1,29 @@ # Monte-Carlo Integration -Most integrals can be estimated with a basic Monte-Carlo estimator: +For uniformly sampled points `x_i` in a domain of volume `\Omega`, a basic Monte-Carlo estimator is ```math -\hat{I}_N = \frac{\Omega}{N}\sum_{i=1}^Nf(x_i) +\hat{I}_N = \frac{\Omega}{N}\sum_{i=1}^N f(x_i). ``` -where `N` is the number of samples, ``\Omega`` is the volume of the domain and ``x_i`` are sampled indpendently and uniformly at random from the domain. This estimator is guranteed to converge (subject to some minor constraints) at a rate independent of the dimension and is embaressingly parallel to compute! -In the example below, we estimate the integral: -```math -I = \int_{-\infty}^{\infty}e^{-x^2}. -``` +This example estimates `\int_{-\infty}^{\infty} e^{-x^2}\,dx` by sampling the finite interval `[-10, 10]`. `@accelerate` frees the non-returned sample arrays after their final use; CUDA may also fuse eligible broadcasts. -Since we cannot uniformly sample form negative to positive infinity, we truncate the domain between -5 and 5. This is ok since the integrand exponentially decays and we won't be off by much in the end. ```julia -# found in examples/integrate.jl +# examples/integrate.jl using cuNumeric -# Note that we do not yet support broadcasting -# custom functions over NDArray, so the broadcasting MUST -# be done inside the function -integrand = (x) -> @. exp(-x^2) - -N = 1_000_000 +integrand(x) = @. exp(-x^2) -x_max = 10.0f0 -domain = [-x_max, x_max] -Ω = domain[2] - domain[1] +@accelerate function monte_carlo(N, x_max) + Ω = 2 * x_max + raw_samples = cuNumeric.rand(N) + samples = @. Ω * raw_samples - x_max + return (Ω / N) * sum(integrand(samples)) +end -samples = Ω * cuNumeric.rand(N) -samples = @. samples - x_max - -# Reductions return 0D NDArrays instead -# of a scalar to avoid blocking runtime -estimate = (Ω / N) * sum(integrand(samples)) - -println("Monte-Carlo Estimate: $(estimate)") -println("Analytical: $(sqrt(pi))") +estimate = monte_carlo(1_000_000, 10.0f0) +println("Monte-Carlo estimate: $(estimate)") +println("Analytical value: $(sqrt(pi))") ``` + +The result is a 0-dimensional `NDArray`, which keeps the reduction asynchronous. Use `unwrap(estimate)` only when a Julia scalar is required. diff --git a/docs/src/index.md b/docs/src/index.md index fe626eb6a..bb266637d 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -5,7 +5,7 @@ ``` -[![Documentation dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://julialegate.github.io/cuNumeric.jl/dev/) [![codecov](https://codecov.io/github/julialegate/cuNumeric.jl/branch/main/graph/badge.svg)](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT) +[![Documentation dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://julialegate.github.io/cuNumeric.jl/dev) [![codecov](https://codecov.io/github/julialegate/cuNumeric.jl/branch/main/graph/badge.svg)](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](https://opensource.org/licenses/MIT) cuNumeric.jl wraps and extends the [cuPyNumeric](https://github.com/nv-legate/cupynumeric) library from NVIDIA to bring distributed array computing on GPUs and CPUs to Julia. The central type is `NDArray`, which behaves like Julia's `Array` or the `CuArray` from [CUDA.jl](https://github.com/juliagpu/cuda.jl), but executes across multiple GPUs/CPUs. We implement array-level operations on `NDArray` which can be composed into larger programs without the need for explicit MPI calls or writing CUDA kernels. @@ -20,7 +20,7 @@ using Pkg Pkg.add(url = "https://github.com/JuliaLegate/cuNumeric.jl", rev = "main") ``` -The first time might take awhile as it has to install multiple large dependencies such as the CUDA SDK (if you have an NVIDIA GPU). To use a local build of cupynumeric.so, see [Build Modes](./install.md). +The first installation can take a while because it includes several large dependencies, such as the CUDA SDK. To use a local cupynumeric build, see [Build Modes](https://julialegate.github.io/cuNumeric.jl/dev/install). ```julia using cuNumeric @@ -30,7 +30,7 @@ cuNumeric.versioninfo() > [!WARNING] > Starting more than one instance of cuNumeric.jl can lead to a hard-crash. The default hardware configuration reserves all available resources. -For more details, see [Hardware](./configuration/hardware.md). +For more details, see [Hardware](https://julialegate.github.io/cuNumeric.jl/dev/configuration/hardware). ### How `NDArray`s work @@ -40,7 +40,7 @@ The semantics of `NDArray` closely mirror Julia's `Array`, and in most cases it **Slices are views.** Indexing an `NDArray` with ranges returns a view onto the same store, not a copy. That differs from Base Julia, where `A[1:n]` allocates a new `Array`. Mutations through an `NDArray` slice are visible through other aliases of the same data. -**Reductions return arrays, not Julia scalars.** Reductions such as `sum(A)` produce a **0D or 1D** `NDArray` (axis reductions produce a lower-rank `NDArray`), rather than a bare `Float64` / `Float32`. That keeps the Legate task graph asynchronous instead of forcing synchronization to communite with the Julia runtime. When you need a plain Julia number, call `unwrap`: +**Reductions return arrays, not Julia scalars.** Reductions such as `sum(A)` produce a **0D or 1D** `NDArray` (axis reductions produce a lower-rank `NDArray`), rather than a bare `Float64` / `Float32`. That keeps the Legate task graph asynchronous instead of forcing synchronization to communicate with the Julia runtime. When you need a plain Julia number, call `unwrap`: ```julia s = sum(A) # NDArray{T,0} @@ -49,7 +49,7 @@ x = unwrap(s) # T, e.g. Float32 **The Legate runtime builds a DAG asynchronously.** Calling `cuNumeric.zeros` or `A .+ B` records work into Legate's task graph rather than blocking until every GPU kernel finishes. Results are materialized when you need them (for example `println`, `unwrap`, or converting with `Array(A)`). Hiding latency enables performant code. -For API details see [Initialization](./api_initialization.md), [Random](./api_random.md), and [NDArray Reference](./api.md). For anti-patterns that kill performance, see [Patterns to Avoid](./perf/patterns_to_avoid.md). +For API details see [Initialization](https://julialegate.github.io/cuNumeric.jl/dev/api_initialization) and [NDArray Reference](https://julialegate.github.io/cuNumeric.jl/dev/api). For common performance pitfalls, see [Patterns to Avoid](https://julialegate.github.io/cuNumeric.jl/dev/perf/patterns_to_avoid). ### Kernel Fusion @@ -59,41 +59,33 @@ Nested broadcast expressions fuse into a single kernel by default when on GPU. P y .= @. -a + b * c ``` -See [Kernel Fusion](./perf/kernel_fusion.md) and [Debugging](./debugging.md) for controls and pretty printers. +See [Kernel Fusion](https://julialegate.github.io/cuNumeric.jl/dev/perf/kernel_fusion) and [Debugging](https://julialegate.github.io/cuNumeric.jl/dev/debugging) for controls and diagnostics. -### Helping the Garbage Collector +### The `@accelerate` macro -Many calls such as array slicing and un-fused broadcasts allocate a new `NDArray`. The Legate runtime keeps track of all references to the underlying data and will not free the memory until Julia's GC frees the `NDArray` handles. Because Julia's GC runs on memory pressure and an `NDArray` only stores a pointer (i.e., Julia's GC does not know the true size), many dead buffers accumulate and can cause out-of-memory errors. +`@accelerate` fuses eligible GPU broadcasts within and across statements, then releases materialized temporary `NDArray`s after their last use on CPU or GPU. See [The `@accelerate` Macro](https://julialegate.github.io/cuNumeric.jl/dev/perf/reduce_allocations) for usage guidance. -`@analyze_lifetimes` performs a **static last-use analysis** at macro-expansion time and inserts eager calls to immediately free unused `NDArrays`. These buffers can then be reused by legate later for same-sized allocations. +### Benchmarks -```julia -@analyze_lifetimes begin - result = @. A[1:end, :] + B[1:end, :] - C .= @. result * 2.0f0 -end -``` - -### Performance at a glance - -A representative benchmark figure will go here (add something like `docs/src/images/benchmarks-overview.png` when ready). - -Numbers, plots, and how to reproduce them live under [Benchmark Results](./benchmarks/results.md) and [How to Benchmark](./benchmarks/howto.md). +Results and reproduction instructions live under [Benchmark Results](https://julialegate.github.io/cuNumeric.jl/dev/benchmarks/results) and [How to Benchmark](https://julialegate.github.io/cuNumeric.jl/dev/benchmarks/howto). ### Try an example ```julia using cuNumeric -integrand = (x) -> @. exp(-x^2) +integrand(x) = @. exp(-x^2) + +@accelerate function monte_carlo(N, x_max) + Ω = 2 * x_max + raw_samples = cuNumeric.rand(N) + samples = @. Ω * raw_samples - x_max + return (Ω / N) * sum(integrand(samples)) +end N = 1_000_000 x_max = 10.0f0 -Ω = 2 * x_max - -samples = Ω .* cuNumeric.rand(N) -samples = samples .- x_max -estimate = (Ω / N) .* sum(integrand(samples)) +estimate = monte_carlo(N, x_max) println("Monte-Carlo Estimate: $(estimate)") ``` diff --git a/docs/src/internals.md b/docs/src/internals.md index 1afa20c03..652433251 100644 --- a/docs/src/internals.md +++ b/docs/src/internals.md @@ -1,72 +1,30 @@ # Internals -This page describes the implementation details of kernel fusion, manual memory management via `@analyze_lifetimes` and automatic memory management via GC heuristics. For docs on how to use these features, see [Kernel Fusion](./perf/kernel_fusion.md) and [Reduce Allocations](./perf/reduce_allocations.md). +This page summarizes the machinery behind [`@accelerate`](./perf/reduce_allocations.md), [kernel fusion](./perf/kernel_fusion.md), and allocation-driven garbage collection. ## Broadcast fusion -We compile nested Julia broadcast expressions on `NDArray` to a single CUDA kernel instead of launching one kernel per operation. +Julia first builds a nested `Broadcasted` tree for a dotted expression. When CUDA fusion is enabled and the tree is eligible, cuNumeric flattens it, generates a CUDA.jl kernel, and launches it through Legate. Otherwise `unravel_broadcast_tree` evaluates the operations individually. CPU execution always uses the unfused path. -### Pipeline +`@accelerate` adds an inter-statement syntax pass. It can merge a single-use broadcast producer into its consumer when doing so preserves mutation and aliasing semantics. The soft-scope `begin` form instead keeps named values materialized and may lower an eligible chain to a multi-output CUDA kernel. -- **Expression:** You write a dotted or `@.` expression such as `y .= @. a * b + c`. -- **Broadcast tree:** Julia builds the `Broadcasted` tree. -- **Fused path:** Flatten the tree, build a kernel with CUDA.jl, and launch it with Legate. -- **Unfused path:** `unravel_broadcast_tree` recursively unravels the tree and executes each operation one at a time. +Relevant source: `src/ndarray/broadcast_fusion.jl` and `src/scoping/`. -## Lifetimes and GC +## Lifetime analysis -Julia's GC sees an `NDArray` as a small handle. The actual data is owned by the Legate runtime. This means that to Julia's GC `NDArrays` do not create memory pressure and GC is never executed. To avoid out-of-memory errors we created the `@analyze_lifetimes` macro so users can manually manage the lifetimes of a code block and also manually track device memory to automatically invoke GC. +An `NDArray` is a small Julia handle to storage owned by Legate, so Julia's heap pressure understates the size of live array data. `@accelerate` therefore performs static last-use analysis: -### Eager last-use freeing with `@analyze_lifetimes` +1. Expand nested `@.` macros and reject non-straight-line code. +2. Apply conservative inter-statement fusion where the selected macro form permits it. +3. Hoist materialized temporary values and find their final uses. +4. Insert `maybe_insert_delete` calls after those uses while protecting caller-owned arguments, named soft-scope bindings, and returned values. -`@analyze_lifetimes` rewrites a block at macro-expansion time: +With fusion enabled, nested dotted nodes remain lazy and are not counted as separate array allocations. With fusion disabled, allocating calls are analyzed individually. `@show_lifetimes` prints the exact expansion without running it. -- Hoist temporary allocations into named temps. -- Find each temp's static last use. -- Insert calls to free temporary `NDArrays` after last-use. +The analysis is statement-linear rather than a control-flow graph pass. Apply it to straight-line function or loop bodies, not to control flow itself. -Under broadcast fusion, intermediate dotted nodes are **not** real `NDArray` allocations. The macro switches to a fusion-aware hoist that keeps dotted trees lazy and only treats slices, broadcast roots, and non-broadcast calls as real allocations. When fusion is off, every call (including dotted ops) is treated as a real allocation. +## Allocation-driven GC -```julia -@analyze_lifetimes begin - result = A[1:end, :] .+ B[1:end, :] - C .= result .* 2 -end -``` - -Use `@show_lifetimes` to print the rewritten block and the free sites without running the code. That is pure AST work and works without a GPU. - -The implementation uses separate passes for inter-statement broadcast fusion, -allocation hoisting, and finalizer insertion. The top-level scoping pass selects -the appropriate lifetime analysis based on whether broadcast fusion is enabled. - -- **Inter-statement broadcast fusion** (`rewrite_scope`): - Merges single-use broadcast statements into their consumer, e.g.: - `p = A .* B; C .= p .+ 1` → `C .= A .* B .+ 1` - This pass operates only on syntax and does not depend on cuNumeric types. - -- **Lifetime analysis**: - - With fusion enabled (`rewrite_broadcast_lifetimes`): keeps dotted trees lazy - and hoists only materialized values (slices, non-broadcast calls, etc.) - - With fusion disabled (`rewrite_eager_lifetimes`): hoists all allocating calls - Both passes rename temps and insert free calls after static last uses. - -- **Finalizer insertion** (`insert_finalizers`): - Traverses the AST generated by either lifetime pass and inserts `delete` calls - at the computed last-use sites. It preserves the final return value of the - block by not freeing it. - -Limits to keep in mind: - -- Analysis is statement-linear. It is not a full control-flow graph pass. Wrap hot loop bodies, not entire programs. -- Some paths free eagerly outside the macro (for example LHS slice views created during indexed assignment). - -Relevant source: `src/scoping/`. - -### Allocation-driven GC heuristics - -Every `NDArray` registers its byte size on construct and free. When predicted live bytes cross soft (~80%) or hard (~90%) fractions of available memory, and enough new growth has accumulated since the last collection, cuNumeric.jl triggers Julia `GC.gc`. - -`@analyze_lifetimes` reduces peak live temps. The heuristics catch cases the macro cannot see. +Each `NDArray` reports its byte size when constructed and freed. When predicted live bytes cross soft and hard fractions of available memory—and enough new growth has accumulated—cuNumeric asks Julia to collect garbage. Eager last-use freeing reduces peak live storage; the GC heuristic covers allocations the static pass cannot prove dead. Relevant source: `src/memory.jl`. diff --git a/docs/src/perf/kernel_fusion.md b/docs/src/perf/kernel_fusion.md index 04c437742..cb715fadf 100644 --- a/docs/src/perf/kernel_fusion.md +++ b/docs/src/perf/kernel_fusion.md @@ -1,61 +1,47 @@ # Kernel Fusion -On CUDA, nested broadcast expressions are fused into a single PTX kernel when fusion is enabled (the default). There is no separate `@fuse` macro. You write ordinary Julia broadcast code, and cuNumeric compiles eligible trees into one kernel instead of launching one op at a time. +When CUDA is available, cuNumeric can compile an eligible nested broadcast tree into one PTX kernel instead of launching one operation at a time. CPU execution follows the normal unfused path. -Prefer Julia's `@.` macro for multi-op elementwise expressions. Placing `.` on every operator by hand is easy to get wrong: missing a dot on unary negation or addition silently changes the meaning, and it can also break fusion by splitting work into the wrong ops. +Prefer Julia's `@.` macro for multi-operation elementwise expressions so every operator is dotted: ```julia -# Easy to miss the dot on negation +# A missing dot on unary negation changes the expression and can prevent fusion. y .= .-a .+ b .* c -# Prefer: @. dots every operator, which is clearer and fusion-friendly +# Prefer this form. y .= @. -a + b * c ``` -## Avoid preallocated intermediate broadcast buffers +Fusion requires array leaves with the same shape and at least `FUSE_BROADCAST_MIN_OPS` broadcast operations (default: 2). Shape-mismatched broadcasts such as `matrix .+ vector` use the unfused path. -Preallocation is useful for a final output or a buffer that must persist across -iterations. It can be counterproductive for a single-use intermediate inside -`@analyze_lifetimes`, however. An in-place `.=` assignment is an observable -mutation, so inter-statement broadcast fusion treats it as a kernel boundary: +## Fuse across statements with `@accelerate` -```julia -tmp = cuNumeric.zeros(Float32, N, N) -result = cuNumeric.zeros(Float32, N, N) +An ordinary assignment lets `@accelerate` substitute a single-use producer into its consumer: -@analyze_lifetimes begin - tmp .= @. A + B +```julia +@accelerate function update!(result, A, B, C) + tmp = @. A + B result .= @. tmp * C + 1.0f0 + return result end ``` -This materializes `tmp` before the second expression and requires separate -kernel launches. Instead, use an ordinary assignment for a single-use -intermediate and keep `.=` for the final destination: +This can become the equivalent of `result .= @. (A + B) * C + 1.0f0`. The rewrite is conservative: the producer must have one use, no intervening statement may invalidate its inputs, and the normal fusion requirements still apply. -```julia -result = cuNumeric.zeros(Float32, N, N) +Do not preallocate a single-use intermediate with `.=` merely to avoid allocation: -@analyze_lifetimes begin - tmp = @. A + B - result .= @. tmp * C + 1.0f0 -end +```julia +tmp .= @. A + B +result .= @. tmp * C + 1.0f0 ``` -The inter-statement pass can substitute `tmp` into its only consumer, producing -the equivalent of `result .= @. (A + B) * C + 1.0f0`. The intermediate is never -materialized, so the full expression can run as one fused kernel. +The mutation of `tmp` is observable, so it remains a kernel boundary. Keep preallocation when the intermediate is reused or its mutation must be visible. -This rewrite is intentionally conservative: the intermediate must have one -use, no intervening statement may invalidate its inputs, and all normal fusion -requirements still apply. Keep preallocation when an intermediate is reused, -must preserve mutation semantics, or cannot be fused. Use -[`BCAST_FUSION_DEBUG`](../debugging.md#inspect-fused-broadcasts-with-bcast_fusion_debug) -to confirm whether the rewrite occurred. +The `begin` form has different ownership semantics: named bindings remain live. On CUDA, an eligible same-shape chain can still be emitted as one multi-output kernel that materializes each binding. See [Accelerate Array Code](./reduce_allocations.md) for all macro forms. -Fusion applies when CUDA is available, the array leaves share the same shape, and the expression has at least `FUSE_BROADCAST_MIN_OPS` ops (default 2). Otherwise cuNumeric falls back to evaluating one op at a time. Shape-mismatched broadcasts such as `matrix .+ vector` use the unfused path. +## Configure and inspect fusion -Toggle fusion through `CNPreferences` (restart Julia after changing these): +Set preferences in one Julia process, then restart Julia: ```julia using CNPreferences @@ -63,14 +49,7 @@ using CNPreferences CNPreferences.enable_broadcast_fusion!() # default CNPreferences.disable_broadcast_fusion!() CNPreferences.set_broadcast_fusion_min_ops!(2) # default -CNPreferences.set_broadcast_fusion_min_ops!(1) # also fuse single-ops +CNPreferences.set_broadcast_fusion_min_ops!(1) # include single-op broadcasts ``` -What `set_broadcast_fusion_min_ops!` controls: - -- **`2` (default):** only trees with two or more ops fuse. Example: `y .= @. a * b + c` can fuse; `y .= cos.(x)` does not. Keeping single-ops on the unfused C-API path avoids PTX compile overhead when there is little to gain. -- **`1`:** every eligible broadcast can fuse, including unary / single-op forms. Prefer this when you want uniform fused behavior (for example in tests) rather than for typical apps. - -The threshold counts `Broadcasted` nodes in the expression tree. Set it through `CNPreferences`, then restart Julia. See [CNPreferences](../api_preferences.md). - -To inspect a fused launch or a lifetime rewrite, see [Debugging](../debugging.md). For the implementation pipeline, see [Internals](../internals.md). +The default threshold avoids PTX compilation overhead when there is little work to combine. See [CNPreferences](../api_preferences.md) for preference details and [`BCAST_FUSION_DEBUG`](../debugging.md#inspect-fused-broadcasts-with-bcast_fusion_debug) to confirm which path ran. diff --git a/docs/src/perf/reduce_allocations.md b/docs/src/perf/reduce_allocations.md index 37b38ca50..f36ab803b 100644 --- a/docs/src/perf/reduce_allocations.md +++ b/docs/src/perf/reduce_allocations.md @@ -1,31 +1,88 @@ -# Reduce Allocations +# The `@accelerate` Macro -Every intermediate `NDArray` (from a slice, broadcast, or function call) allocates a fresh buffer and waits for the Julia GC to free it. Because the GC runs on memory pressure, many dead buffers accumulate and pressure cuNumeric's allocator. +`@accelerate` optimizes *straight-line* array code: a fixed sequence of statements +with no branches, loops, jumps, `try`, or nested functions. Ordinary calls are +opaque boundaries; general control flow is not rewritten. -`@analyze_lifetimes` performs a **static last-use analysis** at macro-expansion time and inserts eager `maybe_insert_delete` calls immediately after each temporary's final use. Freed buffers can then be reused by later same-sized allocations instead of waiting on GC. +Within that restricted body, it performs: -When broadcast fusion is on, intermediate dotted nodes in a broadcast tree are not real `NDArray` allocations. The macro accounts for that automatically. +```@raw html +
    +
  1. Fusion within a broadcast expression. On CUDA, an eligible dotted expression such as @. A + B * C can run as one kernel. CPU execution uses the normal unfused path.
    2. +
  2. Fusion across broadcast statements. A single-use broadcast result can be substituted into its consumer, producing fewer GPU kernel launches.
    3. +
  3. Temporary lifetime analysis. After rewriting the code, the macro releases materialized, non-returned NDArrays after their final use on CPU or GPU.
    4. +
+``` + +These jobs must happen together: an intermediate that fuses into its consumer is never allocated, while an intermediate that cannot fuse is materialized and then released after its last use. + +## Use the function form by default + +Annotate a reusable straight-line function: + +```julia +@accelerate function update!(C, A, B) + combined = @. A + B + C .= @. 2.0f0 * combined + return C +end +``` + +On an eligible GPU path, `combined` can be folded into the second broadcast so the chain runs as one kernel. On CPU, or when fusion is ineligible, `combined` is materialized and released after the update. Function arguments belong to the caller and are never released by `@accelerate`; returned values also remain valid. + +## Choose a form + +The forms differ in which values must remain available, which determines how aggressively the macro may fuse or release intermediates. + +| Form | Use it when | Fusion and lifetime behavior | +| :--- | :--- | :--- | +| `@accelerate function ... end` | Defining reusable array code. This is the recommended default. | Arguments and returned values are protected. Non-returned locals may fuse into consumers or be released after their last use. | +| `@accelerate begin ... end` | Named results must remain in the current scope. | `begin` creates no new Julia scope, so every named binding is protected. An eligible same-shape CUDA chain may still use one multi-output kernel, but each named result is materialized. | +| `@accelerate let ... end` | Writing a one-off multi-statement calculation when only its result is needed. | `let` creates a local scope. Only the result escapes; other locals may fuse away or be released after their last use. | +| `@accelerate expr` | Evaluating one expression without named intermediates. | The result is materialized and returned. Eligible operations fuse within the expression, and transient temporaries are released. | + +For example, nested scope lets a private intermediate feed a value that is also +used by the outer block: ```julia -T = Float32 -A = cuNumeric.ones(T, (N, N)) -B = cuNumeric.ones(T, (N, N)) -C = cuNumeric.zeros(T, (N, N)) - -@analyze_lifetimes begin - result = @. A[1:end, :] + B[1:end, :] - C .= @. result * 2.0f0 +@accelerate begin + shifted = let + product = @. A * B + @. product + 1 + end + x = @. shifted * C end + +consume(shifted, x) ``` -**Benchmark** (Gray-Scott reaction-diffusion, 512×512, 10 000 steps): +Choose `let` when only the final result should escape: +```julia +result = @accelerate let + product = @. A * B + @. product + 1 +end ``` - user system elapsed CPU max RSS -without 106.50 s 23.87 s 58.66 s 222% 3786 MB -with 61.74 s 13.66 s 27.84 s 270% 2999 MB + +For a single unnamed expression, use: + +```julia +result = @accelerate (@. A + B * C) ``` -~2× wall-clock speedup and ~800 MB lower peak memory with no algorithmic changes. +## Writing an accelerated body + +- Apply `@.` to each elementwise right-hand side. Applying it to the entire body would change `x = ...` into `x .= ...` and `f(...)` into `f.(...)`. +- Use ordinary `=` for a disposable intermediate. This allows a single-use producer to fuse into its consumer. +- Use `.=` when the mutation must be visible. The destination write is preserved, although an eligible producer may fuse into it. +- Keep control flow outside the accelerated body. Loops, conditionals, `try`, short-circuit operators, and nested functions are rejected. +- Ordinary function calls run in program order and form rewrite boundaries. Annotate the called function separately if its body should also be accelerated. + +```julia +for _ in 1:nsteps + update!(C, A, B) +end +``` -Use `@show_lifetimes` to print the rewrite without running it ([Debugging](../debugging.md)). For how the rewriter and GC heuristics work, see [Internals](../internals.md). +See [Kernel Fusion](./kernel_fusion.md) for CUDA fusion requirements and [`@show_lifetimes`](../debugging.md#inspect-lifetime-rewrites-with-show_lifetimes) to inspect the exact rewrite without executing it. diff --git a/examples/custom_cuda.jl b/examples/custom_cuda.jl index 6092642bb..c968b29ae 100644 --- a/examples/custom_cuda.jl +++ b/examples/custom_cuda.jl @@ -3,6 +3,8 @@ using cuNumeric using CUDA import CUDA: i32 +cuNumeric.Experimental(true) + function kernel_add(a, b, c, N) i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x if i <= N @@ -19,25 +21,26 @@ function kernel_sin(a, b, N) return nothing end -N = 1024 -threads = 256 -blocks = cld(N, threads) +function run_custom_cuda(N=1024) + threads = 256 + blocks = cld(N, threads) + a = cuNumeric.fill(1.0f0, N) + b = cuNumeric.fill(2.0f0, N) + c = cuNumeric.zeros(Float32, N) + n_scalar = UInt32(N) -a = cuNumeric.fill(1.0f0, N) -b = cuNumeric.fill(2.0f0, N) -c = cuNumeric.ones(Float32, N) + add_task = cuNumeric.@cuda_task kernel_add(a, b, c, n_scalar) + cuNumeric.@launch task=add_task threads=threads blocks=blocks inputs=(a, b) outputs=c scalars=n_scalar -# task = cuNumeric.@cuda_task kernel_add(a, b, c, UInt32(1)) -# cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, b) outputs=c scalars=UInt32(N) -# allowscalar() do -# c_cpu = c[:] -# println("Result of c after kenel launch: ", c_cpu[1]) -# end + sin_task = cuNumeric.@cuda_task kernel_sin(c, b, n_scalar) + cuNumeric.@launch task=sin_task threads=threads blocks=blocks inputs=c outputs=b scalars=n_scalar -task = cuNumeric.@cuda_task kernel_sin(a, b, UInt32(1)) -cuNumeric.@launch task=task threads=threads blocks=blocks inputs=a outputs=b scalars=UInt32(N) + allowscalar() do + return println("sin(1 + 2) = ", b[1]) + end + return b +end -allowscalar() do - b_cpu = b[:] - println("Result of b after kenel launch: ", b_cpu[1]) +if abspath(PROGRAM_FILE) == @__FILE__ + run_custom_cuda() end diff --git a/examples/daxpy.jl b/examples/daxpy.jl index f7983fff0..7ecb82886 100644 --- a/examples/daxpy.jl +++ b/examples/daxpy.jl @@ -6,6 +6,6 @@ arr = cuNumeric.rand(20) α = 1.32f0 b = 2.0f0 -arr2 = @. α * arr + b +arr2 = @accelerate @. α * arr + b println(arr2) diff --git a/examples/gray-scott.jl b/examples/gray-scott.jl index a52617a35..9d2a757e3 100644 --- a/examples/gray-scott.jl +++ b/examples/gray-scott.jl @@ -28,7 +28,7 @@ function bc!(u_new, v_new, u, v) return v_new[end, :] = v[2, :] end -function step!(u, v, u_new, v_new, args::Params) +@accelerate function step!(u, v, u_new, v_new, args::Params) # calculate F_u and F_v functions F_u = ( (-u[2:(end - 1), 2:(end - 1)] .* (v[2:(end - 1), 2:(end - 1)] .^ 2)) .+ diff --git a/examples/gray-scott.py b/examples/gray-scott.py index dce4e6cab..5eab0b8c8 100644 --- a/examples/gray-scott.py +++ b/examples/gray-scott.py @@ -1,82 +1,54 @@ -# python equivalent of gray-scott.jl to test the GC problem +"""cuPyNumeric equivalent of examples/gray-scott.jl.""" import cupynumeric as np -# import matplotlib.animation as animation -# from IPython.display import HTML -# import matplotlib.pyplot as plt - -def greyScottSys(u, v, dx, dt, c_u, c_v, f, k): - # u,v are arrays - # dx,dt are space and time steps - # c_u, c_v, f, k are constant paramaters - - #create new u array +def step(u, v, dx, dt, c_u, c_v, feed, kill): u_new = np.zeros_like(u) v_new = np.zeros_like(v) - #calculate F_u and F_v functions - F_u = (-u[1:-1,1:-1]*(v[1:-1,1:-1]**2)) + f*(1-u[1:-1,1:-1]) - F_v = (u[1:-1,1:-1]*(v[1:-1,1:-1]**2)) - (f+k)*v[1:-1,1:-1] - - # 2-D Laplacian of f using array slicing, excluding boundaries - # For an N x N array f, f_lap is the N-1 x N-1 array in the "middle" - u_lap = (u[2:,1:-1] - 2*u[1:-1,1:-1] + u[:-2,1:-1]) / dx**2\ - + (u[1:-1,2:] - 2*u[1:-1,1:-1] + u[1:-1,:-2]) / dx**2 - v_lap = (v[2:,1:-1] - 2*v[1:-1,1:-1] + v[:-2,1:-1]) / dx**2\ - + (v[1:-1,2:] - 2*v[1:-1,1:-1] + v[1:-1,:-2]) / dx**2 - - # Forward-Euler time step for all points except the boundaries - u_new[1:-1,1:-1] = ((c_u * u_lap) + F_u)*dt + u[1:-1,1:-1] - v_new[1:-1,1:-1] = ((c_v * v_lap) + F_v)*dt + v[1:-1,1:-1] - - # Apply periodic boundary conditions - u_new[:,0] = u[:,-2] - u_new[:,-1] = u[:,1] - u_new[0,:] = u[-2,:] - u_new[-1,:] = u[1,:] - v_new[:,0] = v[:,-2] - v_new[:,-1] = v[:,1] - v_new[0,:] = v[-2,:] - v_new[-1,:] = v[1,:] - + u_mid = u[1:-1, 1:-1] + v_mid = v[1:-1, 1:-1] + reaction = u_mid * v_mid**2 + f_u = -reaction + feed * (1 - u_mid) + f_v = reaction - (feed + kill) * v_mid + + u_lap = ( + u[2:, 1:-1] - 2 * u_mid + u[:-2, 1:-1] + + u[1:-1, 2:] - 2 * u_mid + u[1:-1, :-2] + ) / dx**2 + v_lap = ( + v[2:, 1:-1] - 2 * v_mid + v[:-2, 1:-1] + + v[1:-1, 2:] - 2 * v_mid + v[1:-1, :-2] + ) / dx**2 + + u_new[1:-1, 1:-1] = (c_u * u_lap + f_u) * dt + u_mid + v_new[1:-1, 1:-1] = (c_v * v_lap + f_v) * dt + v_mid + + u_new[:, 0] = u[:, -2] + u_new[:, -1] = u[:, 1] + u_new[0, :] = u[-2, :] + u_new[-1, :] = u[1, :] + v_new[:, 0] = v[:, -2] + v_new[:, -1] = v[:, 1] + v_new[0, :] = v[-2, :] + v_new[-1, :] = v[1, :] return u_new, v_new +def gray_scott(n=4000, n_steps=100): + dx = 1.0 + dt = dx / 5 + u = np.ones((n, n)) + v = np.zeros((n, n)) + seed = min(150, n) + u[:seed, :seed] = np.random.rand(seed, seed) + v[:seed, :seed] = np.random.rand(seed, seed) -# initial conditions and discretizaiton -dx = 1 -dt = dx/5 -u = np.ones((4000,4000)) -v = np.zeros((4000,4000)) -u[:150,:150] = np.random.rand(150,150) -v[:150,:150] = np.random.rand(150,150) - - -# fig = plt.figure() - -c_u = 1 -c_v = 0.3 -f = 0.03 -k = 0.06 - -# t_final = 1000 - -# ims = [] -n_steps = 100 # number of steps to take -frame_interval = 200 # steps to take between making plots - -# build a list of images -for n in range(n_steps) : - - ## This may need to be changed. - u,v = greyScottSys(u, v, dx, dt, c_u, c_v, f, k) + for _ in range(n_steps): + u, v = step(u, v, dx, dt, 1.0, 0.3, 0.03, 0.06) + return u, v - # ## Store frames when n is a multiple of frame_interval - # if n%frame_interval == 0: - # im = plt.imshow(u, vmin=0, vmax=1) # Show a plot of u. - # ims.append([im]) # append single image to the list of images -# anim = animation.ArtistAnimation(fig, ims, interval=100, repeat=False) -# HTML(anim.to_jshtml()) +if __name__ == "__main__": + gray_scott() diff --git a/examples/integrate.jl b/examples/integrate.jl index f27dae031..8d4e9653d 100644 --- a/examples/integrate.jl +++ b/examples/integrate.jl @@ -1,9 +1,4 @@ -using cuNumeric - -# Note that we do not yet support broadcasting -# custom functions over NDArray, so the broadcasting MUST -# be done inside the function -integrand = (x) -> @. exp(-x^2) +integrand = (x) -> exp(-x^2) N = 1_000_000 @@ -11,12 +6,13 @@ x_max = 10.0f0 domain = [-x_max, x_max] Ω = domain[2] - domain[1] -samples = Ω * cuNumeric.rand(N) -samples = @. samples - x_max +estimate = @accelerate begin + samples = @. Ω * cuNumeric.rand(N) - x_max -# Reductions return 0D NDArrays instead -# of a scalar to avoid blocking runtime -estimate = (Ω / N) * sum(integrand(samples)) + # Reductions return 0D NDArrays instead + # of a scalar to avoid blocking runtime + return (Ω / N) * sum(integrand.(samples)) +end println("Monte-Carlo Estimate: $(estimate)") println("Analytical: $(sqrt(pi))") diff --git a/examples/stencil.jl b/examples/stencil.jl index b8d8e3586..ee8c3e6a9 100644 --- a/examples/stencil.jl +++ b/examples/stencil.jl @@ -1,4 +1,3 @@ - using cuNumeric: cuNumeric function initialize(N) @@ -13,7 +12,6 @@ end function run_stencil(N, I, warmup) grid = initialize(N) - println("Running Jacobi stencil...") center = grid[2:(N + 1), 2:(N + 1)] @@ -22,11 +20,14 @@ function run_stencil(N, I, warmup) west = grid[2:(N + 1), 1:N] south = grid[3:(N + 2), 2:(N + 1)] - for i in 1:(I + warmup) + for _ in 1:(I + warmup) average = center .+ north .+ east .+ west .+ south work = 0.2 .* average - center = work + center .= work end + return grid end -run_stencil(1000, 100, 5) +if abspath(PROGRAM_FILE) == @__FILE__ + run_stencil(1000, 100, 5) +end diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 468904acc..bd080578f 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -59,6 +59,10 @@ if !HAS_CUDA @warn "We couldn't find a CUDA-enabled GPU. If you have an NVIDIA GPU something might be wrong." end +# `HAS_CUDA` describes the machine. A CPU-only Legate configuration on a GPU +# machine must still avoid registering or launching GPU tasks. +@inline _has_gpu_target() = HAS_CUDA && Int(Legate.num_gpus()) > 0 + const DEFAULT_FLOAT = Float32 const DEFAULT_INT = Int32 @@ -263,8 +267,6 @@ function __init__() _is_precompiling() && return nothing - _register_scoping_error_hint!() - # Cannot set LEGATE_CONFIG on CI machines used # to register packages. So we will just skip starting # legate/cunumeric when using registry CI machines. diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl index 87962b7fa..44213a47a 100644 --- a/src/cuda/cuda_ptx_task.jl +++ b/src/cuda/cuda_ptx_task.jl @@ -2,7 +2,11 @@ export @cuda_task, @launch, CUDATask struct CUDATask func::String - argtypes::NTuple{N,Type} where {N} #! THIS IS TYPE UNSTABLE + argtypes::Vector{DataType} + + function CUDATask(func, argtypes) + return new(convert(String, func), collect(DataType, argtypes)) + end end #! JUST PASS TYPES HERE INSTEAD OF CALLING typeof() @@ -16,57 +20,59 @@ function to_stdvec(::Type{T}, vec) where {T} return stdvec end -function add_padding(arr::NDArray, dims::Dims{N}; copy=false) where {N} - old_size = size(arr) - - @assert all(dims .>= old_size) "newdims must be ≥ current dims elementwise" - new = zeros(eltype(arr), dims) +@inline _launch_shape(arr::NDArray) = _launch_shape(arr, _padding(arr)) +@inline _launch_shape(arr::NDArray, ::Nothing) = size(arr) +@inline _launch_shape(::NDArray, padding::PaddedStorage) = padding.shape + +@inline _physical_array(arr::NDArray, ::Nothing) = arr +@inline _physical_array(::NDArray, padding::PaddedStorage) = padding.backing + +function _ensure_launch_padding!(arr::NDArray{T,N}, target_shape; copy=false) where {T,N} + padding = _padding(arr) + !isnothing(padding) && padding.shape == target_shape && return arr + isnothing(padding) && size(arr) == target_shape && return arr + @assert all(target_shape .>= size(arr)) "cannot pad $(size(arr)) to $target_shape" + + padded = zeros(T, target_shape) + slices = ntuple(d -> (0, size(arr, d)), N) + logical = nda_get_slice(padded, slice_array(slices...)) + aliases_parent = !isnothing(arr.parent) + copy && !aliases_parent && copyto!(logical, arr) + storage = PaddedStorage{T,N}( + padded, + aliases_parent ? logical : nothing, + target_shape, + ) - if copy # due to being an input. we don't need to copy outputs - slices = ntuple(d -> (0, Int(old_size[d])), length(old_size)) - s = nda_get_slice(new, slice_array(slices...)) - copyto!(s, arr) - destroy!(s) + if aliases_parent + old_padding = _padding(arr) + arr.padding = storage + !isnothing(old_padding) && _destroy_padded_storage!(old_padding) + else + destroy!(arr) + arr.ptr = logical.ptr + arr.nbytes = logical.nbytes + arr.padding = storage + logical.ptr = Ptr{Cvoid}(0) + logical.nbytes = 0 end - - nda_destroy_array(arr.ptr) - register_free!(arr.nbytes) - - # update pointer & update metadata - arr.ptr = new.ptr - arr.nbytes = new.nbytes - arr.padding = old_size # remember the prior (before the padding) - - # julia GC will call finalizer, but we manually cleaned it - new.ptr = Ptr{Cvoid}(0) - new.nbytes = 0 - return new.padding = nothing -end - -function add_padding(arr::NDArray, i::Int64; copy=false) - return add_padding(arr, (i,); copy=copy) + return arr end -function check_sz!(arr, maxshape; copy=false) - sz = cuNumeric.size(arr) - if maxshape != nothing - # currently require all ndarray inputs to be equal - alligned_equal_size = sz == maxshape - if !alligned_equal_size - cuNumeric.add_padding(arr, maxshape; copy=copy) - new_size = padded_shape(arr) - @warn "[Padding Added] $sz output is now $new_size" - end +function _sync_to_launch_padding!(arr::NDArray) + padding = _padding(arr) + if !isnothing(padding) && !isnothing(padding.staging) + copyto!(padding.staging, arr) end + return nothing end -function check_sz(arr, maxshape) - sz = cuNumeric.size(arr) - if maxshape != nothing - # currently require all ndarray inputs to be equal - alligned_equal_size = sz == maxshape - @assert alligned_equal_size +function _sync_from_launch_padding!(arr::NDArray) + padding = _padding(arr) + if !isnothing(padding) && !isnothing(padding.staging) + copyto!(arr, padding.staging) end + return nothing end # `get_store` returns a Julia-owned `LogicalArrayImplAllocated` that shares the @@ -74,42 +80,32 @@ end # array into the task; if we leave the temporary alive until GC, store refcounts # stay elevated and framebuffer reclaim stalls (fusion 1-GPU OOM under load). # Finalize the temporary immediately after the copy into the task. -function _add_task_array!(add_to, task, arr::NDArray) - st = cuNumeric.get_store(arr) - var = add_to(task, st) - finalize(st) - return var +function _add_task_array!(add_to, task, arr::NDArray; physical=false) + task_arr = physical ? _physical_array(arr, _padding(arr)) : arr + st = cuNumeric.get_store(task_arr) + try + return add_to(task, st) + finally + finalize(st) + end end function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}}; - blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing, validate_shapes=true) - max_shape = if validate_shapes - # Generic PTX tasks retain the existing padding/shape behavior. - ndarrays = vcat(inputs..., outputs...) # returns (nbytes, position) - mx = findmax(arr -> arr.nbytes, ndarrays) # first elem nbytes - shape = size(ndarrays[mx[2]]) # second elem max position - @assert !isnothing(shape) - shape - else - # Fused linear broadcast verifies shapes match - nothing - end - + blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing) rt = Legate.get_runtime() lib = cuNumeric.get_lib() task = Legate.create_auto_task(rt, lib, taskid) + physical = taskid == cuNumeric.RUN_PTX input_vars = Vector{Legate.Variable}() for arr in inputs - validate_shapes && check_sz!(arr, max_shape; copy=true) - push!(input_vars, _add_task_array!(Legate.add_input, task, arr)) + push!(input_vars, _add_task_array!(Legate.add_input, task, arr; physical)) end output_vars = Vector{Legate.Variable}() for arr in outputs - validate_shapes && check_sz!(arr, max_shape; copy=false) - push!(output_vars, _add_task_array!(Legate.add_output, task, arr)) + push!(output_vars, _add_task_array!(Legate.add_output, task, arr; physical)) end # Reserved scalars: kernel_name (0), blocks (1,2,3), threads (4,5,6) @@ -136,15 +132,34 @@ function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}}, end function launch(kernel::CUDATask, inputs, outputs, scalars; - blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing, validate_shapes=true) - return Launch(kernel, - isa(inputs, Tuple) ? inputs : (inputs,), - isa(outputs, Tuple) ? outputs : (outputs,), + blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing) + input_tuple = isa(inputs, Tuple) ? inputs : (inputs,) + output_tuple = isa(outputs, Tuple) ? outputs : (outputs,) + + # Custom tasks require equal physical shapes. Keep the padded backing so + # repeated launches do not allocate or copy again. + if taskid == cuNumeric.RUN_PTX + arrays = (input_tuple..., output_tuple...) + if !isempty(arrays) + rank = ndims(first(arrays)) + @assert all(ndims(arr) == rank for arr in arrays) "custom task arrays must have equal ranks" + max_shape = ntuple(d -> maximum(_launch_shape(arr)[d] for arr in arrays), rank) + foreach(arr -> _ensure_launch_padding!(arr, max_shape; copy=true), input_tuple) + foreach(arr -> _ensure_launch_padding!(arr, max_shape), output_tuple) + foreach(_sync_to_launch_padding!, input_tuple) + end + end + + result = Launch(kernel, + input_tuple, + output_tuple, isa(scalars, Tuple) ? scalars : (scalars,); blocks=isa(blocks, Tuple) ? blocks : (blocks,), threads=isa(threads, Tuple) ? threads : (threads,), - taskid=taskid, ctx=ctx, validate_shapes=validate_shapes, + taskid=taskid, ctx=ctx, ) + taskid == cuNumeric.RUN_PTX && foreach(_sync_from_launch_padding!, output_tuple) + return result end function ptx_task(ptx::String, kernel_name) diff --git a/src/ndarray/broadcast.jl b/src/ndarray/broadcast.jl index 8df0b0ee2..6ac971e4c 100644 --- a/src/ndarray/broadcast.jl +++ b/src/ndarray/broadcast.jl @@ -123,11 +123,9 @@ __materialize(x::Base.RefValue{Val{V}}) where {V} = NDArray(V) # Use binary_op P # Catch unknown things... __materialize(x) = error("Unrecognized leaf in broadcast expression: $(x)") -# Scalar-only nested broadcasts (e.g. `s1 .* s2 .+ A`): the inner -# `Broadcasted(*, (s1, s2))` keeps DefaultArrayStyle{0}, not NDArrayStyle. -# Fold to a Number so the parent unravel sees a scalar leaf. +# Use Base for scalar-only broadcasts, including `literal_pow` wrappers. @inline function __materialize(bc::Broadcasted{<:DefaultArrayStyle{0}}) - return bc.f((__materialize.(bc.args))...) + return Base.materialize(bc) end function __materialize(bc::Broadcasted{<:NDArrayStyle}) @@ -135,6 +133,22 @@ function __materialize(bc::Broadcasted{<:NDArrayStyle}) return unravel_broadcast_tree(bc) end +# The C API is binary, so evaluate flattened `+` and `*` chains pairwise. +function _unravel_flattened_associative(f, args::Tuple) + acc = first(args) + owns_acc = false + for arg in Base.tail(args) + next = try + __materialize(Base.broadcasted(f, acc, arg)) + finally + owns_acc && acc isa NDArray && destroy!(acc) + end + acc = next + owns_acc = acc isa NDArray + end + return acc +end + # Destroy promote copies and non-leaf materialized NDArrays (nested results / Val{V}). @inline function _destroy_unfused_arg_temps!(orig, materialized, promoted) if promoted isa NDArray && promoted !== materialized @@ -148,6 +162,9 @@ end # Un-fused implementation of broadcast tree function unravel_broadcast_tree(bc::Broadcasted) + if length(bc.args) > 2 && _is_flattened_associative(bc.f) + return _unravel_flattened_associative(bc.f, bc.args) + end # Recursively materialize/unravel any nested broadcasts # until we reach a Broadcasted expression with only @@ -234,13 +251,11 @@ end ) end - # Fused writes `dest` in place (no post-fuse `nda_move`); promotion is - # checked pre-launch in `fuse_broadcast_tree!`. CPU vs GPU is compile-time - # via `@static if FUSE_BROADCAST_EXPRS && HAS_CUDA`. + # Require an active GPU target so `--gpus 0` stays on the unfused path. # Fusion requires same-shaped NDArray leaves; otherwise fall back. # Single-op exprs (length < `FUSE_BROADCAST_MIN_OPS`) stay unfused by default. @static if FUSE_BROADCAST_EXPRS && HAS_CUDA - if _should_attempt_broadcast_fusion(dest, bc) + if _has_gpu_target() && _should_attempt_broadcast_fusion(dest, bc) return fuse_broadcast_tree!(dest, bc) else return _copyto_unfused!(dest, unravel_broadcast_tree(bc)) diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl index 728deec17..64149a9f7 100644 --- a/src/ndarray/broadcast_fusion.jl +++ b/src/ndarray/broadcast_fusion.jl @@ -619,7 +619,11 @@ end bc::Base.Broadcast.Broadcasted{S,Ax,F,Args} ) where {S,Ax,F,Args} eltypes = _fused_checked_eltypes(bc.args) - T_OUT = __checked_promote_op(bc.f, eltypes) + T_OUT = if length(bc.args) > 2 && _is_flattened_associative(bc.f) + _checked_promote_associative(bc.f, eltypes.parameters...) + else + __checked_promote_op(bc.f, eltypes) + end __my_promote_type(eltypes.parameters...) return T_OUT end @@ -748,10 +752,329 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas threads=fkm.threads, taskid=cuNumeric.RUN_PTX_BROADCAST, ctx=fkm.ctx, - validate_shapes=false, ) end # Fused kernel already wrote `dest` in place; promotion was checked pre-launch. return dest end + +# ============================================================================ +# Multi-output fused broadcast: materialize named intermediates in one launch. +# Segments (dependency order, root last) are flattened independently; a +# `MatRef{K}` leaf reads the K-th segment's per-element local. One kernel +# computes each segment into a local, stores it to that segment's output buffer, +# and chains locals into parents (segmented flatten + chained-local multi-store). +# ============================================================================ + +# Opaque scalar-like leaf that survives `Base.Broadcast.flatten` (never descended +# into, never wrapped in a Ref). +struct MatRef{K} end +MatRef(k::Int) = MatRef{k}() +Base.broadcastable(m::MatRef) = m +Base.Broadcast.BroadcastStyle(::Type{<:MatRef}) = Base.Broadcast.DefaultArrayStyle{0}() +Base.axes(::MatRef) = () +Base.ndims(::Type{<:MatRef}) = 0 + +# Third arg-plan variant (alongside Runtime/Static): read the K-th chained local. +struct LocalBroadcastArg{K} end + +Base.@propagate_inbounds @inline _materialize_ml_arg( + ::RuntimeBroadcastArg{J}, rt, sa, locals, I +) where {J} = _gpu_broadcast_getindex(getfield(rt, J), I) +Base.@propagate_inbounds @inline _materialize_ml_arg( + ::StaticBroadcastArg{J}, rt, sa, locals, I +) where {J} = getfield(sa, J) +Base.@propagate_inbounds @inline _materialize_ml_arg( + ::LocalBroadcastArg{K}, rt, sa, locals, I +) where {K} = getfield(locals, K) + +Base.@propagate_inbounds @inline _materialize_ml_args(::Tuple{}, rt, sa, locals, I) = () +Base.@propagate_inbounds @inline function _materialize_ml_args(plan::Tuple, rt, sa, locals, I) + return ( + @inbounds(_materialize_ml_arg(getfield(plan, 1), rt, sa, locals, I)), + @inbounds(_materialize_ml_args(Base.tail(plan), rt, sa, locals, I))..., + ) +end + +# Device-side: run each segment in order, store to its output, chain the local. +# Generate straight-line code because recursive tuple traversal eventually hits +# Julia's inference limit and leaves a dynamic call in GPU kernels on Julia 1.10. +Base.@propagate_inbounds @inline @generated function _run_segments( + segs::S, outs, rt, sa, locals::L, I +) where {S<:Tuple,L<:Tuple} + body = Expr(:block) + local_values = Any[:(getfield(locals, $k)) for k in 1:fieldcount(L)] + + for k in 1:fieldcount(S) + seg = gensym(:seg) + vals = gensym(:vals) + value = gensym(:value) + local_tuple = Expr(:tuple, local_values...) + push!( + body.args, + quote + $seg = getfield(segs, $k) + $vals = _materialize_ml_args( + getfield($seg, 2), rt, sa, $local_tuple, I + ) + $value = Base.Broadcast._broadcast_getindex_evalf( + getfield($seg, 1), $vals... + ) + @inbounds getfield(outs, $k)[I] = $value + end, + ) + push!(local_values, value) + end + + push!(body.args, :(nothing)) + return body +end + +# Dimension-dispatched (mirrors the single-output linear/cartesian kernels). +# `args` = (outputs[1:NOUT]..., runtime_args...); bounds from the first output. +function make_multi_output_kernel(segs, ::Val{NOUT}, static_args, ::Val{2}) where {NOUT} + @kernel unsafe_indices = true function broadcast_kernel_multi_2d(args...) + I = _broadcast_cartesian_work_id() + dest = getfield(args, 1) + @inbounds if I[1] <= size(dest, 1) && I[2] <= size(dest, 2) + _run_segments(segs, args[1:NOUT], args[(NOUT + 1):end], static_args, (), I) + end + end + return broadcast_kernel_multi_2d +end + +function make_multi_output_kernel(segs, ::Val{NOUT}, static_args, ::Val{3}) where {NOUT} + @kernel unsafe_indices = true function broadcast_kernel_multi_3d(args...) + I = _broadcast_cartesian_work_id_3d() + dest = getfield(args, 1) + @inbounds if I[1] <= size(dest, 1) && I[2] <= size(dest, 2) && I[3] <= size(dest, 3) + _run_segments(segs, args[1:NOUT], args[(NOUT + 1):end], static_args, (), I) + end + end + return broadcast_kernel_multi_3d +end + +# 1-D and any other rank: linear indexing (matches the single-output default). +function make_multi_output_kernel(segs, ::Val{NOUT}, static_args, ::Val) where {NOUT} + @kernel unsafe_indices = true function broadcast_kernel_multi_linear(args...) + I = _broadcast_linear_work_id() + @inbounds if I <= length(getfield(args, 1)) + _run_segments(segs, args[1:NOUT], args[(NOUT + 1):end], static_args, (), I) + end + end + return broadcast_kernel_multi_linear +end + +# Flatten a segment and classify its leaves, deduping NDArrays into shared +# `runtime_args` and static leaves into shared `static_args`. +function _split_segment!(seg_bc, runtime_args, static_args, ndarray_idx) + flat = Base.Broadcast.flatten(seg_bc) + plan = Any[] + for leaf in flat.args + if leaf isa MatRef + push!(plan, LocalBroadcastArg{_matref_k(leaf)}()) + elseif leaf isa Base.RefValue + v = leaf[] + if v isa Number + push!(runtime_args, v) + push!(plan, RuntimeBroadcastArg{length(runtime_args)}()) + else + _push_static_arg!(static_args, plan, v) + end + elseif leaf isa NDArray || leaf isa Base.Broadcast.Extruded + nda = get_ndarray(leaf) + j = get!(() -> (push!(runtime_args, leaf); length(runtime_args)), + ndarray_idx, objectid(nda)) + push!(plan, RuntimeBroadcastArg{j}()) + elseif leaf isa Number + push!(runtime_args, leaf) + push!(plan, RuntimeBroadcastArg{length(runtime_args)}()) + elseif isbits(leaf) + _push_static_arg!(static_args, plan, leaf) + else + throw(ArgumentError("multi-output fusion: cannot lower leaf $(typeof(leaf))")) + end + end + return (flat.f, tuple(plan...)) +end + +_matref_k(::MatRef{K}) where {K} = K + +const _MULTI_PTX_CACHE = Dict{Any,Any}() +const _MULTI_PTX_CACHE_LOCK = ReentrantLock() + +# Compile + register the multi-output kernel -> (ctx, threads, CUDATask). Cached +# by (closure type, arg types) so a repeated fusion signature compiles once. +function get_multi_cuda_task(obj, out_arrs, runtime_args) + arg_types = (map_cuda_type.(typeof.(out_arrs))..., map_cuda_type.(typeof.(runtime_args))...) + key = (typeof(obj), arg_types) + lock(_MULTI_PTX_CACHE_LOCK) do + return get!(_MULTI_PTX_CACHE, key) do + ptx, threads, ctx = get_ptx(obj, arg_types...) + threads == 0 && return (ctx, 0, nothing) + orig = extract_kernel_name(ptx) + uname = orig * "_" * string(hash(ptx); base=16) + ptx = replace(ptx, orig => uname) + ptx_task(ptx, uname) + return (ctx, threads, CUDATask(uname, arg_types)) + end + end +end + +# First NDArray leaf across all segments; used as an allocation template. +function _first_ndarray(seg_bcs::Tuple) + for seg_bc in seg_bcs + for leaf in Base.Broadcast.flatten(seg_bc).args + leaf isa NDArray && return leaf + leaf isa Base.Broadcast.Extruded && return get_ndarray(leaf) + end + end + return throw(ArgumentError("multi-output fusion: no NDArray leaf to size buffers from")) +end + +# Result eltype of a segment, resolving `MatRef{k}` to `eltype(bufs[k])` (earlier +# segments already allocated). Lets each intermediate use its own eltype. +function _segment_eltype(flat, bufs) + ets = map(flat.args) do leaf + if leaf isa MatRef + eltype(bufs[_matref_k(leaf)]) + elseif leaf isa NDArray + eltype(leaf) + elseif leaf isa Base.Broadcast.Extruded + eltype(leaf.x) + else + typeof(leaf) + end + end + T = Base.promote_op(flat.f, ets...) + # Fall back to promoting the leaf eltypes when promote_op can't infer. + return isconcretetype(T) ? T : promote_type(ets...) +end + +# Render a segment's broadcast tree, showing `MatRef{k}` leaves as `seg{k}`. +function _bcast_multi_tree_str(bc) + return _bcast_tree_str(bc) do x + x isa MatRef && return "seg{$(_matref_k(x))}" + x isa NDArray && return "NDArray" + x isa Base.Broadcast.Extruded && return "NDArray" + x isa Number && return repr(x) + x isa Base.RefValue && return string("^", repr(x[])) + return string("<", typeof(x), ">") + end +end + +# Fused multi-output introspection (mirrors `_describe_fused_broadcast`). Enable +# with `cuNumeric.BCAST_FUSION_DEBUG[] = true`. +function _describe_fused_multi( + out_arrs, seg_bcs, input_ndarrays, actual_scalars, argmap, threads, ndrange +) + io = IOBuffer() + field(k, v) = println(io, " ", rpad(k, 8), v) + NOUT = length(out_arrs) + println(io, "\n", "="^40, " fused multi-output broadcast ($NOUT outputs)") + println(io, " segments (each materialized to its own output):") + for (i, seg) in enumerate(seg_bcs) + role = i == NOUT ? "root" : "seg{$i}" + println( + io, " ", rpad(role, 7), _ndarray_debug_summary(out_arrs[i]), + " <- ", _bcast_multi_tree_str(seg), + ) + end + field("inputs", "input{N} ($(length(input_ndarrays)) unique)") + for (i, nd) in enumerate(input_ndarrays) + println(io, " ", rpad(string(i - 1), 4), _ndarray_debug_summary(nd)) + end + isempty(actual_scalars) || field("scalars", join(repr.(actual_scalars), ", ")) + indexing = ndims(out_arrs[1]) in (2, 3) ? "cartesian" : "linear" + field( + "launch", + "host thread budget=$threads, indexing=$indexing, num_outputs=$NOUT, " * + "blocks=device(local tile), global_ndrange=$ndrange", + ) + field("arg_map", string(argmap)) + print(String(take!(io))) + return nothing +end + +# Launch one kernel writing each segment into preallocated `out_arrs[i]` +# (dependency order; `out_arrs[end]` is the root). +function _fused_multi_launch!(out_arrs::Tuple, seg_bcs::Tuple) + NOUT = length(seg_bcs) + runtime_args = Any[] + static_args = Any[] + ndarray_idx = Dict{UInt,Int}() + segs = Any[] + for seg_bc in seg_bcs + push!(segs, _split_segment!(seg_bc, runtime_args, static_args, ndarray_idx)) + end + segs = tuple(segs...) + static_args = tuple(static_args...) + + kernel = make_multi_output_kernel(segs, Val(NOUT), static_args, Val(ndims(out_arrs[1]))) + bck = kernel(CUDACore.CUDAKernels.CUDABackend()) + ctx, threads, task = get_multi_cuda_task(bck, out_arrs, tuple(runtime_args...)) + isnothing(task) && return out_arrs + + # arg_map: kernel args in order (outputs..., runtime_args...). + argmap = Int32[Int32(i) for i in 0:(NOUT - 1)] + input_ndarrays = NDArray[] + ndinput_idx = Dict{UInt,Int}() + actual_scalars = Any[] + for arg in runtime_args + if stores_cudevicearray(map_cuda_type(typeof(arg))) + nda = get_ndarray(arg) + j = get!(() -> (push!(input_ndarrays, nda); length(input_ndarrays) - 1), + ndinput_idx, objectid(nda)) + push!(argmap, Int32(NOUT + j)) + else + push!(argmap, Int32(-1 - length(actual_scalars))) + push!(actual_scalars, arg) + end + end + + if BCAST_FUSION_DEBUG[] + ndrange = ndims(out_arrs[1]) > 0 ? size(out_arrs[1]) : (1,) + _describe_fused_multi( + out_arrs, seg_bcs, input_ndarrays, actual_scalars, argmap, threads, ndrange + ) + end + + launch( + task, tuple(input_ndarrays...), out_arrs, + (Int32(length(argmap)), argmap..., actual_scalars...); + blocks=1, threads=threads, taskid=cuNumeric.RUN_PTX_BROADCAST, ctx=ctx, + ) + return out_arrs +end + +# Allocate a typed tuple so callers retain each segment's concrete NDArray type. +function _alloc_segment_buffers(template::NDArray, seg_bcs::Tuple, dims) + return _alloc_segment_buffers(template, seg_bcs, dims, ()) +end + +@inline _alloc_segment_buffers(template, ::Tuple{}, dims, bufs::Tuple) = bufs + +@inline function _alloc_segment_buffers(template, seg_bcs::Tuple, dims, bufs::Tuple) + flat = Base.Broadcast.flatten(first(seg_bcs)) + buf = similar(template, _segment_eltype(flat, bufs), dims) + return _alloc_segment_buffers(template, Base.tail(seg_bcs), dims, (bufs..., buf)) +end + +# `seg_bcs[1:end-1]` are materialized producers (dependency order); `seg_bcs[end]` +# writes `dest`. Producer buffers are allocated (returned so callers bind names). +function copyto_fused_multi!(dest::NDArray, seg_bcs::Tuple) + bufs = _alloc_segment_buffers(dest, seg_bcs[1:(end - 1)], size(dest)) + outs = (bufs..., dest) + _fused_multi_launch!(outs, seg_bcs) + return outs +end + +# Every segment gets a fresh buffer (all named results stay live). Returns the +# buffers in segment order so callers can bind each user name. +function copyto_fused_multi_alloc!(seg_bcs::Tuple) + tmpl = _first_ndarray(seg_bcs) + outs = _alloc_segment_buffers(tmpl, seg_bcs, size(tmpl)) + _fused_multi_launch!(outs, seg_bcs) + return outs +end diff --git a/src/ndarray/detail/linalg.jl b/src/ndarray/detail/linalg.jl index a15cf9e1a..6722495c1 100644 --- a/src/ndarray/detail/linalg.jl +++ b/src/ndarray/detail/linalg.jl @@ -263,15 +263,12 @@ function _svd(a::NDArray{T,2}, full_matrices::Bool) where {T} ) k = min(m, n) S = real(T) - # cuSolver requires full square buffers regardless of full_matrices - u_buf = cuNumeric.zeros(T, m, m) + + u_buf = full_matrices ? cuNumeric.zeros(T, m, m) : cuNumeric.zeros(T, m, k) s = cuNumeric.zeros(S, k) - vh_buf = cuNumeric.zeros(T, n, n) + vh_buf = full_matrices ? cuNumeric.zeros(T, n, n) : cuNumeric.zeros(T, k, n) svd_single(a, u_buf, s, vh_buf) - # Backend factors are logically ordered; only thin strided views need materialization. - u = full_matrices ? u_buf : copy(u_buf[:, 1:k]) - vh = full_matrices ? vh_buf : copy(vh_buf[1:k, :]) - return u, s, vh + return u_buf, s, vh_buf end # svd runs on float/complex only — no integer backend diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl index 8b7e67671..a8b6f8118 100644 --- a/src/ndarray/detail/ndarray.jl +++ b/src/ndarray/detail/ndarray.jl @@ -43,21 +43,24 @@ get_n_dim(ptr::NDArray_t) = Int(ccall((:nda_array_dim, libnda), Int32, (NDArray_ abstract type AbstractNDArray{T<:SUPPORTED_TYPES,N} <: AbstractArray{T,N} end +# Runtime padding uses an abstract field to break the recursive storage definition. +abstract type AbstractPaddedStorage{T,N} end + @doc""" The NDArray type represents a multi-dimensional array in cuNumeric. It is a wrapper around a Legate array and provides various methods for array manipulation and operations. Finalizer calls `nda_destroy_array` to clean up the underlying Legate array when the NDArray is garbage collected. """ -mutable struct NDArray{T,N,PADDED,P} <: AbstractNDArray{T,N} +mutable struct NDArray{T,N,P} <: AbstractNDArray{T,N} ptr::NDArray_t nbytes::Int64 - padding::Union{Nothing,NTuple{N,Int}} + padding::Union{Nothing,AbstractPaddedStorage{T,N}} parent::P function NDArray(ptr::NDArray_t, ::Type{T}, ::Val{N}) where {T,N} nbytes = cuNumeric.nda_nbytes(ptr) cuNumeric.register_alloc!(nbytes) - handle = new{T,N,false,Nothing}(ptr, nbytes, nothing, nothing) + handle = new{T,N,Nothing}(ptr, nbytes, nothing, nothing) finalizer(_finalize_ndarray!, handle) return handle end @@ -66,26 +69,53 @@ mutable struct NDArray{T,N,PADDED,P} <: AbstractNDArray{T,N} function NDArray(ptr::NDArray_t, ::Type{T}, ::Val{N}, parent::P) where {T,N,P} nbytes = cuNumeric.nda_nbytes(ptr) cuNumeric.register_alloc!(nbytes) - handle = new{T,N,false,P}(ptr, nbytes, nothing, parent) + handle = new{T,N,P}(ptr, nbytes, nothing, parent) finalizer(_finalize_ndarray!, handle) return handle end end +struct PaddedStorage{T,N} <: AbstractPaddedStorage{T,N} + backing::NDArray{T,N,Nothing} + staging::Union{Nothing,NDArray{T,N,NDArray{T,N,Nothing}}} + shape::NTuple{N,Int} +end + +# Narrow the abstract field to its concrete storage type. +@inline _padding(arr::NDArray{T,N}) where {T,N} = + arr.padding::Union{Nothing,PaddedStorage{T,N}} + +function _finalize_padded_storage!(storage::PaddedStorage) + !isnothing(storage.staging) && finalize(storage.staging) + finalize(storage.backing) + return nothing +end + +function _destroy_padded_storage!(storage::PaddedStorage) + !isnothing(storage.staging) && destroy!(storage.staging) + destroy!(storage.backing) + return nothing +end + # May run off the launch thread, so defer the Legate free to drain_pending_frees!. # Accounting is atomic and safe to do here immediately. function _finalize_ndarray!(arr::NDArray) ptr = arr.ptr - ptr == C_NULL && return nothing arr.ptr = Ptr{Cvoid}(0) nbytes = arr.nbytes arr.nbytes = 0 - nbytes > 0 && register_free!(nbytes) - _enqueue_free!(ptr) + padding = _padding(arr) + arr.padding = nothing + + if ptr != C_NULL + nbytes > 0 && register_free!(nbytes) + _enqueue_free!(ptr) + end + !isnothing(padding) && _finalize_padded_storage!(padding) return nothing end -@inline _is_ndarray_slice(arr::NDArray) = arr.parent isa NDArray +@inline _is_ndarray_slice(arr::NDArray) = arr.parent isa NDArray || !isnothing(_padding(arr)) """ destroy!(arr::NDArray) @@ -102,6 +132,9 @@ function destroy!(arr::NDArray) arr.nbytes = 0 nbytes > 0 && register_free!(nbytes) end + padding = _padding(arr) + arr.padding = nothing + !isnothing(padding) && _destroy_padded_storage!(padding) return arr end @@ -591,9 +624,7 @@ end Return the size of the given `NDArray`. """ -shape(arr::NDArray{<:Any,N,true}) where {N} = arr.padding - -function shape(arr::NDArray{<:Any,N,false}) where {N} +function shape(arr::NDArray{<:Any,N}) where {N} shp = cuNumeric.nda_array_shape(arr) return ntuple(i -> Int(shp[i]), Val(N)) end diff --git a/src/ndarray/ndarray.jl b/src/ndarray/ndarray.jl index f42955061..437d69d5c 100644 --- a/src/ndarray/ndarray.jl +++ b/src/ndarray/ndarray.jl @@ -399,7 +399,7 @@ function _setindex!( end #### START OF SLICING #### -# LHS slices from `nda_get_slice` are invisible to `@analyze_lifetimes`; destroy +# LHS slices from `nda_get_slice` are invisible to `@accelerate`; destroy # the view handle after submitting the assign so they cannot pile up under Julia # GC (which sees each NDArray as ~pointer-sized). function _setindex_slice!(lhs::NDArray, rhs::NDArray, slices) diff --git a/src/ndarray/promotion.jl b/src/ndarray/promotion.jl index 20e5ab607..cb1034098 100644 --- a/src/ndarray/promotion.jl +++ b/src/ndarray/promotion.jl @@ -30,9 +30,19 @@ unchecked_promote_scalar(x, ::Type) = x unchecked_promote_arr(::Base.RefValue{typeof(^)}, ::Type{T}) where {T} = typeof(Base.:(^)) unchecked_promote_arr(::Base.RefValue{Val{V}}, ::Type{T}) where {T,V} = Val{V} +@inline _is_flattened_associative(f) = f === (+) || f === (*) + __checked_promote_op(op, ::Type{Tuple{A}}) where {A} = __checked_promote_op(op, A) __checked_promote_op(op, ::Type{Tuple{A,B}}) where {A,B} = __checked_promote_op(op, A, B) +# Julia flattens dotted `+` and `*` chains into n-ary Broadcasted nodes. Fold +# their input types pairwise, matching both the binary C API and fused path. +@inline function __checked_promote_op( + op::Union{typeof(+),typeof(*)}, ::Type{Args} +) where {Args<:Tuple{Any,Any,Any,Vararg{Any}}} + return _checked_promote_associative(op, Args.parameters...) +end + # Path for literal powers @inline function __checked_promote_op( f::typeof(Base.literal_pow), a::Type{Tuple{_,ARR_TYPE,Val{POWER}}} @@ -77,6 +87,16 @@ end return T end +@inline _checked_promote_associative(op, ::Type{A}, ::Type{B}) where {A,B} = + __checked_promote_op(op, A, B) + +@inline function _checked_promote_associative( + op, ::Type{A}, ::Type{B}, ::Type{C}, rest::Type... +) where {A,B,C} + T = __checked_promote_op(op, A, B) + return _checked_promote_associative(op, T, C, rest...) +end + # For literal powers which are often Int64, do not check for promotion to double # The result of promote_op with a literal integer power is always the base type # Base.promote_op(^, Float32, Int64) == Float32 diff --git a/src/ndarray/unary.jl b/src/ndarray/unary.jl index ef234d694..d27a78951 100644 --- a/src/ndarray/unary.jl +++ b/src/ndarray/unary.jl @@ -48,6 +48,13 @@ global const unary_op_map_no_args = Dict{Function,UnaryOpCode}( ### SPECIAL CASES ### +# `dest .= src` lowers to `identity.(src)`. Treat identity like the native +# unary operation it is so ordinary Julia broadcast assignment works for +# NDArrays, including writable slices. +@inline function __broadcast(::typeof(identity), out::NDArray, input::NDArray) + return nda_unary_op!(out, cuNumeric.COPY, input) +end + # Needed to support != Base.:(!)(input::NDArray{Bool,0}) = nda_unary_op!(similar(input), cuNumeric.LOGICAL_NOT, input) Base.:(!)(input::NDArray{Bool,1}) = nda_unary_op!(similar(input), cuNumeric.LOGICAL_NOT, input) @@ -92,7 +99,7 @@ end # Fallbacks for Real types @inline function __broadcast(f::typeof(Base.real), out::NDArray, input::NDArray{<:Real}) # real(real_array) is just the array - return nda_unary_op!(out, cuNumeric.IDENTITY, input) + return nda_unary_op!(out, cuNumeric.COPY, input) end @inline function __broadcast(f::typeof(Base.imag), out::NDArray, input::NDArray{<:Real}) # imag(real_array) is all zeros @@ -100,7 +107,7 @@ end end @inline function __broadcast(f::typeof(Base.conj), out::NDArray, input::NDArray{<:Real}) # conj(real_array) is just the array - return nda_unary_op!(out, cuNumeric.IDENTITY, input) + return nda_unary_op!(out, cuNumeric.COPY, input) end function Base.:(-)(input::NDArray{Bool}) diff --git a/src/scoping/accelerate.jl b/src/scoping/accelerate.jl new file mode 100644 index 000000000..3b0268c69 --- /dev/null +++ b/src/scoping/accelerate.jl @@ -0,0 +1,257 @@ +using MacroTools: MacroTools + +# Rejected everywhere: control flow makes last-use freeing unsound (a temp freed +# after its textual last use could be revived on another path). +const _CONTROL_FLOW_HEADS = (:if, :elseif, :for, :while, :try, :do, :break, :continue, :&&, :||) + +_is_function_lhs(::Any) = false +function _is_function_lhs(lhs::Expr) + lhs.head === :call && return true + lhs.head === :where && return _is_function_lhs(first(lhs.args)) + return false +end + +function _reject_nonstraightline(body) + MacroTools.postwalk(body) do node + node isa Expr || return node + if node.head === :function || node.head === :-> || + (node.head === :(=) && _is_function_lhs(first(node.args))) + error("@accelerate: nested/anonymous function definitions are not supported") + end + if node.head in _CONTROL_FLOW_HEADS + error( + "@accelerate: control flow (`$(node.head)`) is not supported; " * + "only straight-line code can be accelerated", + ) + end + return node + end + return nothing +end + +# Function args are caller-owned: protected roots, never freed or fused away. +function _argument_symbols(def) + names = Set{Symbol}() + for arg in Iterators.flatten((get(def, :args, Any[]), get(def, :kwargs, Any[]))) + name, _, _, _ = MacroTools.splitarg(arg) + name isa Symbol && push!(names, name) + end + return names +end + +# Function form: validate, expand `@.`, normalize trailing `return`, run the +# lifetime/fusion passes protecting `protected_roots` (the args). +function _accelerate_rewrite(body, caller::Module, protected_roots::Set{Symbol}) + _reject_nonstraightline(body) + body = _normalize_return(_expand_dot_macros(body, caller)) + on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing + return process_ndarray_scope(body; on_rewrite, protected_roots) +end + +# `begin`/expr form: 1:1 Julia scope (no `let`) — named bindings stay live. +# On GPU, same-shape chains fuse into one multi-output launch; otherwise only +# anonymous temporaries (slices) are freed. +function _accelerate_block_soft(block, caller::Module) + _reject_nonstraightline(block) + nb = _normalize_return(_expand_dot_macros(block, caller)) + on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing + fallback = process_ndarray_scope( + nb; on_rewrite, protected_roots=_assigned_symbols(nb) + ) + @static if FUSE_BROADCAST_EXPRS && HAS_CUDA + fused = _try_fuse_block_multi(nb) + if !isnothing(fused) + return quote + if cuNumeric._has_gpu_target() + $fused + else + $fallback + end + end + end + end + # Protect named bindings; free only anonymous temps. + return fallback +end + +# Flatten a `let` node's bindings + body into one statement block. +function _let_body(letexpr::Expr) + stmts = Any[] + for part in letexpr.args + if part isa Expr && part.head === :block + append!(stmts, part.args) + elseif part isa Expr && part.head === :(=) + push!(stmts, part) + elseif !isnothing(part) + push!(stmts, part) + end + end + return Expr(:block, stmts...) +end + +# `let` form: hard scope. Full analysis — combine single-use producers, free +# every non-returned temp — re-wrapped in a `let` so only the result escapes. +function _accelerate_block_hard(letexpr, caller::Module) + body = _let_body(letexpr) + _reject_nonstraightline(body) + nb = _normalize_return(_expand_dot_macros(body, caller)) + on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing + rewritten = process_ndarray_scope(nb; on_rewrite, protected_roots=Set{Symbol}()) + bindings = union(_assigned_symbols(nb), _assigned_symbols(rewritten)) + return _lexical_scope(rewritten, bindings) +end + +# Drop the leading dot: `.+` -> `+`, `.^` -> `^`. +_undot(op::Symbol) = Symbol(chop(string(op); head=1, tail=0)) + +# Dotted RHS -> lazy `Base.broadcasted(...)`: chain vars become `MatRef{k}`, +# slice leaves are hoisted into `hoisted` (temp => slice) to free post-launch. +# `nothing` when not lowerable (caller falls back). +function _to_broadcasted(expr, idx::AbstractDict{Symbol,Int}, hoisted::Vector) + if expr isa Symbol + haskey(idx, expr) && return :(cuNumeric.MatRef($(idx[expr]))) + return expr + end + expr isa Expr || return expr + if expr.head === :ref + # Slice temp: bail if it indexes a chain var, else hoist to free later. + any(s -> haskey(idx, s), walk_symbols(expr)) && return nothing + tmp = gensym(:slice) + push!(hoisted, tmp => expr) + return tmp + end + if expr.head === :call && expr.args[1] isa Symbol && _is_broadcast_op(expr.args[1]) + cargs = map(a -> _to_broadcasted(a, idx, hoisted), expr.args[2:end]) + any(isnothing, cargs) && return nothing + return Expr(:call, :(Base.broadcasted), _undot(expr.args[1]), cargs...) + end + if expr.head === :. && length(expr.args) == 2 && + expr.args[2] isa Expr && expr.args[2].head === :tuple + cargs = map(a -> _to_broadcasted(a, idx, hoisted), expr.args[2].args) + any(isnothing, cargs) && return nothing + return Expr(:call, :(Base.broadcasted), expr.args[1], cargs...) + end + # Non-dotted scalar leaf; unsafe if it reads a chain var as a scalar. + any(s -> haskey(idx, s), walk_symbols(expr)) && return nothing + return expr +end + +function _is_top_broadcast(rhs) + return ( + rhs isa Expr && rhs.head === :call && rhs.args[1] isa Symbol && + _is_broadcast_op(rhs.args[1]) + ) || + ( + rhs isa Expr && rhs.head === :. && length(rhs.args) == 2 && + rhs.args[2] isa Expr && rhs.args[2].head === :tuple + ) +end + +# SSA chain of `sym = ` (+ optional trailing return) -> one +# multi-output launch materializing each result. `nothing` -> caller falls back. +function _try_fuse_block_multi(block) + stmts = _scope_statements(block) + isnothing(stmts) && return nothing + stmts = filter(s -> !(s isa LineNumberNode), stmts) + isempty(stmts) && return nothing + + assigns = stmts + ret = nothing + if isnothing(_assignment(last(stmts))) + ret = last(stmts) + assigns = stmts[1:(end - 1)] + end + length(assigns) >= 2 || return nothing + + syms = Symbol[] + idx = Dict{Symbol,Int}() + seg_exprs = Any[] + hoisted = Pair{Symbol,Any}[] # slice temp => slice expr + for stmt in assigns + a = _assignment(stmt) + isnothing(a) && return nothing + a.lhs isa Symbol || return nothing # no indexed-assign in this path + a.lhs in syms && return nothing # SSA: no reassignment + _is_top_broadcast(a.rhs) || return nothing # must be a real broadcast + seg = _to_broadcasted(a.rhs, idx, hoisted) + isnothing(seg) && return nothing + push!(seg_exprs, seg) + push!(syms, a.lhs) + idx[a.lhs] = length(syms) + end + outs = gensym(:outs) + slice_binds = [:($t = $e) for (t, e) in hoisted] + slice_frees = [:(cuNumeric.maybe_insert_delete($t)) for (t, _) in hoisted] + binds = [:($(syms[i]) = $outs[$i]) for i in eachindex(syms)] + value = isnothing(ret) ? last(syms) : ret + return quote + $(slice_binds...) # materialize slice views + $outs = cuNumeric.copyto_fused_multi_alloc!(($(seg_exprs...),)) + $(slice_frees...) # free them after the launch + $(binds...) + $value + end +end + +# AST `@accelerate` emits (pre-`esc`); shared with `@show_lifetimes`. Dispatch: +# function def / `let` (hard scope) / `begin`-expr (soft, 1:1 Julia scope). +function _accelerate_expand(input, caller::Module) + if MacroTools.isdef(input) + def = MacroTools.splitdef(input) + def[:body] = _accelerate_rewrite(def[:body], caller, _argument_symbols(def)) + return MacroTools.combinedef(def) + elseif input isa Expr && input.head === :let + return _accelerate_block_hard(input, caller) + end + return _accelerate_block_soft(input, caller) +end + +@doc""" + @accelerate function f(args...) ... end + @accelerate begin ... end + @accelerate let ... end + @accelerate expr + +Optimize straight-line array code by coordinating CUDA broadcast fusion within +expressions, fusion across broadcast statements, and scope-aware cleanup of +materialized temporaries. Control flow and nested/anonymous functions are +rejected. Four forms determine which values must remain valid: + + * **function** (preferred): arguments and returned values are protected; + non-returned locals may fuse into consumers or be freed after their last use. + * **`begin`**: creates no new Julia scope, so every named binding stays live; + eligible GPU chains may use one multi-output kernel that materializes them. + * **`let`**: creates a local scope; only the result escapes, so other locals may + fuse away or be freed after their last use. + * **expression**: materializes and returns one expression; eligible operations + fuse within it and transient temporaries are released. + +```julia +@accelerate function step(u, v) # c may fuse away; the result is returned + c = u .* v + return c .^ 2 +end +a, b = @accelerate begin # a and b both stay live, one GPU launch + a = x .* y + b = a .+ 1 + (a, b) +end +result = @accelerate (x .+ y .* z) +``` +""" +macro accelerate(input) + return esc(_accelerate_expand(input, __module__)) +end + +@doc""" + @show_lifetimes function f(args...) ... end + @show_lifetimes begin ... end + @show_lifetimes let ... end + +Print the exact expansion [`@accelerate`](@ref) produces for the same input +(all forms), without running it; inserted frees are highlighted. Pure AST work. +""" +macro show_lifetimes(input) + expansion = _accelerate_expand(input, __module__) + return :(print_lifetime_analysis($(QuoteNode(expansion)))) +end diff --git a/src/scoping/broadcast_lifetimes.jl b/src/scoping/broadcast_lifetimes.jl index 47bb6b301..f1228f58c 100644 --- a/src/scoping/broadcast_lifetimes.jl +++ b/src/scoping/broadcast_lifetimes.jl @@ -115,9 +115,11 @@ function rewrite_broadcast_lifetimes(scope) return _prepend_statements(rewritten, temps), assigned_vars end -function process_broadcast_lifetime_scope(scope; on_rewrite=nothing) - # Returned producers must stay materialized, so exempt them from fusion. - protected = _returned_symbols(scope) +function process_broadcast_lifetime_scope( + scope; on_rewrite=nothing, protected_roots=Set{Symbol}() +) + # Returned producers and caller-owned roots stay materialized: exempt from fusion. + protected = union(_returned_symbols(scope), protected_roots) scope = InterBroadcastFusion.rewrite_scope(scope; on_rewrite, protected) - return _process_lifetime_scope(scope, rewrite_broadcast_lifetimes) + return _process_lifetime_scope(scope, rewrite_broadcast_lifetimes; protected_roots) end diff --git a/src/scoping/lifetimes.jl b/src/scoping/lifetimes.jl index b5b88469c..7f3d97db4 100644 --- a/src/scoping/lifetimes.jl +++ b/src/scoping/lifetimes.jl @@ -70,6 +70,6 @@ function rewrite_eager_lifetimes(scope) return _prepend_statements(rewritten, temps), assigned_vars end -function process_lifetime_scope(scope) - return _process_lifetime_scope(scope, rewrite_eager_lifetimes) +function process_lifetime_scope(scope; protected_roots=Set{Symbol}()) + return _process_lifetime_scope(scope, rewrite_eager_lifetimes; protected_roots) end diff --git a/src/scoping/scoping.jl b/src/scoping/scoping.jl index 676708b30..dec17913f 100644 --- a/src/scoping/scoping.jl +++ b/src/scoping/scoping.jl @@ -1,4 +1,4 @@ -export @analyze_lifetimes, @show_lifetimes +export @accelerate, @show_lifetimes # Include generic syntax layers before the cuNumeric-specific lifetime passes. include("util.jl") @@ -44,9 +44,7 @@ function _normalize_return(block) for (i, stmt) in enumerate(stmts) stmt isa Expr && stmt.head === :return || continue i == length(stmts) || throw( - ArgumentError( - "@analyze_lifetimes: `return` is only allowed as the block's final statement" - ), + ArgumentError("`return` is only allowed as the final statement") ) value = isempty(stmt.args) ? :nothing : only(stmt.args) return Expr(block.head, stmts[1:(end - 1)]..., value) @@ -54,47 +52,6 @@ function _normalize_return(block) return block end -@doc""" - @analyze_lifetimes expr - -Wraps a block of code so that all temporary `NDArray` allocations -(e.g. from slicing or function calls) are tracked and safely freed -at the end of the block. Ensures proper cleanup of GPU memory by -inserting `maybe_insert_delete` calls automatically. - -Assignments created inside the macro are scoped to its lexical region. Existing -arrays can still be mutated in place, and the final value of the block is -returned, but internal bindings do not leak into the surrounding scope. - -The block's final statement determines what leaves the region. Any binding it -returns (a bare name or the elements of a returned tuple) is both protected from -the automatic free and, under fusion, kept materialized rather than inlined into -its consumer, so a real `NDArray` escapes rather than a lazy broadcast tree: - - x, y = @analyze_lifetimes begin - x = e1 .+ e2 - c = x .* e1 # not returned, single-use -> fused into y - y = c .^ 2 - (x, y) # returned -> x and y stay materialized - end - -A trailing `return expr` is accepted as an explicit spelling of the final -statement (`return (x, y)` above); a `return` anywhere else is an error. - -When broadcast fusion is enabled (`FUSE_BROADCAST_EXPRS`), dotted operators -(`.+`, `.*`, etc.) form a lazy `Base.Broadcast.Broadcasted` tree compiled into -a single PTX kernel; intermediate nodes are not real `NDArray` allocations and -are not individually hoisted. The macro automatically selects the -broadcast-aware analysis in that case and the plain analysis otherwise. -""" -macro analyze_lifetimes(block) - block = _normalize_return(_expand_dot_macros(block, __module__)) - on_rewrite = BCAST_FUSION_DEBUG[] ? InterBroadcastFusion.log_rewrite : nothing - rewritten = process_ndarray_scope(block; on_rewrite) - bindings = union(_assigned_symbols(block), _assigned_symbols(rewritten)) - return esc(_lexical_scope(rewritten, bindings)) -end - const counter = Ref(0) function maybe_insert_delete(var::NDArray) @@ -103,10 +60,8 @@ end maybe_insert_delete(x) = x -# `@analyze_lifetimes` is an ownership region, analogous to a C++ `{ ... }` -# block. Bind every source and generated assignment explicitly so it cannot -# accidentally reuse or leak a caller local with the same name. Indexed and -# broadcast assignments are mutations, not new bindings, and remain visible. +# Symbols bound by an assignment anywhere in `expr` (let-form locals; soft-form +# protected bindings). function _assigned_symbols(expr) assigned = Set{Symbol}() @@ -124,9 +79,7 @@ function _assigned_symbols(expr) function visit(node) node isa Expr || return nothing assignment = _assignment(node) - if !isnothing(assignment) - collect_binding(assignment.lhs) - end + isnothing(assignment) || collect_binding(assignment.lhs) foreach(visit, node.args) return nothing end @@ -140,20 +93,6 @@ function _lexical_scope(body, bindings::Set{Symbol}) return Expr(:let, Expr(:block, ordered...), body) end -function _register_scoping_error_hint!() - isdefined(Base.Experimental, :register_error_hint) || return nothing - Base.Experimental.register_error_hint(UndefVarError) do io, exc - return print( - io, - "\nHint: bindings assigned inside `@analyze_lifetimes` are local to its " * - "block. If `", - exc.var, - "` was created there, return it from the block to use it afterward.", - ) - end - return nothing -end - function _hoist_temporary(expr, assigned_vars) counter[] += 1 temporary = Symbol(:tmp, counter[]) @@ -202,7 +141,9 @@ end insert_finalizers(stmts::Vector) Insert `cuNumeric.maybe_insert_delete(var)` after the last use of each temporary variable. """ -function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) +function insert_finalizers( + exprs::Vector, assigned_vars::Set{Symbol}; protected_roots::Set{Symbol}=Set{Symbol}() +) last_use = Dict{Symbol,Int}() alias_map = Dict{Symbol,Symbol}() @@ -261,7 +202,8 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) # return `nothing` rather than leak it or hand back a dangling handle. terminal_indexed = n > 0 && is_indexed_assign(stmts[n]) - protected = Set{Symbol}() + # Roots (function args) are protected regardless of the terminal statement. + protected = Set{Symbol}(canon(root) for root in protected_roots) if n > 0 && !terminal_indexed for result in _result_symbols(stmts[n]) push!(protected, canon(result)) @@ -315,16 +257,20 @@ end insert_finalizers(block::Expr) Apply finalizer insertion to a `begin ... end` or `:block` expression. """ -function insert_finalizers(block::Expr, assigned_vars::Set{Symbol}) +function insert_finalizers( + block::Expr, assigned_vars::Set{Symbol}; protected_roots::Set{Symbol}=Set{Symbol}() +) stmts = _scope_statements(block) isnothing(stmts) && error("Expected a begin/block expression") - return Expr(:block, insert_finalizers(stmts, assigned_vars)...) + return Expr(:block, insert_finalizers(stmts, assigned_vars; protected_roots)...) end -function _process_lifetime_scope(scope, rewrite_lifetimes) +function _process_lifetime_scope( + scope, rewrite_lifetimes; protected_roots::Set{Symbol}=Set{Symbol}() +) try rewritten, assigned_vars = rewrite_lifetimes(scope) - return insert_finalizers(rewritten, assigned_vars) + return insert_finalizers(rewritten, assigned_vars; protected_roots) finally counter[] = 0 end @@ -335,13 +281,15 @@ end include("lifetimes.jl") include("broadcast_lifetimes.jl") -function process_ndarray_scope(scope; on_rewrite=nothing) +function process_ndarray_scope( + scope; on_rewrite=nothing, protected_roots::Set{Symbol}=Set{Symbol}() +) # Broadcast expressions stay lazy only when fusion is enabled; otherwise # every call is analyzed as an eager allocation. @static if FUSE_BROADCAST_EXPRS - return process_broadcast_lifetime_scope(scope; on_rewrite) + return process_broadcast_lifetime_scope(scope; on_rewrite, protected_roots) end - return process_lifetime_scope(scope) + return process_lifetime_scope(scope; protected_roots) end # Return the deleted value for a generated finalizer call. @@ -354,12 +302,27 @@ function _delete_argument(expr) return only(call.args) end -function print_lifetime_analysis(block; io::IO=stdout) +# Header + body statements per form, so the printout mirrors the real expansion. +function _analysis_parts(ex) + ex = _strip_lines(ex) + if ex isa Expr && ex.head === :function + return "function " * string(first(ex.args)), _flatten_statements(ex.args[2]) + elseif ex isa Expr && ex.head === :let + binds = _strip_lines(first(ex.args)) + bindstr = binds isa Expr ? join(binds.args, ", ") : string(binds) + return "let " * bindstr, _flatten_statements(ex.args[2]) + end + return nothing, _flatten_statements(ex) +end + +# Pretty-print `expansion` (the exact `@accelerate` output), highlighting frees. +function print_lifetime_analysis(expansion; io::IO=stdout) rule = "-"^60 - stmts = _flatten_statements(process_ndarray_scope(block)) + header, stmts = _analysis_parts(expansion) mode = FUSE_BROADCAST_EXPRS ? "fusion-aware" : "plain" - println(io, "@analyze_lifetimes expansion ($mode analysis)\n", rule) + println(io, "@accelerate expansion ($mode)\n", rule) + isnothing(header) || println(io, header) n = 0 for s in stmts @@ -368,24 +331,13 @@ function print_lifetime_analysis(block; io::IO=stdout) printstyled(io, lpad("✗ free ", 11), deleted, "\n"; color=:red) else n += 1 - println(io, lpad(n, 4), " ", s) + println(io, lpad(n, 4), " ", _strip_lines(s)) end end + isnothing(header) || println(io, "end") println(io, rule) return nothing end -@doc""" - @show_lifetimes expr - -Print the lifetime-analysis rewrite of `expr` — the same transformation -[`@analyze_lifetimes`](@ref) applies — without running it. Every statement is -shown in source order and each inserted `maybe_insert_delete` is highlighted so -you can see exactly where each temporary is freed. Pure AST work, so it runs on -CPU-only checkouts. -""" -macro show_lifetimes(block) - block = _normalize_return(_expand_dot_macros(block, __module__)) - return :(print_lifetime_analysis($(QuoteNode(block)))) -end +include("accelerate.jl") diff --git a/test/Project.toml b/test/Project.toml index 73f4b047f..6e9f7cd1a 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,7 +1,7 @@ [deps] CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" -CUDA_Driver_jll = "4ee394cb-3365-5eb0-8335-949819d2adfc" +CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" diff --git a/test/analysis/accelerate.jl b/test/analysis/accelerate.jl new file mode 100644 index 000000000..28eee8897 --- /dev/null +++ b/test/analysis/accelerate.jl @@ -0,0 +1,190 @@ +#= Copyright 2026 Northwestern University, + * Carnegie Mellon University University + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * Author(s): David Krasowska + * Ethan Meitz +=# + +# Coverage of the four `@accelerate` forms and their scope contracts: +# @accelerate function f(...) ... end -> function scope, frees non-returned +# @accelerate begin ... end -> 1:1 Julia scope, bindings stay alive +# @accelerate let ... end -> hard scope, combine + free non-returned +# @accelerate expr -> materialized result, temps freed + +using InteractiveUtils: code_typed + +@testset "@accelerate — four forms" begin + T = Float32 + N = 64 + _nd(v) = @allowscalar NDArray(v) + approx(x, ref) = isapprox(Array(x), ref; rtol=1.0f-4) + ja = my_rand(T, N) + jb = my_rand(T, N) + + @testset "1. function form" begin + @accelerate function _acc_fsq(a, b) + c = a .* b + return c .^ 2 + end + a = _nd(ja) + b = _nd(jb) + @test approx(_acc_fsq(a, b), (ja .* jb) .^ 2) + # Arguments are caller-owned: a second call on the same inputs still works. + @test approx(_acc_fsq(a, b), (ja .* jb) .^ 2) + end + + @testset "3. let form (hard scope)" begin + function _acc_let(a, b) + s = @accelerate let + r = a .+ b + s = r .* T(2) + s + end + return s, @isdefined(r) + end + s, r_leaked = _acc_let(_nd(ja), _nd(jb)) + @test approx(s, (ja .+ jb) .* T(2)) + @test r_leaked == false # `r` must not escape the let scope + end + + @testset "4. expr form" begin + a = _nd(ja) + b = _nd(jb) + res = @accelerate (a .+ b) .^ 2 + @test res isa NDArray + @test approx(res, (ja .+ jb) .^ 2) + end + + @testset "2. begin form (bindings stay alive)" begin + function _acc_begin(a, b) + q = @accelerate begin + p = a .* b + q = p .+ one(T) + q + end + return p, q # both must be defined in this scope + end + p, q = _acc_begin(_nd(ja), _nd(jb)) + @test approx(p, ja .* jb) + @test approx(q, (ja .* jb) .+ one(T)) + + # A nested `let` keeps its intermediate private while the outer block + # can consume and return the value it produces. + a = _nd(ja) + b = _nd(jb) + one_t = one(T) + shifted, x = @accelerate begin + shifted = let + product = @. a * b + @. product + one_t + end + x = @. shifted * 2 + (shifted, x) + end + @test approx(shifted, (ja .* jb) .+ one(T)) + @test approx(x, ((ja .* jb) .+ one(T)) .* 2) + end + + @testset "expansion contracts (white-box)" begin + expand(ex) = cuNumeric._accelerate_expand(ex, @__MODULE__) + hasfree(ex) = occursin("maybe_insert_delete", string(expand(ex))) + + # Scope shape per form. + @test expand(:(function f(a) + ;c = a .* a; + c .^ 2; + end)).head === :function + @test expand(:( + begin + C .= a[2:end] .+ b[2:end] + end + )).head === :block + @test expand(:( + let + r = a .+ b; + r .* 2 + end + )).head === :let + + # Slices are freed in every non-`let` form (uniform cleanup). + @test hasfree(:(function f(a) + ;s = a[2:end]; + s .+ 1; + end)) + @test hasfree(:( + begin + C .= a[2:end] .+ b[2:end] + end + )) + + if cuNumeric.FUSE_BROADCAST_EXPRS && cuNumeric.HAS_CUDA + # A same-shape chain fuses into one multi-output launch and still + # frees the hoisted slice temporaries. + mo = string(expand(:( + begin + p = a[2:end] .* b[2:end] + q = p .+ 1 + q + end + ))) + @test occursin("copyto_fused_multi_alloc!", mo) + @test occursin("maybe_insert_delete", mo) + else + # Multi-output fusion is GPU-only; CPU expansion must use the + # ordinary broadcast path even when fusion is enabled in preferences. + cpu = string(expand(:( + begin + p = a .* b + q = p .+ 1 + q + end + ))) + @test !occursin("copyto_fused_multi_alloc!", cpu) + end + end + + @testset "multi-output segment runner is fully unrolled" begin + # GPU compilation requires every chained segment call to be statically + # dispatched. This three-segment shape crossed Julia 1.10's recursive + # inference limit when `_run_segments` recursed over `Base.tail`. + segs = ( + (+, (cuNumeric.RuntimeBroadcastArg{1}(), cuNumeric.RuntimeBroadcastArg{2}())), + (*, (cuNumeric.LocalBroadcastArg{1}(), cuNumeric.RuntimeBroadcastArg{1}())), + (^, (cuNumeric.LocalBroadcastArg{2}(), cuNumeric.RuntimeBroadcastArg{3}())), + ) + outs = ntuple(_ -> zeros(T, 2, 2), 3) + runtime_args = (ones(T, 2, 2), ones(T, 2, 2), 2) + + @test @inferred( + cuNumeric._run_segments( + segs, outs, runtime_args, (), (), CartesianIndex(1, 1) + ) + ) === nothing + @test getindex.(outs, Ref(CartesianIndex(1, 1))) == (T(2), T(2), T(4)) + + argtypes = ( + typeof(segs), + typeof(outs), + typeof(runtime_args), + Tuple{}, + Tuple{}, + CartesianIndex{2}, + ) + typed = only(code_typed(cuNumeric._run_segments, argtypes; optimize=true)).first + @test !occursin( + "_run_segments", sprint(show, MIME("text/plain"), typed) + ) + end +end diff --git a/test/analysis/promotion.jl b/test/analysis/promotion.jl index 7b1fb39a7..422e5766a 100644 --- a/test/analysis/promotion.jl +++ b/test/analysis/promotion.jl @@ -19,3 +19,8 @@ @test safe_compare(r1, r2, atol(Float64), rtol(Float64)) end end + +@testset "Flattened associative broadcast promotion" begin + @test @inferred(cuNumeric.__checked_promote_op(+, NTuple{5,Float64})) === Float64 + @test @inferred(cuNumeric.__checked_promote_op(*, NTuple{4,Int32})) === Int32 +end diff --git a/test/analysis/type_stability.jl b/test/analysis/type_stability.jl index 9d0240749..ce81c1cb0 100644 --- a/test/analysis/type_stability.jl +++ b/test/analysis/type_stability.jl @@ -17,6 +17,32 @@ * Ethan Meitz =# +@accelerate function _type_stable_accelerate_function(a, b) + intermediate = @. a + b + return @. intermediate * 2.0f0 +end + +function _type_stable_accelerate_begin(a, b) + return @accelerate begin + intermediate = @. a + b + result = @. intermediate * 2.0f0 + (intermediate, result) + end +end + +function _type_stable_accelerate_let(a, b) + return @accelerate let + intermediate = @. a + b + @. intermediate * 2.0f0 + end +end + +function _type_stable_accelerate_expr(a, b) + return @accelerate (@. (a + b) * 2.0f0) +end + +_type_stable_cuda_argtypes(task::cuNumeric.CUDATask) = task.argtypes + @testset verbose = true "core" begin a = cuNumeric.zeros(5) b = cuNumeric.zeros(Float64, 3, 4) @@ -66,6 +92,18 @@ end @test @inferred(cuNumeric.NDArray(rand(Float32, 3, 3))) !== nothing end +@testset verbose = true "custom CUDA metadata" begin + task = cuNumeric.CUDATask("kernel", (Float32, Int32)) + @test isconcretetype(typeof(task)) + @test all(isconcretetype, fieldtypes(typeof(task))) + @test @inferred(_type_stable_cuda_argtypes(task)) == DataType[Float32, Int32] + + storage_type = cuNumeric.PaddedStorage{Float32,1} + @test all(isconcretetype, Base.uniontypes(fieldtype(storage_type, :backing))) + @test all(isconcretetype, Base.uniontypes(fieldtype(storage_type, :staging))) + @test all(isconcretetype, Base.uniontypes(fieldtype(storage_type, :shape))) +end + @testset verbose = true "conversion" begin # cast to array, as_type a = cuNumeric.zeros(Float64, 5, 5) @@ -104,6 +142,23 @@ end @test @inferred(((a .* b) .+ a) .* 2.0f0) !== nothing end +@testset verbose = true "@accelerate forms" begin + a = cuNumeric.ones(Float32, 3, 3) + b = cuNumeric.ones(Float32, 3, 3) + + function_result = @inferred _type_stable_accelerate_function(a, b) + @test function_result isa NDArray{Float32,2} + + begin_result = @inferred _type_stable_accelerate_begin(a, b) + @test begin_result isa Tuple{NDArray{Float32,2},NDArray{Float32,2}} + + let_result = @inferred _type_stable_accelerate_let(a, b) + @test let_result isa NDArray{Float32,2} + + expr_result = @inferred _type_stable_accelerate_expr(a, b) + @test expr_result isa NDArray{Float32,2} +end + @testset verbose = true "solve" begin # native float/complex, 2D and 1D rhs @testset "$(T)" for T in Base.uniontypes(cuNumeric.SUPPORTED_SOLVE_TYPES) diff --git a/test/array/broadcast_basic.jl b/test/array/broadcast_basic.jl index e9eb95c04..16b7f710b 100644 --- a/test/array/broadcast_basic.jl +++ b/test/array/broadcast_basic.jl @@ -68,6 +68,18 @@ end result_cpu = zeros(dims) @test result == result_cpu + # Plain broadcast assignment lowers to identity.(source). It must copy + # into both dense NDArrays and writable views. + result .= arrA + @test result == arrA_cpu + + parent = cuNumeric.zeros(Float64, N + 2, N + 2) + center = parent[2:(N + 1), 2:(N + 1)] + center .= arrA + expected_parent = zeros(N + 2, N + 2) + expected_parent[2:(N + 1), 2:(N + 1)] .= arrA_cpu + @test parent == expected_parent + # where the real testing starts arrA = 13.74 .- arrA arrA_cpu = 13.74 .- arrA_cpu @@ -125,6 +137,21 @@ end result_cpu = arrA_cpu .* arrB_cpu @test result == result_cpu + # `@.` lowers associative chains to n-ary Broadcasted nodes. Both the + # fused GPU path and pairwise unfused CPU path must accept them. + result = @. arrA + arrB + arrA + arrB + arrA + result_cpu = @. arrA_cpu + arrB_cpu + arrA_cpu + arrB_cpu + arrA_cpu + @test result == result_cpu + + result = @. 0.5 * arrA * arrB * arrA + result_cpu = @. 0.5 * arrA_cpu * arrB_cpu * arrA_cpu + @test result == result_cpu + + dx = 0.1 + result = @. arrA / dx^2 + result_cpu = @. arrA_cpu / dx^2 + @test result == result_cpu + operator(arrA, arrB) operator(arrA_cpu, arrB_cpu) @test arrA == arrA_cpu diff --git a/test/defunct/fusion_compare.jl b/test/cuda.jl/fusion_compare.jl similarity index 87% rename from test/defunct/fusion_compare.jl rename to test/cuda.jl/fusion_compare.jl index dd4934fca..be0d5090e 100644 --- a/test/defunct/fusion_compare.jl +++ b/test/cuda.jl/fusion_compare.jl @@ -1,3 +1,9 @@ +using CUDA: CUDA, @cuda +using CUDACore: blockDim, blockIdx, threadIdx +import CUDACore: i32 + +cuNumeric.Experimental(true) + function unfused_cunumeric(u, v, f, k) F_u = ( ( @@ -101,8 +107,8 @@ function run_unfused_baseline(N, u, v) end function fusion_test(; N=1024, atol=1.0f-6, rtol=1.0f-6) - u = cuNumeric.as_type(cuNumeric.random(Float32, (N, N)), Float32) - v = cuNumeric.as_type(cuNumeric.random(Float32, (N, N)), Float32) + u = cuNumeric.rand(Float32, (N, N)) + v = cuNumeric.rand(Float32, (N, N)) # using CUDA u_base = CUDA.rand(Float32, (N, N)) @@ -117,8 +123,14 @@ function fusion_test(; N=1024, atol=1.0f-6, rtol=1.0f-6) Fu_fused, Fv_fused = run_fused_cunumeric(N, u, v) Fu_unfused, Fv_unfused = run_unfused_cunumeric(N, u, v) - @test isapprox(Fu_fused, Fu_unfused; atol=atol, rtol=rtol) - @test isapprox(Fv_fused, Fv_unfused; atol=atol, rtol=rtol) + @test isapprox(Array(Fu_fused), Array(Fu_unfused); atol=atol, rtol=rtol) + @test isapprox(Array(Fv_fused), Array(Fv_unfused); atol=atol, rtol=rtol) end -fusion_test() +try + @testset "2D fusion comparison" begin + fusion_test() + end +finally + cuNumeric.Experimental(false) +end diff --git a/test/defunct/fusion_compare_1d.jl b/test/cuda.jl/fusion_compare_1d.jl similarity index 85% rename from test/defunct/fusion_compare_1d.jl rename to test/cuda.jl/fusion_compare_1d.jl index 8163a5db1..8efab2989 100644 --- a/test/defunct/fusion_compare_1d.jl +++ b/test/cuda.jl/fusion_compare_1d.jl @@ -1,4 +1,10 @@ +using CUDA: CUDA, @cuda +using CUDACore: blockDim, blockIdx, threadIdx +import CUDACore: i32 + +cuNumeric.Experimental(true) + function unfused_cunumeric(u, v, f, k) F_u = ( ( @@ -102,8 +108,8 @@ function run_unfused_baseline(N, u, v) end function fusion_test(; N=1024*1024, atol=1.0f-6, rtol=1.0f-6) - u = cuNumeric.as_type(cuNumeric.rand(NDArray, N), Float32) - v = cuNumeric.as_type(cuNumeric.rand(NDArray, N), Float32) + u = cuNumeric.rand(Float32, N) + v = cuNumeric.rand(Float32, N) # using CUDA u_base = CUDA.rand(Float32, N) @@ -116,8 +122,14 @@ function fusion_test(; N=1024*1024, atol=1.0f-6, rtol=1.0f-6) # using cuNumeric Fu_fused, Fv_fused = run_fused_cunumeric(N, u, v) Fu_unfused, Fv_unfused = run_unfused_cunumeric(N, u, v) - @test isapprox(Fu_fused, Fu_unfused; atol=atol, rtol=rtol) - @test isapprox(Fv_fused, Fv_unfused; atol=atol, rtol=rtol) + @test isapprox(Array(Fu_fused), Array(Fu_unfused); atol=atol, rtol=rtol) + @test isapprox(Array(Fv_fused), Array(Fv_unfused); atol=atol, rtol=rtol) end -fusion_test() +try + @testset "1D fusion comparison" begin + fusion_test() + end +finally + cuNumeric.Experimental(false) +end diff --git a/test/cuda.jl/padding.jl b/test/cuda.jl/padding.jl new file mode 100644 index 000000000..f38e875a8 --- /dev/null +++ b/test/cuda.jl/padding.jl @@ -0,0 +1,224 @@ +#= Copyright 2025 Northwestern University, + * Carnegie Mellon University University + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + * + * Author(s): David Krasowska + * Ethan Meitz +=# + +#= Purpose of test: cuda + -- Validate custom-kernel padding, synchronization, and lifetime management +=# + +using CUDACore: blockDim, blockIdx, threadIdx +import CUDACore: i32 + +cuNumeric.Experimental(true) + +function padding_add(a, b, c, N) + i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x + if i <= N + @inbounds c[i] = a[i] + b[i] + end + return nothing +end + +function padding_mul(a, c, b, N) + i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x + if i <= N + @inbounds b[i] = a[i] * c[i] + end + return nothing +end + +function cuda_padding_lifetime() + N = 1_000_000 + M = N - 2 + threads = 256 + blocks = cld(M, threads) + initial_bytes = cuNumeric.current_device_bytes[] + + a = cuNumeric.ones(Float32, N) + b = cuNumeric.ones(Float32, N) + c = cuNumeric.zeros(Float32, M) + task = cuNumeric.@cuda_task padding_add(a, b, c, UInt32(M)) + unpadded_bytes = cuNumeric.current_device_bytes[] + + try + @test @inferred( + cuNumeric.launch( + task, (a, b), c, UInt32(M); threads=threads, blocks=blocks + ) + ) === nothing + padded_bytes = cuNumeric.current_device_bytes[] + + @test @inferred(cuNumeric._launch_shape(c)) == (N,) + @test @inferred(cuNumeric._sync_to_launch_padding!(c)) === nothing + @test @inferred(cuNumeric._sync_from_launch_padding!(c)) === nothing + + accounting_ok = true + for _ in 1:15 + cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, b) outputs=c scalars=UInt32( + M + ) + accounting_ok &= cuNumeric.current_device_bytes[] == padded_bytes + end + + @test padded_bytes > unpadded_bytes + @test accounting_ok + @test size(c) == (M,) + @test all(Array(c) .== 2.0f0) + finally + cuNumeric.destroy!(a) + cuNumeric.destroy!(b) + cuNumeric.destroy!(c) + end + @test cuNumeric.current_device_bytes[] == initial_bytes +end + +function cuda_padding_api_interop() + N = 4096 + M = N - 2 + threads = 256 + blocks = cld(M, threads) + initial_bytes = cuNumeric.current_device_bytes[] + + a = cuNumeric.ones(Float32, N) + b = cuNumeric.ones(Float32, N) + c = cuNumeric.zeros(Float32, M) + library_result = nothing + + try + task = cuNumeric.@cuda_task padding_add(a, b, c, UInt32(M)) + cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, b) outputs=c scalars=UInt32( + M + ) + + # The broadcast writes through c's logical view into its padded backing. + c .= c .* 2.0f0 .+ 0.0f0 + @test all(Array(c) .== 4.0f0) + + # Non-broadcasted operators also consume the logical shape. + library_result = c + c + @test size(library_result) == (M,) + @test all(Array(library_result) .== 8.0f0) + + # A later custom launch sees the values written by the regular API. + task = cuNumeric.@cuda_task padding_mul(a, c, b, UInt32(M)) + cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, c) outputs=b scalars=UInt32( + M + ) + result = Array(b) + @test all(result[1:M] .== 4.0f0) + @test all(result[(M + 1):N] .== 1.0f0) + finally + !isnothing(library_result) && cuNumeric.destroy!(library_result) + cuNumeric.destroy!(a) + cuNumeric.destroy!(b) + cuNumeric.destroy!(c) + end + @test cuNumeric.current_device_bytes[] == initial_bytes +end + +function cuda_padding_slice_output() + N = 4096 + M = N - 2 + threads = 256 + blocks = cld(M, threads) + initial_bytes = cuNumeric.current_device_bytes[] + + a = cuNumeric.ones(Float32, N) + b = cuNumeric.ones(Float32, N) + parent = cuNumeric.zeros(Float32, N) + output = parent[1:M] + task = cuNumeric.@cuda_task padding_add(a, b, output, UInt32(M)) + unpadded_bytes = cuNumeric.current_device_bytes[] + + try + @test @inferred( + cuNumeric.launch( + task, (a, b), output, UInt32(M); threads=threads, blocks=blocks + ) + ) === nothing + padded_bytes = cuNumeric.current_device_bytes[] + @test @inferred(cuNumeric._launch_shape(output)) == (N,) + @test @inferred(cuNumeric._sync_from_launch_padding!(output)) === nothing + + # Mutate the logical parent view before reusing it as a custom-kernel input. + output .= output .* 1.0f0 .+ 1.0f0 + library_result = output + output + @test all(Array(library_result) .== 6.0f0) + cuNumeric.destroy!(library_result) + + task = cuNumeric.@cuda_task padding_mul(a, output, b, UInt32(M)) + @test @inferred( + cuNumeric.launch( + task, (a, output), b, UInt32(M); threads=threads, blocks=blocks + ) + ) === nothing + values = Array(parent) + product = Array(b) + + @test padded_bytes > unpadded_bytes + @test cuNumeric.current_device_bytes[] == padded_bytes + @test all(values[1:M] .== 3.0f0) + @test all(product[1:M] .== 3.0f0) + @test values[end] == 0.0f0 + finally + cuNumeric.destroy!(output) + cuNumeric.destroy!(parent) + cuNumeric.destroy!(a) + cuNumeric.destroy!(b) + end + @test cuNumeric.current_device_bytes[] == initial_bytes +end + +Base.@noinline function drop_padded_arrays() + N = 4096 + M = N - 2 + a = cuNumeric.ones(Float32, N) + b = cuNumeric.ones(Float32, N) + c = cuNumeric.zeros(Float32, M) + task = cuNumeric.@cuda_task padding_add(a, b, c, UInt32(M)) + cuNumeric.@launch task=task threads=256 blocks=cld(M, 256) inputs=(a, b) outputs=c scalars=UInt32( + M + ) + return nothing +end + +function cuda_padding_finalizer() + GC.gc(true) + cuNumeric.drain_pending_frees!() + baseline = cuNumeric.current_device_bytes[] + + drop_padded_arrays() + allocated = cuNumeric.current_device_bytes[] + GC.gc(true) + GC.gc(true) + cuNumeric.drain_pending_frees!() + + @test allocated > baseline + @test cuNumeric.current_device_bytes[] == baseline +end + +try + @testset "Custom CUDA padding" begin + cuda_padding_lifetime() + cuda_padding_api_interop() + cuda_padding_slice_output() + cuda_padding_finalizer() + end +finally + cuNumeric.Experimental(false) +end diff --git a/test/defunct/vecadd.jl b/test/cuda.jl/vecadd.jl similarity index 81% rename from test/defunct/vecadd.jl rename to test/cuda.jl/vecadd.jl index ea01e7ab3..e2b5de552 100644 --- a/test/defunct/vecadd.jl +++ b/test/cuda.jl/vecadd.jl @@ -21,6 +21,11 @@ -- Register various custom kernels using CUDA.jl =# +using CUDACore: blockDim, blockIdx, threadIdx +import CUDACore: i32 + +cuNumeric.Experimental(true) + function kernel_add(a, b, c, N) i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x if i <= N @@ -29,9 +34,8 @@ function kernel_add(a, b, c, N) return nothing end -# testing a second kernel -# on purpose switching inputs and outputs -function kernel_mul(a, b, c, N) +# Test a second kernel with `c` as an input and `b` as the output. +function kernel_mul(a, c, b, N) i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x if i <= N @inbounds b[i] = a[i] * c[i] @@ -70,18 +74,16 @@ function cuda_binaryop(max_diff) N ) - @test @allowscalar cuNumeric.compare(c, c_cpu, atol(Float32), rtol(Float32)) + @test @allowscalar cuNumeric.compare(c, c_cpu, max_diff, max_diff) - for i in 1:N - @allowscalar b[i] = a[i] * c[i] - end + b_cpu .= a_cpu .* c_cpu - task = cuNumeric.@cuda_task kernel_mul(a, b, c, UInt32(1)) + task = cuNumeric.@cuda_task kernel_mul(a, c, b, UInt32(1)) cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, c) outputs=b scalars=UInt32( N ) - @test @allowscalar cuNumeric.compare(b, b_cpu, atol(Float32), rtol(Float32)) + @test @allowscalar cuNumeric.compare(b, b_cpu, max_diff, max_diff) end function kernel_sin(a, b, N) @@ -119,5 +121,14 @@ function cuda_unaryop(max_diff) # TODO explore getting inplace ops working. cuNumeric.@launch task=task threads=threads blocks=blocks inputs=a outputs=b scalars=UInt32(N) - @test @allowscalar cuNumeric.compare(b, b_cpu, atol(Float32), rtol(Float32)) + @test @allowscalar cuNumeric.compare(b, b_cpu, max_diff, max_diff) +end + +try + @testset "Custom CUDA kernels" begin + cuda_binaryop(1.0f-5) + cuda_unaryop(1.0f-5) + end +finally + cuNumeric.Experimental(false) end diff --git a/test/gpu_only/broadcast_fusion.jl b/test/gpu_only/broadcast_fusion.jl index 5c92911c8..8ab3afb3d 100644 --- a/test/gpu_only/broadcast_fusion.jl +++ b/test/gpu_only/broadcast_fusion.jl @@ -203,7 +203,7 @@ _broadcast_fusion_user_add(x, y) = x + y @testset "z .= scalar * f.(A, B)" begin expected = T(2.0) .* (julia_a .+ julia_b) z = cuNumeric.zeros(T, (N,)) - @analyze_lifetimes begin + @accelerate begin z .= T(2.0) .* _broadcast_fusion_user_add.(a, b) end @allowscalar @test safe_compare(expected, z, atol, rtol) @@ -566,7 +566,7 @@ end ja = reshape(T.(1:(N * N)), N, N) a = @allowscalar NDArray(ja) out = cuNumeric.zeros(T, (N + 2, N + 2)) - @analyze_lifetimes begin + @accelerate begin producer = a .* s1 out[2:(end - 1), 2:(end - 1)] = producer .+ s2 end diff --git a/test/runtests.jl b/test/runtests.jl index da06059c3..a7ca7bbb7 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,28 +1,18 @@ using cuNumeric -using CUDA: CUDA +using CUDACore: CUDACore using ParallelTestRunner using Pkg using InteractiveUtils: versioninfo -run_gpu_tests = CUDA.functional() +run_gpu_tests = CUDACore.functional() @info "Julia information:\n" * sprint(io -> versioninfo(io)) -run_gpu_tests && @info "CUDA information:\n" * sprint(io -> CUDA.versioninfo(io)) @info "cuNumeric information:\n" * sprint(io -> cuNumeric.versioninfo(io)) # Forcibly precompile the current environment in parallel: Pkg sometimes ignores # dependencies pointed through via `[sources]` Pkg.precompile() -cuda_init = if run_gpu_tests - quote - using CUDA - import CUDA: i32 - end -else - :() -end - const init_code = quote using LinearAlgebra using Random @@ -32,8 +22,6 @@ const init_code = quote ENV["LEGATE_SKIP_RUNTIME"] = "false" using cuNumeric - $cuda_init - include("util.jl") end @@ -44,16 +32,23 @@ delete!(testsuite, "util") delete!(testsuite, "array/unary/tests") delete!(testsuite, "array/binary/tests") -if !run_gpu_tests - @warn "CUDA GPU not available, skipping GPU-only tests" - filter!(test -> !startswith(first(test), "gpu_only/"), testsuite) -end +test_args = parse_args(ARGS) +if filter_tests!(testsuite, test_args) + if !run_gpu_tests + @warn "CUDA GPU not available, skipping GPU-only tests" + filter!( + test -> + !startswith(first(test), "gpu_only/") && + !startswith(first(test), "cuda.jl/"), + testsuite, + ) + end -if !run_gpu_tests || !cuNumeric.FUSE_BROADCAST_EXPRS - @warn "Broadcast fusion is disabled, skipping fusion tests" - filter!(test -> !startswith(first(test), "gpu_only/broadcast_fusion"), testsuite) + if !run_gpu_tests || !cuNumeric.FUSE_BROADCAST_EXPRS + @warn "Broadcast fusion is disabled, skipping fusion tests" + filter!(test -> !startswith(first(test), "gpu_only/broadcast_fusion"), testsuite) + end end -filter!(test -> !startswith(first(test), "defunct/"), testsuite) - -runtests(cuNumeric, ARGS; testsuite, init_code) +cuda_tests = filter(test -> startswith(test, "cuda.jl/"), collect(keys(testsuite))) +runtests(cuNumeric, test_args; testsuite, init_code, serial=cuda_tests) diff --git a/test/workflows/grayscott.jl b/test/workflows/grayscott.jl index e81c364f6..7be2662ba 100644 --- a/test/workflows/grayscott.jl +++ b/test/workflows/grayscott.jl @@ -33,64 +33,62 @@ struct ParamsGS{T<:AbstractFloat} end end -function step(u, v, u_new, v_new, args::ParamsGS) - @analyze_lifetimes begin - # calculate F_u and F_v functions - # currently we don't have NDArray^x working yet. - F_u = ( - ( - -u[2:(end - 1), 2:(end - 1)] .* - (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) - ) + args.f*(1 .- u[2:(end - 1), 2:(end - 1)]) - ) - F_v = ( - ( - u[2:(end - 1), 2:(end - 1)] .* - (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) - ) - (args.f+args.k)*v[2:(end - 1), 2:(end - 1)] - ) - # 2-D Laplacian of f using array slicing, excluding boundaries - # For an N x N array f, f_lap is the Nend x Nend array in the "middle" - u_lap = ( - ( - u[3:end, 2:(end - 1)] - 2*u[2:(end - 1), 2:(end - 1)] + - u[1:(end - 2), 2:(end - 1)] - ) ./ args.dx^2 + - ( - u[2:(end - 1), 3:end] - 2*u[2:(end - 1), 2:(end - 1)] + - u[2:(end - 1), 1:(end - 2)] - ) ./ args.dx^2 - ) - v_lap = ( - ( - v[3:end, 2:(end - 1)] - 2*v[2:(end - 1), 2:(end - 1)] + - v[1:(end - 2), 2:(end - 1)] - ) ./ args.dx^2 + - ( - v[2:(end - 1), 3:end] - 2*v[2:(end - 1), 2:(end - 1)] + - v[2:(end - 1), 1:(end - 2)] - ) ./ args.dx^2 - ) +@accelerate function step(u, v, u_new, v_new, args::ParamsGS) + # calculate F_u and F_v functions + # currently we don't have NDArray^x working yet. + F_u = ( + ( + -u[2:(end - 1), 2:(end - 1)] .* + (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) + ) + args.f*(1 .- u[2:(end - 1), 2:(end - 1)]) + ) + F_v = ( + ( + u[2:(end - 1), 2:(end - 1)] .* + (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)]) + ) - (args.f+args.k)*v[2:(end - 1), 2:(end - 1)] + ) + # 2-D Laplacian of f using array slicing, excluding boundaries + # For an N x N array f, f_lap is the Nend x Nend array in the "middle" + u_lap = ( + ( + u[3:end, 2:(end - 1)] - 2*u[2:(end - 1), 2:(end - 1)] + + u[1:(end - 2), 2:(end - 1)] + ) ./ args.dx^2 + + ( + u[2:(end - 1), 3:end] - 2*u[2:(end - 1), 2:(end - 1)] + + u[2:(end - 1), 1:(end - 2)] + ) ./ args.dx^2 + ) + v_lap = ( + ( + v[3:end, 2:(end - 1)] - 2*v[2:(end - 1), 2:(end - 1)] + + v[1:(end - 2), 2:(end - 1)] + ) ./ args.dx^2 + + ( + v[2:(end - 1), 3:end] - 2*v[2:(end - 1), 2:(end - 1)] + + v[2:(end - 1), 1:(end - 2)] + ) ./ args.dx^2 + ) - # # Forward-Euler time step for all points except the boundaries - u_new[2:(end - 1), 2:(end - 1)] = - ((args.c_u * u_lap) + F_u) * args.dt + u[2:(end - 1), 2:(end - 1)] - v_new[2:(end - 1), 2:(end - 1)] = - ((args.c_v * v_lap) + F_v) * args.dt + v[2:(end - 1), 2:(end - 1)] - - # Apply periodic boundary conditions - u_new[:, 1] = u[:, end - 1] - u_new[:, end] = u[:, 2] - u_new[1, :] = u[end - 1, :] - u_new[end, :] = u[2, :] - v_new[:, 1] = v[:, end - 1] - v_new[:, end] = v[:, 2] - v_new[1, :] = v[end - 1, :] - v_new[end, :] = v[2, :] - end + # # Forward-Euler time step for all points except the boundaries + u_new[2:(end - 1), 2:(end - 1)] = + ((args.c_u * u_lap) + F_u) * args.dt + u[2:(end - 1), 2:(end - 1)] + v_new[2:(end - 1), 2:(end - 1)] = + ((args.c_v * v_lap) + F_v) * args.dt + v[2:(end - 1), 2:(end - 1)] + + # Apply periodic boundary conditions + u_new[:, 1] = u[:, end - 1] + u_new[:, end] = u[:, 2] + u_new[1, :] = u[end - 1, :] + u_new[end, :] = u[2, :] + v_new[:, 1] = v[:, end - 1] + v_new[:, end] = v[:, 2] + v_new[1, :] = v[end - 1, :] + v_new[end, :] = v[2, :] end -# same as above but without @analyze_lifetimes macro +# same as above but without the @accelerate macro function step_base(u, v, u_new, v_new, args::ParamsGS) # calculate F_u and F_v functions # currently we don't have NDArray^x working yet. @@ -239,8 +237,8 @@ function run_slice_test(op, op_scoped, FT, N; f=0.04, k=0.06, dx=1.0) return base, scoped end -binary_scope(op) = (a, b, out) -> @analyze_lifetimes out[:, :] = op(a, b) -slice_scope(op) = (u, v, out, args) -> @analyze_lifetimes out[:, :] = op(u, v, args) +binary_scope(op) = (a, b, out) -> @accelerate out[:, :] = op(a, b) +slice_scope(op) = (u, v, out, args) -> @accelerate out[:, :] = op(u, v, args) const OPS = Dict( :add => (+), @@ -436,68 +434,30 @@ function test_scoping_rewrite_pipeline() @test isempty(freed) end - @testset "Lexical lifetime scope" begin - function hidden_binding() - @analyze_lifetimes begin - internal_result = 41 - nothing - end - return internal_result - end - - function shadowed_binding() - internal_result = :outer - @analyze_lifetimes begin - internal_result = :inner - nothing - end - return internal_result - end - - function hidden_destructured_bindings() - @analyze_lifetimes begin - internal_first, internal_second = (1, 2) + @testset "Block form keeps Julia scope" begin + # Block form adds no scope: bindings stay live in the enclosing scope (1:1 Julia). + function visible_binding() + @accelerate begin + internal_result = 42 nothing end - return internal_first, internal_second - end - - function unrelated_undefined_binding() - return unrelated_result - end - - function rendered_error(f) - try - f() - catch exc - return sprint(io -> showerror(io, exc, catch_backtrace())) - end - return "" + return internal_result # would be UndefVar under a `let` end + @test visible_binding() == 42 output = [0] - returned = @analyze_lifetimes begin + returned = @accelerate begin internal_result = 42 output[1] = internal_result internal_result end - - @test_throws UndefVarError hidden_binding() - @test_throws UndefVarError hidden_destructured_bindings() - @test occursin( - "If `internal_result` was created there", rendered_error(hidden_binding) - ) - @test occursin( - "If `unrelated_result` was created there", - rendered_error(unrelated_undefined_binding), - ) - @test shadowed_binding() == :outer @test output == [42] @test returned == 42 if cuNumeric.FUSE_BROADCAST_EXPRS - function hidden_fused_binding(a, b, destination) - @analyze_lifetimes begin + # A fused intermediate is materialized and also stays live. + function fused_binding(a, b, destination) + @accelerate begin fused_result = a .* b destination .= fused_result .+ 1 end @@ -505,9 +465,7 @@ function test_scoping_rewrite_pipeline() end destination = zeros(Int, 2) - @test_throws UndefVarError hidden_fused_binding( - [2, 3], [4, 5], destination - ) + @test fused_binding([2, 3], [4, 5], destination) == [8, 15] @test destination == [9, 16] end end @@ -519,7 +477,7 @@ function test_scoping_regressions(T, N) C = cuNumeric.zeros(T, (N, N)) @testset "In-place assignment" begin - @analyze_lifetimes begin + @accelerate begin result = A[1:end, :] .+ B[1:end, :] C .= result .* T(2.0) end @@ -529,7 +487,7 @@ function test_scoping_regressions(T, N) @testset "Macro as RHS" begin # Test values: (1+1)^2 = 4 - res = @analyze_lifetimes (A .+ B) .^ 2 + res = @accelerate (A .+ B) .^ 2 @test res isa cuNumeric.NDArray @test all(Array(res) .== T(4.0)) end @@ -537,7 +495,7 @@ function test_scoping_regressions(T, N) @testset "Returned bindings stay materialized" begin # A returned producer must come back as a materialized NDArray, not a # lazy broadcast tree; `c` stays a private intermediate that fuses away. - x, y = @analyze_lifetimes begin + x, y = @accelerate begin x = A .+ B c = x .* A y = c .^ 2 @@ -551,14 +509,14 @@ function test_scoping_regressions(T, N) @testset "Return forms yield materialized bindings" begin # `x = y` alias, tuple, and trailing `return` all return real NDArrays. - aliased = @analyze_lifetimes begin + aliased = @accelerate begin y = A .+ B x = y end @test aliased isa cuNumeric.NDArray @test all(Array(aliased) .== T(2)) - rx, ry = @analyze_lifetimes begin + rx, ry = @accelerate begin rx = A .+ B ry = rx .^ 2 return (rx, ry) @@ -570,7 +528,7 @@ function test_scoping_regressions(T, N) if cuNumeric.FUSE_BROADCAST_EXPRS @testset "Indexed fused assignment writes through NDArray slices" begin out = cuNumeric.zeros(T, (N + 2, N + 2)) - @analyze_lifetimes begin + @accelerate begin producer = A .* T(2) out[2:(end - 1), 2:(end - 1)] = producer .+ T(1) end @@ -582,7 +540,7 @@ function test_scoping_regressions(T, N) @testset "Nested @. macros fuse before lifetime analysis" begin multiplier = cuNumeric.ones(T, (N, N)) result = cuNumeric.zeros(T, (N, N)) - @analyze_lifetimes begin + @accelerate begin tmp = @. A + B result .= @. tmp * multiplier + T(1.0) end