diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 7d413a0fe..e379130af 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -8,12 +8,23 @@ steps: - group: ":julia: Julia" key: "julia" steps: - - label: "GPU - Julia {{matrix.julia}}" + # Fusion prefs are compile-time (CNPreferences). Write LocalPreferences.toml and + # set both values before instantiate/test so each job precompiles cuNumeric + # with the intended fusion configuration. + # - fusion=on: FUSE_BROADCAST_EXPRS=true, MIN_OPS=1 (fuse single-op exprs) + # - fusion=off: FUSE_BROADCAST_EXPRS=false (MIN_OPS irrelevant) + - label: "GPU - Julia {{matrix.julia}} (fusion {{matrix.fusion}})" plugins: - JuliaCI/julia#v1: version: "{{matrix.julia}}" - jquick/pre-hook#v1.2.0: - command: "julia --project -e 'using Pkg; Pkg.resolve(); Pkg.instantiate()'" + command: | + if [ "{{matrix.fusion}}" = "on" ]; then + printf '%s\n' '[CNPreferences]' 'FUSE_BROADCAST_EXPRS = true' 'FUSE_BROADCAST_MIN_OPS = 1' > LocalPreferences.toml + else + printf '%s\n' '[CNPreferences]' 'FUSE_BROADCAST_EXPRS = false' > LocalPreferences.toml + fi + julia --project -e 'using Pkg; Pkg.resolve(); Pkg.instantiate()' - JuliaCI/julia-test#v1: test_args: "--quickfail" - JuliaCI/julia-coverage#v1: @@ -42,6 +53,9 @@ steps: - "1.10" - "1.11" - "1.12" + fusion: + - "on" + - "off" arch: - "x64" # adjustments: @@ -49,6 +63,37 @@ steps: # julia: "1.12" # soft_fail : true + - group: ":hammer: Developer" + key: "developer" + steps: + - label: "GPU - Julia 1.11 (source wrapper, fusion {{matrix.fusion}})" + plugins: + - JuliaCI/julia#v1: + version: "1.11" + command: ".buildkite/run_developer_ci.sh" + artifact_paths: + - "deps/build.log" + - "deps/*.log" + - "deps/*.err" + - "*.log" + agents: + queue: "cuda" + if: build.message !~ /\[skip tests\]/ + timeout_in_minutes: 60 + env: + LD_LIBRARY_PATH: "" + LEGATE_AUTO_CONFIG: "0" + LEGATE_SHOW_CONFIG: "1" + GPUTESTS: "1" + LEGATE_TEST: "1" + LEGATE_CONFIG: "--cpus 1 --gpus 1 --utility 1 --fbmem 500 --logging legate=info,level=2 --log-to-file 1" + CUNUMERIC_FUSION: "{{matrix.fusion}}" + matrix: + setup: + fusion: + - "on" + - "off" + # - group: ":computer: CUDA Versions (Julia 1.11)" # key: "cuda" # depends_on: "julia" diff --git a/.buildkite/run_developer_ci.sh b/.buildkite/run_developer_ci.sh new file mode 100755 index 000000000..1a91c0b15 --- /dev/null +++ b/.buildkite/run_developer_ci.sh @@ -0,0 +1,81 @@ +#!/usr/bin/env bash + +set -euo pipefail + +CI_PROJECT="$(mktemp -d)" + +case "${CUNUMERIC_FUSION:-}" in + on | off) ;; + *) + echo "CUNUMERIC_FUSION must be 'on' or 'off'" >&2 + exit 2 + ;; +esac + +CMAKE_VERSION="3.30.7" +CMAKE_ROOT="$(mktemp -d)" +CMAKE_INSTALLER="$CMAKE_ROOT/cmake-installer.sh" +curl --fail --silent --show-error --location \ + --output "$CMAKE_INSTALLER" \ + "https://github.com/Kitware/CMake/releases/download/v$CMAKE_VERSION/cmake-$CMAKE_VERSION-linux-x86_64.sh" +sh "$CMAKE_INSTALLER" --skip-license --prefix="$CMAKE_ROOT" +export PATH="$CMAKE_ROOT/bin:$PATH" +cmake --version + +LEGATE_BRANCH_INPUT="${BUILDKITE_MESSAGE:-}" +if [[ "${BUILDKITE_PULL_REQUEST:-false}" =~ ^[0-9]+$ ]]; then + echo "Reading Legate branch override from pull request #$BUILDKITE_PULL_REQUEST" + PR_BODY="$( + curl --fail --silent --show-error --location \ + --header "Accept: application/vnd.github+json" \ + "https://api.github.com/repos/JuliaLegate/cuNumeric.jl/pulls/$BUILDKITE_PULL_REQUEST" | + python3 -c 'import json, sys; print(json.load(sys.stdin).get("body") or "")' + )" + LEGATE_BRANCH_INPUT+=$'\n'"$PR_BODY" +fi + +shopt -s nocasematch +LEGATE_BRANCH="" +if [[ "$LEGATE_BRANCH_INPUT" =~ legate[-_]branch:[[:space:]]*([A-Za-z0-9._/-]+) ]]; then + LEGATE_BRANCH="${BASH_REMATCH[1]}" +fi +shopt -u nocasematch + +if [[ -n "$LEGATE_BRANCH" ]]; then + LEGATE_CHECKOUT="$(mktemp -d)/Legate.jl" + echo "Using Legate.jl branch override: $LEGATE_BRANCH" + git clone --depth 1 --branch "$LEGATE_BRANCH" \ + https://github.com/JuliaLegate/Legate.jl.git "$LEGATE_CHECKOUT" + git -C "$LEGATE_CHECKOUT" log -1 --format="Legate checkout: %D (%H)" + ( + cd "$LEGATE_CHECKOUT" + julia --color=yes --project="$CI_PROJECT" -e ' + using Pkg + Pkg.develop(PackageSpec(path = "lib/LegatePreferences")) + Pkg.develop(PackageSpec(path = ".")) + using LegatePreferences + LegatePreferences.use_developer_mode() + Pkg.build("Legate") + ' + ) +else + echo "No Legate.jl branch override - using JLL" +fi + +julia --color=yes --project="$CI_PROJECT" -e ' + using Pkg + Pkg.develop(PackageSpec(path = "lib/CNPreferences")) + using CNPreferences + CNPreferences.use_developer_mode() + CNPreferences.set_broadcast_fusion!(ENV["CUNUMERIC_FUSION"] == "on") + CNPreferences.set_broadcast_fusion_min_ops!(1) + Pkg.develop(PackageSpec(path = ".")) + Pkg.build("cuNumeric") +' + +cp "$CI_PROJECT/LocalPreferences.toml" test/LocalPreferences.toml + +julia --color=yes --project="$CI_PROJECT" -e ' + using Pkg + Pkg.test("cuNumeric"; test_args = ["--quickfail"]) +' diff --git a/.github/workflows/TagBot.yml b/.github/workflows/TagBot.yml index a3de268c2..c114fd58d 100644 --- a/.github/workflows/TagBot.yml +++ b/.github/workflows/TagBot.yml @@ -28,3 +28,5 @@ jobs: - uses: JuliaRegistries/TagBot@v1 with: token: ${{ secrets.GITHUB_TOKEN }} + # SSH deploy key so the tag push can trigger Docs (GITHUB_TOKEN pushes cannot). + ssh: ${{ secrets.DOCUMENTER_KEY }} diff --git a/.github/workflows/docs-tags.yml b/.github/workflows/docs-tags.yml new file mode 100644 index 000000000..ead34e481 --- /dev/null +++ b/.github/workflows/docs-tags.yml @@ -0,0 +1,43 @@ +# Deploy docs for version tags. Kept separate from docs.yml so tag pushes are +# not subject to the docs/README path filter (a tag tip often only bumps +# Project.toml). +name: Docs (tags) +on: + push: + tags: + - 'v*' +jobs: + docs: + name: Documentation + permissions: + actions: write + contents: write + pull-requests: read + statuses: write + runs-on: ubuntu-latest + env: + LEGATE_AUTO_CONFIG: 0 + steps: + - uses: actions/checkout@v4 + - uses: julia-actions/setup-julia@v2 + with: + version: '1.11' + - uses: julia-actions/cache@v2 + + - name: Setup cuNumeric.jl + run: | + julia --color=yes --project=docs -e ' + using Pkg + Pkg.develop(PackageSpec(path = "lib/CNPreferences")) + Pkg.develop(PackageSpec(path = ".")) + ' + - name: Instantiate docs environment + run: | + julia --color=yes --project=docs -e ' + using Pkg + Pkg.instantiate()' + - name: Build documentation + run: julia --color=yes --project=docs docs/make.jl + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }} diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index a688ae083..3317cea6a 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -1,44 +1,21 @@ name: Docs on: push: - paths: - - '.github/workflows/docs.yml' - - 'src/**' - - 'docs/src/**' - - 'docs/make.jl' - - 'docs/Project.toml' - - 'README.md' - tags: - - 'v*' + # After merge: deploy live docs. Feature branches only run via pull_request. branches: - main pull_request: paths: - - '.github/workflows/docs.yml' - - 'src/**' - - 'docs/src/**' - - 'docs/make.jl' - - 'docs/Project.toml' - - 'README.md' + - '.github/workflows/docs.yml' + - '.github/workflows/docs-tags.yml' + - 'src/**' + - 'docs/**' + - 'README.md' + - 'lib/CNPreferences/**' + workflow_dispatch: jobs: - check_changes: - name: Check for wrapper changes - runs-on: ubuntu-latest - outputs: - wrapper_changed: ${{ steps.filter.outputs.wrapper }} - steps: - - uses: actions/checkout@v4 - - uses: dorny/paths-filter@v3 - id: filter - with: - filters: | - wrapper: - - 'lib/cunumeric_jl_wrapper/**' - docs: - name : Documentation - needs: check_changes - if: ${{ github.base_ref == 'main' || needs.check_changes.outputs.wrapper_changed != 'true' }} + name: Documentation permissions: actions: write contents: write diff --git a/.gitignore b/.gitignore index 29f3d39d6..0311f9e68 100644 --- a/.gitignore +++ b/.gitignore @@ -18,8 +18,9 @@ debug debug/* # benchmark outputs -benchmark/results +benchmark/results** benchmark/results/* +benchmark/plots** compile_wrapper.sh @@ -28,8 +29,8 @@ compile_wrapper.sh build build/* build.log -*.log -*.err +*.log** +*.err** kernel.ptx *.perf *.folded @@ -94,3 +95,8 @@ node_modules *.out *.app .githash + +# VitePress caches/output (keep docs/src/.vitepress/config.mts + theme/) +docs/.vitepress/ +docs/src/.vitepress/cache/ +docs/src/.vitepress/dist/ diff --git a/Project.toml b/Project.toml index feca11f03..53e722549 100644 --- a/Project.toml +++ b/Project.toml @@ -4,8 +4,12 @@ version = "0.1.1" [deps] CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" +CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" +CUDATools = "9ec180c6-1c07-47c7-9e6e-ebefa4d1f6d0" CxxWrap = "1f15a43c-97ca-5a2a-ae31-89f07a497df4" +ExpressionExplorer = "21656369-7473-754a-2065-74616d696c43" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Legate = "1238f2cf-6593-4d60-9aca-2f5364e49909" LegatePreferences = "8028f36a-2b64-49e9-aa04-2d0933fd2ed9" Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb" @@ -20,17 +24,14 @@ cunumeric_jl_wrapper_jll = "49048992-29d2-5fd1-994f-9cecf112d624" cupynumeric_jll = "2862d674-414d-5b0b-a494-b21f8deca547" libcxxwrap_julia_jll = "3eaa8342-bff7-56a5-9981-c04077f7cee7" -[weakdeps] -CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" - -[extensions] -CUDAExt = "CUDA" - [compat] CNPreferences = "0.1.2" -CUDA = "5.9" +CUDACore = "6.0.0" +CUDATools = "6.0.0" CxxWrap = "0.17" +ExpressionExplorer = "1.1.4" JuliaFormatter = "2.3.0" +KernelAbstractions = "0.9.41" Legate = "0.1.2" LegatePreferences = "0.1.6" MacroTools = "0.5.16" diff --git a/README.md b/README.md index ad1c5660d..b51c551b3 100644 --- a/README.md +++ b/README.md @@ -1,69 +1,102 @@ -# 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) +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. -The cuNumeric.jl package wraps the [cuPyNumeric](https://github.com/nv-legate/cupynumeric) C++ API from NVIDIA to bring simple distributed computing on GPUs and CPUs to Julia! We provide a simple array abstraction, the `NDArray`, which supports most of the operations you would expect from a normal Julia array. - -> [!WARNING] -> cuNumeric.jl is under active development. This is an alpha API and is subject to change. Stability is not guaranteed until the first official release. We are actively working to improve the build experience to be more seamless and Julia-friendly. +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. From the Julia REPL, type `]` to enter the Pkg REPL mode and run: -```julia -pkg> add cuNumeric -``` -Or, using the `Pkg` API: + +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 run might take awhile as it has to install multiple large dependencies such as the CUDA SDK (if you have an NVIDIA GPU). For more install instructions, please visit out install guide in the documentation. -To see information about your cuNumeric install run the `versioninfo` function. +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). ```julia +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, please visit our hardware configuration documentation. +> Starting more than one instance of cuNumeric.jl can lead to a hard-crash. The default hardware configuration reserves all available resources. -### Monte-Carlo Example -```julia -using cuNumeric +For more details, see [Hardware](./configuration/hardware.md). -integrand = (x) -> exp.(-x.^2) +### How `NDArray`s work -N = 1_000_000 +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. -x_max = 10.0f0 -domain = [-x_max, x_max] -Ω = domain[2] - domain[1] +**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. -samples = Ω*cuNumeric.rand(N) .- x_max -estimate = (Ω/N) * sum(integrand(samples)) +**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. -println("Monte-Carlo Estimate: $(estimate)") +**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`: + +```julia +s = sum(A) # NDArray{T,0} +x = unwrap(s) # T, e.g. Float32 ``` -### Helping the Garbage Collector +**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). -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. +### Kernel Fusion -`@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 are returned to cuNumeric's pool and recycled by the next same-sized allocation, skipping new buffer allocation. +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. -This macro can improve runtime and reduce memory overheads. +```julia +y .= @. -a + b * c +``` + +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. + +`@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. ```julia @analyze_lifetimes begin - result = A[1:end, :] .+ B[1:end, :] - C .= result .* 2.0 + result = @. A[1:end, :] + B[1:end, :] + C .= @. result * 2.0f0 end ``` -### Requirements +### 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). + +### Try an example + +```julia +using cuNumeric + +integrand = (x) -> @. exp(-x^2) + +N = 1_000_000 +x_max = 10.0f0 +Ω = 2 * x_max + +samples = Ω .* cuNumeric.rand(N) +samples = samples .- x_max +estimate = (Ω / N) .* sum(integrand(samples)) + +println("Monte-Carlo Estimate: $(estimate)") +``` +More worked examples (initialization, Gray-Scott, …) are in the documentation sidebar under **Examples**. + +### Known Limitations -We require an x86 Linux platform and Julia >=1.10. For GPU support we require an NVIDIA GPU and a CUDA driver which supports CUDA 13.0. ARM support is theoretically possible, but we do not make binaries or test on ARM. Please open an issue if ARM support is of interest. +- There is no support for `Float16` or `ComplexF16` diff --git a/benchmark/Project.toml b/benchmark/Project.toml index 62eb4c276..805db8aad 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -1,5 +1,6 @@ [deps] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -8,5 +9,4 @@ TOML = "fa267f1f-6049-4f14-aa54-33bafae1ed76" cuNumeric = "0fd9ffd4-7e84-4cd0-b8f8-645bd8c73620" [extras] -CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" LegatePreferences = "8028f36a-2b64-49e9-aa04-2d0933fd2ed9" diff --git a/benchmark/README.md b/benchmark/README.md index 753a04166..205cd08ca 100644 --- a/benchmark/README.md +++ b/benchmark/README.md @@ -41,18 +41,24 @@ gpus = 1 cpus = 2 N = 150 M = 150 # optional, defaults to 1 +fusion = true # optional, defaults to true; toggles cuNumeric broadcast fusion ``` Repeat a `[[name]]` block to add independent configs. ## Lists -Any of `T`, `gpus`, `cpus`, `N`, `M` may be a list. They expand along +Any of `T`, `fusion`, `gpus`, `cpus`, `N`, `M` may be a list. They expand along two axes: -- **`T` multiply.** The whole sweep runs once per type. +- **`T` and `fusion` multiply.** The whole sweep runs once per type and once per + fusion setting (`fusion = [true, false]` sweeps both). - **`gpus`, `cpus`, `N`, `M` zip** into a single lockstep sweep — element `i` of each is paired together. +`fusion` toggles cuNumeric broadcast fusion (`true`/`false` or `"on"`/`"off"`, +default `true`); it only affects cuNumeric, so comparison backends run once, not +per variant. + Each zipped field must be one of: - a scalar or single-element list (`cpus = 2` or `[2]`) -> broadcast to every config diff --git a/benchmark/benchmarks.toml b/benchmark/benchmarks.toml index 688ae9e2f..70d5a47a6 100644 --- a/benchmark/benchmarks.toml +++ b/benchmark/benchmarks.toml @@ -4,37 +4,48 @@ n_iter = 1000 n_trial = 5 cupynumeric = true # (needs install_cupynumeric.sh) cuda = false # compare against CUDA.jl (single-GPU configs only) +# One CPU-reference check per config (not per timed iter). Written to CSV. +check_correctness = true +n_correctness_iter = 5 #################################### # GEMM # -# Work ~ 2*N^2*M. Hold N, scale M. # +# Weak scaling, restricted to NxN. # +# Work ~ 2*N^2*M = 2*N^3. # +# N^3 / P --> constant. # +# N = baseline * P^(1/3). # #################################### [[gemm]] -T = ["Float32", "Float64"] +T = ["Float32"] gpus = [1, 2, 4, 8] -cpus = 2 -N = 4096 -M = [4096, 8192, 16384, 32768] +cpus = 16 +N = [20000, 25200, 31752, 40000] +M = [20000, 25200, 31752, 40000] ################################# # Gray-Scott # -# Work ~ N*M. Hold N, scale M. # +# Weak scaling, square NxN grid.# +# Work ~ N*M = N^2. # +# N^2 / P --> constant. # +# N = baseline * P^(1/2). # ################################# [[grayscott_baseline]] T = "Float32" gpus = [1, 2, 4, 8] -cpus = 2 -N = 1024 -M = [1024, 2048, 4096, 8192] +cpus = 16 +fusion = [true, false] +N = [2000, 2832, 4000, 5656] +M = [2000, 2832, 4000, 5656] [[grayscott_lifetimes]] T = "Float32" gpus = [1, 2, 4, 8] -cpus = 2 -N = 1024 -M = [1024, 2048, 4096, 8192] +cpus = 16 +fusion = false +N = [2000, 2832, 4000, 5656] +M = [2000, 2832, 4000, 5656] ################################# # Monte-Carlo Integration # @@ -44,5 +55,5 @@ M = [1024, 2048, 4096, 8192] [[montecarlo]] T = "Float32" gpus = [1, 2, 4, 8] -cpus = 2 +cpus = 16 N = [1_000_000, 2_000_000, 4_000_000, 8_000_000] diff --git a/benchmark/install_cupynumeric.sh b/benchmark/install_cupynumeric.sh index 541a654c1..6140cf64a 100755 --- a/benchmark/install_cupynumeric.sh +++ b/benchmark/install_cupynumeric.sh @@ -52,9 +52,12 @@ fi echo "cupynumeric_jll major.minor: $VER" SPEC="cupynumeric=$VER.*" +# numpy 2.3 dropped the private numpy.linalg.linalg path that cupynumeric 25.10 imports. +NUMPY_SPEC="numpy<2.3" + if [[ -n "$INTO_ENV" ]]; then echo "Installing $SPEC into existing env '$INTO_ENV'..." - conda install -y -n "$INTO_ENV" -c conda-forge -c legate "$SPEC" + conda install -y -n "$INTO_ENV" -c conda-forge -c legate "$SPEC" "$NUMPY_SPEC" echo "Done. Activate with: conda activate $INTO_ENV" exit 0 fi @@ -68,6 +71,6 @@ if conda env list | awk '{print $1}' | grep -qx "$ENV_NAME"; then fi echo "Creating env '$ENV_NAME' with $SPEC..." -conda create -y -n "$ENV_NAME" -c conda-forge -c legate "$SPEC" +conda create -y -n "$ENV_NAME" -c conda-forge -c legate "$SPEC" "$NUMPY_SPEC" echo "Done. Activate with: conda activate $ENV_NAME" diff --git a/benchmark/run.jl b/benchmark/run.jl index 7d6c3bb4a..574bb4ea7 100644 --- a/benchmark/run.jl +++ b/benchmark/run.jl @@ -2,7 +2,7 @@ # dispatches it; the script sets LEGATE_CONFIG (from --gpus/--cpus) before # launching the worker (single.jl) that actually runs the benchmark. # no args -> one command per benchmarks.toml entry -# with args -> one command from +# with args -> one command from [fusion] # Orchestrator stays off the GPU: it only needs GlobalSettings + parse_config, # both cuNumeric-free. The worker (single.jl) loads cuNumeric and the kernels. @@ -16,14 +16,24 @@ const RUNNER = joinpath(@__DIR__, "run_benchmark.sh") const WORKER = joinpath(@__DIR__, "src/single.jl") const PY_WORKER = joinpath(@__DIR__, "src_py/single.py") +# quiet by default (banner + results only); -v/--verbose shows the plumbing +const VERBOSE_FLAGS = ("-v", "--verbose") +const VERBOSE = any(in(ARGS), VERBOSE_FLAGS) +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") -# ensure things are resolved and devlop'd properly +const LAST_FUSION_TOGGLE = Ref{Union{Nothing,Bool}}(nothing) + +# dev CNPreferences && cuNumeric function ensure_project_ready() - Pkg.develop(; path=joinpath(@__DIR__, "..")) + Pkg.develop([ + Pkg.PackageSpec(; path=joinpath(@__DIR__, "..", "lib", "CNPreferences")), + Pkg.PackageSpec(; path=joinpath(@__DIR__, "..")), + ]) Pkg.instantiate() end @@ -40,21 +50,40 @@ function cupynumeric_env_name() end function dispatch(; gpus, cpus, name, T, N, M, n_iter, n_warmup, n_trial, - cupynumeric=false, cudajl=false) + fusion=true, cupynumeric=false, cudajl=false, + check_correctness=false, n_correctness_iter=5) + fstr = fusion ? "enabled" : "disabled" banner( - "$(name): T=$(T) gpus=$(gpus) cpus=$(cpus) N=$(N) M=$(M) " * + "$(name): T=$(T) gpus=$(gpus) cpus=$(cpus) N=$(N) M=$(M) fusion=$(fstr) " * "n_iter=$(n_iter) n_warmup=$(n_warmup) n_trial=$(n_trial)", ) + # precompile in the orchestrator so the worker loads a warm cache quietly + CNPreferences.set_broadcast_fusion!(fusion) + if LAST_FUSION_TOGGLE[] != fusion + VERBOSE && println("Precompiling cuNumeric (fusion=$(fstr))") + Pkg.precompile("cuNumeric"; io=devnull) + LAST_FUSION_TOGGLE[] = fusion + end + # each backend runs in its own worker process + vflag = VERBOSE ? `--verbose` : `` args = `--gpus $gpus --cpus $cpus $name $T $N $M $n_iter $n_warmup $n_trial` - cmds = [`bash $RUNNER $WORKER $args cunumeric`] - # CUDA.jl is single-GPU only - if cudajl && gpus == 1 && !cunumeric_only(name) - push!(cmds, `bash $RUNNER $WORKER $args cudajl`) - end - if cupynumeric && !cunumeric_only(name) - push!(cmds, `bash $RUNNER $PY_WORKER --pyenv $(cupynumeric_env_name()) $args`) + # trailing: backend check_correctness n_correctness_iter + corr_args = `$check_correctness $n_correctness_iter` + cmds = [`bash $RUNNER $WORKER $vflag $args cunumeric $corr_args`] + + # comparison backends have no fusion knob, so run them once instead of per + # fusion variant; the fused pass (the default) is that single run + run_comparison_backends = fusion + if run_comparison_backends + # CUDA.jl is single-GPU only + if cudajl && gpus == 1 && !cunumeric_only(name) + push!(cmds, `bash $RUNNER $WORKER $vflag $args cudajl $corr_args`) + end + if cupynumeric && !cunumeric_only(name) + push!(cmds, `bash $RUNNER $PY_WORKER $vflag --pyenv $(cupynumeric_env_name()) $args`) + end end for cmd in cmds @@ -76,28 +105,35 @@ function run_all_benchmarks(config="benchmarks.toml") name=spec.name, T=spec.T, N=N, M=M, + fusion=spec.fusion, n_iter=gs.n_iter, n_warmup=gs.n_warmup, n_trial=gs.n_trial, cupynumeric=gs.cupynumeric, cudajl=gs.cuda, + check_correctness=gs.check_correctness, + n_correctness_iter=gs.n_correctness_iter, ) end end ensure_project_ready() -if isempty(ARGS) +using CNPreferences: CNPreferences +if isempty(POSARGS) run_all_benchmarks() else # dispatch on args dispatch(; - gpus=parse(Int, ARGS[1]), - cpus=parse(Int, ARGS[2]), - name=ARGS[3], - T=ARGS[4], - N=parse(Int, ARGS[5]), - M=parse(Int, ARGS[6]), - n_iter=parse(Int, ARGS[7]), - n_warmup=parse(Int, ARGS[8]), - n_trial=parse(Int, ARGS[9]), + gpus=parse(Int, POSARGS[1]), + cpus=parse(Int, POSARGS[2]), + name=POSARGS[3], + T=POSARGS[4], + N=parse(Int, POSARGS[5]), + M=parse(Int, POSARGS[6]), + n_iter=parse(Int, POSARGS[7]), + n_warmup=parse(Int, POSARGS[8]), + n_trial=parse(Int, POSARGS[9]), + fusion=length(POSARGS) >= 10 ? parse_fusion(POSARGS[10]) : true, + check_correctness=length(POSARGS) >= 11 ? parse(Bool, POSARGS[11]) : false, + n_correctness_iter=length(POSARGS) >= 12 ? parse(Int, POSARGS[12]) : 5, ) end diff --git a/benchmark/run_benchmark.sh b/benchmark/run_benchmark.sh index b802f7bce..8fc47aa8d 100755 --- a/benchmark/run_benchmark.sh +++ b/benchmark/run_benchmark.sh @@ -12,6 +12,7 @@ shift GPUS=0 CPUS=1 PYENV="" +VERBOSE=0 while [[ $# -gt 0 ]]; do case $1 in @@ -27,6 +28,10 @@ while [[ $# -gt 0 ]]; do PYENV=$2 shift 2 ;; + --verbose) + VERBOSE=1 + shift + ;; *) # Collect all other arguments as extra arguments EXTRA_ARGS+=("$1") @@ -54,11 +59,11 @@ fi export LEGATE_AUTO_CONFIG=1 export LEGATE_CONFIG="--cpus=$CPUS --gpus=$GPUS" -export LEGATE_SHOW_CONFIG=1 +export LEGATE_SHOW_CONFIG=$VERBOSE export LD_LIBRARY_PATH="" -echo "Running $FILENAME with $CPUS CPUs and $GPUS GPUs" +[[ $VERBOSE == 1 ]] && echo "Running $FILENAME with $CPUS CPUs and $GPUS GPUs" # Python (cupynumeric) workers run in the conda env built by install_cupynumeric.sh; # Julia (cuNumeric) workers run against the local project. @@ -72,5 +77,5 @@ else CMD="julia --project $FILENAME $GPUS ${EXTRA_ARGS[@]}" fi -printf "Running: %s\n" "$CMD" +[[ $VERBOSE == 1 ]] && printf "Running: %s\n" "$CMD" eval "$CMD" diff --git a/benchmark/src/benchmarks/grayscott.jl b/benchmark/src/benchmarks/grayscott.jl index a2d51315a..29f4151c6 100644 --- a/benchmark/src/benchmarks/grayscott.jl +++ b/benchmark/src/benchmarks/grayscott.jl @@ -8,7 +8,7 @@ struct GSParams{T} end function GSParams{T}(; dx=1, c_u=1.0, c_v=0.3, f=0.03, k=0.06) where {T} - GSParams{T}(T(dx), T(dx / 5), T(c_u), T(c_v), T(f), T(k)) + return GSParams{T}(T(dx), T(dx / 5), T(c_u), T(c_v), T(f), T(k)) end abstract type AbstractGrayScott{T} <: AbstractBenchmark{T} end @@ -30,7 +30,7 @@ allowed_types(::Type{AbstractGrayScott}) = cuNumeric.SUPPORTED_FLOAT_TYPES total_flops(b::AbstractGrayScott) = b.N * b.M # grid points updated per step function build_benchmark(::Type{A}, ::Type{T}, N, M) where {A<:AbstractGrayScott,T} - A{T}(; N=N, M=M) + return A{T}(; N=N, M=M) end mutable struct GrayScottState{A,P} @@ -41,7 +41,7 @@ mutable struct GrayScottState{A,P} params::P end -function initialize(b::AbstractGrayScott{T}; mod=cuNumeric) where {T} +function initialize(b::AbstractGrayScott{T}; mod=cuNumeric, deterministic::Bool=false) where {T} d = (b.N, b.M) u = mod.ones(T, d) v = mod.zeros(T, d) @@ -49,56 +49,95 @@ function initialize(b::AbstractGrayScott{T}; mod=cuNumeric) where {T} v_new = mod.zeros(T, d) seed = min(150, b.N, b.M) - u[1:seed, 1:seed] = mod.rand(T, (seed, seed)) - v[1:seed, 1:seed] = mod.rand(T, (seed, seed)) + if deterministic + # Fixed host pattern so CPU and GPU (any GPU count) share the same IC. + # Avoids Random streams differing across array backends. + host_u = T[ + T(0.5) + T(0.5) * sin(T(i)) * cos(T(j)) for i in 1:seed, j in 1:seed + ] + host_v = T[ + T(0.25) + T(0.25) * cos(T(i)) * sin(T(j)) for i in 1:seed, j in 1:seed + ] + u[1:seed, 1:seed] = mod === cuNumeric ? NDArray(host_u) : host_u + v[1:seed, 1:seed] = mod === cuNumeric ? NDArray(host_v) : host_v + else + u[1:seed, 1:seed] = mod.rand(T, (seed, seed)) + v[1:seed, 1:seed] = mod.rand(T, (seed, seed)) + end return (GrayScottState(u, v, u_new, v_new, GSParams{T}()),) end +correctness_supported(::AbstractGrayScott) = true + +function check_benchmark_correctness( + b::AbstractGrayScott{T}, gs::GlobalSettings; mod=cuNumeric, atol=1e-4, rtol=1e-4 +) where {T} + # CPU reference compares via cuNumeric.compare (scalar gather). Other backends skip. + mod === cuNumeric || return "skipped" + + n = gs.n_correctness_iter + st_gpu = only(initialize(b; mod=mod, deterministic=true)) + st_cpu = only(initialize(b; mod=Base, deterministic=true)) + + for _ in 1:n + run!(b, st_gpu) + run!(b, st_cpu) + end + + # Element-wise NDArray indexing gathers across tiles — do not use Array(NDArray) + # for multi-GPU (get_ptr is local-tile only). + u_ok = @allowscalar cuNumeric.compare(st_cpu.u, st_gpu.u, atol, rtol) + v_ok = @allowscalar cuNumeric.compare(st_cpu.v, st_gpu.v, atol, rtol) + 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. + # 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)] .* + .-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)]) + ) .+ 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)] + ) .- (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[3:end, 2:(end - 1)] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ u[1:(end - 2), 2:(end - 1)] - ) ./ args.dx^2 + + ) ./ args.dx^2 .+ ( - u[2:(end - 1), 3:end] - 2 * u[2:(end - 1), 2:(end - 1)] + + 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[3:end, 2:(end - 1)] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+ v[1:(end - 2), 2:(end - 1)] - ) ./ args.dx^2 + + ) ./ args.dx^2 .+ ( - v[2:(end - 1), 3:end] - 2 * v[2:(end - 1), 2:(end - 1)] + + 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)] + ((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)] + ((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] diff --git a/benchmark/src/core.jl b/benchmark/src/core.jl index db452c6ec..5d9b3abbb 100644 --- a/benchmark/src/core.jl +++ b/benchmark/src/core.jl @@ -11,6 +11,9 @@ using Statistics standard deviations/errors. - `n_gpu::Int` : The number of GPUs used by legate. Set through the LEGATE_CONFIG, this value is just bookkeeping. +- `check_correctness::Bool` : If true, run one CPU-reference check per config + (not per timed iteration) before timing; result is recorded in the CSV. +- `n_correctness_iter::Int` : Steps to run for that single correctness check. """ Base.@kwdef struct GlobalSettings n_warmup::Int # Number of warmup steps, where timing is not done. @@ -19,6 +22,8 @@ Base.@kwdef struct GlobalSettings n_gpu::Int = 0 cupynumeric::Bool = false # also run baselines under cupynumeric for comparison cuda::Bool = false # also run under CUDA.jl for comparison (single-GPU only) + check_correctness::Bool = false + n_correctness_iter::Int = 5 end ######################################### @@ -49,10 +54,19 @@ end # Per-trial timings for one benchmark. `times_ms[i]`/`gflops[i]` are the mean # over `n_iter` iterations for trial `i`; the spread across trials gives stddev. +# `correctness` is one of "pass", "fail", "skipped" — checked once per config. struct BenchmarkResult{B<:AbstractBenchmark} times_ms::Vector{Float64} gflops::Vector{Float64} benchmark::B + correctness::String +end + +# Optional per-benchmark correctness vs a CPU/`Array` reference. +# Return "pass", "fail", or "skipped". Default: no check implemented. +correctness_supported(::AbstractBenchmark) = false +function check_benchmark_correctness(b::AbstractBenchmark, gs::GlobalSettings; mod=cuNumeric) + return "skipped" end # One timed trial: warmup, then time `n_iter` iterations of `run!`. @@ -75,7 +89,17 @@ function _trial(b::AbstractBenchmark, gs::GlobalSettings; mod=cuNumeric) end # Run `n_trial` independent trials and collect their per-trial measurements. +# Correctness (if enabled) runs once before timing, not per trial/iteration. function run_benchmark(b::AbstractBenchmark, gs::GlobalSettings; mod=cuNumeric) + correctness = "skipped" + if gs.check_correctness + if correctness_supported(b) + correctness = check_benchmark_correctness(b, gs; mod=mod) + else + correctness = "skipped" + end + end + times_ms = Float64[] gflops = Float64[] for _ in 1:gs.n_trial @@ -83,7 +107,7 @@ function run_benchmark(b::AbstractBenchmark, gs::GlobalSettings; mod=cuNumeric) push!(times_ms, t) push!(gflops, g) end - return BenchmarkResult(times_ms, gflops, b) + return BenchmarkResult(times_ms, gflops, b, correctness) end _std(x) = length(x) > 1 ? std(x) : 0.0 @@ -94,10 +118,11 @@ function save_result(br::BenchmarkResult, gpus; mod::String="cunumeric") mkpath(dirname(path)) open(path, "a") do io for trial in eachindex(br.times_ms) + # correctness is per-config; repeated on each trial row for CSV joins @printf( - io, "%s,%d,%d,%d,%d,%.6f,%.6f\n", + io, "%s,%d,%d,%d,%d,%.6f,%.6f,%s\n", mod, gpus, N, M, trial, - br.times_ms[trial], br.gflops[trial], + br.times_ms[trial], br.gflops[trial], br.correctness, ) end end @@ -126,5 +151,5 @@ end # end # register_variant("baseline") -# register_variant("fusion_off", cuNumeric.CNPreferences.disable_broadcast_fusion!) -# register_variant("fusion_on", cuNumeric.CNPreferences.enable_broadcast_fusion!) +# register_variant("fusion_off", cuNumeric.disable_broadcast_fusion!) +# register_variant("fusion_on", cuNumeric.enable_broadcast_fusion!) diff --git a/benchmark/src/parse_benchmarks.jl b/benchmark/src/parse_benchmarks.jl index 605c5002b..cff2eba45 100644 --- a/benchmark/src/parse_benchmarks.jl +++ b/benchmark/src/parse_benchmarks.jl @@ -10,12 +10,22 @@ struct BenchmarkSpec T::String gpus::Int cpus::Int + fusion::Bool args::Vector{Int} end # A field may be a scalar or a list. aslist(x) = x isa AbstractVector ? collect(x) : [x] +# `fusion` accepts a bool or "on"/"off" (or a list of these). +function parse_fusion(x) + x isa Bool && return x + s = lowercase(string(x)) + s in ("on", "true") && return true + s in ("off", "false") && return false + error("fusion must be on/off (or true/false); got $(repr(x))") +end + # Value of a zipped field for sweep position `i`. length==1 field broadcasts. sweep_value(field, i) = length(field) == 1 ? field[1] : field[i] @@ -52,6 +62,8 @@ function parse_config(path) n_warmup=g["n_warmup"], n_iter=g["n_iter"], n_trial=get(g, "n_trial", 1), cupynumeric=get(g, "cupynumeric", false), cuda=get(g, "cuda", false), + check_correctness=get(g, "check_correctness", false), + n_correctness_iter=get(g, "n_correctness_iter", 5), ) specs = BenchmarkSpec[] @@ -61,13 +73,13 @@ function parse_config(path) types = aslist(get(e, "T", "Float32")) gpus = aslist(e["gpus"]) cpus = aslist(e["cpus"]) - # fusion = get(e, "fusion", true) + fusion = aslist(get(e, "fusion", true)) N = aslist(e["N"]) M = aslist(get(e, "M", 1)) n = sweep_length(name, ["gpus" => gpus, "cpus" => cpus, "N" => N, "M" => M]) - for T in types, i in 1:n + for T in types, fuse in fusion, i in 1:n push!( specs, BenchmarkSpec( @@ -75,6 +87,7 @@ function parse_config(path) T, sweep_value(gpus, i), sweep_value(cpus, i), + parse_fusion(fuse), [sweep_value(N, i), sweep_value(M, i)], ), ) diff --git a/benchmark/src/single.jl b/benchmark/src/single.jl index 5b2fff549..e37f1d399 100644 --- a/benchmark/src/single.jl +++ b/benchmark/src/single.jl @@ -1,7 +1,9 @@ # single.jl: worker that runs exactly one benchmark under one backend. Launched by # run_benchmark.sh (dispatched from run.jl), which sets LEGATE_CONFIG before julia starts. # Args: +# [check_correctness] [n_correctness_iter] # backend is "cunumeric" or "cudajl"; run.jl launches one worker per backend. +# run.jl sets the compile-time fusion pref before launch; we read it back to label results. using cuNumeric using CUDACore @@ -20,7 +22,10 @@ const BACKENDS = Dict( "cudajl" => (mod=CUDACore, label="CUDA.jl", save_as="CUDA.jl"), ) -function run_single(gpus, name, T_str, N, M, n_iter, n_warmup, n_trial, backend) +function run_single( + gpus, name, T_str, N, M, n_iter, n_warmup, n_trial, backend; + check_correctness=false, n_correctness_iter=5, +) haskey(BENCHMARKS, name) || error( "No benchmark registered for '$(name)'. Known: $(join(sort(collect(keys(BENCHMARKS))), ", "))" ) @@ -28,18 +33,31 @@ function run_single(gpus, name, T_str, N, M, n_iter, n_warmup, n_trial, backend) "Unknown backend '$(backend)'. Known: $(join(sort(collect(keys(BACKENDS))), ", "))" ) bk = BACKENDS[backend] + + # unfused cuNumeric runs land in their own CSV so they stay a distinct series + fused = cuNumeric.FUSE_BROADCAST_EXPRS + save_as = fused ? bk.save_as : "$(bk.save_as)_nofusion" + label = fused ? bk.label : "$(bk.label) (no fusion)" + T = parse_type(T_str) b = build_benchmark(BENCHMARKS[name], T, N, M) - gs = GlobalSettings(; n_warmup=n_warmup, n_iter=n_iter, n_trial=n_trial) + gs = GlobalSettings(; + n_warmup=n_warmup, + n_iter=n_iter, + n_trial=n_trial, + check_correctness=check_correctness, + n_correctness_iter=n_correctness_iter, + ) println( - "[$(bk.label)] $(name) benchmark ($(T)) on $(N)x$(M) for $(n_iter) " * + "[$(label)] $(name) benchmark ($(T)) on $(N)x$(M) for $(n_iter) " * "iterations ($(n_warmup) warmup) x $(n_trial) trials", ) br = run_benchmark(b, gs; mod=bk.mod) - @printf("[%s] Mean Run Time: %.5f ± %.5f ms\n", bk.label, mean(br.times_ms), _std(br.times_ms)) - @printf("[%s] FLOPS: %.5f ± %.5f GFLOPS\n", bk.label, mean(br.gflops), _std(br.gflops)) - save_result(br, gpus; mod=bk.save_as) + @printf("[%s] Mean Run Time: %.5f ± %.5f ms\n", label, mean(br.times_ms), _std(br.times_ms)) + @printf("[%s] FLOPS: %.5f ± %.5f GFLOPS\n", label, mean(br.gflops), _std(br.gflops)) + println("[$(label)] Correctness: $(br.correctness)") + save_result(br, gpus; mod=save_as) end gpus = parse(Int, ARGS[1]) @@ -51,4 +69,9 @@ n_iter = parse(Int, ARGS[6]) n_warmup = parse(Int, ARGS[7]) n_trial = parse(Int, ARGS[8]) backend = ARGS[9] -run_single(gpus, bench_name, T_str, N, M, n_iter, n_warmup, n_trial, backend) +check_correctness = length(ARGS) >= 10 ? parse(Bool, ARGS[10]) : false +n_correctness_iter = length(ARGS) >= 11 ? parse(Int, ARGS[11]) : 5 +run_single( + gpus, bench_name, T_str, N, M, n_iter, n_warmup, n_trial, backend; + check_correctness=check_correctness, n_correctness_iter=n_correctness_iter, +) diff --git a/benchmark/src_py/core.py b/benchmark/src_py/core.py index f32a4fc31..1632e5999 100644 --- a/benchmark/src_py/core.py +++ b/benchmark/src_py/core.py @@ -54,4 +54,4 @@ def save_result(name, dims, gpus, times_ms, gflops): path = os.path.join(RESULTS_DIR, f"{name}_{MOD}.csv") with open(path, "a") as io: for i, (t, g) in enumerate(zip(times_ms, gflops), start=1): - io.write(f"{MOD},{gpus},{N},{M},{i},{t:.6f},{g:.6f}\n") + io.write(f"{MOD},{gpus},{N},{M},{i},{t:.6f},{g:.6f},skipped\n") diff --git a/deps/build.jl b/deps/build.jl index bacfa699e..9bcdd2ad0 100644 --- a/deps/build.jl +++ b/deps/build.jl @@ -21,6 +21,11 @@ using Pkg using Preferences using Legate using CNPreferences +using CUDACore: CUDACore + +# Maybe needed as build deps +using cupynumeric_jll: cupynumeric_jll +using OpenBLAS32_jll: OpenBLAS32_jll const BuildTools = Legate.BuildTools @@ -72,9 +77,17 @@ function build(::CNPreferences.Conda) pkg_root = BuildTools.start_build("cuNumeric.jl", @__DIR__) cupynumeric_root = load_preference(CNPreferences, "cunumeric_conda_env", nothing) + cuda_toolkit_root = load_preference(CNPreferences, "CUDA_TOOLKIT_ROOT", nothing) if isnothing(cupynumeric_root) error("This shouldn't happen. cunumeric_conda_env = nothing?") end + if isnothing(cuda_toolkit_root) + error( + "CUDA_TOOLKIT_ROOT must be set by CNPreferences to point to the CUDA linked in your cupynumeric build." + ) + end + + #!TODO SET LocalPreferences.toml to use local CUDA libraries is_cupynumeric_installed(cupynumeric_root; throw_errors=true) build_deps(pkg_root, cupynumeric_root, cupynumeric_root) diff --git a/docs/Project.toml b/docs/Project.toml index 92966fdf7..289148278 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,15 +1,20 @@ +[compat] +CNPreferences = "0.1.2" +Documenter = "1.5" +cuNumeric = "0.1" + [deps] -cuNumeric = "0fd9ffd4-7e84-4cd0-b8f8-645bd8c73620" CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365" LiveServer = "16fef848-5104-11e9-1b77-fb7a48bbb589" +cuNumeric = "0fd9ffd4-7e84-4cd0-b8f8-645bd8c73620" -[compat] -Documenter = "1.5" -cuNumeric = "0.1" -CNPreferences = "0.1.2" +[extras] +LegatePreferences = "8028f36a-2b64-49e9-aa04-2d0933fd2ed9" + +[sources.CNPreferences] +path = "../lib/CNPreferences" -[sources] -cuNumeric = {path = ".."} -CNPreferences = {path = "../lib/CNPreferences"} +[sources.cuNumeric] +path = ".." diff --git a/docs/make.jl b/docs/make.jl index e1a7f0ddf..bb6f78260 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -1,9 +1,29 @@ using Documenter, DocumenterVitepress using cuNumeric using CNPreferences +using Random +using TOML ci = get(ENV, "CI", "") == "true" +# Read cupynumeric_jll compat from the package Project.toml so conda docs stay in sync. +const CUPYNUMERIC_JLL_COMPAT, CUPYNUMERIC_MAJOR_MINOR = let + project_toml = joinpath(@__DIR__, "..", "Project.toml") + compat = String(TOML.parsefile(project_toml)["compat"]["cupynumeric_jll"]) + m = match(r"(\d+)\.(\d+)", compat) + m === nothing && error("Could not parse major.minor from cupynumeric_jll compat '$compat'") + compat, "$(m.captures[1]).$(m.captures[2])" +end +@info "Docs using cupynumeric_jll compat=$CUPYNUMERIC_JLL_COMPAT (conda line $CUPYNUMERIC_MAJOR_MINOR)" + +# DocumenterVitepress hardcodes nested sidebar groups as collapsed: false. +# Prefer closed-by-default dropdowns in the left nav. +@eval DocumenterVitepress function pagelist2str(doc, name_contents::Pair{String,<:AbstractVector}) + name, contents = name_contents + rendered_contents = pagelist2str.((doc,), contents) + return "{ text: '$(replace(name, "'" => "\\'"))', collapsed: true, items: [\n$(join(rendered_contents, ",\n"))]\n }" +end + makedocs(; sitename="cuNumeric.jl", authors="Ethan Meitz and David Krasowska", @@ -14,11 +34,41 @@ makedocs(; ), pages=[ "Home" => "index.md", - "Examples" => "examples.md", - "Performance Tips" => "perf.md", - "Custom Installalation" => "install.md", - "Benchmarks" => "benchmark.md", - "Public API" => "api.md", + "Examples" => [ + "Initialization" => "examples/initialization.md", + "Monte-Carlo" => "examples/montecarlo.md", + "Gray-Scott" => "examples/grayscott.md", + "HDF5 I/O" => "examples/hdf5.md", + ], + "Performance Tips" => [ + "Kernel Fusion" => "perf/kernel_fusion.md", + "Reduce Allocations" => "perf/reduce_allocations.md", + "Patterns to Avoid" => "perf/patterns_to_avoid.md", + ], + "Configuration" => [ + "Hardware" => "configuration/hardware.md", + "Build Modes" => "install.md", + "CNPreferences" => "api_preferences.md", + ], + "Benchmarks" => [ + "Results" => "benchmarks/results.md", + "How to Benchmark" => "benchmarks/howto.md", + ], + "Developer" => [ + "Developer Mode" => "developer_mode.md", + "Debugging" => "debugging.md", + "Internals" => "internals.md", + ], + "Public API" => [ + "Initialization" => "api_initialization.md", + "Unary Operations" => "api_unary.md", + "Binary Operations" => "api_binary.md", + "Linear Algebra" => "linalg.md", + "HDF5" => "api_hdf5.md", + "NDArray Reference" => "api.md", + "CUDA.jl Tasking" => "api_cuda.md", + "Internal API" => "api_internal.md", + ], ], ) diff --git a/docs/src/.vitepress/config.mts b/docs/src/.vitepress/config.mts index ee0f97044..fb1fcbc4d 100644 --- a/docs/src/.vitepress/config.mts +++ b/docs/src/.vitepress/config.mts @@ -15,16 +15,32 @@ const baseTemp = { } const navTemp = { + // DocumenterVitepress still fills this token; we intentionally omit page links + // from the top bar (sidebar has the TOC). Keep only the version picker. nav: 'REPLACE_ME_DOCUMENTER_VITEPRESS', } +const sidebarTemp = { + sidebar: 'REPLACE_ME_DOCUMENTER_VITEPRESS', +} + const nav = [ - ...navTemp.nav, { component: 'VersionPicker' } ] +// VitePress packs root-level leaves into anonymous groups (muted level-1 text). +// Attach an empty `items` array so Home stays a bold, clickable level-0 title. +function keepRootLeavesAsGroups(items: T[]): T[] { + return items.map((item) => { + if (item.link && !item.items) { + return { ...item, items: [] } + } + return item + }) +} + // https://vitepress.dev/reference/site-config export default defineConfig({ base: 'REPLACE_ME_DOCUMENTER_VITEPRESS',// TODO: replace this in makedocs! @@ -77,7 +93,9 @@ export default defineConfig({ }, themeConfig: { outline: 'deep', - logo: 'REPLACE_ME_DOCUMENTER_VITEPRESS', + // Keep logo.png in assets/ for the README/home page, but do not show it + // in the VitePress navbar (DocumenterVitepress would inject /logo.png here). + logo: false, search: { provider: 'local', options: { @@ -85,7 +103,7 @@ export default defineConfig({ } }, nav, - sidebar: 'REPLACE_ME_DOCUMENTER_VITEPRESS', + sidebar: keepRootLeavesAsGroups(sidebarTemp.sidebar), editLink: 'REPLACE_ME_DOCUMENTER_VITEPRESS', socialLinks: [ { icon: 'github', link: 'REPLACE_ME_DOCUMENTER_VITEPRESS' } diff --git a/docs/src/.vitepress/theme/style.css b/docs/src/.vitepress/theme/style.css index 936322d13..a4582ff8b 100644 --- a/docs/src/.vitepress/theme/style.css +++ b/docs/src/.vitepress/theme/style.css @@ -173,3 +173,64 @@ mjx-container > svg { margin: auto; display: inline-block; } + +/* Top-level leaf sidebar titles (e.g. Home with empty items[]) match group weight. + VitePress otherwise packs lone root links into muted level-1 entries. */ +.VPSidebarItem.level-0.is-link > .item .text { + font-weight: 700; + color: var(--vp-c-text-1); +} + +/* Home is level-0 but not `.collapsed`, so VitePress gives it 24px bottom padding + vs 10px for collapsed section headers. Match the collapsed spacing. */ +.VPSidebarItem.level-0.is-link { + padding-bottom: 10px; +} + +/* VitePress defaults doc images to display:block, which stacks README badges. */ +.vp-doc p > a > img { + display: inline-block; + vertical-align: middle; +} + +/* Keep the home title logo inline with the heading text (width comes from the + HTML width="50" attribute). */ +.vp-doc h1 img { + display: inline-block; + vertical-align: middle; + margin: 0 0.25em 0 0; +} + +/* Pull search to the right, beside the version picker (menu), and give it a + real width. VitePress sets flex-grow:1 on search, which parks a tiny button + on the left of the content bar. */ +@media (min-width: 768px) { + .VPNavBarSearch { + flex-grow: 0 !important; + margin-left: auto; + padding-left: 0 !important; + min-width: 16rem; + max-width: 22rem; + } + + .VPNavBarSearch .DocSearch-Button { + width: 100%; + min-width: 16rem; + justify-content: flex-start; + } +} + +@media (min-width: 960px) { + .VPNavBarSearch { + min-width: 18rem; + max-width: 24rem; + } + + .VPNavBarSearch .DocSearch-Button { + min-width: 18rem; + } +} + +.VPNavBar .search + .menu { + margin-left: 0.5rem; +} diff --git a/docs/src/api.md b/docs/src/api.md index ba651bb4f..fd7ba459e 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -1,118 +1,9 @@ +# NDArray Reference -# Public API - -User facing functions supported by cuNumeric.jl - -```@contents -Pages = ["api.md"] -Depth = 2:2 -``` - -### Supported Unary Operations -The following unary operations are supported and can be broadcast over `NDArray`: - - • `-`, `!`, `abs`, `acos`, `acosh`, `asin`, `asinh`, `atan`, `atanh`, `cbrt`, `conj`, `cos`, `cosh`, `deg2rad`, `exp`, `exp2`, `expm1`, `floor`, `imag`, `isfinite`, `log`, `log10`, `log1p`, `log2`, `rad2deg`, `real`, `sign`, `signbit`, `sin`, `sinh`, `sqrt`, `tan`, `tanh`, `^2`, `^-1` or `inv`, - -##### Differences -- The `acosh` function in Julia will error on inputs outside of the domain (x >= 1) - but cuNumeric.jl will return NaN. - - - -### Supported Binary Operations -The following binary operations are supported and can be applied elementwise to pairs of `NDArray` values: - - • `+`, `-`, `*`, `/`, `^`, `<`, `<=`, `>`, `>=`, `==`, `!=`, `atan`, `hypot`, `max`, `min`, `lcm`, `gcd` - -These operations are applied elementwise by default and follow standard Julia semantics. - - -```@autodocs -Modules = [cuNumeric] -Pages = ["ndarray/ndarray.jl", "ndarray/unary.jl", "ndarray/binary.jl", "cuNumeric.jl", "warnings.jl", "util.jl", "memory.jl", "scoping.jl"] -``` - -# CNPreferences - -This section details how to set custom build configuration options. To see more details visit our install guide [here](./install.md). - -```@autodocs -Modules = [CNPreferences] -Pages = ["CNPreferences.jl"] -``` - -# CUDA.jl Tasking - -Write custom GPU kernels in Julia using CUDA.jl and execute them through the Legate distributed runtime. Your kernels automatically benefit from Legate's data partitioning, dependency tracking, and multi-GPU scheduling. - -!!! warning "Experimental Feature" - CUDA.jl tasking is experimental. You must opt in before using `@cuda_task` or `@launch`: - ```julia - cuNumeric.Experimental(true) - ``` - -The interface has two steps: -1. **Compile & Register** — [`@cuda_task`](@ref) JIT-compiles a kernel to PTX and registers it with Legate. -2. **Launch** — [`@launch`](@ref) submits the kernel with grid dimensions, inputs, outputs, and scalars. - -`NDArray` arguments are automatically mapped to their CUDA equivalents (`NDArray{T,1}` → `CuDeviceVector{T,1}`, etc.). Scalar arguments are passed through by copy. - -!!! warning "Inputs vs. outputs" - Correctly separating `inputs` and `outputs` is critical for Legate's - dependency analysis. If an array is both read and written, list it as an `output`. - -!!! warning "Array sizes" - Mismatched array sizes are automatically padded to the largest shape. To address this, we plan to add support for other Legate constraints in the future (more information [here](https://docs.nvidia.com/legate/latest/api/cpp/generated/group/group__partitioning.html)). - -## Example - -```julia -using cuNumeric -using CUDA -import CUDA: i32 - -# Enable experimental features -cuNumeric.Experimental(true) - -# 1. Write a standard CUDA.jl kernel -function kernel_sin(a, b, N) - i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x - if i <= N - @inbounds b[i] = sin(a[i]) - end - return nothing -end - -N = 1024 -threads = 256 -blocks = cld(N, threads) - -a = cuNumeric.fill(1.0f0, N) -b = cuNumeric.zeros(Float32, N) - -# 2. Compile & register — args are used only for type inference -task = cuNumeric.@cuda_task kernel_sin(a, b, UInt32(1)) - -# 3. Launch through Legate -cuNumeric.@launch task=task threads=threads blocks=blocks inputs=a outputs=b scalars=UInt32(N) - -allowscalar() do - println("sin(1) = ", b[:][1]) # ≈ 0.8414709 -end -``` - -See `examples/custom_cuda.jl` for a more complete example with multiple kernels. - -## API Reference - -```@autodocs -Modules = [cuNumeric] -Pages = ["utilities/cuda_stubs.jl"] -``` - -# Internal API +Indexing, reshaping, reductions, comparisons, memory helpers, lifetime macros, and related utilities. For constructors (`zeros`, `ones`, `rand`, …) see [Initialization](./api_initialization.md). ```@autodocs Modules = [cuNumeric] -Pages = ["ndarray/detail/ndarray.jl"] +Pages = ["ndarray/ndarray.jl", "ndarray/linalg.jl", "cuNumeric.jl", "warnings.jl", "util.jl", "memory.jl", "scoping.jl"] +Filter = t -> !(t isa Function && nameof(t) in (:zeros, :ones, :fill, :trues, :falses, :eye, :rand, :rand!)) ``` diff --git a/docs/src/api_binary.md b/docs/src/api_binary.md new file mode 100644 index 000000000..7285ae15e --- /dev/null +++ b/docs/src/api_binary.md @@ -0,0 +1,15 @@ +# Binary Operations + +>[!NOTE] +> Prefer `@.` for multi-op elementwise expressions so every operator is dotted. Missing a dot silently changes the meaning and can prevent broadcast fusion. See [Kernel Fusion](./perf/kernel_fusion.md). + + +The following binary operations are supported and can be applied elementwise to pairs of `NDArray` values: + +- `+`, `-`, `*`, `/`, `^`, `<`, `<=`, `>`, `>=`, `==`, `!=`, `atan`, `hypot`, `max`, `min`, `lcm`, `gcd` + +```@autodocs +Modules = [cuNumeric] +Pages = ["ndarray/binary.jl"] +Filter = t -> !(t isa Function && nameof(t) === :mul!) +``` diff --git a/docs/src/api_cuda.md b/docs/src/api_cuda.md new file mode 100644 index 000000000..6e912cd8a --- /dev/null +++ b/docs/src/api_cuda.md @@ -0,0 +1,68 @@ +# CUDA.jl Tasking + +Write custom GPU kernels in Julia using CUDA.jl and execute them through the Legate distributed runtime. Your kernels automatically benefit from Legate's data partitioning, dependency tracking, and multi-GPU scheduling. + +!!! warning "Experimental Feature" + CUDA.jl tasking is experimental. You must opt in before using `@cuda_task` or `@launch`: + ```julia + cuNumeric.Experimental(true) + ``` + +The interface has two steps: +1. **Compile & Register** - `@cuda_task` JIT-compiles a kernel to PTX and registers it with Legate. +2. **Launch** - `@launch` submits the kernel with grid dimensions, inputs, outputs, and scalars. + +`NDArray` arguments are automatically mapped to their CUDA equivalents (`NDArray{T,1}` → `CuDeviceVector{T,1}`, etc.). Scalar arguments are passed through by copy. + +!!! warning "Inputs vs. outputs" + Correctly separating `inputs` and `outputs` is critical for Legate's + dependency analysis. If an array is both read and written, list it as an `output`. + +!!! warning "Array sizes" + Mismatched array sizes are automatically padded to the largest shape. To address this, we plan to add support for other Legate constraints in the future (more information [here](https://docs.nvidia.com/legate/latest/api/cpp/generated/group/group__partitioning.html)). + +## Example + +```julia +using cuNumeric +using CUDA +import CUDA: i32 + +# Enable experimental features +cuNumeric.Experimental(true) + +# 1. Write a standard CUDA.jl kernel +function kernel_sin(a, b, N) + i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x + if i <= N + @inbounds b[i] = sin(a[i]) + end + return nothing +end + +N = 1024 +threads = 256 +blocks = cld(N, threads) + +a = cuNumeric.fill(1.0f0, N) +b = cuNumeric.zeros(Float32, N) + +# 2. Compile and register (args are used only for type inference) +task = cuNumeric.@cuda_task kernel_sin(a, b, UInt32(1)) + +# 3. Launch through Legate +cuNumeric.@launch task=task threads=threads blocks=blocks inputs=a outputs=b scalars=UInt32(N) + +allowscalar() do + println("sin(1) = ", b[:][1]) # ≈ 0.8414709 +end +``` + +See `examples/custom_cuda.jl` for a more complete example with multiple kernels. + +## API Reference + +```@autodocs +Modules = [cuNumeric] +Pages = ["cuda/cuda_ptx_task.jl"] +``` diff --git a/docs/src/api_hdf5.md b/docs/src/api_hdf5.md new file mode 100644 index 000000000..27b97f3f1 --- /dev/null +++ b/docs/src/api_hdf5.md @@ -0,0 +1,29 @@ +# HDF5 + +> [!NOTE] +> HDF5 support is planned. Signatures below are placeholders and will be replaced with `@docs` blocks when the API is implemented. + +I/O helpers for reading and writing `NDArray`s via HDF5. Prefer these over host-side gather + HDF5.jl when arrays are large or distributed. + +## h5read + +```julia +# Planned: +# cuNumeric.h5read(path, dataset) -> NDArray +``` + +Load a dataset from an HDF5 file into an `NDArray`. + +## h5write + +```julia +# Planned: +# cuNumeric.h5write(path, dataset, arr::NDArray) +``` + +Write an `NDArray` to an HDF5 dataset. + +## Related + +- Example sketch: [HDF5 I/O](./examples/hdf5.md) +- Host conversion when you must leave the runtime: `Array(arr)` (see [NDArray Reference](./api.md)) diff --git a/docs/src/api_initialization.md b/docs/src/api_initialization.md new file mode 100644 index 000000000..b6b08ef05 --- /dev/null +++ b/docs/src/api_initialization.md @@ -0,0 +1,53 @@ +# Initialization + +Constructors for new `NDArray`s. Default floating-point type is `Float32`. + +## zeros + +```@docs +cuNumeric.zeros +``` + +## ones + +```@docs +cuNumeric.ones +``` + +## fill + +```@docs +cuNumeric.fill +``` + +## trues + +```@docs +cuNumeric.trues +``` + +## falses + +```@docs +cuNumeric.falses +``` + +## eye + +```@docs +cuNumeric.eye +``` + +## rand + +```@docs +cuNumeric.rand +``` + +## rand! + +```@docs +Random.rand!(::NDArray{Float64}) +``` + +The backend currently draws `Float64` uniforms. `cuNumeric.rand(Float32, dims...)` converts for you. `rand!` on `NDArray` currently requires `Float64` storage. diff --git a/docs/src/api_internal.md b/docs/src/api_internal.md new file mode 100644 index 000000000..6d554377c --- /dev/null +++ b/docs/src/api_internal.md @@ -0,0 +1,6 @@ +# Internal API + +```@autodocs +Modules = [cuNumeric] +Pages = ["ndarray/detail/ndarray.jl"] +``` diff --git a/docs/src/api_preferences.md b/docs/src/api_preferences.md new file mode 100644 index 000000000..12892687a --- /dev/null +++ b/docs/src/api_preferences.md @@ -0,0 +1,67 @@ +# CNPreferences + +Function reference for [`CNPreferences`](https://github.com/JuliaLegate/cuNumeric.jl/tree/main/lib/CNPreferences). Preferences write `LocalPreferences.toml` and generally require a **fresh Julia process**. + +Out of the box (no `LocalPreferences.toml` changes): + +| Setting | Default | +|---|---| +| Binary / build mode | **JLL** prebuilt binaries | +| Broadcast fusion | **on** | +| `FUSE_BROADCAST_MIN_OPS` | **2** (single-op broadcasts stay unfused) | +| Task scope names | **off** | + +Build-mode setup (JLL / conda / developer) is documented under [Build Modes](./install.md). Fusion usage tips live under [Kernel Fusion](./perf/kernel_fusion.md). + +## Build mode + +```@docs +CNPreferences.use_jll_binary +CNPreferences.use_conda +CNPreferences.use_developer_mode +``` + +## Broadcast fusion + +Defaults: fusion **on**, `FUSE_BROADCAST_MIN_OPS == 2`. + +```julia +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 +``` + +`set_broadcast_fusion_min_ops!` counts `Broadcasted` nodes (ops) in the tree: + +- **`2` (default):** fuse multi-op trees such as `y .= @. a * b + c`. Single-ops like `y .= cos.(x)` stay on the unfused C-API path. +- **`1`:** fuse every eligible expression, including single-ops. + +Set the preference in one Julia process, then start a fresh process to use it. + +```@docs +CNPreferences.set_broadcast_fusion! +CNPreferences.enable_broadcast_fusion! +CNPreferences.disable_broadcast_fusion! +CNPreferences.set_broadcast_fusion_min_ops! +``` + +## Task scope names + +Default: **off**. Optional Legate task-scope naming for debugging. When on, cuNumeric wraps many ops in `Legate.with_scope` so provenance strings (for example `matmul`, `zeros`, or fused `broadcast.`) appear in Legate logs and profiles. Pair this with `--logging legate=debug --log-to-file` (or `--profile`) in `LEGATE_CONFIG`; see [Debugging](./debugging.md#trace-legate-runtime-work). + +```julia +using CNPreferences +CNPreferences.enable_task_scope_names!() +CNPreferences.disable_task_scope_names!() # default +``` + +Restart Julia after changing this preference (it is compile-time in cuNumeric.jl). + +```@docs +CNPreferences.set_task_scope_names! +CNPreferences.enable_task_scope_names! +CNPreferences.disable_task_scope_names! +``` diff --git a/docs/src/api_unary.md b/docs/src/api_unary.md new file mode 100644 index 000000000..e2099a911 --- /dev/null +++ b/docs/src/api_unary.md @@ -0,0 +1,18 @@ +# 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). + + +The following unary operations are supported and can be broadcast over `NDArray`: + +- `-`, `!`, `abs`, `acos`, `acosh`, `asin`, `asinh`, `atan`, `atanh`, `cbrt`, `conj`, `cos`, `cosh`, `deg2rad`, `exp`, `exp2`, `expm1`, `floor`, `imag`, `isfinite`, `log`, `log10`, `log1p`, `log2`, `rad2deg`, `real`, `sign`, `signbit`, `sin`, `sinh`, `sqrt`, `tan`, `tanh`, `^2`, `^-1` or `inv` + +## Differences from Base Julia + +- The `acosh` function in Julia will error on inputs outside of the domain (`x >= 1`), but cuNumeric.jl will return `NaN`. + +```@autodocs +Modules = [cuNumeric] +Pages = ["ndarray/unary.jl"] +``` diff --git a/docs/src/assets/logo.png b/docs/src/assets/logo.png new file mode 100644 index 000000000..2b4668b4a Binary files /dev/null and b/docs/src/assets/logo.png differ diff --git a/docs/src/benchmark.md b/docs/src/benchmark.md deleted file mode 100644 index ac082766d..000000000 --- a/docs/src/benchmark.md +++ /dev/null @@ -1,174 +0,0 @@ -# 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! - - -```@contents -Pages = ["benchmark_results.md"] -Depth = 2:2 -``` - -## SGEMM - -Code Outline: -```julia -mul!(C, A, B) -``` - -```@raw html - - - - - - - - - -
GEMM EfficiencyGEMM GFLOPS
GEMM EfficiencyGEMM GFLOPS
-``` - -## 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. - -Code Outline: -```julia -integrand = (x) -> exp.(-x.^2) -val = (V/N) * sum(integrand(x)) -``` - -```@raw html - - - - - - - - - -
MC EfficiencyMC GFLOPS
MC EfficiencyMC GFLOPS
-``` - - -## Gray-Scott (2D) - -Solving a PDE requires halo-exchanges and lots of data movement. In this benchmark we fall an order of magnitude short of the `ImplicitGlobalGrid.jl` library which specifically targets multi-node, multi-GPU halo exchanges. We attribute this to the lack of kernel fusion in cuNumeric.jl - -```@raw html - - - - - - - -
GS GFLOPS
GS GFLOPS
-``` - - -# Benchmarking cuNumeric.jl Programs - -Since there is no programatic way to set the hardware configuration (as of 24.11) benchmarking cuNumeric.jl code is a bit tedious. As an introduction, we walk through a benchmark of matrix multiplication (SGEMM). All the code for this benchmark can be found in the `cuNumeric.jl/pkg/benchmark` directory. - -> [!WARNING] -> We do not commit to maintaining the benchmark scripts, due to difficulty programatically configuring legate and API overturn as we work on cuNumeric v1.0. The general principles used should work, even if the code does not. - - -## Weak Scaling of SGEMM - -In this benchmark we will try to understand the weak scaling behavior of the SGEMM kernel (Float32 MatMul). To get started we need to decide our initial problem size, `N` and create some arrays. Note depending on the choice of `N`, Legate may decide to schedule your tasks on the CPU or GPU or both. We will also define two functions, `total_flops` and `total_space` which will help us calculate some useful benchmark metrics. - -```julia -using cuNumeric - -function initialize_cunumeric(N, M) - A = cuNumeric.rand(Float32, N, M) - B = cuNumeric.rand(Float32, M, N) - C = cuNumeric.zeros(Float32, N, N) - GC.gc() # remove any intermediate arrays - return A, B, C -end - -function total_flops(N, M) - return N * N * ((2*M) - 1) -end - -function total_space(N, M) - return 2 * (N*M) * sizeof(Float32) + (N*N) * sizeof(Float32) -end -``` - -We cannot use rely on common benchmark tools in Julia like [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl) or [ChairMarks.jl](https://github.com/LilithHafner/Chairmarks.jl) or even the built in `Base.@time` macro. The asynchronous nature of operations on NDArrays means that function calls will execute almost immediately and program execution must be blocked to properly time a kernel. It is technically possible to time NDArray operations with something like [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl) by adding a blocking operation (e.g., accessing the result), but the allocations reported by these tools will never be correct and it is safer to use the timing functionality from CuPyNumeric. - -The timer built into CuPyNumeric blocks execution until all Legate operations preceding the call that generated the timing object complete. We provide two timing utilities: `get_time_microseconds` and `get_time_nanoseconds`. - -Now we can write the benchmark code. There are two more parameters we need to set: the number of samples, `n_samples`, and the number of warm-up samples, `n_warnup`. With all this the benchmark loop is: - -```julia -using LinearAlgebra -using cuNumeric - -function gemm_cunumeric(N, M, n_samples, n_warmup) - A, B, C = initialize_cunumeric(N, M) - - start_time = nothing - for idx in range(1, n_samples + n_warmup) - if idx == n_warmup + 1 - start_time = get_time_microseconds() - end - - mul!(C, A, B) - end - total_time_μs = get_time_microseconds() - start_time - mean_time_ms = total_time_μs / (n_samples * 1e3) - gflops = total_flops(N, M) / (mean_time_ms * 1e6) # GFLOP is 1e9 - - return mean_time_ms, gflops -end - -N = 100 -n_samples = 10 -n_warmup = 2 - -mean_time_ms, gflops = gemm_cunumeric(N, N, n_samples, n_warmup) -``` - -Since there is no programatic way to set the hardware configuration we must manipulate the environment variables described in [Setting Hardware Configuration](@ref) through shell scripts to make a weak scaling plot. These variables must be set before we launch the Julia runtime where we will run our benchmark. Therefore, I do not recommend generating scaling plots from the REPL because you would have to start and stop the REPL each time to re-configure the hardware settings. To make benchmarking easier, we provide a small shell script, `run_benchmark.sh`, located in `cuNumeric.jl/pkg/benchmark`. This script will automatically set the `LEGATE_CONFIG` according to the passed flags and run the specified benchmark file. - -The first time using this script it is important to set the memory configuration which is currently hard coded in `run_benchmark.sh` where `LEGATE_CONFIG` is exported. We recommend using the memory configuration which is automatically chosen by Legate. It is also important to note that `run_benchmark.sh` assumes the environment containing cuNumeric.jl is at the relative path `..` as hardcoded in the variable `CMD` (this is correct if you did not move files around). If your cuNumeric.jl environment is different be sure to update the path in `run_benchmark.sh`. - -We have placed the SGEMM example from above into a Julia file called `sgemm.jl` located in the `cuNumeric.jl/pkg/benchmark` directory. We also added basic argument parsing to get settings from the command line so that we do not have to open the file to edit the settings each time. -```julia -gpus = parse(Int, ARGS[1]) -N = parse(Int, ARGS[2]) -M = parse(Int, ARGS[3]) -n_samples = parse(Int, ARGS[4]) -n_warmup = parse(Int, ARGS[5]) -``` - -To run the benchmark simply run the following command with your settings. Be sure to update the legate config in this file as described above! -```sh -./run_benchmark.sh sgemm.jl --cpus --gpus -``` - -Successful completion of one run should look like: - -```bash -./run_benchmark.sh sgemm.jl --cpus 1 --gpus 1 10000 10 2 -Running sgemm.jl with 1 CPUs and 1 GPUs -Running: julia --project='..' sgemm.jl 10000 10 2 -[ Info: Starting Legate -Legate hardware configuration: --cpus=1 --gpus=1 --omps=1 --ompthreads=3 --utility=2 --sysmem=256 --numamem=19029 --fbmem=7569 --zcmem=128 --regmem=0 -[ Info: Started Legate successfully -[ Info: Running MATMUL benchmark on 10000x10000 matricies for 10 iterations, 2 warmups -cuNumeric Mean Run Time: 310.1302 ms -cuNumeric FLOPS: 6448.581918175012 GFLOPS -[ Info: Cleaning Up Legate -``` - -To generate a weak scaling plot, you must increment the problem size in proportion to the number of GPUs. This helps reveal any communication overhead in our SGEMM implementation since data may be transfered between GPUs or even across nodes in a server. - - -As part of a more complete benchmark we ran our code on up to 8 A100 GPUs (single-node) and compared it to the Python library cuPyNumeric as well as a custom implementation using CUDA.jl. From these resutls we can see that cuNumeric.jl is capable of scaling and saturating the GPU memory bandwidth for matrix multiplication. diff --git a/docs/src/benchmarks/howto.md b/docs/src/benchmarks/howto.md new file mode 100644 index 000000000..6e6e8e6e4 --- /dev/null +++ b/docs/src/benchmarks/howto.md @@ -0,0 +1,118 @@ +# How to Benchmark + +The benchmark harness lives in `benchmark/` at the repo root. Configs are declared in `benchmarks.toml`. `run.jl` expands those configs and launches one worker process per run. Workers never share a GPU runtime within a measurement. + +> [!WARNING] +> We do not commit to maintaining the benchmark scripts forever. The harness evolves with the package. The ideas here (declare configs in TOML, one process per run, time with Legate fences) should still apply even if file names move. + +## Why not BenchmarkTools.jl? + +cuNumeric ops are asynchronous. A Julia call usually returns before the GPU work finishes, so tools like [BenchmarkTools.jl](https://github.com/JuliaCI/BenchmarkTools.jl), [Chairmarks.jl](https://github.com/LilithHafner/Chairmarks.jl), or even a bare `Base.@time` will under-report time unless you force a fence. The harness uses `get_time_microseconds` / `get_time_nanoseconds`, which block until preceding Legate work completes. Allocations reported by BenchmarkTools will also not reflect Legion / CUDA buffers. + +## Quick start + +From the repo: + +```bash +cd benchmark +julia --project=. run.jl +``` + +With no extra args, `run.jl` reads `benchmarks.toml` and runs every expanded config. It develops `CNPreferences` and `cuNumeric` from the parent checkout, then for each config: + +1. Sets the broadcast-fusion preference if needed and precompiles +2. Calls `run_benchmark.sh`, which exports `LEGATE_CONFIG` from `--gpus` / `--cpus` **before** Julia starts +3. Launches `src/single.jl` for the cuNumeric backend (and optional comparison backends) + +Add `-v` / `--verbose` for more plumbing output. + +## `benchmarks.toml` + +`[Global]` sets defaults shared by every run: + +```toml +[Global] +n_warmup = 5 +n_iter = 1000 +n_trial = 5 +cupynumeric = true # also run Python cupynumeric (needs install_cupynumeric.sh) +cuda = false # also run CUDA.jl (single-GPU configs only) +check_correctness = true +n_correctness_iter = 5 +``` + +- `n_warmup`: untimed iterations (hide compile / first-touch cost) +- `n_iter`: timed iterations per trial (build task queue depth) +- `n_trial`: independent trials; mean ± stddev across trials is what gets printed / saved +- `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`, `grayscott_baseline`, `grayscott_lifetimes`, …). Names must match what `src/benchmarks/*.jl` registers. + +```toml +[[gemm]] +T = ["Float32"] +gpus = [1, 2, 4, 8] +cpus = 16 +N = [20000, 25200, 31752, 40000] +M = [20000, 25200, 31752, 40000] +``` + +### How lists expand + +Any of `T`, `fusion`, `gpus`, `cpus`, `N`, `M` may be a scalar or a list. + +- **`T` and `fusion` multiply.** The sweep runs once per type and once per fusion setting (`fusion = [true, false]` sweeps both). +- **`gpus`, `cpus`, `N`, `M` zip** in lockstep. Element `i` of each is paired together. A scalar broadcasts to every position. + +```toml +[[gemm]] +T = ["Float64", "Float32"] # multiplies +gpus = [1, 2, 4] +cpus = 2 # zip -> (1,2,150,150), (2,2,300,300), (4,2,600,600) +N = [150, 300, 600] +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. + +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. + +## One-off runs + +You can dispatch a single config without editing the TOML: + +```bash +julia --project=. run.jl [fusion] +``` + +Example: + +```bash +julia --project=. run.jl 1 16 gemm Float32 20000 20000 1000 5 5 true +``` + +`run.jl` still goes through `run_benchmark.sh` so Legate sees the right GPU/CPU count at process start. + +## Comparison backends + +- **CUDA.jl:** set `cuda = true` in `[Global]`. Only runs when `gpus == 1`. +- **cupynumeric (Python):** set `cupynumeric = true`, then build a matching conda env once: + +```bash +./install_cupynumeric.sh # creates cupynumeric-bench- +``` + +`run.jl` picks the env from the resolved `cupynumeric_jll` version. Override with `CUPYNUMERIC_ENV`. + +## Results and timing + +Each worker prints mean ± stddev run time (ms) and GFLOPS, plus a correctness tag (`pass` / `fail` / `skipped`). CSVs append under `benchmark/results/`. + +Unfused cuNumeric runs are labeled and saved separately (for example `cunumeric_nofusion`) so they stay a distinct series from fused runs. + +## Hardware notes + +`LEGATE_CONFIG` must be set before Julia / Legate starts. The harness does that for you via `run_benchmark.sh`. For manual REPL experiments, see [Hardware Configuration](../configuration/hardware.md). Do not expect to change GPU count mid-session without restarting Julia. diff --git a/docs/src/benchmarks/results.md b/docs/src/benchmarks/results.md new file mode 100644 index 000000000..fd62aedcf --- /dev/null +++ b/docs/src/benchmarks/results.md @@ -0,0 +1,61 @@ +# 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! + +## SGEMM + +Code Outline: +```julia +mul!(C, A, B) +``` + +```@raw html + + + + + + + + + +
GEMM EfficiencyGEMM GFLOPS
GEMM EfficiencyGEMM GFLOPS
+``` + +## 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. + +Code Outline: +```julia +integrand = (x) -> exp.(-x.^2) +val = (V/N) * sum(integrand(x)) +``` + +```@raw html + + + + + + + + + +
MC EfficiencyMC GFLOPS
MC EfficiencyMC GFLOPS
+``` + +## Gray-Scott (2D) + +Solving a PDE requires halo exchanges and lots of data movement. In this benchmark we fall an order of magnitude short of the `ImplicitGlobalGrid.jl` library which specifically targets multi-node, multi-GPU halo exchanges. Broadcast fusion helps on the elementwise update, but communication and stencil data movement still dominate the gap. + +```@raw html + + + + + + + +
GS GFLOPS
GS GFLOPS
+``` diff --git a/docs/src/configuration/hardware.md b/docs/src/configuration/hardware.md new file mode 100644 index 000000000..0102ee183 --- /dev/null +++ b/docs/src/configuration/hardware.md @@ -0,0 +1,15 @@ +# 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_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 +- `LEGATE_CONFIG`: A string representing the hardware configuration to set + +These variables must be set before launching the Julia instance running cuNumeric.jl. We recommend setting `export LEGATE_SHOW_CONFIG=1` so that the hardware configuration will be printed when Legate starts. This output is automatically captured and relayed to the user. + +To manually set the hardware configuration, `export LEGATE_AUTO_CONFIG=0`, and then define your own config with something like `export LEGATE_CONFIG="--gpus 1 --cpus 10"`. We recommend using the default memory configuration for your machine and only setting the `gpus`, `cpus`. More details about the Legate configuration can be found in the [NVIDIA Legate documentation](https://docs.nvidia.com/legate/latest/usage.html#resource-allocation). + +The same `LEGATE_CONFIG` string can carry logging / profiling flags (for example `--logging legate=debug --log-to-file`). Those are covered under [Debugging](../debugging.md#trace-legate-runtime-work), including how they pair with [CNPreferences](../api_preferences.md#task-scope-names) task-scope naming. + +The benchmark harness (`benchmark/run.jl` via `run_benchmark.sh`) sets `LEGATE_CONFIG` from `--gpus` / `--cpus` before each worker starts. See [How to Benchmark](../benchmarks/howto.md). diff --git a/docs/src/debugging.md b/docs/src/debugging.md new file mode 100644 index 000000000..8686e304b --- /dev/null +++ b/docs/src/debugging.md @@ -0,0 +1,156 @@ +# Debugging + +Debug the layer that matches the problem: + +| Question | Tool | +|---|---| +| Which operations did Legate submit, and when did they run? | [Legate logs and profiles](#trace-legate-runtime-work) | +| Did a broadcast become one fused kernel? | [`BCAST_FUSION_DEBUG`](#inspect-fused-broadcasts-with-bcast_fusion_debug) | +| Where does `@analyze_lifetimes` free temporaries? | [`@show_lifetimes`](#inspect-lifetime-rewrites-with-show_lifetimes) | + +## Trace Legate runtime work + +Legate records provenance on submitted operations. cuNumeric can provide that provenance automatically for individual operations, or you can add broader phase names manually. + +Choose one output format for a run: + +| Goal | `LEGATE_CONFIG` flags | Output | +|---|---|---| +| Read runtime decisions as text | `--logging legate=debug --log-to-file` | `legate_*.log` | +| Inspect an execution timeline | `--profile` | `legate_*.prof` | + +Set `LEGATE_CONFIG` before Julia starts. Add `--logdir ` to keep the generated files outside the working directory. Logging and profiling add overhead, so collect them in separate runs when measuring performance. + +### Label individual operations + +Task-scope naming is off by default. Enable the compile-time preference in one Julia process: + +```bash +julia --project=. -e \ + 'using CNPreferences; CNPreferences.enable_task_scope_names!()' +``` + +Then start a fresh process with logging enabled: + +```bash +export LEGATE_AUTO_CONFIG=0 +export LEGATE_CONFIG="--gpus 1 --cpus 4 --logging legate=debug --log-to-file" +julia --project=. workload.jl +``` + +cuNumeric operations now carry short labels such as `zeros`, `matmul`, and `broadcast.+(*(input0, input1), scalar0)`. Search for those labels in `legate_*.log` to connect runtime messages to calls in `workload.jl`. + +For a timeline instead, replace the logging flags with `--profile`, run the same workload, and process the resulting `legate_*.prof` files with `legate_prof`. + +Disable the preference after debugging, then restart Julia: + +```bash +julia --project=. -e \ + 'using CNPreferences; CNPreferences.disable_task_scope_names!()' +``` + +See [Task scope names](./api_preferences.md#task-scope-names) for the preference API. + +### Label larger phases + +Use `Legate.with_scope` when phase names such as `initialize` and `update` are more useful than per-operation names: + +```julia +using cuNumeric +import Legate + +A, B = Legate.with_scope("initialize") do + A = cuNumeric.ones(Float32, 64, 64) + B = cuNumeric.ones(Float32, 64, 64) + (A, B) +end + +D = Legate.with_scope("update") do + @. A * B + 2.0f0 +end +``` + +cuNumeric already supplies names for individual operations when task-scope naming is enabled. Scopes can be nested, so you can wrap those operations in your own phase-level scopes, as shown above. + +## Inspect fused broadcasts with `BCAST_FUSION_DEBUG` + +When broadcast fusion is on, set `cuNumeric.BCAST_FUSION_DEBUG[] = true` to print each fused kernel before launch: the expression tree, inputs, scalars, arg map, and launch geometry. + +```julia +using cuNumeric + +cuNumeric.BCAST_FUSION_DEBUG[] = true + +N = 8 +A = cuNumeric.ones(Float32, N, N) +B = cuNumeric.ones(Float32, N, N) +C = cuNumeric.zeros(Float32, N, N) + +C .= @. A * B + 2.0f0 + +cuNumeric.BCAST_FUSION_DEBUG[] = false +``` + +Example output: + +```text +======================================== fused broadcast kernel + expr +(*(NDArray, NDArray), 2.0f0) + output NDArray{Float32, 2, false, Nothing} (8, 8) + inputs 2 unique NDArray(s) + [0] NDArray{Float32, 2, false, Nothing} (8, 8) + [1] NDArray{Float32, 2, false, Nothing} (8, 8) + scalars 2.0f0 + arg_map [0, 1, 2, -1] (0=output, >=1=input idx+1, <0=scalar) + launch host thread budget=1024, indexing=linear, blocks=device(local tile), global_ndrange=(8, 8) + call broadcast.gpu_broadcast_kernel_linear_splat(input0, input1, scalar0) +``` + +How to read it: + +- `expr` is the fused op tree. +- `inputs` / `scalars` are the runtime arguments passed into the kernel. +- `arg_map` encodes how kernel slots map to the output (`0`), inputs (`>= 1`), and scalars (`< 0`). +- If nothing prints, the expression took the unfused path (for example shape-mismatched leaves, fusion disabled, or below the min-ops threshold). + +Turn the flag off when you are done. It prints on every fused launch and can be noisy in loops. + +See [Kernel Fusion](./perf/kernel_fusion.md) for `@.` / fusion usage, and [Internals](./internals.md) for how these features work. + +## 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. + +```julia +using cuNumeric + +@show_lifetimes begin + result = A[1:end, :] .+ B[1:end, :] + C .= result .* 2.0 +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. +- 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. + +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 new file mode 100644 index 000000000..5cac1631d --- /dev/null +++ b/docs/src/developer_mode.md @@ -0,0 +1,72 @@ +# Developer Mode + +Developer mode builds the Julia C++ wrapper from the sources in this repo under `lib/cunumeric_jl_wrapper`, instead of downloading a prebuilt `cunumeric_jl_wrapper_jll`. Use it when you change the wrapper, debug the C-API boundary, or point at a custom cupynumeric install. + +## When to use it + +- You edited C++ or CxxWrap code under `lib/cunumeric_jl_wrapper/` and need a new `.so`. +- You are pairing cuNumeric.jl against a cupynumeric build that is not the JLL artifact (for example a source tree). +- You want faster iterate-rebuild cycles on the wrapper without waiting on a JLL release. + +Ordinary Julia-only changes under `src/` do **not** require developer mode. Restart or re-`using` as usual. + +## Enable developer mode + +From an environment that can load `CNPreferences` (the docs project or a develop checkout of cuNumeric.jl): + +```julia +using CNPreferences +CNPreferences.use_developer_mode(; use_jll=true, path=nothing) +``` + +- `use_jll=true` (default): still use the `cupynumeric_jll` binary, but compile the wrapper in-tree against it. +- `use_jll=false, path="..."`: compile the wrapper against a cupynumeric root you already installed. See [Build Modes](./install.md) for compiler and CMake requirements. + +Preferences are written to `LocalPreferences.toml`. **Restart Julia** after changing the mode. + +## Rebuild the wrapper after changes + +After editing files under `lib/cunumeric_jl_wrapper/` (or after switching into developer mode for the first time): + +```julia +using Pkg +Pkg.build("cuNumeric") +``` + +That runs `deps/build.jl`, which in developer mode: + +1. Resolves cupynumeric (JLL or your `path`) +2. Builds / refreshes the CxxWrap pieces as needed +3. Compiles `lib/cunumeric_jl_wrapper` via `scripts/build_cpp_wrapper.sh` into `lib/cunumeric_jl_wrapper/build/` +4. Points the package at that local library (artifact override) + +Then restart Julia (or at least reload cuNumeric) so the new shared library is picked up. + +If the build fails, check CMake / g++ (C++20) / CUDA toolkit availability as described on [Build Modes](./install.md). Build logs from the helper scripts are written under `deps/`. + +### Typical edit loop + +```text +1. Edit lib/cunumeric_jl_wrapper/... +2. Pkg.build("cuNumeric") +3. Restart Julia +4. using cuNumeric +``` + +## Switch back to JLLs + +When you no longer need a local wrapper: + +```julia +using CNPreferences +CNPreferences.use_jll_binary() +``` + +Restart Julia. You do not need `Pkg.build` for pure JLL mode (the build script exits early). + +## Related pages + +- [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 diff --git a/docs/src/errors.md b/docs/src/errors.md index 61c65e2dd..ff924ef35 100644 --- a/docs/src/errors.md +++ b/docs/src/errors.md @@ -1,4 +1,4 @@ # Common Errors ## OOM on Startup -If you have other processes using GPU RAM (e.g. another instance of cuNumeric.jl) then cuNumeric.jl will fail to start and will segfault. The first symbol is typically something like `_ZN5Realm4CudaL22allocate_device_memoryEPNS0_3GPUEm`. You can fix this by killing the other jobs or modifying the amount of GPU RAM requested in `LEGATE_CONFIG`. See the [performance](./perf.md) documentation for examples on how to set the `LEGATE_CONFIG` environment variable. +If you have other processes using GPU RAM (e.g. another instance of cuNumeric.jl) then cuNumeric.jl will fail to start and will segfault. The first symbol is typically something like `_ZN5Realm4CudaL22allocate_device_memoryEPNS0_3GPUEm`. You can fix this by killing the other jobs or modifying the amount of GPU RAM requested in `LEGATE_CONFIG`. See the [Hardware](./configuration/hardware.md) documentation for examples on how to set the `LEGATE_CONFIG` environment variable. diff --git a/docs/src/examples.md b/docs/src/examples.md deleted file mode 100644 index 005aa4ef9..000000000 --- a/docs/src/examples.md +++ /dev/null @@ -1,145 +0,0 @@ -# Examples - - -## DAXPY -```julia -# found in examples/daxpy.jl -using cuNumeric - -arr = cuNumeric.rand(20) - -α = 1.32f0 -b = 2.0f0 - -arr2 = α .* arr .+ b -``` -## Monte-Carlo Integration - -Most integrals can be estimated with a basic Monte-Carlo estimator: - -```math -\hat{I}_N = \frac{\Omega}{N}\sum_{i=1}^Nf(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}. -``` - -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 -using cuNumeric - -# Note that we do not yet support broadcasting -# custom functions, so the braodcasting MUST -# be done inside the function -integrand = (x) -> exp.(-x.^2) - -N = 1_000_000 - -x_max = 10.0f0 -domain = [-x_max, x_max] -Ω = domain[2] - domain[1] - -samples = Ω*cuNumeric.rand(N) .- 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))") -``` -## Gray Scott Reaction Diffusion -```julia -# found in examples/gray-scott.jl -using cuNumeric -using Plots - -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) - # calculate F_u and F_v functions - 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]) - # 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 - bc!(u_new, v_new, u, v) -end - -function gray_scott() - #anim = Animation() - - N = 100 - dims = (N, N) - - args = Params() - - n_steps = 2000 # number of steps to take - frame_interval = 200 # steps to take between making plots - - u = cuNumeric.ones(dims) - v = cuNumeric.zeros(dims) - u_new = cuNumeric.zeros(dims) - v_new = cuNumeric.zeros(dims) - - u[1:15,1:15] = cuNumeric.rand(15,15) - v[1:15,1:15] = cuNumeric.rand(15,15) - - for n in 1:n_steps - step!(u, v, u_new, v_new, args) - # update u and v - # this doesn't copy, this switching references - u, u_new = u_new, u - v, v_new = v_new, v - - if n%frame_interval == 0 - u_cpu = u[:, :] - heatmap(u_cpu, clims=(0, 1)) - frame(anim) - end - end - gif(anim, "gray-scott.gif", fps=10) - return u, v - -end - -u, v = gray_scott() -``` -![Simulation Output](./gray-scott.gif) diff --git a/docs/src/examples/grayscott.md b/docs/src/examples/grayscott.md new file mode 100644 index 000000000..dc0fca426 --- /dev/null +++ b/docs/src/examples/grayscott.md @@ -0,0 +1,95 @@ +# Gray-Scott Reaction Diffusion + +```julia +# found in examples/gray-scott.jl +using cuNumeric +using Plots + +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 + + bc!(u_new, v_new, u, v) +end + +function gray_scott() + #anim = Animation() + + N = 100 + dims = (N, N) + + args = Params() + + n_steps = 2000 # number of steps to take + frame_interval = 200 # steps to take between making plots + + u = cuNumeric.ones(dims) + v = cuNumeric.zeros(dims) + u_new = cuNumeric.zeros(dims) + v_new = cuNumeric.zeros(dims) + + u[1:15,1:15] = cuNumeric.rand(15,15) + v[1:15,1:15] = cuNumeric.rand(15,15) + + for n in 1:n_steps + step!(u, v, u_new, v_new, args) + # update u and v + # this doesn't copy, this switching references + u, u_new = u_new, u + v, v_new = v_new, v + + if n%frame_interval == 0 + u_cpu = u[:, :] + heatmap(u_cpu, clims=(0, 1)) + frame(anim) + end + end + gif(anim, "gray-scott.gif", fps=10) + return u, v + +end + +u, v = gray_scott() +``` +![Simulation Output](../gray-scott.gif) diff --git a/docs/src/examples/hdf5.md b/docs/src/examples/hdf5.md new file mode 100644 index 000000000..7c46104b2 --- /dev/null +++ b/docs/src/examples/hdf5.md @@ -0,0 +1,22 @@ +# HDF5 I/O + +> [!NOTE] +> HDF5 support is planned. This page is a placeholder for an end-to-end example once the API lands. + +Reading and writing `NDArray`s through HDF5 will let you checkpoint distributed arrays and exchange data with NumPy / cuPyNumeric workflows without gathering everything to the host first. + +## Planned sketch + +```julia +using cuNumeric + +# Write (API names TBD) +# cuNumeric.h5write("checkpoint.h5", "fields/u", u) + +# Read into an NDArray (API names TBD) +# u = cuNumeric.h5read("checkpoint.h5", "fields/u") +``` + +When available, prefer the cuNumeric HDF5 entry points over collecting to a Julia `Array` and using HDF5.jl alone, so large arrays can stay partitioned across devices. + +See [HDF5](../api_hdf5.md) in the Public API for the (forthcoming) function reference. diff --git a/docs/src/examples/initialization.md b/docs/src/examples/initialization.md new file mode 100644 index 000000000..6a1ddbe4a --- /dev/null +++ b/docs/src/examples/initialization.md @@ -0,0 +1,36 @@ +# Initialization + +Create `NDArray`s with the usual Julia-style constructors. The default element type is `Float32` unless you pass one. + +```julia +using cuNumeric + +# Zeros / ones / fill +Z = cuNumeric.zeros(4, 4) # Float32 by default +Z64 = cuNumeric.zeros(Float64, 4, 4) +O = cuNumeric.ones(3, 3) +F = cuNumeric.fill(7.5f0, (2, 3)) + +# Boolean arrays +T = cuNumeric.trues(2, 3) +Fbool = cuNumeric.falses(2, 3) + +# Identity +I = cuNumeric.eye(5) +I16 = cuNumeric.eye(Float32, 5) + +# Uniform random values (default Float32; backend draws Float64 then converts) +R = cuNumeric.rand(4, 4) +R64 = cuNumeric.rand(Float64, 1000) +cuNumeric.rand!(R64) # fill an existing Float64 array +``` + +Shapes can be passed as separate `Int`s or as a `Tuple` / `Dims`: + +```julia +cuNumeric.zeros(2, 3) +cuNumeric.zeros((2, 3)) +cuNumeric.ones(Float64, (10, 10)) +``` + +For signatures and more detail, see [Initialization](../api_initialization.md) in the Public API. diff --git a/docs/src/examples/montecarlo.md b/docs/src/examples/montecarlo.md new file mode 100644 index 000000000..5a66572db --- /dev/null +++ b/docs/src/examples/montecarlo.md @@ -0,0 +1,40 @@ +# Monte-Carlo Integration + +Most integrals can be estimated with a basic Monte-Carlo estimator: + +```math +\hat{I}_N = \frac{\Omega}{N}\sum_{i=1}^Nf(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}. +``` + +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 +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 + +x_max = 10.0f0 +domain = [-x_max, x_max] +Ω = domain[2] - domain[1] + +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))") +``` diff --git a/docs/src/index.md b/docs/src/index.md deleted file mode 120000 index fe8400541..000000000 --- a/docs/src/index.md +++ /dev/null @@ -1 +0,0 @@ -../../README.md \ No newline at end of file diff --git a/docs/src/index.md b/docs/src/index.md new file mode 100644 index 000000000..932cc5b52 --- /dev/null +++ b/docs/src/index.md @@ -0,0 +1,104 @@ +```@raw html +

+ 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) + +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. + +### Quick Start + +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") +``` + +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). + +```julia +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). + +### 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. + +**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`: + +```julia +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. + +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). + +### 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. + +```julia +y .= @. -a + b * c +``` + +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. + +`@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. + +```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). + +### Try an example + +```julia +using cuNumeric + +integrand = (x) -> @. exp(-x^2) + +N = 1_000_000 +x_max = 10.0f0 +Ω = 2 * x_max + +samples = Ω .* cuNumeric.rand(N) +samples = samples .- x_max +estimate = (Ω / N) .* sum(integrand(samples)) + +println("Monte-Carlo Estimate: $(estimate)") +``` +More worked examples (initialization, Gray-Scott, …) are in the documentation sidebar under **Examples**. + +### Known Limitations + +- There is no support for `Float16` or `ComplexF16` diff --git a/docs/src/install.md b/docs/src/install.md index 8a2b2dadd..a4dc9ef37 100644 --- a/docs/src/install.md +++ b/docs/src/install.md @@ -1,47 +1,40 @@ -# Build Options +# Build Modes -To make customization of the build options easier we have the `CNPreferences.jl` package to generate the `LocalPreferences.toml` which is read by the build script to determine which build option to use. CNPreferences.jl will also enforce that Julia is restarted for changes to take effect. +cuNumeric.jl gets its cupynumeric / Legate binaries from one of three providers, chosen through `CNPreferences` (writes `LocalPreferences.toml`; **restart Julia** after changing mode): +| Mode | When to use | +|---|---| +| **JLL (default)** | Normal installs; prebuilt artifacts from the Julia package server | +| **Developer** | Building or hacking the in-tree C++ wrapper, or a custom cupynumeric tree | +| **Conda** | Linking against an existing conda env that already has cupynumeric | -## Julia Installation +Install `CNPreferences` on its own if you want to set the mode before adding `cuNumeric`: -cuNumeric supports Julia 1.10 and 1.11. We recommend installing Julia with [juliaup](https://github.com/JuliaLang/juliaup): - -``` -curl -fsSL https://install.julialang.org | sh -s -- --default-channel 1.11 -``` - -This will install version 1.11 by default since that is what we have tested against. To verify 1.11 is the default run either of the following (you may need to source bashrc): -```bash -juliaup status -julia --version -``` - -If 1.11 is not your default, please set it to be the default. Other versions of Julia are untested. -```bash -juliaup default 1.11 +```julia +using Pkg +Pkg.add("CNPreferences") ``` ## Default Build (jlls) ```julia -pkg> add cuNumeric +using Pkg +Pkg.add("cuNumeric") ``` -If you previously used a custom build or conda build and would like to revert back to using prebuilt JLLs, run the following command in the directory containing the Project.toml of your environment. -```julia -using CNPreferences; CNPreferences.use_jll_binary() -``` - -`CNPreferences` is a separate module so that it can be used to configure the build settings before `cuNumeric.jl` is added to your environment. To install it separately run +If you previously used a custom build or conda build and would like to revert back to using prebuilt JLLs: ```julia -pkg> add CNPreferences +using CNPreferences +CNPreferences.use_jll_binary() ``` +Then restart Julia. Run `Pkg.build("cuNumeric")` if you left a non-JLL mode. + ## Developer mode > [!TIP] > This gives the most flexibility in installs. It is meant for developing on cuNumeric.jl. +> For rebuilding `lib/cunumeric_jl_wrapper` after C++ changes, see [Developer Mode](./developer_mode.md). We support using a custom install version of cupynumeric. See https://docs.nvidia.com/cupynumeric/latest/installation.html for details about different install configurations, or building cupynumeric from source. @@ -51,12 +44,14 @@ To use developer mode, ```julia using CNPreferences; CNPreferences.use_developer_mode(; use_jll=true, path=nothing) ``` -By default `use_cunumeric_jll` will be set to true. However, you can set a custom branch and/or use a custom path of cupynumeric. By setting `use_jll=false`, you can set `path` to your custom install. +By default `use_jll` will be set to true. However, you can use a custom path of cupynumeric. By setting `use_jll=false`, you can set `path` to your custom install. ```julia using CNPreferences; CNPreferences.use_developer_mode(;use_jll=false, path="/path/to/cupynumeric/root") ``` +After enabling developer mode (and after any wrapper edits), rebuild with `Pkg.build("cuNumeric")` and restart Julia. Details are on [Developer Mode](./developer_mode.md). + ## Link Against Existing Conda Environment > [!WARNING] @@ -64,12 +59,21 @@ using CNPreferences; CNPreferences.use_developer_mode(;use_jll=false, path="/pat Note, you need conda >= 24.1 to install the conda package. More installation details are found [here](https://docs.nvidia.com/cupynumeric/latest/installation.html). +````@eval +using Markdown +mm = Main.CUPYNUMERIC_MAJOR_MINOR +compat = Main.CUPYNUMERIC_JLL_COMPAT +Markdown.parse(""" +Supported cupynumeric versions match the `cupynumeric_jll` major.minor in this repo's `Project.toml`. Currently that pin is `cupynumeric_jll = "$(compat)"`, so use the **$(mm)** conda line: + ```bash # with a new environment -conda create -n myenv -c conda-forge -c cupynumeric +conda create -n myenv -c conda-forge -c cupynumeric cupynumeric=$(mm) # into an existing environment -conda install -c conda-forge -c cupynumerice +conda install -c conda-forge -c cupynumeric cupynumeric=$(mm) ``` +""") +```` Once you have the conda package installed, you can activate here. ```bash conda activate [conda-env-with-cupynumeric] @@ -77,6 +81,11 @@ conda activate [conda-env-with-cupynumeric] To update `LocalPreferences.toml` so that a local conda environment is used as the binary provider for cupynumeric run the following command. `conda_env` should be the absolute path to the conda environment (e.g., the value of CONDA_PREFIX when your environment is active). For example, this path is: `/home/JuliaLegate/.conda/envs/cupynumeric-gpu`. ```julia -using CNPreferences; CNPreferences.use_conda("conda-env-with-legate"); -Pkg.build() +using CNPreferences +using Pkg + +CNPreferences.use_conda(ENV["CONDA_PREFIX"]) # absolute path, e.g. from CONDA_PREFIX +Pkg.build("cuNumeric") ``` + +Then **restart Julia** so the new mode loads. diff --git a/docs/src/internals.md b/docs/src/internals.md new file mode 100644 index 000000000..10f655faf --- /dev/null +++ b/docs/src/internals.md @@ -0,0 +1,52 @@ +# 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). + +## Broadcast fusion + +We compile nested Julia broadcast expressions on `NDArray` to a single CUDA kernel instead of launching one kernel per operation. + +### Pipeline + +- **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. + +## Lifetimes and GC + +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_lifetiems` macro so users can manually manage the lifetimes of a code block and also manually track device memory to automatically invoke GC. + +### Eager last-use freeing with `@analyze_lifetimes` + +`@analyze_lifetimes` rewrites a block at macro-expansion time: + +- Hoist temporary allocations into named temps. +- Find each temp's static last use. +- Insert calls to free temporary `NDArrays` after last-use. + +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. + +```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. + +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.jl`. + +### 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. + +Relevant source: `src/memory.jl`. diff --git a/docs/src/linalg.md b/docs/src/linalg.md new file mode 100644 index 000000000..12742fad4 --- /dev/null +++ b/docs/src/linalg.md @@ -0,0 +1,77 @@ +# Linear Algebra + +cuNumeric.jl supports a small set of linear algebra operations on `NDArray`. This page covers matrix multiply, batched solve, and related helpers. Related autodocs also appear under [NDArray Reference](./api.md). + +## Matrix multiply + +For two 2D arrays, `*` is matrix multiplication (GEMM), not elementwise multiply. Use `.*` when you want an elementwise product of matrices. + +```julia +using LinearAlgebra +using cuNumeric + +A = cuNumeric.rand(Float32, 128, 128) +B = cuNumeric.rand(Float32, 128, 128) + +C = A * B # allocates a new result +mul!(similar(C), A, B) # in-place GEMM into an existing array +``` + +```@autodocs +Modules = [cuNumeric] +Pages = ["ndarray/binary.jl"] +Filter = t -> t isa Function && nameof(t) === :mul! +``` + +## Solve (batched) + +`cuNumeric.solve(A, b)` solves linear systems. It is not Julia's `\`. + +```@docs +cuNumeric.solve +``` + +Shapes follow the batched signature: + +- `A` is `(..., m, m)` (last two dims square) +- `b` is `(..., m)` or `(..., m, n)` +- result is `(..., m)` or `(..., m, n)` + +A 1D right-hand side is reshaped internally to a single column, then reshaped back. + +```julia +using cuNumeric + +# Single system: (m, m) and (m,) +A = cuNumeric.rand(Float32, 64, 64) +b = cuNumeric.rand(Float32, 64) +x = cuNumeric.solve(A, b) + +# Several right-hand sides: (m, m) and (m, n) +B = cuNumeric.rand(Float32, 64, 4) +X = cuNumeric.solve(A, B) + +# Batched systems: (batch, m, m) and (batch, m, n) +As = cuNumeric.rand(Float32, 8, 32, 32) +Bs = cuNumeric.rand(Float32, 8, 32, 2) +Xs = cuNumeric.solve(As, Bs) +``` + +Notes: + +- Accepted types: `Float32`, `Float64`, `ComplexF32`, `ComplexF64`. Integer or `Bool` inputs promote to `Float64` only when promotion is allowed (`@allowpromotion` / `allowpromotion`). +- The implementation always goes through a batched Legate `SOLVE` task, including the 2D case. +- Batch dimensions are supported in the API. Coverage for higher-rank batches in the test suite is still thin, so start with 2D and small batches when validating new code. + +## Helpers + +These helpers live on `NDArray` and are also listed in the Public API: + +- `cuNumeric.transpose` +- `cuNumeric.eye` +- `cuNumeric.diag` (2D to 1D) +- `cuNumeric.trace` + +## Not available yet + +There is no public `svd`, `qr`, `cholesky`, `eig`, `lu`, matrix `inv`, or `ldiv!` in cuNumeric.jl yet. Elementwise `inv` / `^-1` exist as unary ops; those are not matrix inverse. diff --git a/docs/src/perf.md b/docs/src/perf.md deleted file mode 100644 index 83a37e205..000000000 --- a/docs/src/perf.md +++ /dev/null @@ -1,50 +0,0 @@ -# Performance Tips - -## Avoid Scalar Indexing -Accessing elements of an NDArray one at a time (e.g., `arr[5]`) is slow and should be avoided. Indexing like this requires data to be trasfered between device and host and maybe even communicated across nodes. Scalar indexing will emit an error which can be opted out of with `@allowscalar` or `allwoscalar() do ... end`. Several functions in the existing API invoke scalar indexing and are intended for testing (e.g., the `==` operator). - -## Avoid Implicit Promotion -Mixing integral types of different size (e.g., `Float64` and `Float32`) will result in implicit promotion of the smaller type to the larger types. This creates a copy of the data and hurts performance. Implicit promotion from a smaller integral type to a larger integral type will emit an error which can be opted out of with `@allowpromotion` or `allowpromotion() do ... end`. This error is common when mixing literals with `NDArrays`. By default a floating point literal (i.e., 1.0) is `Float64` but the default type of an `NDArray` is `Float32`. - -## Setting Hardware Configuration - -There is no programatic 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_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 -- `LEGATE_CONFIG`: A string representing the hardware configuration to set - -These variables must be set before launching the Julia instance running cuNumeric.jl. We recommend setting `export LEGATE_SHOW_CONFIG=1` so that the hardware configuration will be printed when Legate starts. This output is automatically captured and relayed to the user. - -To manually set the hardware configuration, `export LEGATE_AUTO_CONFIG=0`, and then define your own config with something like `export LEGATE_CONFIG="--gpus 1 --cpus 10 --ompthreads 10"`. We recommend using the default memory configuration for your machine and only settings the `gpus`, `cpus` and `ompthreads`. More details about the Legate configuration can be found in the [NVIDIA Legate documentation](https://docs.nvidia.com/legate/latest/usage.html#resource-allocation). If you know where Legate is installed on your computer you can also run `legate --help` for more detailed information. - -## Reduce Allocations with `@analyze_lifetimes` - -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. - -`@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 are returned to cuNumeric's pool and recycled by the next same-sized allocation, skipping new buffer allocation. - -```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.0 -end -``` - -**Benchmark** (Gray–Scott reaction–diffusion, 512×512, 10 000 steps): - -``` - 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 -``` - -~2× wall-clock speedup and ~800 MB lower peak memory with no algorithmic changes. - -## Kernel Fusion -cuPyNumeric does not fuse independent operations automatically, even in broadcast expressions. This is a priority for a future release. diff --git a/docs/src/perf/kernel_fusion.md b/docs/src/perf/kernel_fusion.md new file mode 100644 index 000000000..daed74284 --- /dev/null +++ b/docs/src/perf/kernel_fusion.md @@ -0,0 +1,35 @@ +# 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. + +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. + +```julia +# Easy to miss the dot on negation +y .= .-a .+ b .* c + +# Prefer: @. dots every operator, which is clearer and fusion-friendly +y .= @. -a + b * c +``` + +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. + +Toggle fusion through `CNPreferences` (restart Julia after changing these): + +```julia +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 +``` + +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). diff --git a/docs/src/perf/patterns_to_avoid.md b/docs/src/perf/patterns_to_avoid.md new file mode 100644 index 000000000..3fe864acc --- /dev/null +++ b/docs/src/perf/patterns_to_avoid.md @@ -0,0 +1,9 @@ +# Patterns to Avoid + +## Scalar indexing + +Accessing elements of an NDArray one at a time (e.g., `arr[5]`) is slow and should be avoided. Indexing like this requires data to be transferred between device and host and maybe even communicated across nodes. Scalar indexing will emit an error which can be opted out of with `@allowscalar` or `allowscalar() do ... end`. Several functions in the existing API invoke scalar indexing and are intended for testing (e.g., the `==` operator). + +## Implicit promotion + +Mixing integral types of different size (e.g., `Float64` and `Float32`) will result in implicit promotion of the smaller type to the larger types. This creates a copy of the data and hurts performance. Implicit promotion from a smaller integral type to a larger integral type will emit an error which can be opted out of with `@allowpromotion` or `allowpromotion() do ... end`. This error is common when mixing literals with `NDArrays`. By default a floating point literal (i.e., 1.0) is `Float64` but the default type of an `NDArray` is `Float32`. diff --git a/docs/src/perf/reduce_allocations.md b/docs/src/perf/reduce_allocations.md new file mode 100644 index 000000000..37b38ca50 --- /dev/null +++ b/docs/src/perf/reduce_allocations.md @@ -0,0 +1,31 @@ +# Reduce Allocations + +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. + +`@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. + +When broadcast fusion is on, intermediate dotted nodes in a broadcast tree are not real `NDArray` allocations. The macro accounts for that automatically. + +```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 +end +``` + +**Benchmark** (Gray-Scott reaction-diffusion, 512×512, 10 000 steps): + +``` + 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 +``` + +~2× wall-clock speedup and ~800 MB lower peak memory with no algorithmic changes. + +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). diff --git a/examples/daxpy.jl b/examples/daxpy.jl index db2617f9d..f7983fff0 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 = @. α * arr + b println(arr2) diff --git a/examples/integrate.jl b/examples/integrate.jl index fd03436e3..f27dae031 100644 --- a/examples/integrate.jl +++ b/examples/integrate.jl @@ -1,9 +1,9 @@ using cuNumeric # Note that we do not yet support broadcasting -# custom functions, so the braodcasting MUST +# 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,11 +11,12 @@ x_max = 10.0f0 domain = [-x_max, x_max] Ω = domain[2] - domain[1] -samples = Ω*cuNumeric.rand(N) .- x_max +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)) +estimate = (Ω / N) * sum(integrand(samples)) println("Monte-Carlo Estimate: $(estimate)") println("Analytical: $(sqrt(pi))") diff --git a/ext/CUDAExt/CUDAExt.jl b/ext/CUDAExt/CUDAExt.jl deleted file mode 100644 index ed9c07c63..000000000 --- a/ext/CUDAExt/CUDAExt.jl +++ /dev/null @@ -1,26 +0,0 @@ -module CUDAExt - -using Random -using CUDA -using Legate: Legate -using CxxWrap: CxxWrap -using cuNumeric: cuNumeric -import cuNumeric: - @cuda_task, @launch, NDArray, assert_experimental - -const KERNEL_OFFSET = sizeof(CUDA.KernelState) - -include("cuda.jl") - -function __init__() - if CUDA.functional() - # in cuda.jl to notify /wrapper/src/cuda.cpp about CUDA.jl kernel state size - cuNumeric.register_kernel_state_size(UInt64(KERNEL_OFFSET)) - # in /wrapper/src/cuda.cpp - cuNumeric.register_tasks(); - else - @warn "CUDA.jl is not functional; skipping CUDA kernel registration." - end -end - -end # module CUDAExt diff --git a/ext/CUDAExt/cuda.jl b/ext/CUDAExt/cuda.jl deleted file mode 100644 index 33421640c..000000000 --- a/ext/CUDAExt/cuda.jl +++ /dev/null @@ -1,210 +0,0 @@ -function ndarray_cuda_type(A::NDArray{T,N}) where {T,N} - if N == 1 - CuDeviceVector{T,1} - elseif N == 2 - CuDeviceMatrix{T,1} - else - CuDeviceArray{T,N,1} - end -end - -function ndarray_cuda_type(arg::T) where {T} - Base.isbits(arg) || throw(ArgumentError("Unsupported argument type: $(typeof(arg))")) - typeof(arg) -end - -cuNumeric.map_ndarray_cuda_types(args...) = tuple(map(ndarray_cuda_type, args)...) - -function to_stdvec(::Type{T}, vec) where {T} - stdvec = CxxWrap.StdVector{T}() - for x in vec - push!(stdvec, T(x)) - end - 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) - - if copy # due to being an input. we don't need to copy outputs - indices = ntuple(d -> 1:old_size[d], length(old_size)) - assign(new[indices...], arr) - 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 - new.padding = nothing -end - -function add_padding(arr::NDArray, i::Int64; copy=false) - add_padding(arr, (i,); copy=copy) -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 - end -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 - end -end - -function nda_to_logical_array(arr::NDArray{T,N}) where {T,N} - st_handle = cuNumeric.get_store(arr) - return Legate.LogicalArray{T,N}(st_handle, size(arr)) -end - -function Launch(kernel::cuNumeric.CUDATask, inputs::Tuple{Vararg{NDArray}}, - outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}}; blocks, threads) - - # we find the largest input/output. - ndarrays = vcat(inputs..., outputs...) - mx = findmax(arr -> arr.nbytes, ndarrays) # returns (nbytes, position) - max_size = mx[1] # first elem nbytes - max_shape = size(ndarrays[mx[2]]) # second elem max position - @assert !isnothing(max_shape) - - rt = Legate.get_runtime() - lib = cuNumeric.get_lib() - taskid = cuNumeric.RUN_PTX - task = Legate.create_auto_task(rt, lib, taskid) - - input_vars = Vector{Legate.Variable}() - for arr in inputs - check_sz!(arr, max_shape; copy=true) - la = nda_to_logical_array(arr) - p = Legate.add_input(task, la) - push!(input_vars, p) - end - - output_vars = Vector{Legate.Variable}() - for arr in outputs - check_sz!(arr, max_shape; copy=false) - la = nda_to_logical_array(arr) - p = Legate.add_output(task, la) - push!(output_vars, p) - end - - # next 3 lines are reserved scalars in the RUN_PTX task - Legate.add_scalar(task, Legate.string_to_scalar(kernel.func)) # 0 - cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, blocks)) # bx,by,bz 1,2,3 - cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, threads)) # tx,ty,tz 4,5,6 - - # any user defined scalars in the launch macro - for s in scalars - Legate.add_scalar(task, Legate.Scalar(s)) # 7+ -> ARG_OFFSET - end - - # all inputs are aligned with all outputs - Legate.default_alignment(task, input_vars, output_vars) - Legate.submit_auto_task(rt, task) -end - -function cuNumeric.launch(kernel::cuNumeric.CUDATask, inputs, outputs, scalars; blocks, threads) - Launch(kernel, - isa(inputs, Tuple) ? inputs : (inputs,), - isa(outputs, Tuple) ? outputs : (outputs,), - isa(scalars, Tuple) ? scalars : (scalars,); - blocks=isa(blocks, Tuple) ? blocks : (blocks,), - threads=isa(threads, Tuple) ? threads : (threads,), - ) -end - -function cuNumeric.ptx_task(ptx::String, kernel_name) - rt = Legate.get_runtime() - lib = cuNumeric.get_lib() # grab lib of legate app - # this taskid is directly tied to cpp code in our setup - taskid = cuNumeric.LOAD_PTX - task = Legate.create_auto_task(rt, lib, taskid) - # assign task arguments - Legate.add_scalar(task, Legate.string_to_scalar(ptx)) - Legate.add_scalar(task, Legate.string_to_scalar(kernel_name)) - Legate.submit_auto_task(rt, task) -end - -macro cuda_task(call_expr) - cuNumeric.assert_experimental() - - fname = call_expr.args[1] - fargs = call_expr.args[2:end] - - esc(quote - local _buf = IOBuffer() - local _types = cuNumeric.map_ndarray_cuda_types($(fargs...)) - # generate ptx using CUDA.jl - CUDA.code_ptx(_buf, $fname, _types; raw=false, kernel=true) - - local _ptx = String(take!(_buf)) - local _func_name = cuNumeric.extract_kernel_name(_ptx) - - # issue ptx_task within legate runtime to register cufunction ptr with cucontext - cuNumeric.ptx_task(_ptx, _func_name) - - # create a cuNumeric.CUDAtask that stores some info for a launch config - cuNumeric.CUDATask(_func_name, _types) - end) -end - -macro launch(args...) - cuNumeric.assert_experimental() - - allowed_keys = Set([:task, :blocks, :threads, :inputs, :outputs, :scalars]) - kwargs = Dict{Symbol,Any}() - - for ex in args - if !(ex isa Expr && ex.head == :(=)) - error("All arguments must be keyword assignments, e.g. task=..., threads=...") - end - key = ex.args[1] - val = ex.args[2] - - if !(key in allowed_keys) - error("@launch macro received unexpected keyword: $(key)") - end - - kwargs[key] = val - end - - if !haskey(kwargs, :task) - error("@launch macro requires 'task=...' to be provided.") - end - task = kwargs[:task] - blocks = get(kwargs, :blocks, :((1))) - threads = get(kwargs, :threads, :((256))) - inputs = get(kwargs, :inputs, :(())) - outputs = get(kwargs, :outputs, :(())) - scalars = get(kwargs, :scalars, :(())) - - esc( - quote - cuNumeric.launch( - $task, $inputs, $outputs, $scalars; blocks=($blocks), threads=($threads) - ) - end, - ) -end diff --git a/gray-scott.gif b/gray-scott.gif deleted file mode 100644 index 4811410d2..000000000 Binary files a/gray-scott.gif and /dev/null differ diff --git a/lib/CNPreferences/Project.toml b/lib/CNPreferences/Project.toml index 1231671e6..e612a4c5b 100644 --- a/lib/CNPreferences/Project.toml +++ b/lib/CNPreferences/Project.toml @@ -1,7 +1,7 @@ name = "CNPreferences" uuid = "3e078157-ea10-49d5-bf32-908f777cd46f" authors = ["""David Krasowska and Ethan Meitz """] -version = "0.1.2" +version = "0.1.3" [deps] LegatePreferences = "8028f36a-2b64-49e9-aa04-2d0933fd2ed9" diff --git a/lib/CNPreferences/src/CNPreferences.jl b/lib/CNPreferences/src/CNPreferences.jl index 60138ed65..558e47713 100644 --- a/lib/CNPreferences/src/CNPreferences.jl +++ b/lib/CNPreferences/src/CNPreferences.jl @@ -6,4 +6,82 @@ const DEVEL_DEFAULT_WRAPPER_BRANCH = "main" LegatePreferences.@make_preferences("cunumeric_") +# Compile-time: flipping it recompiles cuNumeric, so set it then load in a fresh process. +const FUSE_BROADCAST = @load_preference("FUSE_BROADCAST_EXPRS", true) +# Fuse when the broadcast tree has at least this many `Broadcasted` nodes (ops). +# Default 2 → single ops like `y .= cos.(x)` stay on the unfused path (no PTX compile). +# Set to 1 to fuse every eligible expression (e.g. in tests). +const FUSE_BROADCAST_MIN_OPS = @load_preference("FUSE_BROADCAST_MIN_OPS", 2) +const TASK_SCOPE_NAMES = @load_preference("TASK_SCOPE_NAMES", false) + +""" + set_broadcast_fusion!(enabled::Bool; export_prefs=false, force=true) + +Enable or disable broadcast fusion. When enabled (the default), eligible nested +broadcast expressions compile to a single CUDA PTX kernel. + +Restart Julia after changing this preference. +""" +function set_broadcast_fusion!(enabled::Bool; export_prefs=false, force=true) + return set_preferences!(@__MODULE__, "FUSE_BROADCAST_EXPRS" => enabled; export_prefs, force) +end + +""" + enable_broadcast_fusion!(; export_prefs=false, force=true) + +Enable broadcast fusion. This is the default. +""" +enable_broadcast_fusion!(; kwargs...) = set_broadcast_fusion!(true; kwargs...) + +""" + disable_broadcast_fusion!(; export_prefs=false, force=true) + +Disable broadcast fusion so each broadcast node runs as a separate cuNumeric op. +""" +disable_broadcast_fusion!(; kwargs...) = set_broadcast_fusion!(false; kwargs...) + +""" + set_broadcast_fusion_min_ops!(n::Integer; export_prefs=false, force=true) + +Fuse only when a broadcast tree has at least `n` ops. Default is `2`, so single +ops such as `y .= cos.(x)` stay on the unfused path. Set `n = 1` to fuse those +too. Values must be `>= 1`. + +Restart Julia after changing this preference. +""" +function set_broadcast_fusion_min_ops!(n::Integer; export_prefs=false, force=true) + n >= 1 || throw(ArgumentError("FUSE_BROADCAST_MIN_OPS must be >= 1, got $n")) + return set_preferences!( + @__MODULE__, "FUSE_BROADCAST_MIN_OPS" => Int(n); export_prefs, force + ) +end + +""" + set_task_scope_names!(enabled::Bool; export_prefs=false, force=true) + +Enable or disable named Legate task scopes for debugging. Default is off. + +When enabled, cuNumeric wraps ops in `Legate.with_scope` so provenance labels +appear in Legate logs/profiles. Pair with `LEGATE_CONFIG` flags such as +`--logging legate=debug --log-to-file` (set before Julia starts). Requires a +fresh Julia process after changing the preference. +""" +function set_task_scope_names!(enabled::Bool; export_prefs=false, force=true) + return set_preferences!(@__MODULE__, "TASK_SCOPE_NAMES" => enabled; export_prefs, force) +end + +""" + enable_task_scope_names!(; export_prefs=false, force=true) + +Enable named Legate task scopes for debugging. See [`set_task_scope_names!`](@ref). +""" +enable_task_scope_names!(; kwargs...) = set_task_scope_names!(true; kwargs...) + +""" + disable_task_scope_names!(; export_prefs=false, force=true) + +Disable named Legate task scopes. This is the default. +""" +disable_task_scope_names!(; kwargs...) = set_task_scope_names!(false; kwargs...) + end # module CNPreferences diff --git a/lib/cunumeric_jl_wrapper/CMakeLists.txt b/lib/cunumeric_jl_wrapper/CMakeLists.txt index dd0841a1a..b4eb3c551 100644 --- a/lib/cunumeric_jl_wrapper/CMakeLists.txt +++ b/lib/cunumeric_jl_wrapper/CMakeLists.txt @@ -4,6 +4,7 @@ project(cuNumericWrapper) set(cuNumericWrapperVersion 0.0.1) message(STATUS "Project version: v${cuNumericWrapperVersion}") + set(CXX_CUNUMERICJL_WRAPPER cunumeric_jl_wrapper) set(C_INTERFACE_LIB cunumeric_c_wrapper) @@ -20,19 +21,21 @@ find_package(cupynumeric REQUIRED) # CxxWrap Stuff if(NOT BINARYBUILDER) -execute_process( - COMMAND julia -e "println(DEPOT_PATH[1])" - OUTPUT_VARIABLE JULIA_DEP_PATH - OUTPUT_STRIP_TRAILING_WHITESPACE -) - -set(JlCxx_DIR "${JULIA_DEP_PATH}/dev/libcxxwrap_julia_jll/override") -message(STATUS "Setting JlCxx_DIR to ${JlCxx_DIR} (not using BinaryBuilder)") + execute_process( + COMMAND julia -e "println(DEPOT_PATH[1])" + OUTPUT_VARIABLE JULIA_DEP_PATH + OUTPUT_STRIP_TRAILING_WHITESPACE + ) + + set(JlCxx_DIR "${JULIA_DEP_PATH}/dev/libcxxwrap_julia_jll/override") + message(STATUS "Setting JlCxx_DIR to ${JlCxx_DIR} (not using BinaryBuilder)") endif() find_package(JlCxx REQUIRED) + get_target_property(JlCxx_location JlCxx::cxxwrap_julia LOCATION) get_filename_component(JlCxx_location ${JlCxx_location} DIRECTORY) + set(CMAKE_INSTALL_RPATH "${CMAKE_INSTALL_PREFIX}/lib;${JlCxx_location}") message(STATUS "Found JlCxx at ${JlCxx_location}") @@ -54,7 +57,9 @@ endif() # Library: C++ wrapper add_library(${CXX_CUNUMERICJL_WRAPPER} SHARED ${SOURCES}) -set_target_properties(${CXX_CUNUMERICJL_WRAPPER} PROPERTIES VERSION ${cuNumericWrapperVersion}) +set_target_properties(${CXX_CUNUMERICJL_WRAPPER} PROPERTIES + VERSION ${cuNumericWrapperVersion} +) target_link_libraries(${CXX_CUNUMERICJL_WRAPPER} PRIVATE cupynumeric::cupynumeric @@ -77,7 +82,9 @@ set(C_SOURCES ) add_library(${C_INTERFACE_LIB} SHARED ${C_SOURCES}) -set_target_properties(${C_INTERFACE_LIB} PROPERTIES VERSION ${cuNumericWrapperVersion}) +set_target_properties(${C_INTERFACE_LIB} PROPERTIES + VERSION ${cuNumericWrapperVersion} +) target_link_libraries(${C_INTERFACE_LIB} PRIVATE cupynumeric::cupynumeric diff --git a/lib/cunumeric_jl_wrapper/include/cuda_macros.h b/lib/cunumeric_jl_wrapper/include/cuda_macros.h new file mode 100644 index 000000000..de9a9f2fa --- /dev/null +++ b/lib/cunumeric_jl_wrapper/include/cuda_macros.h @@ -0,0 +1,74 @@ +#pragma once + +#define ERROR_CHECK(x) \ + { \ + cudaError_t status = x; \ + if (status != cudaSuccess) { \ + fprintf(stderr, "CUDA Error at %s:%d: %s\n", __FILE__, __LINE__, \ + cudaGetErrorString(status)); \ + if (stream_) cudaStreamDestroy(stream_); \ + exit(-1); \ + } \ + } + +#define DRIVER_ERROR_CHECK(x) \ + { \ + CUresult status = x; \ + if (status != CUDA_SUCCESS) { \ + const char *err_str = nullptr; \ + cuGetErrorString(status, &err_str); \ + fprintf(stderr, "CUDA Driver Error at %s:%d: %s\n", __FILE__, __LINE__, \ + err_str); \ + if (stream_) cudaStreamDestroy(stream_); \ + exit(-1); \ + } \ + } + +#define TEST_PRINT_DEBUG(dev_ptr, N, T, format, stream, message) \ + { \ + std::vector host_arr(N); \ + ERROR_CHECK(cudaMemcpy(host_arr.data(), \ + reinterpret_cast(dev_ptr), \ + sizeof(T) * N, cudaMemcpyDeviceToHost)); \ + ERROR_CHECK(cudaStreamSynchronize(stream)); \ + fprintf(stderr, "[TEST_PRINT] %s: " format "\n", message, host_arr[0]); \ + } + +#ifdef CUDA_DEBUG +#define CUDA_DEBUG_PRINT(x) \ + do { \ + x; \ + } while (0) +#else +#define CUDA_DEBUG_PRINT(x) \ + do { \ + } while (0) +#endif + +#define CUDA_DEVICE_ARRAY_ARG(MODE, ACCESSOR_CALL) \ + template < \ + typename T, int D, \ + typename std::enable_if<(D >= 1 && D <= REALM_MAX_DIM), int>::type = 0> \ + void cuda_device_array_arg_##MODE(char *&p, \ + const legate::PhysicalArray &rf) { \ + auto shp = rf.shape(); \ + auto acc = rf.data().ACCESSOR_CALL(); \ + CUDA_DEBUG_PRINT(std::cerr << "[RunPTXTask] " #MODE " accessor shape: " \ + << shp.lo << " - " << shp.hi << ", dim: " << D \ + << std::endl; \ + std::cerr << "[RunPTXTask] " #MODE " accessor strides: " \ + << acc.accessor.strides << std::endl;); \ + void *dev_ptr = const_cast(/*.lo to ensure multiple GPU support*/ \ + static_cast( \ + acc.ptr(Realm::Point(shp.lo)))); \ + auto extents = shp.hi - shp.lo + legate::Point::ONES(); \ + CuDeviceArray desc; \ + desc.ptr = dev_ptr; \ + desc.maxsize = shp.volume() * sizeof(T); \ + for (size_t i = 0; i < D; ++i) { \ + desc.dims[i] = extents[i]; \ + } \ + desc.length = shp.volume(); \ + memcpy(p, &desc, sizeof(CuDeviceArray)); \ + p += sizeof(CuDeviceArray); \ + } diff --git a/lib/cunumeric_jl_wrapper/include/ufi.h b/lib/cunumeric_jl_wrapper/include/ufi.h index b132dd661..983337494 100644 --- a/lib/cunumeric_jl_wrapper/include/ufi.h +++ b/lib/cunumeric_jl_wrapper/include/ufi.h @@ -1,6 +1,6 @@ /* 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 @@ -28,6 +28,7 @@ namespace ufi { enum TaskIDs { LOAD_PTX_TASK = 143432, RUN_PTX_TASK = 143433, + RUN_PTX_BROADCAST_TASK = 143434, }; class LoadPTXTask : public legate::LegateTask { @@ -46,6 +47,14 @@ class RunPTXTask : public legate::LegateTask { static void gpu_variant(legate::TaskContext context); }; +class RunPTXBroadcastTask : public legate::LegateTask { + public: + static inline const auto TASK_CONFIG = + legate::TaskConfig{legate::LocalTaskID{ufi::RUN_PTX_BROADCAST_TASK}}; + + static void gpu_variant(legate::TaskContext context); +}; + } // namespace ufi void wrap_cuda_methods(jlcxx::Module& mod); diff --git a/lib/cunumeric_jl_wrapper/src/cuda.cpp b/lib/cunumeric_jl_wrapper/src/cuda.cpp index d5e779d0f..e2eb5ae7f 100644 --- a/lib/cunumeric_jl_wrapper/src/cuda.cpp +++ b/lib/cunumeric_jl_wrapper/src/cuda.cpp @@ -20,6 +20,7 @@ #include "cuda.h" +#include #include #include @@ -29,7 +30,7 @@ #include "types.h" #include "ufi.h" -#define CUDA_DEBUG 0 +// #define CUDA_DEBUG #define BLOCK_START 1 #define THREAD_START 4 @@ -125,6 +126,8 @@ enum class AccessMode { WRITE, }; +// Dense — MUST match CUDA.jl CuDeviceArray bit layout (RunPTXTask / +// @cuda_task). template struct CuDeviceArray { void *ptr; // Pointer to device memory @@ -133,6 +136,18 @@ struct CuDeviceArray { uint64_t length; // Number of elements (at the end) }; +// Strided — matches Julia cuNumeric.CuStridedDeviceArray (RunPTXBroadcastTask +// only). +template +struct CuStridedDeviceArray { + void *ptr; + uint64_t maxsize; + std::array dims; + std::array + strides; // element strides (byte strides / sizeof(T)) + uint64_t length; +}; + #define CUDA_DEVICE_ARRAY_ARG(MODE, ACCESSOR_CALL) \ template < \ typename T, int D, \ @@ -161,8 +176,48 @@ struct CuDeviceArray { p += sizeof(CuDeviceArray); \ } +#define CUDA_STRIDED_DEVICE_ARRAY_ARG(MODE, ACCESSOR_CALL) \ + template < \ + typename T, int D, \ + typename std::enable_if<(D >= 1 && D <= REALM_MAX_DIM), int>::type = 0> \ + void cuda_strided_device_array_arg_##MODE(char *&p, \ + const legate::PhysicalArray &rf) { \ + auto shp = rf.shape(); \ + auto acc = rf.data().ACCESSOR_CALL(); \ + CUDA_DEBUG_PRINT( \ + std::cerr << "[RunPTXBroadcastTask] " #MODE " accessor shape: " \ + << shp.lo << " - " << shp.hi << ", dim: " << D << std::endl; \ + std::cerr << "[RunPTXBroadcastTask] " #MODE \ + << " accessor byte strides: " << acc.accessor.strides \ + << std::endl;); \ + void *dev_ptr = const_cast( \ + static_cast(acc.ptr(Realm::Point(shp.lo)))); \ + auto extents = shp.hi - shp.lo + legate::Point::ONES(); \ + CuStridedDeviceArray desc; \ + desc.ptr = dev_ptr; \ + desc.maxsize = shp.volume() * sizeof(T); \ + for (size_t i = 0; i < D; ++i) { \ + desc.dims[i] = extents[i]; \ + /* Legion AffineAccessor::strides are in bytes */ \ + desc.strides[i] = acc.accessor.strides[i] / sizeof(T); \ + } \ + desc.length = shp.volume(); \ + CUDA_DEBUG_PRINT(std::cerr << "[RunPTXBroadcastTask] " #MODE \ + << " packed dims="; \ + for (size_t i = 0; i < D; ++i) std::cerr \ + << desc.dims[i] << (i + 1 < D ? "," : ""); \ + std::cerr << " elem_strides="; \ + for (size_t i = 0; i < D; ++i) std::cerr \ + << desc.strides[i] << (i + 1 < D ? "," : ""); \ + std::cerr << " length=" << desc.length << std::endl;); \ + memcpy(p, &desc, sizeof(CuStridedDeviceArray)); \ + p += sizeof(CuStridedDeviceArray); \ + } + CUDA_DEVICE_ARRAY_ARG(read, read_accessor); // cuda_device_array_arg_read CUDA_DEVICE_ARRAY_ARG(write, write_accessor); // cuda_device_array_arg_write +CUDA_STRIDED_DEVICE_ARRAY_ARG(read, read_accessor); +CUDA_STRIDED_DEVICE_ARRAY_ARG(write, write_accessor); struct ufiFunctor { template @@ -175,37 +230,51 @@ struct ufiFunctor { } }; -// https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc -/*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { - cudaStream_t stream_ = context.get_task_stream(); - std::string kernel_name = context.scalar(0).value(); // 0 - - std::uint32_t bx = - context.scalar(BLOCK_START + 0).value(); // 1 - std::uint32_t by = - context.scalar(BLOCK_START + 1).value(); // 2 - std::uint32_t bz = - context.scalar(BLOCK_START + 2).value(); // 3 - - std::uint32_t tx = - context.scalar(THREAD_START + 0).value(); // 4 - std::uint32_t ty = - context.scalar(THREAD_START + 1).value(); // 5 - std::uint32_t tz = - context.scalar(THREAD_START + 2).value(); // 6 +struct ufiStridedFunctor { + template + void operator()(AccessMode mode, char *&p, const legate::PhysicalArray &arr) { + using CppT = typename legate_util::code_to_cxx::type; + if (mode == AccessMode::READ) + cuda_strided_device_array_arg_read(p, arr); + else + cuda_strided_device_array_arg_write(p, arr); + } +}; + +struct PTXLaunchParams { + cudaStream_t stream; + CUstream custream; + CUfunction func; + std::string kernel_name; + std::uint32_t bx, by, bz; + std::uint32_t tx, ty, tz; +}; + +// Reads common scalars (kernel_name, blocks, threads) and looks up the +// compiled CUfunction. Shared by RunPTXTask and RunPTXBroadcastTask. +static PTXLaunchParams read_launch_params(legate::TaskContext &context) { + PTXLaunchParams p; + p.stream = context.get_task_stream(); + p.kernel_name = context.scalar(0).value(); + + p.bx = context.scalar(BLOCK_START + 0).value(); + p.by = context.scalar(BLOCK_START + 1).value(); + p.bz = context.scalar(BLOCK_START + 2).value(); + + p.tx = context.scalar(THREAD_START + 0).value(); + p.ty = context.scalar(THREAD_START + 1).value(); + p.tz = context.scalar(THREAD_START + 2).value(); CUcontext ctx; - cuStreamGetCtx(stream_, &ctx); + cuStreamGetCtx(p.stream, &ctx); - FunctionKey key = {ctx, kernel_name}; + FunctionKey key = {ctx, p.kernel_name}; assert(cufunction_ptr.has_value()); FunctionMap &fmap = cufunction_ptr.get(); - auto it = fmap.find(key); #ifdef CUDA_DEBUG if (it == fmap.end()) { - // for DEBUG output std::cerr << "[RunPTXTask] Could not find key: " << key_to_string(key) << std::endl; for (const auto &[k, v] : fmap) { @@ -216,17 +285,51 @@ struct ufiFunctor { #endif assert(it != fmap.end()); - CUfunction func = it->second; + p.func = it->second; + p.custream = reinterpret_cast(p.stream); + return p; +} + +// Launch the kernel with the filled arg_buffer. +static void launch_kernel(const PTXLaunchParams &lp, + std::vector &arg_buffer, + std::size_t buffer_size) { + cudaStream_t stream_ = lp.stream; // alias for DRIVER_ERROR_CHECK macro + void *config[] = { + CU_LAUNCH_PARAM_BUFFER_POINTER, + static_cast(arg_buffer.data()), + CU_LAUNCH_PARAM_BUFFER_SIZE, + &buffer_size, + CU_LAUNCH_PARAM_END, + }; + +#ifdef CUDA_DEBUG + std::cerr << "[RunPTXTask] Launching kernel " << lp.kernel_name + << " with blocks (" << lp.bx << "," << lp.by << "," << lp.bz + << ") and threads (" << lp.tx << "," << lp.ty << "," << lp.tz << ")" + << std::endl; +#endif + + DRIVER_ERROR_CHECK(cuLaunchKernel(lp.func, lp.bx, lp.by, lp.bz, lp.tx, lp.ty, + lp.tz, 0, lp.custream, nullptr, config)); +} + +// Helper: align pointer to 8-byte boundary. +static inline void align8(char *&ptr) { + std::uintptr_t addr = reinterpret_cast(ptr); + ptr = reinterpret_cast((addr + 7) & ~std::uintptr_t(7)); +} + +// RunPTXTask: user-defined @cuda_task kernels +// Arg buffer: [kernel_state | inputs... | outputs... | scalars...] +// https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/eval_udf_gpu.cc +/*static*/ void RunPTXTask::gpu_variant(legate::TaskContext context) { + auto lp = read_launch_params(context); const std::size_t num_inputs = context.num_inputs(); const std::size_t num_outputs = context.num_outputs(); const std::size_t num_scalars = context.num_scalars(); - const std::size_t num_reductions = - context.num_reductions(); // unused for now - // compute total size: all device arrays + all scalars - // skip scalar 0-2 (kernel_name, threads, blocks) - // we allocate extra to decrease looping and dynamic dispatching on dim std::size_t max_buffer_size = padded_bytes_kernel_state + (num_inputs + num_outputs) * sizeof(CuDeviceArray); @@ -238,26 +341,13 @@ struct ufiFunctor { for (std::size_t i = 0; i < num_inputs; ++i) { auto ps = context.input(i); - auto code = ps.type().code(); - auto dim = ps.dim(); -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Input " << i << " type: " << code - << ", dim: " << dim << std::endl; -#endif - // dispatch on dim and code with ufiFunctor operator() - legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::READ, p, - ps); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiFunctor{}, + ufi::AccessMode::READ, p, ps); } for (std::size_t i = 0; i < num_outputs; ++i) { auto ps = context.output(i); - auto code = ps.type().code(); - auto dim = ps.dim(); -#ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Output " << i << " type: " << code - << ", dim: " << dim << std::endl; -#endif - legate::double_dispatch(dim, code, ufiFunctor{}, ufi::AccessMode::WRITE, p, - ps); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiFunctor{}, + ufi::AccessMode::WRITE, p, ps); } for (std::size_t i = ARG_OFFSET; i < num_scalars; ++i) { const auto &scalar = context.scalar(i); @@ -265,29 +355,163 @@ struct ufiFunctor { p += scalar.size(); } - std::size_t buffer_size = p - arg_buffer.data(); // calc used buffer + launch_kernel(lp, arg_buffer, p - arg_buffer.data()); +} - void *config[] = { - CU_LAUNCH_PARAM_BUFFER_POINTER, - static_cast(arg_buffer.data()), - CU_LAUNCH_PARAM_BUFFER_SIZE, - &buffer_size, - CU_LAUNCH_PARAM_END, - }; +// RunPTXBroadcastTask: broadcast fusion kernels +// Arg buffer: [kernel_state | ctx | arg_map-driven args...] +// +// Scalars after ARG_OFFSET: +// [7] = ctx (CompilerMetadata, raw bytes) +// [8] = num_kernel_args (Int32) +// [9..8+N] = arg_map entries (Int32 each) +// [9+N..] = actual scalar values +// +// Host passes an occupancy thread *budget* in tx (ty/tz unused) and a +// placeholder bx. This task overwrites bx/tx from the local output tile +// (linear: threads=min(budget,volume), blocks=cld(volume,threads)). +// +// arg_map encoding: +// val >= 0, val < num_outputs → output[val] (write CuStridedDeviceArray) +// val >= num_outputs → input[val - num_outputs] (read +// CuStridedDeviceArray) val < 0 → scalar at index -(val + 1) in +// trailing scalars + +static void broadcast_launch_dims_from_tile(PTXLaunchParams &lp, + const legate::PhysicalArray &out) { + const std::uint32_t budget = std::max(lp.tx, 1u); + const int dim = out.dim(); + + if (dim <= 0) { + lp.bx = 1; + lp.by = 1; + lp.bz = 1; + lp.tx = 1; + lp.ty = 1; + lp.tz = 1; + return; + } + + std::uint64_t volume = 0; + +#define CU_BCAST_FILL_VOLUME(D) \ + do { \ + volume = out.shape().volume(); \ + } while (0) + + switch (dim) { + case 1: + CU_BCAST_FILL_VOLUME(1); + break; + case 2: + CU_BCAST_FILL_VOLUME(2); + break; + case 3: + CU_BCAST_FILL_VOLUME(3); + break; + case 4: + CU_BCAST_FILL_VOLUME(4); + break; + case 5: + CU_BCAST_FILL_VOLUME(5); + break; + case 6: + CU_BCAST_FILL_VOLUME(6); + break; + default: + assert(0 && "broadcast launch: unsupported array dim"); + return; + } +#undef CU_BCAST_FILL_VOLUME + + if (volume == 0) { + lp.bx = 1; + lp.by = 1; + lp.bz = 1; + lp.tx = 1; + lp.ty = 1; + lp.tz = 1; + return; + } + + const std::uint32_t threads = + static_cast(std::min(budget, volume)); + const std::uint32_t blocks = + static_cast((volume + threads - 1) / threads); - CUstream custream_ = reinterpret_cast(stream_); + lp.bx = blocks; + lp.by = 1; + lp.bz = 1; + lp.tx = threads; + lp.ty = 1; + lp.tz = 1; #ifdef CUDA_DEBUG - std::cerr << "[RunPTXTask] Launching kernel " << kernel_name - << " with blocks (" << bx << "," << by << "," << bz - << ") and threads (" << tx << "," << ty << "," << tz << ")" - << " on CUcontext " << context_to_string(ctx) << std::endl; + std::cerr << "[RunPTXBroadcastTask] local volume=" << volume << " dim=" << dim + << " -> blocks=" << lp.bx << " threads=" << lp.tx + << " (budget=" << budget << ")" << std::endl; #endif - // Launch the kernel - DRIVER_ERROR_CHECK(cuLaunchKernel(func, bx, by, bz, tx, ty, tz, 0, custream_, - nullptr, config)); +} + +/*static*/ void RunPTXBroadcastTask::gpu_variant(legate::TaskContext context) { + auto lp = read_launch_params(context); + + const std::size_t num_inputs = context.num_inputs(); + const std::size_t num_outputs = context.num_outputs(); + const std::size_t num_scalars = context.num_scalars(); + + assert(num_outputs >= 1); + broadcast_launch_dims_from_tile(lp, context.output(0)); + + // Read num_kernel_args first so we can size the buffer precisely + std::int32_t num_kernel_args = + context.scalar(ARG_OFFSET + 1).value(); + std::size_t map_start = ARG_OFFSET + 2; + std::size_t scalar_values_start = map_start + num_kernel_args; + + std::size_t max_buffer_size = + padded_bytes_kernel_state + + context.scalar(ARG_OFFSET).size() + // ctx (CompilerMetadata) + num_kernel_args * + (sizeof(CuStridedDeviceArray) + + 8); // worst case: all args are CuStridedDeviceArrays + alignment + for (std::size_t i = scalar_values_start; i < num_scalars; ++i) { + max_buffer_size += context.scalar(i).size(); + } + + std::vector arg_buffer(max_buffer_size); + char *p = arg_buffer.data() + padded_bytes_kernel_state; + + // 1. Write ctx (CompilerMetadata) + if (num_scalars > ARG_OFFSET) { + const auto &ctx_scalar = context.scalar(ARG_OFFSET); + memcpy(p, ctx_scalar.ptr(), ctx_scalar.size()); + p += ctx_scalar.size(); + } - // DRIVER_ERROR_CHECK(cuStreamSynchronize(stream_)); + // 2. Read arg_map and reconstruct arg buffer (strided descriptors) + for (std::int32_t i = 0; i < num_kernel_args; ++i) { + std::int32_t val = context.scalar(map_start + i).value(); + + if (val >= 0 && val < static_cast(num_outputs)) { + align8(p); + auto ps = context.output(val); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiStridedFunctor{}, + ufi::AccessMode::WRITE, p, ps); + } else if (val >= static_cast(num_outputs)) { + align8(p); + auto ps = context.input(val - num_outputs); + legate::double_dispatch(ps.dim(), ps.type().code(), ufiStridedFunctor{}, + ufi::AccessMode::READ, p, ps); + } else { + std::size_t scalar_idx = static_cast(-(val + 1)); + const auto &scalar = context.scalar(scalar_values_start + scalar_idx); + memcpy(p, scalar.ptr(), scalar.size()); + p += scalar.size(); + } + } + + launch_kernel(lp, arg_buffer, p - arg_buffer.data()); } // https://github.com/nv-legate/legate.pandas/blob/branch-22.01/src/udf/load_ptx.cc @@ -389,6 +613,13 @@ inline void add_xyz_scalars(legate::AutoTask &task, task.add_scalar_arg(legate::Scalar(xyz[2])); } +inline void add_scalar_from_ptr(legate::AutoTask &task, void *ptr, + size_t size) { + uint8_t *byte_ptr = static_cast(ptr); + std::vector vec(byte_ptr, byte_ptr + size); + task.add_scalar_arg(legate::Scalar(vec)); +} + void gpu_sync() { cudaStream_t stream_ = nullptr; ERROR_CHECK(cudaDeviceSynchronize()); @@ -412,9 +643,12 @@ void register_kernel_state_size(uint64_t st_size) { void wrap_cuda_methods(jlcxx::Module &mod) { mod.method("add_xyz_scalars", &add_xyz_scalars); + mod.method("add_scalar_from_ptr", &add_scalar_from_ptr); mod.method("register_kernel_state_size", ®ister_kernel_state_size); mod.method("gpu_sync", &gpu_sync); mod.method("extract_kernel_name", &extract_kernel_name); mod.set_const("LOAD_PTX", legate::LocalTaskID{ufi::TaskIDs::LOAD_PTX_TASK}); mod.set_const("RUN_PTX", legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_TASK}); + mod.set_const("RUN_PTX_BROADCAST", + legate::LocalTaskID{ufi::TaskIDs::RUN_PTX_BROADCAST_TASK}); } diff --git a/lib/cunumeric_jl_wrapper/src/ndarray.cpp b/lib/cunumeric_jl_wrapper/src/ndarray.cpp index ed157b077..de5f3f0f2 100644 --- a/lib/cunumeric_jl_wrapper/src/ndarray.cpp +++ b/lib/cunumeric_jl_wrapper/src/ndarray.cpp @@ -237,21 +237,20 @@ void nda_unary_reduction(CN_NDArray* out, CuPyNumericUnaryRedCode op_code, out->obj.unary_reduction(op_code, input->obj); } -CN_NDArray* nda_unary_reduction_axes(CuPyNumericUnaryRedCode op_code, CN_NDArray* input, const int32_t* axes, int32_t num_axes, bool keepdims) { - std::vector axis_vec(axes, axes + num_axes); - NDArray result = input->obj._perform_unary_reduction( - static_cast(op_code), - input->obj, - axis_vec, - std::nullopt, // dtype - std::nullopt, // res_dtype - std::nullopt, // out - keepdims, - {}, // args - std::nullopt, // initial - std::nullopt // where - ); - return new CN_NDArray{NDArray(std::move(result))}; +CN_NDArray* nda_unary_reduction_axes(CuPyNumericUnaryRedCode op_code, + CN_NDArray* input, const int32_t* axes, + int32_t num_axes, bool keepdims) { + std::vector axis_vec(axes, axes + num_axes); + NDArray result = input->obj._perform_unary_reduction( + static_cast(op_code), input->obj, axis_vec, + std::nullopt, // dtype + std::nullopt, // res_dtype + std::nullopt, // out + keepdims, {}, // args + std::nullopt, // initial + std::nullopt // where + ); + return new CN_NDArray{NDArray(std::move(result))}; } NDArray get_slice(NDArray arr, std::vector slices) { diff --git a/lib/cunumeric_jl_wrapper/src/types.cpp b/lib/cunumeric_jl_wrapper/src/types.cpp index f181dc2e1..f51369d6d 100644 --- a/lib/cunumeric_jl_wrapper/src/types.cpp +++ b/lib/cunumeric_jl_wrapper/src/types.cpp @@ -164,6 +164,8 @@ void wrap_binary_ops(jlcxx::Module& mod) { } void wrap_linalg_ops(jlcxx::Module& mod) { - mod.set_const("SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SOLVE}); - mod.set_const("MP_SOLVE", legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_SOLVE}); -} \ No newline at end of file + mod.set_const("SOLVE", + legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_SOLVE}); + mod.set_const("MP_SOLVE", + legate::LocalTaskID{CuPyNumericOpCode::CUPYNUMERIC_MP_SOLVE}); +} diff --git a/lib/cunumeric_jl_wrapper/src/wrapper.cpp b/lib/cunumeric_jl_wrapper/src/wrapper.cpp index c5c792faa..03434a514 100644 --- a/lib/cunumeric_jl_wrapper/src/wrapper.cpp +++ b/lib/cunumeric_jl_wrapper/src/wrapper.cpp @@ -58,6 +58,7 @@ void register_tasks() { auto library = get_lib(); ufi::LoadPTXTask::register_variants(library); ufi::RunPTXTask::register_variants(library); + ufi::RunPTXBroadcastTask::register_variants(library); } #endif diff --git a/scripts/install_cxxwrap.sh b/scripts/install_cxxwrap.sh index ba430d507..8335858cb 100755 --- a/scripts/install_cxxwrap.sh +++ b/scripts/install_cxxwrap.sh @@ -69,7 +69,7 @@ JULIA_CXXWRAP=$JULIA_CXXWRAP_DEV/override cd $CUNUMERIC_ROOT_DIR [ -f Manifest.toml ] && rm Manifest.toml rm -rf $JULIA_CXXWRAP_DEV -julia -e 'using Pkg; Pkg.activate("."); Pkg.add(url="https://github.com/JuliaLegate/Legate.jl")' +julia -e 'using Pkg; Pkg.activate("."); Pkg.add("Legate")' julia -e 'using Pkg; Pkg.activate("."); Pkg.precompile(["CxxWrap"])' # https://github.com/JuliaInterop/libcxxwrap-julia/tree/v0.13.3?tab=readme-ov-file#preparing-the-install-location diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl index 0da59605c..498948382 100644 --- a/src/cuNumeric.jl +++ b/src/cuNumeric.jl @@ -27,6 +27,12 @@ using Legate using Libdl using CxxWrap +using CUDATools: CUDATools +using CUDACore: CUDACore +import CUDACore: CuArray +import KernelAbstractions: @kernel, @index +import KernelAbstractions as KA + using cupynumeric_jll using cunumeric_jl_wrapper_jll @@ -139,25 +145,35 @@ include("memory.jl") # allowscalar and allowpromotion include("warnings.jl") +# Compile-time so task scope instrumentation is fully elided when disabled. +const TASK_SCOPE_NAMES = CNPreferences.TASK_SCOPE_NAMES + # NDArray internal include("ndarray/detail/ndarray.jl") -# NDArray interface +# Utilities +include("cuda/strided_device_array.jl") +include("cuda/cuda_util.jl") +include("utilities/version.jl") +include("util.jl") + +# Compile-time so the fusion branch is elided; flip via CNPreferences before loading. +const FUSE_BROADCAST_EXPRS = CNPreferences.FUSE_BROADCAST +# Fuse when broadcast tree length is at least this (see `_broadcast_tree_length`). +# Default 2 skips single-op exprs like `y .= cos.(x)`. Use 1 to fuse everything. +const FUSE_BROADCAST_MIN_OPS = CNPreferences.FUSE_BROADCAST_MIN_OPS + +# Functionality include("ndarray/promotion.jl") +include("cuda/cuda_ptx_task.jl") +include("ndarray/broadcast_fusion.jl") include("ndarray/broadcast.jl") include("ndarray/ndarray.jl") include("ndarray/unary.jl") include("ndarray/binary.jl") include("ndarray/linalg.jl") - -# scoping macro include("scoping.jl") -# Utilities -include("utilities/version.jl") -include("utilities/cuda_stubs.jl") -include("util.jl") - # From https://github.com/JuliaGraphics/QML.jl/blob/dca239404135d85fe5d4afe34ed3dc5f61736c63/src/QML.jl#L147 mutable struct ArgcArgv argv @@ -239,7 +255,11 @@ function __init__() # legate/cunumeric when using registry CI machines. get(ENV, "JULIA_REGISTRYCI_AUTOMERGE", false) == "true" && return nothing + # Start runtime, but only if not pre-compiling ensure_runtime!() + + # Requries runtime to be started + return _setup_cuda_tasking() end end #module cuNumeric diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl new file mode 100644 index 000000000..8e34fcb45 --- /dev/null +++ b/src/cuda/cuda_ptx_task.jl @@ -0,0 +1,307 @@ +export @cuda_task, @launch, CUDATask + +struct CUDATask + func::String + argtypes::NTuple{N,Type} where {N} #! THIS IS TYPE UNSTABLE +end + +#! JUST PASS TYPES HERE INSTEAD OF CALLING typeof() +map_ndarray_cuda_types(args...) = tuple(map(ndarray_cuda_type, typeof.(args))...) + +function to_stdvec(::Type{T}, vec) where {T} + stdvec = CxxWrap.StdVector{T}() + for x in vec + push!(stdvec, T(x)) + end + 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) + + 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) + 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) +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 + end +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 + end +end + +# Unused by Launch (which uses `_add_task_array!` + eager finalize). If revived, +# callers must finalize the returned LogicalArray's handle after add_input/output. +function nda_to_logical_array(arr::NDArray{T,N}) where {T,N} + st_handle = cuNumeric.get_store(arr) + return Legate.LogicalArray{T,N}(st_handle, size(arr)) +end + +# `get_store` returns a Julia-owned `LogicalArrayImplAllocated` that shares the +# underlying Legate store with the NDArray. `add_input`/`add_output` copy that +# 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 +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) + + # we find the largest input/output. + ndarrays = vcat(inputs..., outputs...) + mx = findmax(arr -> arr.nbytes, ndarrays) # returns (nbytes, position) + max_size = mx[1] # first elem nbytes + max_shape = size(ndarrays[mx[2]]) # second elem max position + @assert !isnothing(max_shape) + + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + task = Legate.create_auto_task(rt, lib, taskid) + + input_vars = Vector{Legate.Variable}() + for arr in inputs + check_sz!(arr, max_shape; copy=true) + push!(input_vars, _add_task_array!(Legate.add_input, task, arr)) + end + + output_vars = Vector{Legate.Variable}() + for arr in outputs + check_sz!(arr, max_shape; copy=false) + push!(output_vars, _add_task_array!(Legate.add_output, task, arr)) + end + + # Reserved scalars: kernel_name (0), blocks (1,2,3), threads (4,5,6) + Legate.add_scalar(task, Legate.string_to_scalar(kernel.func)) # 0 + cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, blocks)) # 1,2,3 + cuNumeric.add_xyz_scalars(task, to_stdvec(UInt32, threads)) # 4,5,6 + + # CompilerMetadata ctx for broadcast tasks (scalar 7, raw bytes) + if !isnothing(ctx) + ref = Ref(ctx) + GC.@preserve ref begin + cuNumeric.add_scalar_from_ptr(task, Base.unsafe_convert(Ptr{Cvoid}, ref), sizeof(ctx)) + end + end + + # User-defined scalars + for s in scalars + Legate.add_scalar(task, Legate.Scalar(s)) + end + + # all inputs are aligned with all outputs + Legate.default_alignment(task, input_vars, output_vars) + return Legate.submit_auto_task(rt, task) +end + +function launch(kernel::CUDATask, inputs, outputs, scalars; + blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing) + return Launch(kernel, + isa(inputs, Tuple) ? inputs : (inputs,), + isa(outputs, Tuple) ? outputs : (outputs,), + isa(scalars, Tuple) ? scalars : (scalars,); + blocks=isa(blocks, Tuple) ? blocks : (blocks,), + threads=isa(threads, Tuple) ? threads : (threads,), + taskid=taskid, + ctx=ctx, + ) +end + +function ptx_task(ptx::String, kernel_name) + rt = Legate.get_runtime() + lib = cuNumeric.get_lib() + taskid = cuNumeric.LOAD_PTX + + # One point task per GPU so every GPU compiles the module. + ngpus = max(Int(Legate.num_gpus()), 1) + domain = Legate.domain_from_shape(Legate.Shape(Legate.to_cxx_vector((ngpus,)))) + task = Legate.create_manual_task(rt, lib, taskid, domain) + Legate.add_scalar(task, Legate.string_to_scalar(ptx)) + Legate.add_scalar(task, Legate.string_to_scalar(kernel_name)) + Legate.submit_manual_task(rt, task) + + # Fence so every load finishes before any launch reads the cache. + return issue_execution_fence(; block=false) +end + +""" + @cuda_task(f(args...)) + +Compile a Julia GPU kernel to PTX, register it with the Legate runtime, +and return a `CUDATask` object for later launch. + +# Arguments +- `f` — The name of the Julia CUDA.jl GPU kernel function to compile. +- `args...` — Example arguments to the kernel, used to determine the + argument type signature when generating PTX. + +# Description +This macro automates the process of: +1. Inferring the CUDA argument types for the given `args` using + `map_ndarray_cuda_types`. +2. Using `CUDA.code_ptx` to compile the specified GPU kernel + (`f`) into raw PTX text for the inferred types. +3. Extracting the kernel's function symbol name from the PTX using + `extract_kernel_name`. +4. Registering the compiled PTX and kernel name with the Legate runtime + via `ptx_task`, making it available for GPU execution. +5. Returning a `CUDATask` struct that stores the kernel name and type signature, + which can be used to configure and launch the kernel later. + +# Notes +- The `args...` are not executed; they are used solely for type inference. +- This macro is intended for use with the Legate runtime and + assumes a CUDA context is available. +- Make sure your kernel code is GPU-compatible and does not rely on + unsupported Julia features. + +# Example +```julia +mytask = @cuda_task my_kernel(A, B, C) +``` +""" +macro cuda_task(call_expr) + cuNumeric.assert_experimental() + + fname = call_expr.args[1] + fargs = call_expr.args[2:end] + + return esc( + quote + local _buf = IOBuffer() + local _types = map_ndarray_cuda_types($(fargs...)) + # generate ptx using CUDA.jl + CUDATools.code_ptx(_buf, $fname, _types; raw=false, dump_module=true, kernel=true) + + local _ptx = String(take!(_buf)) + local _func_name = extract_kernel_name(_ptx) + + # issue ptx_task within legate runtime to register cufunction ptr with cucontext + ptx_task(_ptx, _func_name) + + # create a cuNumeric.CUDAtask that stores some info for a launch config + CUDATask(_func_name, _types) + end, + ) +end + +""" + @launch(; task, blocks=(1,), threads=(256,), inputs=(), outputs=(), scalars=()) + +Launch a GPU kernel (previously registered via [`@cuda_task`](@ref)) through the Legate runtime. + +# Keywords +- `task` — A `CUDATask` object, typically returned by [`@cuda_task`](@ref). +- `blocks` — Tuple or single element specifying the CUDA grid dimensions. Defaults to `(1,)`. +- `threads` — Tuple or single element specifying the CUDA block dimensions. Defaults to `(256,)`. +- `inputs` — Tuple or single element of input NDArray objects. +- `outputs` — Tuple or single element of output NDArray objects. +- `scalars` — Tuple or single element of scalar values. + +# Description +The `@launch` macro validates the provided keywords, ensuring only +the allowed set (`:task`, `:blocks`, `:threads`, `:inputs`, `:outputs`, `:scalars`) +are present. It then expands to a call to `cuNumeric.launch`, +passing the given arguments to the Legate runtime for execution. + +This macro is meant to provide a concise, declarative syntax for +launching GPU kernels, separating kernel compilation (via `@cuda_task`) +from execution configuration. + +# Notes +- `task` **must** be a kernel registered with the runtime, usually from `@cuda_task`. +- All keyword arguments must be specified as assignments, e.g. `blocks=(2,2)` not positional arguments. +- Defaults are chosen for single-block, 256-thread 1D launches. +- The macro escapes its body so that the values of inputs/outputs/scalars are captured + from the surrounding scope at macro expansion time. + +# Example +```julia +mytask = @cuda_task my_kernel(A, B, C) + +@launch task=mytask blocks=(8,8) threads=(32,32) inputs=(A, B) outputs=(C) +``` +""" +macro launch(args...) + cuNumeric.assert_experimental() + + allowed_keys = Set([:task, :blocks, :threads, :inputs, :outputs, :scalars]) + kwargs = Dict{Symbol,Any}() + + for ex in args + if !(ex isa Expr && ex.head == :(=)) + error("All arguments must be keyword assignments, e.g. task=..., threads=...") + end + key = ex.args[1] + val = ex.args[2] + + if !(key in allowed_keys) + error("@launch macro received unexpected keyword: $(key)") + end + + kwargs[key] = val + end + + if !haskey(kwargs, :task) + error("@launch macro requires 'task=...' to be provided.") + end + task = kwargs[:task] + blocks = get(kwargs, :blocks, :((1))) + threads = get(kwargs, :threads, :((256))) + inputs = get(kwargs, :inputs, :(())) + outputs = get(kwargs, :outputs, :(())) + scalars = get(kwargs, :scalars, :(())) + + return esc( + quote + cuNumeric.launch( + $task, $inputs, $outputs, $scalars; + blocks=($blocks), threads=($threads), + ) + end, + ) +end diff --git a/src/cuda/cuda_util.jl b/src/cuda/cuda_util.jl new file mode 100644 index 000000000..68b141c77 --- /dev/null +++ b/src/cuda/cuda_util.jl @@ -0,0 +1,46 @@ +const KERNEL_OFFSET = sizeof(CUDACore.KernelState) + +function _setup_cuda_tasking() + if CUDACore.functional() + # in cuda.jl to notify /wrapper/src/cuda.cpp about CUDA.jl kernel state size + register_kernel_state_size(UInt64(KERNEL_OFFSET)) + # in /wrapper/src/cuda.cpp + register_tasks() + else + @warn "CUDA.jl is not functional; skipping CUDA kernel registration." + end +end + +# Dense @cuda_task / RunPTXTask — MUST match CUDA.jl CuDeviceArray layout. +# Other memory types: https://github.com/JuliaGPU/CUDA.jl/blob/345c1600ebd561135148bb04ee2657f521a40e25/CUDACore/src/device/pointer.jl#L7 +function ndarray_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} + CUDACore.CuDeviceArray{T,N,CUDACore.AS.Global} +end + +function ndarray_cuda_type(::Type{T}) where {T} + Base.isbitstype(T) || throw(ArgumentError("Unsupported argument type: $(T)")) + return T +end + +""" + map_cuda_type(::Type{T})::Type + +Recursively rewrite cuNumeric broadcast-related types for fused-broadcast PTX +(e.g. mapping `NDArray{...}` to `CuStridedDeviceArray{...}`). Dense `@cuda_task` +uses `ndarray_cuda_type` → CUDA.jl `CuDeviceArray` instead. +""" +map_cuda_type(::Type{T}) where {T} = T + +map_cuda_type(::Type{<:NDArray{T,N}}) where {T,N} = CuStridedDeviceArray{T,N,CUDACore.AS.Global} + +function map_cuda_type(::Type{T}) where {T<:Tuple} + return Tuple{map_cuda_type.(T.parameters)...} +end + +function map_cuda_type(::Type{Base.Broadcast.Broadcasted{S,Ax,F,Args}}) where {S,Ax,F,Args} + return Base.Broadcast.Broadcasted{map_cuda_type(S),Ax,F,map_cuda_type(Args)} +end + +function map_cuda_type(::Type{Base.Broadcast.Extruded{X,K,D}}) where {X,K,D} + return Base.Broadcast.Extruded{map_cuda_type(X),K,D} +end diff --git a/src/cuda/strided_device_array.jl b/src/cuda/strided_device_array.jl new file mode 100644 index 000000000..58efd836e --- /dev/null +++ b/src/cuda/strided_device_array.jl @@ -0,0 +1,117 @@ +#= 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. +=# + +# Device-side strided array packed by RunPTXBroadcastTask only. +# Layout must match C++ `CuStridedDeviceArray` in +# lib/cunumeric_jl_wrapper/src/cuda.cpp: +# ptr, maxsize, dims[N], strides[N] (element strides), length +# +# Dense RunPTXTask / @cuda_task still uses CUDA.jl CuDeviceArray (unchanged). + +struct CuStridedDeviceArray{T,N,A} <: AbstractArray{T,N} + ptr::CUDACore.LLVMPtr{T,A} + maxsize::Int + dims::Dims{N} + strides::Dims{N} + len::Int +end + +Base.elsize(::Type{<:CuStridedDeviceArray{T}}) where {T} = sizeof(T) +Base.size(a::CuStridedDeviceArray) = a.dims +Base.size(a::CuStridedDeviceArray{<:Any,1}) = (a.len,) +Base.length(a::CuStridedDeviceArray) = a.len +Base.IndexStyle(::Type{<:CuStridedDeviceArray}) = IndexLinear() + +function Base.pointer(a::CuStridedDeviceArray{T,<:Any,A}) where {T,A} + Base.unsafe_convert(CUDACore.LLVMPtr{T,A}, a) +end +function Base.unsafe_convert( + ::Type{CUDACore.LLVMPtr{T,A}}, a::CuStridedDeviceArray{T,<:Any,A} +) where {T,A} + a.ptr +end + +# 0-based element offset from a 1-based linear index in Julia column-major order +# over `dims`, using Legate element `strides`. +# +# Must not use checked rem/div or signed↔unsigned converts — those emit +# DivideError / InexactError → gpu_report_exception and break LoadPTX. +@inline _bitcast_uint(x::Int) = reinterpret(UInt, x) +@inline _bitcast_int(x::UInt) = reinterpret(Int, x) + +@inline function _strided_elem_offset(dims::Dims{N}, strides::Dims{N}, I::Integer) where {N} + idx = _bitcast_uint(Int(I) - 1) + off = 0 + @inbounds for d in 1:N + dlen = _bitcast_uint(Int(dims[d])) + c = _bitcast_int(Core.Intrinsics.urem_int(idx, dlen)) + idx = Core.Intrinsics.udiv_int(idx, dlen) + off += c * Int(strides[d]) + end + return off +end + +@inline function _strided_elem_offset(strides::Dims{N}, I::CartesianIndex{N}) where {N} + off = 0 + @inbounds for d in 1:N + off += (Int(I[d]) - 1) * Int(strides[d]) + end + return off +end + +@inline _strided_align(::CuStridedDeviceArray{T}) where {T} = Base.datatype_alignment(T) + +# No @boundscheck / throw — those emit gpu_report_exception and break LoadPTX. +CUDACore.@device_function @inline function _strided_arrayref( + A::CuStridedDeviceArray{T}, index::Integer +) where {T} + off = _strided_elem_offset(A.dims, A.strides, index) + return unsafe_load(pointer(A), off + 1, Val(_strided_align(A))) +end + +CUDACore.@device_function @inline function _strided_arrayset( + A::CuStridedDeviceArray{T}, x::T, index::Integer +) where {T} + off = _strided_elem_offset(A.dims, A.strides, index) + unsafe_store!(pointer(A), x, off + 1, Val(_strided_align(A))) + return A +end + +Base.@propagate_inbounds Base.getindex(A::CuStridedDeviceArray{T}, i::Integer) where {T} = _strided_arrayref( + A, i +) +Base.@propagate_inbounds function Base.setindex!( + A::CuStridedDeviceArray{T}, x, i::Integer +) where {T} + return _strided_arrayset(A, convert(T, x)::T, i) +end + +Base.to_index(::CuStridedDeviceArray, i::Integer) = i + +Base.@propagate_inbounds function Base.getindex( + A::CuStridedDeviceArray{T,N}, I::CartesianIndex{N} +) where {T,N} + off = _strided_elem_offset(A.strides, I) + return unsafe_load(pointer(A), off + 1, Val(_strided_align(A))) +end + +Base.@propagate_inbounds function Base.setindex!( + A::CuStridedDeviceArray{T,N}, x, I::CartesianIndex{N} +) where {T,N} + off = _strided_elem_offset(A.strides, I) + unsafe_store!(pointer(A), convert(T, x)::T, off + 1, Val(_strided_align(A))) + return A +end diff --git a/src/memory.jl b/src/memory.jl index 2bf447dcc..4a3442940 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -16,6 +16,11 @@ const current_host_bytes = Atomic{Int64}(0) # predicted host allocations const soft_frac = Ref{Float64}(0.80) const hard_frac = Ref{Float64}(0.90) const AUTO_GC_ENABLE = Ref{Bool}(false) +# memory measured right after the last GC +const post_gc_device_bytes = Atomic{Int64}(0) +const post_gc_host_bytes = Atomic{Int64}(0) +# how much new memory must accumulate before GC fires again +const gc_hysteresis_frac = Ref{Float64}(0.05) # memory measured right after the last GC const post_gc_device_bytes = Atomic{Int64}(0) diff --git a/src/ndarray/binary.jl b/src/ndarray/binary.jl index c04b515bc..07745f483 100644 --- a/src/ndarray/binary.jl +++ b/src/ndarray/binary.jl @@ -41,17 +41,91 @@ global const floaty_binary_op_map = Dict{Function,BinaryOpCode}( ) ## SPECIAL CASES ## +# Promote into out's eltype, then destroy any new temps (dispatch; no runtime !==). +@inline function _nda_binary_op_promoted!( + out::NDArray{T}, op, rhs1::NDArray{T}, rhs2::NDArray{T} +) where {T} + return nda_binary_op!(out, op, rhs1, rhs2) +end +function _nda_binary_op_promoted!(out::NDArray{T}, op, rhs1::NDArray, rhs2::NDArray{T}) where {T} + p1 = unchecked_promote_arr(rhs1, T) # always new when eltype ≠ T + result = nda_binary_op!(out, op, p1, rhs2) + destroy!(p1) + return result +end +function _nda_binary_op_promoted!(out::NDArray{T}, op, rhs1::NDArray{T}, rhs2::NDArray) where {T} + p2 = unchecked_promote_arr(rhs2, T) + result = nda_binary_op!(out, op, rhs1, p2) + destroy!(p2) + return result +end +function _nda_binary_op_promoted!(out::NDArray{T}, op, rhs1::NDArray, rhs2::NDArray) where {T} + p1 = unchecked_promote_arr(rhs1, T) + p2 = unchecked_promote_arr(rhs2, T) + result = nda_binary_op!(out, op, p1, p2) + destroy!(p1) + destroy!(p2) + return result +end + +@inline function _nda_three_dot_promoted!( + rhs1::NDArray{T}, rhs2::NDArray{T}, out::NDArray{T} +) where {T} + return nda_three_dot_arg(rhs1, rhs2, out) +end +function _nda_three_dot_promoted!(rhs1::NDArray, rhs2::NDArray{T}, out::NDArray{T}) where {T} + p1 = unchecked_promote_arr(rhs1, T) + result = nda_three_dot_arg(p1, rhs2, out) + destroy!(p1) + return result +end +function _nda_three_dot_promoted!(rhs1::NDArray{T}, rhs2::NDArray, out::NDArray{T}) where {T} + p2 = unchecked_promote_arr(rhs2, T) + result = nda_three_dot_arg(rhs1, p2, out) + destroy!(p2) + return result +end +function _nda_three_dot_promoted!(rhs1::NDArray, rhs2::NDArray, out::NDArray{T}) where {T} + p1 = unchecked_promote_arr(rhs1, T) + p2 = unchecked_promote_arr(rhs2, T) + result = nda_three_dot_arg(p1, p2, out) + destroy!(p1) + destroy!(p2) + return result +end + +@inline function _nda_three_dot_checked!( + rhs1::NDArray{T}, rhs2::NDArray{T}, out::NDArray{T} +) where {T} + return nda_three_dot_arg(rhs1, rhs2, out) +end +function _nda_three_dot_checked!(rhs1::NDArray, rhs2::NDArray{T}, out::NDArray{T}) where {T} + p1 = checked_promote_arr(rhs1, T) + result = nda_three_dot_arg(p1, rhs2, out) + destroy!(p1) + return result +end +function _nda_three_dot_checked!(rhs1::NDArray{T}, rhs2::NDArray, out::NDArray{T}) where {T} + p2 = checked_promote_arr(rhs2, T) + result = nda_three_dot_arg(rhs1, p2, out) + destroy!(p2) + return result +end +function _nda_three_dot_checked!(rhs1::NDArray, rhs2::NDArray, out::NDArray{T}) where {T} + p1 = checked_promote_arr(rhs1, T) + p2 = checked_promote_arr(rhs2, T) + result = nda_three_dot_arg(p1, p2, out) + destroy!(p1) + destroy!(p2) + return result +end + # Do not need broadcast operation when same shape function Base.:(-)(rhs1::NDArray{A,N}, rhs2::NDArray{B,N}) where {A,B,N} promote_shape(size(rhs1), size(rhs2)) T_OUT = __checked_promote_op(-, A, B) out = cuNumeric.zeros(T_OUT, size(rhs1)) - return nda_binary_op!( - out, - cuNumeric.SUBTRACT, - unchecked_promote_arr(rhs1, T_OUT), - unchecked_promote_arr(rhs2, T_OUT), - ) + return _nda_binary_op_promoted!(out, cuNumeric.SUBTRACT, rhs1, rhs2) end # Do not need broadcast operation when same shape @@ -59,19 +133,18 @@ function Base.:(+)(rhs1::NDArray{A,N}, rhs2::NDArray{B,N}) where {A,B,N} promote_shape(size(rhs1), size(rhs2)) T_OUT = __checked_promote_op(+, A, B) out = cuNumeric.zeros(T_OUT, size(rhs1)) - return nda_binary_op!( - out, cuNumeric.ADD, unchecked_promote_arr(rhs1, T_OUT), unchecked_promote_arr(rhs2, T_OUT) - ) + return _nda_binary_op_promoted!(out, cuNumeric.ADD, rhs1, rhs2) end -function Base.:(*)(val::V, arr::NDArray{A}) where {A,V} - T = __my_promote_type(A, V) - out = cuNumeric.zeros(T, size(arr)) - return nda_binary_op!(out, cuNumeric.MULTIPLY, NDArray(T(val)), unchecked_promote_arr(arr, T)) -end +Base.:(*)(val::V, arr::NDArray{A}) where {A,V} = _mul_scalar(__my_promote_type(A, V), val, arr) +Base.:(*)(arr::NDArray{A}, val::V) where {A,V} = val * arr -function Base.:(*)(arr::NDArray{A}, val::V) where {A,V} - val * arr +_mul_scalar(::Type{T}, val, arr::NDArray{T}) where {T} = nda_multiply_scalar(arr, T(val)) +function _mul_scalar(::Type{U}, val, arr::NDArray) where {U} + promoted = unchecked_promote_arr(arr, U) # always a new array when U ≠ eltype + out = nda_multiply_scalar(promoted, U(val)) + destroy!(promoted) + return out end function Base.:(*)(rhs1::NDArray{A,2}, rhs2::NDArray{B,2}) where {A,B} @@ -79,7 +152,7 @@ function Base.:(*)(rhs1::NDArray{A,2}, rhs2::NDArray{B,2}) where {A,B} throw(DimensionMismatch("Matrix dimensions incompatible: $(size(rhs1)) × $(size(rhs2))")) T = __my_promote_type(A, B) out = cuNumeric.zeros(T, (size(rhs1, 1), size(rhs2, 2))) - return nda_three_dot_arg(unchecked_promote_arr(rhs1, T), unchecked_promote_arr(rhs2, T), out) + return _nda_three_dot_promoted!(rhs1, rhs2, out) end function Base.:(*)(rhs1::NDArray{Bool,2}, rhs2::NDArray{Bool,2}) @@ -143,7 +216,7 @@ function LinearAlgebra.mul!( ), ) end - return nda_three_dot_arg(checked_promote_arr(rhs1, T), checked_promote_arr(rhs2, T), out) + return _nda_three_dot_checked!(rhs1, rhs2, out) end function LinearAlgebra.mul!(out::NDArray, rhs1::NDArray{Bool,2}, rhs2::NDArray{Bool,2}) @@ -182,11 +255,16 @@ for (julia_fn, op_code) in floaty_binary_op_map return nda_binary_op!(out, $(op_code), rhs1, rhs2) end - # If input is not already float, promote to that + # If input is not already float, promote to that (temps always new → always destroy) @inline function __broadcast( f::typeof($(julia_fn)), out::NDArray{A}, rhs1::NDArray{B}, rhs2::NDArray{B} ) where {A<:SUPPORTED_FLOAT_TYPES,B<:Union{SUPPORTED_INT_TYPES,Bool}} - return __broadcast(f, out, checked_promote_arr(rhs1, A), checked_promote_arr(rhs2, A)) + p1 = checked_promote_arr(rhs1, A) + p2 = checked_promote_arr(rhs2, A) + result = __broadcast(f, out, p1, p2) + destroy!(p1) + destroy!(p2) + return result end end end @@ -195,18 +273,24 @@ end f::typeof(Base.:(+)), out::NDArray{O}, rhs1::NDArray{Bool}, rhs2::NDArray{Bool} ) where {O<:Integer} assertpromotion(".+", Bool, O) - return nda_binary_op!( - out, cuNumeric.ADD, unchecked_promote_arr(rhs1, O), unchecked_promote_arr(rhs2, O) - ) + p1 = unchecked_promote_arr(rhs1, O) # always new (Bool → O) + p2 = unchecked_promote_arr(rhs2, O) + result = nda_binary_op!(out, cuNumeric.ADD, p1, p2) + destroy!(p1) + destroy!(p2) + return result end @inline function __broadcast( f::typeof(Base.:(-)), out::NDArray{O}, rhs1::NDArray{Bool}, rhs2::NDArray{Bool} ) where {O<:Integer} assertpromotion(".-", Bool, O) - return nda_binary_op!( - out, cuNumeric.SUBTRACT, unchecked_promote_arr(rhs1, O), unchecked_promote_arr(rhs2, O) - ) + p1 = unchecked_promote_arr(rhs1, O) + p2 = unchecked_promote_arr(rhs2, O) + result = nda_binary_op!(out, cuNumeric.SUBTRACT, p1, p2) + destroy!(p1) + destroy!(p2) + return result end # function Base.:(==)(lhs::NDArray{A}, rhs::NDArray{B}) where {A,B} diff --git a/src/ndarray/broadcast.jl b/src/ndarray/broadcast.jl index fcf0a3e52..ec931b290 100644 --- a/src/ndarray/broadcast.jl +++ b/src/ndarray/broadcast.jl @@ -4,6 +4,11 @@ struct NDArrayStyle{N} <: AbstractArrayStyle{N} end Base.BroadcastStyle(::Type{<:NDArray{<:Any,N}}) where {N} = NDArrayStyle{N}() Base.BroadcastStyle(::NDArrayStyle{N}, ::NDArrayStyle{M}) where {N,M} = NDArrayStyle{max(N, M)}() +# Some other functions in cuda_util.jl +function map_cuda_type(::Type{cuNumeric.NDArrayStyle{N}}) where {N} + return CUDACore.CuArrayStyle{N,CUDACore.DeviceMemory} +end # Also can be HostMemory or UnifiedMemory + function _nd_forbid_mix() throw( ArgumentError( @@ -34,11 +39,12 @@ Base.similar(arr::NDArray, ::Type{T}) where {T} = similar(arr, T, size(arr)) #* IS THERE A BETTER WAY TO ALLOCATE THE NEW ARRAY??? Base.similar(::Type{NDArray{T}}, axes) where {T} = cuNumeric.zeros(T, Base.to_shape.(axes)) function Base.similar(bc::Broadcasted{NDArrayStyle{N}}, ::Type{ElType}) where {N,ElType} - similar(NDArray{ElType}, axes(bc)) + return similar(NDArray{ElType}, axes(bc)) end function __broadcast(f::Function, _, args...) - error( + #! WITH FUSION I THINK WE CAN SUPPORT THIS BY JUST CALLING MAP or MAP! + return error( """ Tried to broadcast $(f). cuNumeric.jl does not support broadcasting user-defined functions yet. Please re-define \ functions to match supported patterns. For example g(x) = x + 1 could be re-defined as \ @@ -52,23 +58,33 @@ end bcast_depth(bc::Base.Broadcast.Broadcasted) = maximum(bcast_depth, bc.args, init=0) + 1; bcast_depth(::Any) = 0 -function Base.Broadcast.materialize(bc::Broadcasted{<:NDArrayStyle}) +struct BrokenBroadcast{T} end +Base.convert(::Type{BrokenBroadcast{T}}, x) where {T} = BrokenBroadcast{T}() +Base.convert(::Type{BrokenBroadcast{T}}, x::BrokenBroadcast{T}) where {T} = x +Base.eltype(::Type{BrokenBroadcast{T}}) where {T} = T + +function Broadcast.copy(bc::Broadcasted{<:NDArrayStyle{0}}) ElType = Broadcast.combine_eltypes(bc.f, bc.args) - if ElType == Union{} || !Base.allocatedinline(ElType) - error("Cannot broadcast $(bc.f) over NDArrays with eltypes: $(eltype.(bc.args))") + if ElType == Union{} + ElType = Nothing end + dest = copyto!(similar(bc, ElType), bc) + #! CHECK THIS DOESNT CAUSE ISSUES DUE TO BLOCKING NATURE + return @allowscalar dest[CartesianIndex()] +end - #* This be the place to inject kernel fusion via CUDA.jl - #* Use the function in Base.Broadcast.flatten(bc). - #* How can we check all the funcs in this expr - #* are supported by CUDA? - - return unravel_broadcast_tree(bc) +@inline function Broadcast.copy(bc::Broadcasted{<:NDArrayStyle}) + ElType = Broadcast.combine_eltypes(bc.f, bc.args) + if ElType == Union{} || !Base.allocatedinline(ElType) + ElType = BrokenBroadcast{ElType} + end + return copyto!(similar(bc, ElType), bc) end # Recursion base cases __materialize(x::NDArray) = x -__materialize(x::Number) = NDArray(x) +# Keep Numbers as scalars; unchecked_promote_arr builds the 0-d NDArray once. +__materialize(x::Number) = x # These are necessary to handle integer powers __materialize(x::Base.RefValue{typeof(^)}) = x @@ -79,11 +95,30 @@ __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. +@inline function __materialize(bc::Broadcasted{<:DefaultArrayStyle{0}}) + return bc.f((__materialize.(bc.args))...) +end + function __materialize(bc::Broadcasted{<:NDArrayStyle}) bc = Base.Broadcast.instantiate(bc) - unravel_broadcast_tree(bc) + return unravel_broadcast_tree(bc) +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 + destroy!(promoted) + end + if materialized isa NDArray && !(orig isa NDArray) + destroy!(materialized) + end + return nothing end +# Un-fused implementation of broadcast tree function unravel_broadcast_tree(bc::Broadcasted) # Recursively materialize/unravel any nested broadcasts @@ -106,13 +141,76 @@ function unravel_broadcast_tree(bc::Broadcasted) # If not it falls back to a pass-through that just calls # the Julia function and assumes the user defined a function # composed of supported operations. - return __broadcast(bc.f, out, in_args...) + result = __broadcast(bc.f, out, in_args...) + for i in eachindex(materialized_args) + _destroy_unfused_arg_temps!(bc.args[i], materialized_args[i], in_args[i]) + end + return result end -# Support .= -function Base.copyto!(dest::NDArray{T,N}, bc::Broadcasted{<:NDArrayStyle{N}}) where {T,N} - # Moves result from broadcast (src) to dest. src array is no longer valid - #! THIS ENABLES FOOT GUN IF USER SPECIFIES INTEGER ARRAY AT OUTPUT - nda_move(dest, checked_promote_arr(Base.Broadcast.materialize(bc), T)) +@inline function _copyto_unfused!(dest::NDArray{T}, temp_result::NDArray{T}) where {T} + nda_move(dest, temp_result) return dest end + +@inline function _copyto_unfused!(dest::NDArray{T}, temp_result::NDArray) where {T} + promoted = checked_promote_arr(temp_result, T) + nda_move(dest, promoted) + destroy!(temp_result) + return dest +end + +# Number of nested `Broadcasted` nodes (ops) in the pre-flatten tree. +@inline _broadcast_tree_length(@nospecialize(_)) = 0 +@inline _broadcast_tree_length(bc::Broadcasted) = + 1 + _broadcast_tree_length_args(bc.args) +@inline _broadcast_tree_length_args(::Tuple{}) = 0 +@inline function _broadcast_tree_length_args(args::Tuple) + return _broadcast_tree_length(getfield(args, 1)) + + _broadcast_tree_length_args(Base.tail(args)) +end + +# Prefer fusion only when the tree has at least `FUSE_BROADCAST_MIN_OPS` ops. +# When that const is <= 1, every Broadcasted qualifies and the length check +# compiles out (`@static`). +@inline function _should_attempt_broadcast_fusion(dest::NDArray, bc::Broadcasted) + @static if FUSE_BROADCAST_MIN_OPS <= 1 + return can_fuse_linear_broadcast(dest, bc) + else + return _broadcast_tree_length(bc) >= FUSE_BROADCAST_MIN_OPS && + can_fuse_linear_broadcast(dest, bc) + end +end + +@inline function _copyto!(dest::NDArray, bc::Broadcasted) + axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc)) + isempty(dest) && return dest + if eltype(dest) <: BrokenBroadcast + throw( + ArgumentError( + "Broadcast operation resulting in $(eltype(eltype(dest))) is not NDArray compatible" + ), + ) + 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`. + # Linear-only 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) + return fuse_broadcast_tree!(dest, bc) + else + return _copyto_unfused!(dest, unravel_broadcast_tree(bc)) + end + else + return _copyto_unfused!(dest, unravel_broadcast_tree(bc)) + end +end + +# Support .= +@inline Base.copyto!(dest::NDArray, bc::Broadcasted{Nothing}) = _copyto!(dest, bc) +@inline Base.copyto!(dest::NDArray, bc::Broadcasted{<:NDArrayStyle}) = _copyto!(dest, bc) + +#! TODO ADD MAP FUSED IMPLEMENTATIONS diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl new file mode 100644 index 000000000..1de900eef --- /dev/null +++ b/src/ndarray/broadcast_fusion.jl @@ -0,0 +1,619 @@ + +struct RuntimeBroadcastArg{J} end +struct StaticBroadcastArg{J} end + +Base.@propagate_inbounds @inline function _gpu_broadcast_getindex(x, I) + return @inbounds Base.Broadcast._broadcast_getindex(x, I) +end + +Base.@propagate_inbounds _gpu_broadcast_getindex(x::Number, I) = x + +# Bypass Broadcast.newindex: for N-D arrays, `_broadcast_getindex(A, I::Int)` becomes +# `A[CartesianIndex(I, 1, 1, ...)]`, which is wrong for our linear work ids. +# Dest already uses linear `dest[I]`; reads must use the same strided linear path. +Base.@propagate_inbounds @inline function _gpu_broadcast_getindex( + x::CuStridedDeviceArray, I::Integer +) + return @inbounds x[I] +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_arg( + ::RuntimeBroadcastArg{J}, + runtime_args, + static_args, + I, +) where {J} + arg = getfield(runtime_args, J) + return @inbounds _gpu_broadcast_getindex(arg, I) +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_arg( + ::StaticBroadcastArg{J}, + runtime_args, + static_args, + I, +) where {J} + return getfield(static_args, J) +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_args( + ::Tuple{}, + runtime_args, + static_args, + I, +) + return () +end + +Base.@propagate_inbounds @inline function _materialize_broadcast_args( + plan::Tuple, + runtime_args, + static_args, + I, +) + head = getfield(plan, 1) + tail = Base.tail(plan) + + return ( + @inbounds(_materialize_broadcast_arg(head, runtime_args, static_args, I)), + @inbounds(_materialize_broadcast_args(tail, runtime_args, static_args, I))..., + ) +end + +_is_runtime_broadcast_arg(x::NDArray) = true +_is_runtime_broadcast_arg(x::Base.Broadcast.Extruded) = true +_is_runtime_broadcast_arg(x::Number) = true +_is_runtime_broadcast_arg(x) = false + +function _push_runtime_arg!(runtime_args, arg_plan, x) + push!(runtime_args, x) + push!(arg_plan, RuntimeBroadcastArg{length(runtime_args)}()) + return nothing +end + +function _push_static_arg!(static_args, arg_plan, x) + isbits(x) || throw( + ArgumentError( + "Broadcast fusion cannot statically capture non-isbits broadcast leaf " * + "$(repr(x)) of type $(typeof(x))", + ), + ) + + push!(static_args, x) + push!(arg_plan, StaticBroadcastArg{length(static_args)}()) + return nothing +end + +function split_broadcast_args_for_kernel(args::Tuple) + runtime_args = Any[] + static_args = Any[] + arg_plan = Any[] + + for arg in args + if arg isa Base.RefValue + value = arg[] + + # Keep ordinary numeric broadcast scalars dynamic so scalar values + # do not cause recompilation. + if value isa Number + _push_runtime_arg!(runtime_args, arg_plan, value) + else + # Function singletons, Val{N}(), etc. + _push_static_arg!(static_args, arg_plan, value) + end + + elseif _is_runtime_broadcast_arg(arg) + _push_runtime_arg!(runtime_args, arg_plan, arg) + + elseif isbits(arg) + # Conservative fallback for singleton/isbits scalar broadcast leaves. + _push_static_arg!(static_args, arg_plan, arg) + + else + throw( + ArgumentError( + "Broadcast fusion does not know how to lower argument " * + "$(repr(arg)) of type $(typeof(arg))", + ), + ) + end + end + + return tuple(runtime_args...), tuple(static_args...), tuple(arg_plan...) +end + +############## + +@inline function _broadcast_linear_work_id() + return (Int(CUDACore.blockIdx().x) - 1) * Int(CUDACore.blockDim().x) + + Int(CUDACore.threadIdx().x) +end + +function make_linear_kernel(dest, bc::Base.Broadcast.Broadcasted, arg_plan, static_args) + f = bc.f + + @kernel unsafe_indices = true function broadcast_kernel_linear_splat(dest, runtime_args...) + I = _broadcast_linear_work_id() + if I <= length(dest) + @inbounds args_modified = _materialize_broadcast_args( + arg_plan, runtime_args, static_args, I + ) + @inbounds dest[I] = Base.Broadcast._broadcast_getindex_evalf(f, args_modified...) + end + end + + return broadcast_kernel_linear_splat +end + +struct FusedBroadcastMetadata + ctx::Any # KA.CompilerMetadata (compilation / arg layout; not global ndrange) + threads::Int # occupancy thread budget; device chooses final launch dims + cuda_task::CUDATask +end + +# Cache by kernel identity + arg types. Launch geometry is derived per-GPU from +# the local PhysicalArray, so global ndrange is not a key. +const _BCAST_PTX_CACHE = Dict{Tuple{Any,DataType,DataType},FusedBroadcastMetadata}() +const _BCAST_PTX_CACHE_LOCK = ReentrantLock() + +# Linear-only fusion: every NDArray leaf must match `dest` shape. Shape-mismatched +# broadcasts (e.g. matrix .+ vector) need cartesian / Extruded indexing and are +# handled by the unfused path instead. +# +# Slice views are allowed: RunPTXBroadcastTask packs element strides so linear +# `I` indexes through CuStridedDeviceArray correctly. +@inline _can_fuse_linear_broadcast_leaf(dest, ::Number) = true +@inline _can_fuse_linear_broadcast_leaf(dest, ::Base.RefValue) = true +@inline function _can_fuse_linear_broadcast_leaf(dest, x::NDArray) + return size(x) == size(dest) +end +@inline function _can_fuse_linear_broadcast_leaf(dest, x::Base.Broadcast.Extruded) + return _can_fuse_linear_broadcast_leaf(dest, x.x) +end +@inline function _can_fuse_linear_broadcast_leaf(dest, bc::Base.Broadcast.Broadcasted) + return _can_fuse_linear_broadcast_args(dest, bc.args) +end +@inline _can_fuse_linear_broadcast_leaf(dest, @nospecialize(x)) = false + +@inline _can_fuse_linear_broadcast_args(dest, ::Tuple{}) = true +@inline function _can_fuse_linear_broadcast_args(dest, args::Tuple) + return _can_fuse_linear_broadcast_leaf(dest, getfield(args, 1)) && + _can_fuse_linear_broadcast_args(dest, Base.tail(args)) +end + +""" +Return true when fused linear broadcast is safe for `bc` into `dest`. + +Requires every NDArray leaf to have the same size as `dest`. Scalars / RefValues +are allowed. Unknown leaf types refuse fusion (fall back to unfused). + +Also refuses 0-d destinations: `RunPTXBroadcastTask` only supports dims in +`[1, 6]`; 0-d falls back to the unfused path. +""" +@inline function can_fuse_linear_broadcast(dest::NDArray, bc::Base.Broadcast.Broadcasted) + # Device-side launch dims require at least one dimension. + ndims(dest) >= 1 || return false + return _can_fuse_linear_broadcast_leaf(dest, bc) +end + +# After Broadcast.preprocess, same-shape arrays are wrapped in Extruded with all +# keeps=true. Unwrap those so the kernel indexes CuStridedDeviceArray with linear `I` +# (avoids Extruded/CartesianIndices paths that emit gpu_report_exception in PTX). +@inline function _unwrap_linear_fusion_arg(x::Base.Broadcast.Extruded) + if all(x.keeps) + return x.x + end + throw( + ArgumentError( + "Broadcast fusion (linear-only) does not support shape-mismatched " * + "Extruded arguments; use the unfused broadcast path", + ), + ) +end +@inline _unwrap_linear_fusion_arg(x) = x +@inline _unwrap_linear_fusion_args(args::Tuple) = _unwrap_linear_fusion_arg.(args) + +# Host-side scalar alignment for fusion — same as unfused +# `T_IN = __my_promote_type(...); unchecked_promote_arr.(args, T_IN)`, but +# `unchecked_promote_scalar` keeps Numbers as scalars for the PTX arg buffer. +# +# Must run on the host *before* `get_cuda_task` so the PTX cache key sees the +# promoted scalar types. `a .^ 2` is unaffected: `Val{N}` is static; op-aware +# checks stay in pre-flatten `_assert_fused_broadcast_promotion`. +function _align_fused_runtime_args(runtime_args::Tuple) + isempty(runtime_args) && return runtime_args + T_IN = __my_promote_type(map(eltype, runtime_args)...) + return map(a -> unchecked_promote_scalar(a, T_IN), runtime_args) +end + +cudevice_array_offset(::Type{T}) where {T<:CuStridedDeviceArray} = 0 +cudevice_array_offset(::Type{T}) where {T<:Base.Broadcast.Extruded} = Int(fieldoffset(T, 1)) + +stores_cudevicearray(::Type{T}) where {T<:CuStridedDeviceArray} = true +stores_cudevicearray(::Type{T}) where {T<:Base.Broadcast.Extruded} = true +stores_cudevicearray(::Type{T}) where {T<:Number} = false +stores_cudevicearray(::Type{T}) where {T<:Bool} = false + +function stores_cudevicearray(::Type{T}) where {T} + throw(error("Broadcast fusion. Don't know what to do with type: $T")) +end + +function find_cudevicearray_offsets_and_indices(::Type{BC_ARGS}) where {BC_ARGS} + offsets = Vector{Int}() + indices = Vector{Int}() + for (i, T) in enumerate(fieldtypes(BC_ARGS)) + if stores_cudevicearray(T) + offset = fieldoffset(BC_ARGS, i) + cudevice_array_offset(T) + push!(offsets, offset) + push!(indices, i) + end + end + return tuple(offsets...), tuple(indices...) +end + +function find_scalar_offsets_and_indices(::Type{BC_ARGS}) where {BC_ARGS} + offsets = Vector{Int}() + indices = Vector{Int}() + for (i, T) in enumerate(fieldtypes(BC_ARGS)) + if (T <: Number) + offset = fieldoffset(BC_ARGS, i) + push!(offsets, offset) + push!(indices, i) + end + end + return tuple(offsets...), tuple(indices...) +end + +get_ndarray(x::T) where {T<:NDArray} = x +get_ndarray(x::T) where {T<:Base.Broadcast.Extruded} = x.x +get_ndarray(x::T) where {T} = throw(error("Broadcast fusion. Don't know what to do with type: $T")) + +""" +Host-side occupancy probe: return a thread *budget* and a minimal KA `ctx` for +PTX compilation. Final blocks/threads are chosen in `RunPTXBroadcastTask` from +each GPU's local `PhysicalArray` shape. +""" +function _threads_from_occupancy( + obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, + ::Type{DEST_T}, + ARG_TYPES...; + ndrange=(1024,), +) where {DEST_T} + backend = KA.backend(obj) + + # Compile-time iterspace only — not used for multi-GPU launch coverage. + ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, nothing) + ctx = KA.mkcontext(obj, ndrange, iterspace) + + maxthreads = + if KA.workgroupsize(obj) <: KA.StaticSize + prod(KA.get(KA.workgroupsize(obj))) + else + nothing + end + + tt = Base.to_tuple_type((typeof(ctx), DEST_T, ARG_TYPES...)) + host_kernel = CUDACore.cufunction( + obj.f, + tt; + kernel=true, + maxthreads=maxthreads, + always_inline=backend.always_inline, + ) + config = CUDACore.launch_configuration(host_kernel.fun; max_threads=prod(ndrange)) + threads = Int(config.threads) + + # Bake ctx workitems to the occupancy budget so KA metadata stays consistent + # with the thread count we pass to the device as a budget. + workgroupsize = CUDACore.CUDAKernels.threads_to_workgroupsize(threads, ndrange) + iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize) + ctx = KA.mkcontext(obj, ndrange, iterspace) + threads = length(KA.workitems(iterspace)) + + return threads, ctx +end + +""" + get_ptx(obj, DEST_T, arg_types...) -> (ptx, threads, ctx) + +Compile a KA CUDA kernel using types only and choose an occupancy thread budget. +""" +function get_ptx( + obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, + ::Type{DEST_T}, + arg_types...; +) where {DEST_T} + threads, ctx = _threads_from_occupancy(obj, DEST_T, arg_types...) + threads == 0 && return "", 0, ctx + + buf = IOBuffer() + # dump_module=true keeps only_entry=false so linked libdevice helpers (e.g. cos + # slowpaths) are not emptied into unresolved .externs before cuModuleLoad. + CUDATools.code_ptx(buf, obj.f, (typeof(ctx), DEST_T, arg_types...); + raw=false, dump_module=true, kernel=true) + + return String(take!(buf)), threads, ctx +end + +function get_cuda_task( + obj::KA.Kernel{CUDACore.CUDAKernels.CUDABackend}, + dest::D, + runtime_args::RT, +) where {D<:NDArray,RT<:Tuple} + DEST_T = map_cuda_type(D) + ARG_TYPES = map_cuda_type.(typeof.(runtime_args)) + + key = (obj, D, RT) + + lock(_BCAST_PTX_CACHE_LOCK) do + return get!(_BCAST_PTX_CACHE, key) do + ptx, threads, ctx = get_ptx(obj, DEST_T, ARG_TYPES...) + + orig_name = extract_kernel_name(ptx) + unique_name = orig_name * "_" * string(hash(ptx); base=16) + ptx = replace(ptx, orig_name => unique_name) + + ptx_task(ptx, unique_name) + cuda_task = CUDATask(unique_name, (DEST_T, ARG_TYPES...)) + + return FusedBroadcastMetadata(ctx, threads, cuda_task) + end + end +end + +# Fused-broadcast introspection. Set `cuNumeric.BCAST_FUSION_DEBUG[] = true` to +# dump each kernel's expr/inputs/scalars/launch geometry before launch. +const BCAST_FUSION_DEBUG = Ref(false) + +_fname(f::Function) = string(nameof(f)) +_fname(@nospecialize(f)) = string(f) + +# Reconstruct the op tree; call before `flatten`, while nesting mirrors the source. +function _bcast_tree_str(leaf_name, bc::Base.Broadcast.Broadcasted) + args = (_bcast_tree_str(leaf_name, arg) for arg in bc.args) + return string(_fname(bc.f), "(", join(args, ", "), ")") +end +_bcast_tree_str(leaf_name, x::Base.Broadcast.Extruded) = _bcast_tree_str(leaf_name, x.x) +_bcast_tree_str(leaf_name, x) = leaf_name(x) + +function _bcast_tree_str(bc::Base.Broadcast.Broadcasted) + return _bcast_tree_str(bc) do x + x isa NDArray && return "NDArray" + x isa Number && return repr(x) + x isa Base.RefValue && return string("^", repr(x[])) + return string("<", typeof(x), ">") + end +end + +function _bcast_scope_name(bc::Base.Broadcast.Broadcasted, ndarray_to_input_idx) + scalar_idx = 0 + tree = _bcast_tree_str(bc) do x + if x isa NDArray + input_idx = get(ndarray_to_input_idx, objectid(x), nothing) + return input_idx === nothing ? "NDArray" : string("input", input_idx) + end + + if x isa Number + idx = scalar_idx + scalar_idx += 1 + return string("scalar", idx) + end + + if x isa Base.RefValue + value = x[] + if value isa Number + idx = scalar_idx + scalar_idx += 1 + return string("scalar", idx) + end + return repr(value) + end + + return string("<", typeof(x), ">") + end + return string("broadcast.", tree) +end + +# Recover the kernel's plain name +function _demangle_head(s::AbstractString) + m = match(r"^_Z([0-9]+)(.*)$", s) + m === nothing && return s + return first(m.captures[2], parse(Int, m.captures[1])) +end + +# `arg_map` records the kernel's argument order (0=output, ≥1=input idx+1, +# <0=scalar idx); reconstruct it as a readable `output -> name(args...)` call. +function _kernel_signature( + dest::NDArray, + unique_ndarrays::AbstractVector{<:NDArray}, + actual_scalars::AbstractVector, + arg_map::AbstractVector{<:Integer}, + func::AbstractString, +) + token(a::Integer) = + if a == 0 + "output" + elseif a > 0 + string("input", a - 1) + else + string("scalar", -a - 1) + end + call_args = [token(a) for a in arg_map if a != 0] + return string("broadcast.", _demangle_head(func), "(", join(call_args, ", "), ")") +end + +function _describe_fused_broadcast( + dest, tree_str, unique_ndarrays, actual_scalars, static_args, arg_map, fkm, ndrange +) + io = IOBuffer() + field(k, v) = println(io, " ", rpad(k, 8), v) + println(io, "\n", "="^40, " fused broadcast kernel") + field("expr", tree_str) + field("output", "$(typeof(dest)) $(size(dest))") + field("inputs", "$(length(unique_ndarrays)) unique NDArray(s)") + for (i, nd) in enumerate(unique_ndarrays) + alias = objectid(nd) == objectid(dest) ? " (aliases output)" : "" + println(io, " [", i - 1, "] ", typeof(nd), " ", size(nd), alias) + end + isempty(actual_scalars) || field("scalars", join(repr.(actual_scalars), ", ")) + isempty(static_args) || field("static", join(repr.(static_args), ", ")) + field("arg_map", "$(Int.(arg_map)) (0=output, >=1=input idx+1, <0=scalar)") + field( + "launch", + "host thread budget=$(fkm.threads), indexing=linear, " * + "blocks=device(local tile), global_ndrange=$ndrange", + ) + field( + "call", + _kernel_signature(dest, unique_ndarrays, actual_scalars, arg_map, fkm.cuda_task.func), + ) + print(String(take!(io))) + return nothing +end + +# Pre-launch promotion policy for fused broadcast. Mirrors unfused +# `unravel_broadcast_tree` (`__checked_promote_op` / `__my_promote_type` per +# tree node) and the dest-side widen check of `checked_promote_arr` in +# `_copyto_unfused!`. Fused writes `dest` in place, so there is no post-fuse +# promote / `nda_move` — this must run before kernel launch. +# +# The tree walk is typed on `Broadcasted{S,Ax,F,Args}` and the concrete +# `Args` Tuple, so Julia specializes/inlines per broadcast shape. Same-eltype +# trees (e.g. all Float32) constant-fold `is_wider_type` to false and DCE the +# `assertpromotion` calls; residual cost is only the specialize/inline frame. +@inline function _assert_fused_broadcast_promotion( + dest::NDArray{DT}, bc::B +) where {DT,B<:Base.Broadcast.Broadcasted} + T_OUT = _assert_fused_broadcast_tree(bc) + is_wider_type(DT, T_OUT) && assertpromotion(promote_type, T_OUT, DT) + return nothing +end + +# Leaf NDArray / Number / RefValue / etc. +@inline _assert_fused_broadcast_tree(x) = eltype(x) + +# Nested node: recurse through typed Args, then op/input promote checks. +@inline function _assert_fused_broadcast_tree( + 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) + __my_promote_type(eltypes.parameters...) + return T_OUT +end + +# Typed Tuple walk — method selection replaces runtime `isa` / Any iteration. +# Returns `Type{Tuple{...}}` for `__checked_promote_op(f, ::Type{Tuple{...}})`. +@inline _fused_checked_eltypes(::Tuple{}) = Tuple{} + +@inline function _fused_checked_eltypes(args::Tuple{A}) where {A} + T1 = _assert_fused_broadcast_tree(getfield(args, 1)) + return Tuple{T1} +end + +@inline function _fused_checked_eltypes(args::Tuple{A,B}) where {A,B} + T1 = _assert_fused_broadcast_tree(getfield(args, 1)) + T2 = _assert_fused_broadcast_tree(getfield(args, 2)) + return Tuple{T1,T2} +end + +# literal_pow and other ternary broadcast args +@inline function _fused_checked_eltypes(args::Tuple{A,B,C}) where {A,B,C} + T1 = _assert_fused_broadcast_tree(getfield(args, 1)) + T2 = _assert_fused_broadcast_tree(getfield(args, 2)) + T3 = _assert_fused_broadcast_tree(getfield(args, 3)) + return Tuple{T1,T2,T3} +end + +@inline function _fused_checked_eltypes(args::Tuple) + T1 = _assert_fused_broadcast_tree(getfield(args, 1)) + rest = _fused_checked_eltypes(Base.tail(args)) + return Tuple{T1,rest.parameters...} +end + +function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcast.Broadcasted} + # Promotion checks use the pre-flatten tree (same shape as unfused unravel). + _assert_fused_broadcast_promotion(dest, bc) + + # Capture the readable tree before flatten collapses the nesting. + bc_scope = bc + tree_str = BCAST_FUSION_DEBUG[] ? _bcast_tree_str(bc) : "" + + bc = Base.Broadcast.preprocess(dest, bc) + bc = Base.Broadcast.instantiate(bc) + bc = Base.Broadcast.flatten(bc) + + # Things like exponentiation generate arguments like Base.RefValue + # which do not work with our pattern for making CUDA kernels as they are + # not is-bits types. We split these out manually into static args and handle + # them separately from runtime args (i.e., arrays, scalars) + runtime_args, static_args, arg_plan = split_broadcast_args_for_kernel( + _unwrap_linear_fusion_args(bc.args) + ) + # Host-only, before PTX cache: same `__my_promote_type` + Number convert as + # unfused `T_IN` / `unchecked_promote_arr`. Kernel sees already-aligned types. + runtime_args = _align_fused_runtime_args(runtime_args) + + broadcast_kernel = make_linear_kernel(dest, bc, arg_plan, static_args) + + ndrange = ndims(dest) > 0 ? size(dest) : (1,) + + bck_cuda = broadcast_kernel(CUDACore.CUDAKernels.CUDABackend()) + + fkm = get_cuda_task(bck_cuda, dest, runtime_args) + + num_outputs = 1 + + unique_ndarrays = NDArray[] + ndarray_to_input_idx = Dict{UInt,Int}() + + arg_map = Int32[] + actual_scalars = Any[] + + # First PTX argument after ctx is dest. + push!(arg_map, Int32(0)) + + # Now map only runtime args, not original bc.args. + for arg in runtime_args + if stores_cudevicearray(map_cuda_type(typeof(arg))) + nda = get_ndarray(arg) + oid = objectid(nda) + + if !haskey(ndarray_to_input_idx, oid) + push!(unique_ndarrays, nda) + ndarray_to_input_idx[oid] = length(unique_ndarrays) - 1 + end + + input_idx = ndarray_to_input_idx[oid] + push!(arg_map, Int32(num_outputs + input_idx)) + else + push!(arg_map, Int32(-1 - length(actual_scalars))) + push!(actual_scalars, arg) + end + end + + input_ndarrays = tuple(unique_ndarrays...) + + BCAST_FUSION_DEBUG[] && _describe_fused_broadcast( + dest, tree_str, unique_ndarrays, actual_scalars, static_args, arg_map, fkm, ndrange + ) + + @task_scope _bcast_scope_name(bc_scope, ndarray_to_input_idx) begin + # `blocks=1` is a placeholder; RunPTXBroadcastTask overwrites grid dims + # from the local PhysicalArray. `threads` is only the occupancy budget (tx). + # Scalars after ctx: num_kernel_args, arg_map... + launch( + fkm.cuda_task, + input_ndarrays, + (dest,), + (Int32(length(arg_map)), arg_map..., actual_scalars...); + blocks=1, + threads=fkm.threads, + taskid=cuNumeric.RUN_PTX_BROADCAST, + ctx=fkm.ctx, + ) + end + + # Fused kernel already wrote `dest` in place; promotion was checked pre-launch. + return dest +end diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl index 9d630fcc0..dfae001fc 100644 --- a/src/ndarray/detail/ndarray.jl +++ b/src/ndarray/detail/ndarray.jl @@ -8,6 +8,15 @@ struct Slice stop::Int64 end +macro task_scope(scope_name, body) + TASK_SCOPE_NAMES || return esc(body) + return quote + Legate.with_scope($(esc(scope_name))) do + return $(esc(body)) + end + end +end + # Opaque pointer const NDArray_t = Ptr{Cvoid} const CN_Store_t = Ptr{Cvoid} @@ -43,10 +52,7 @@ mutable struct NDArray{T,N,PADDED,P} <: AbstractNDArray{T,N} nbytes = cuNumeric.nda_nbytes(ptr) cuNumeric.register_alloc!(nbytes) handle = new{T,N,false,Nothing}(ptr, nbytes, nothing, nothing) - finalizer(handle) do h - cuNumeric.nda_destroy_array(h.ptr) - cuNumeric.register_free!(h.nbytes) - end + finalizer(destroy!, handle) return handle end @@ -55,13 +61,29 @@ mutable struct NDArray{T,N,PADDED,P} <: AbstractNDArray{T,N} nbytes = cuNumeric.nda_nbytes(ptr) cuNumeric.register_alloc!(nbytes) handle = new{T,N,false,P}(ptr, nbytes, nothing, parent) - finalizer(handle) do h - cuNumeric.nda_destroy_array(h.ptr) - cuNumeric.register_free!(h.nbytes) - end + finalizer(destroy!, handle) return handle end end + +""" + destroy!(arr::NDArray) + +Eagerly drop the underlying cuPyNumeric/Legate handle and update allocation +counters. Safe to call more than once. +""" +function destroy!(arr::NDArray) + ptr = arr.ptr + if ptr != C_NULL + nbytes = arr.nbytes + nda_destroy_array(ptr) + arr.ptr = Ptr{Cvoid}(0) + arr.nbytes = 0 + nbytes > 0 && register_free!(nbytes) + end + return arr +end + # this here is to avoid if else patterns @inline _NDArray(ptr, T, v, ::Nothing) = NDArray(ptr, T, v) @inline _NDArray(ptr, T, v, parent) = NDArray(ptr, T, v, parent) @@ -71,6 +93,8 @@ function NDArray(ptr::NDArray_t; T=get_julia_type(ptr), N::Integer=get_n_dim(ptr return _NDArray(ptr, T, Val(N), parent) end +_scope_op(kind, op_code) = string(kind, "#", Int32(op_code)) + #! JUST USE FULL TO MAKE a 0D? # $ cuNumeric.nda_full_array(UInt64[], 2.0f0) # TODO DAVID TEST THIS HERE @@ -88,9 +112,11 @@ NDArray(value::T) where {T<:SUPPORTED_TYPES} = nda_full_array((), value) function nda_zeros_array(dims::Dims{N}, ::Type{T}) where {T,N} shape = collect(UInt64, dims) legate_type = Legate.to_legate_type(T) - ptr = ccall((:nda_zeros_array, libnda), - NDArray_t, (Int32, Ptr{UInt64}, Legate.LegateTypeAllocated), - Int32(N), shape, legate_type) + ptr = @task_scope "zeros" begin + ccall((:nda_zeros_array, libnda), + NDArray_t, (Int32, Ptr{UInt64}, Legate.LegateTypeAllocated), + Int32(N), shape, legate_type) + end return NDArray(ptr, T, Val(N)) end @@ -98,33 +124,42 @@ function nda_full_array(dims::Dims{N}, value::T) where {T,N} shape = collect(UInt64, dims) type = Legate.to_legate_type(T) - ptr = ccall((:nda_full_array, libnda), - NDArray_t, - (Int32, Ptr{UInt64}, Legate.LegateTypeAllocated, Ptr{Cvoid}), - Int32(N), shape, type, Ref(value)) + ptr = @task_scope "full" begin + ccall((:nda_full_array, libnda), + NDArray_t, + (Int32, Ptr{UInt64}, Legate.LegateTypeAllocated, Ptr{Cvoid}), + Int32(N), shape, type, Ref(value)) + end return NDArray(ptr, T, Val(N)) end function nda_random(arr::NDArray, gen_code) - ccall((:nda_random, libnda), - Cvoid, (NDArray_t, Int32), - arr.ptr, Int32(gen_code)) + @task_scope "rand!" begin + ccall((:nda_random, libnda), + Cvoid, (NDArray_t, Int32), + arr.ptr, Int32(gen_code)) + end end function nda_random_array(dims::Dims{N}) where {N} shape = collect(UInt64, dims) - ptr = ccall((:nda_random_array, libnda), - NDArray_t, (Int32, Ptr{UInt64}), - Int32(N), shape) + ptr = @task_scope "rand" begin + ccall((:nda_random_array, libnda), + NDArray_t, (Int32, Ptr{UInt64}), + Int32(N), shape) + end return NDArray(ptr, Float64, Val(N)) #* T is always Float64 cause of cupynumeric end function nda_get_slice(arr::NDArray{T,N}, slices::Vector{Slice}) where {T,N} - ptr = ccall((:nda_get_slice, libnda), - NDArray_t, (NDArray_t, Ptr{Slice}, Cint), - arr.ptr, pointer(slices), length(slices)) - return NDArray(ptr, T, Val(N)) + ptr = @task_scope "slice" begin + ccall((:nda_get_slice, libnda), + NDArray_t, (NDArray_t, Ptr{Slice}, Cint), + arr.ptr, pointer(slices), length(slices)) + end + # Keep parent so callers can detect views (slices share the parent store). + return NDArray(ptr, T, Val(N), arr) end # queries @@ -133,7 +168,7 @@ nda_array_dim(arr::NDArray) = ccall((:nda_array_dim, libnda), nda_array_size(arr::NDArray) = ccall((:nda_array_size, libnda), Int32, (NDArray_t,), arr.ptr) function nda_array_type_code(arr::NDArray) - ccall((:nda_array_type_code, libnda), + return ccall((:nda_array_type_code, libnda), Int32, (NDArray_t,), arr.ptr) end @@ -149,73 +184,91 @@ end # modify function nda_reshape_array(arr::NDArray{T}, newdims::Dims{N}) where {T,N} newshape = collect(UInt64, newdims) - ptr = ccall((:nda_reshape_array, libnda), - NDArray_t, (NDArray_t, Int32, Ptr{UInt64}), - arr.ptr, Int32(N), newshape) + ptr = @task_scope "reshape" begin + ccall((:nda_reshape_array, libnda), + NDArray_t, (NDArray_t, Int32, Ptr{UInt64}), + arr.ptr, Int32(N), newshape) + end return NDArray(ptr, T, Val(N)) end function nda_astype(arr::NDArray{OLD_T,N}, ::Type{NEW_T}) where {OLD_T,NEW_T,N} type = Legate.to_legate_type(NEW_T) - ptr = ccall((:nda_astype, libnda), - NDArray_t, - (NDArray_t, Legate.LegateTypeAllocated), - arr.ptr, type) + ptr = @task_scope "astype" begin + ccall((:nda_astype, libnda), + NDArray_t, + (NDArray_t, Legate.LegateTypeAllocated), + arr.ptr, type) + end return NDArray(ptr, NEW_T, Val(N)) end function nda_fill_array(arr::NDArray{T}, value::T) where {T} type = Legate.to_legate_type(T) val = Ref(value) - ccall((:nda_fill_array, libnda), - Cvoid, (NDArray_t, Legate.LegateTypeAllocated, Ptr{Cvoid}), - arr.ptr, type, val) + @task_scope "fill!" begin + ccall((:nda_fill_array, libnda), + Cvoid, (NDArray_t, Legate.LegateTypeAllocated, Ptr{Cvoid}), + arr.ptr, type, val) + end return nothing end function nda_assign(arr::NDArray{T}, other::NDArray{T}) where {T} - ccall((:nda_assign, libnda), - Cvoid, (NDArray_t, NDArray_t), - arr.ptr, other.ptr) + @task_scope "copyto!" begin + ccall((:nda_assign, libnda), + Cvoid, (NDArray_t, NDArray_t), + arr.ptr, other.ptr) + end end -function nda_copy(arr::NDArray) - ptr = ccall((:nda_copy, libnda), - NDArray_t, (NDArray_t,), - arr.ptr) - return NDArray(ptr) +function nda_copy(arr::NDArray{T,N}) where {T,N} + ptr = @task_scope "copy" begin + ccall((:nda_copy, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + end + return NDArray(ptr, T, Val(N)) end # src will be unused after this function nda_move(dst::NDArray{T,N}, src::NDArray{T,N}) where {T,N} - ccall((:nda_move, libnda), - Cvoid, (NDArray_t, NDArray_t), - dst.ptr, src.ptr) + @task_scope "move!" begin + ccall((:nda_move, libnda), + Cvoid, (NDArray_t, NDArray_t), + dst.ptr, src.ptr) + end src.ptr = Ptr{Cvoid}(0) src.nbytes = 0 - register_free!(dst.nbytes) + return register_free!(dst.nbytes) end # operations function nda_binary_op!(out::NDArray, op_code::BinaryOpCode, rhs1::NDArray, rhs2::NDArray) - ccall((:nda_binary_op, libnda), - Cvoid, (NDArray_t, BinaryOpCode, NDArray_t, NDArray_t), - out.ptr, op_code, rhs1.ptr, rhs2.ptr) + @task_scope _scope_op("binary", op_code) begin + ccall((:nda_binary_op, libnda), + Cvoid, (NDArray_t, BinaryOpCode, NDArray_t, NDArray_t), + out.ptr, op_code, rhs1.ptr, rhs2.ptr) + end return out end function nda_unary_op!(out::NDArray, op_code::UnaryOpCode, input::NDArray) - ccall((:nda_unary_op, libnda), - Cvoid, (NDArray_t, UnaryOpCode, NDArray_t), - out.ptr, op_code, input.ptr) + @task_scope _scope_op("unary", op_code) begin + ccall((:nda_unary_op, libnda), + Cvoid, (NDArray_t, UnaryOpCode, NDArray_t), + out.ptr, op_code, input.ptr) + end return out end function nda_unary_reduction(out::NDArray, op_code::UnaryRedCode, input::NDArray) - ccall((:nda_unary_reduction, libnda), - Cvoid, (NDArray_t, UnaryRedCode, NDArray_t), - out.ptr, op_code, input.ptr) + @task_scope _scope_op("reduce", op_code) begin + ccall((:nda_unary_reduction, libnda), + Cvoid, (NDArray_t, UnaryRedCode, NDArray_t), + out.ptr, op_code, input.ptr) + end return out end @@ -223,87 +276,109 @@ function nda_unary_reduction_axes( op_code::UnaryRedCode, input::NDArray{T,N}, axes::Vector{Int32}, keepdims::Bool ) where {T,N} axes_c = collect(Int32, axes) - ptr = ccall((:nda_unary_reduction_axes, libnda), - NDArray_t, (UnaryRedCode, NDArray_t, Ptr{Int32}, Int32, Cint), - op_code, input.ptr, axes_c, Int32(length(axes_c)), keepdims) + ptr = @task_scope _scope_op("reduce_axes", op_code) begin + ccall((:nda_unary_reduction_axes, libnda), + NDArray_t, (UnaryRedCode, NDArray_t, Ptr{Int32}, Int32, Cint), + op_code, input.ptr, axes_c, Int32(length(axes_c)), keepdims) + end return NDArray(ptr) end function nda_array_equal(rhs1::NDArray{T,N}, rhs2::NDArray{T,N}) where {T,N} - ptr = ccall((:nda_array_equal, libnda), - NDArray_t, (NDArray_t, NDArray_t), - rhs1.ptr, rhs2.ptr) + ptr = @task_scope "array_equal" begin + ccall((:nda_array_equal, libnda), + NDArray_t, (NDArray_t, NDArray_t), + rhs1.ptr, rhs2.ptr) + end return NDArray(ptr, Bool, Val(1)) end # 2D -> 1D: extract the k-th diagonal. Backend only supports the 2D case # (1D-construct and >2D both abort), so non-2D input is a MethodError. function nda_diag(arr::NDArray{T,2}, k::Int32) where {T} - ptr = ccall((:nda_diag, libnda), - NDArray_t, (NDArray_t, Int32), - arr.ptr, k) + ptr = @task_scope "diag" begin + ccall((:nda_diag, libnda), + NDArray_t, (NDArray_t, Int32), + arr.ptr, k) + end return NDArray(ptr, T, Val(1)) end # unique always returns a flat 1D array of the input's element type function nda_unique(arr::NDArray{T}) where {T} - ptr = ccall((:nda_unique, libnda), - NDArray_t, (NDArray_t,), - arr.ptr) + ptr = @task_scope "unique" begin + ccall((:nda_unique, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + end return NDArray(ptr, T, Val(1)) end function nda_ravel(arr::NDArray) - ptr = ccall((:nda_ravel, libnda), - NDArray_t, (NDArray_t,), - arr.ptr) + ptr = @task_scope "ravel" begin + ccall((:nda_ravel, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + end return NDArray(ptr) end function nda_add(rhs1::NDArray, rhs2::NDArray, out::NDArray) - ccall((:nda_add, libnda), - Cvoid, (NDArray_t, NDArray_t, NDArray_t), - rhs1.ptr, rhs2.ptr, out.ptr) + @task_scope "add" begin + ccall((:nda_add, libnda), + Cvoid, (NDArray_t, NDArray_t, NDArray_t), + rhs1.ptr, rhs2.ptr, out.ptr) + end return out end function nda_multiply_scalar(rhs1::NDArray{T,N}, value::T) where {T,N} type = Legate.to_legate_type(T) - ptr = ccall((:nda_multiply_scalar, libnda), - NDArray_t, (NDArray_t, Legate.LegateTypeAllocated, Ptr{Cvoid}), - rhs1.ptr, type, Ref(value)) + ptr = @task_scope "multiply_scalar" begin + ccall((:nda_multiply_scalar, libnda), + NDArray_t, (NDArray_t, Legate.LegateTypeAllocated, Ptr{Cvoid}), + rhs1.ptr, type, Ref(value)) + end return NDArray(ptr, T, Val(N)) end function nda_add_scalar(rhs1::NDArray{T,N}, value::T) where {T,N} type = Legate.to_legate_type(T) - ptr = ccall((:nda_add_scalar, libnda), - NDArray_t, (NDArray_t, Legate.LegateTypeAllocated, Ptr{Cvoid}), - rhs1.ptr, type, Ref(value)) + ptr = @task_scope "add_scalar" begin + ccall((:nda_add_scalar, libnda), + NDArray_t, (NDArray_t, Legate.LegateTypeAllocated, Ptr{Cvoid}), + rhs1.ptr, type, Ref(value)) + end return NDArray(ptr, T, Val(N)) end function nda_three_dot_arg(rhs1::NDArray{T}, rhs2::NDArray{T}, out::NDArray{T}) where {T} - ccall((:nda_three_dot_arg, libnda), - Cvoid, (NDArray_t, NDArray_t, NDArray_t), - rhs1.ptr, rhs2.ptr, out.ptr) + @task_scope "matmul" begin + ccall((:nda_three_dot_arg, libnda), + Cvoid, (NDArray_t, NDArray_t, NDArray_t), + rhs1.ptr, rhs2.ptr, out.ptr) + end return out end function nda_dot(rhs1::NDArray, rhs2::NDArray) - ptr = ccall((:nda_dot, libnda), - NDArray_t, (NDArray_t, NDArray_t), - rhs1.ptr, rhs2.ptr) + ptr = @task_scope "dot" begin + ccall((:nda_dot, libnda), + NDArray_t, (NDArray_t, NDArray_t), + rhs1.ptr, rhs2.ptr) + end return NDArray(ptr) end function nda_eye(rows::Int32, ::Type{T}) where {T} legate_type = Legate.to_legate_type(T) - ptr = ccall((:nda_eye, libnda), - NDArray_t, (Int32, Legate.LegateTypeAllocated), - rows, legate_type) + ptr = @task_scope "eye" begin + ccall((:nda_eye, libnda), + NDArray_t, (Int32, Legate.LegateTypeAllocated), + rows, legate_type) + end return NDArray(ptr, T, Val(2)) end @@ -311,26 +386,33 @@ function nda_trace( arr::NDArray, offset::Int32, a1::Int32, a2::Int32, ::Type{T} ) where {T} legate_type = Legate.to_legate_type(T) - ptr = ccall((:nda_trace, libnda), - NDArray_t, - (NDArray_t, Int32, Int32, Int32, Legate.LegateTypeAllocated), - arr.ptr, offset, a1, a2, legate_type) + ptr = @task_scope "trace" begin + ccall((:nda_trace, libnda), + NDArray_t, + (NDArray_t, Int32, Int32, Int32, Legate.LegateTypeAllocated), + arr.ptr, offset, a1, a2, legate_type) + end return NDArray(ptr, T, Val(1)) end # transpose reverses the axes: element type and rank are preserved function nda_transpose(arr::NDArray{T,N}) where {T,N} - ptr = ccall((:nda_transpose, libnda), - NDArray_t, (NDArray_t,), - arr.ptr) + ptr = @task_scope "transpose" begin + ccall((:nda_transpose, libnda), + NDArray_t, (NDArray_t,), + arr.ptr) + end return NDArray(ptr, T, Val(N)) end -function nda_attach_external(arr::AbstractArray{T,N}) where {T,N} - st = Legate.attach_external(arr) +function nda_attach_external(arr::Array{T,N}; shape::Dims{N}=size(arr)) where {T,N} + st = Legate.attach_external_row_major(arr; shape) # Use the CxxWrap method for type-safe interaction # This returns a raw pointer compatible with the NDArray constructor + # `nda_store_to_ndarray` takes the store by value; drop the Julia-owned + # LogicalStoreImpl so it does not pin alongside the NDArray until GC. nda_ptr = cuNumeric.nda_store_to_ndarray(st.handle) + finalize(st.handle) return NDArray(nda_ptr, T, Val(N), arr) end @@ -342,9 +424,13 @@ function get_store(arr::NDArray) end function get_ptr(arr::NDArray{T,N}) where {T,N} + # `get_store` returns a Julia-owned LogicalArrayImplAllocated that shares the + # store with the NDArray; finalize after use (same pin class as `_add_task_array!`). st_handle = get_store(arr) # LogicalArrayImplAllocated (returned by value) la = Legate.LogicalArray{T,N}(st_handle, size(arr)) - return Legate.get_ptr(la) + ptr = Legate.get_ptr(la) + finalize(st_handle) + return ptr end @doc""" @@ -357,7 +443,7 @@ Converts a Julia 1-based index tuple `idx` to a zero-based C++ style index wrapp Each element of `idx` is decremented by 1 to adjust from Julia’s 1-based indexing to C++ 0-based indexing. """ function to_cpp_index(idx::Dims{N}, (::Type{T})=UInt64) where {N,T<:Integer} - StdVector(T.([e - 1 for e in idx])) + return StdVector(T.([e - 1 for e in idx])) end @doc""" @@ -390,7 +476,7 @@ Constructs a `cuNumeric.Slice` object representing a slice with optional start a """ function slice(start::Union{Nothing,Integer}, stop::Union{Nothing,Integer}) - cuNumeric.Slice( + return cuNumeric.Slice( isnothing(start) ? 0 : 1, isnothing(start) ? 0 : Int64(start), isnothing(stop) ? 0 : 1, @@ -509,5 +595,7 @@ end function nda_to_logical_store(arr::NDArray{T,N}) where {T,N} la_handle = cuNumeric.get_store(arr) # LogicalArrayImplAllocated (returned by value) st_handle = Legate.data(Legate.LogicalArray{T,N}(la_handle, size(arr))) + # Drop temp LogicalArray owner after extracting the store. + finalize(la_handle) return Legate.LogicalStore{T,N}(st_handle, size(arr)) end diff --git a/src/ndarray/linalg.jl b/src/ndarray/linalg.jl index 1c01b44a4..6893b0ca5 100644 --- a/src/ndarray/linalg.jl +++ b/src/ndarray/linalg.jl @@ -39,17 +39,28 @@ function solve_batched(a::NDArray{T,N}, b::NDArray, x::NDArray) where {T,N} tiled_a = Legate.partition_by_tiling(store_a, collect(tilesize_a)) tiled_b = Legate.partition_by_tiling(store_b, collect(tilesize_b)) tiled_x = Legate.partition_by_tiling(store_x, collect(tilesize_b)) + # Same Legate Julia-wrapper pin class as Launch `_add_task_array!` / get_store + # temps: one could finalize store_/tiled_ handles here after partition/add_* + # copies ownership into the task. Not enabled yet — weak linalg test coverage. + # finalize(store_a.handle) + # finalize(store_b.handle) + # finalize(store_x.handle) - rt = Legate.get_runtime() - domain = Legate.domain_from_shape(Legate.Shape(Legate.to_cxx_vector(color_shape))) - lib = cuNumeric.get_lib() - task = Legate.create_manual_task(rt, lib, cuNumeric.SOLVE, domain) + @task_scope "solve" begin + rt = Legate.get_runtime() + domain = Legate.domain_from_shape(Legate.Shape(Legate.to_cxx_vector(color_shape))) + lib = cuNumeric.get_lib() + task = Legate.create_manual_task(rt, lib, cuNumeric.SOLVE, domain) - Legate.add_input(task, tiled_a) - Legate.add_input(task, tiled_b) - Legate.add_output(task, tiled_x) + Legate.add_input(task, tiled_a) + # finalize(tiled_a.handle) + Legate.add_input(task, tiled_b) + # finalize(tiled_b.handle) + Legate.add_output(task, tiled_x) + # finalize(tiled_x.handle) - Legate.submit_manual_task(rt, task) + Legate.submit_manual_task(rt, task) + end end # solve runs in floating point: @@ -60,6 +71,18 @@ _solve_eltype(::Type{T}) where {T<:_SOLVE_PROMOTABLE} = Float64 _solve_eltype(::Type{T}) where {T<:SUPPORTED_SOLVE_TYPES} = T # Type/dim guards dispatch on one argument at a time, then forward to `_solve`. +""" + cuNumeric.solve(A, b) + +Solve linear system(s) `A * x = b`. + +`A` must have shape `(..., m, m)`. `b` must have shape `(..., m)` or `(..., m, n)`. +The result has the same shape as `b`. Batch dimensions are supported; the +implementation always uses the batched Legate `SOLVE` path. + +Accepted element types are `Float32`, `Float64`, `ComplexF32`, and `ComplexF64`. +Integer or `Bool` inputs promote to `Float64` only when promotion is allowed. +""" function solve(a::NDArray{<:_SOLVE_ACCEPTED}, b::NDArray{<:_SOLVE_ACCEPTED}) A, B = eltype(a), eltype(b) O = promote_type(_solve_eltype(A), _solve_eltype(B)) diff --git a/src/ndarray/ndarray.jl b/src/ndarray/ndarray.jl index 086a60caa..d06a65bef 100644 --- a/src/ndarray/ndarray.jl +++ b/src/ndarray/ndarray.jl @@ -30,10 +30,10 @@ function transpose(arr::NDArray) end @doc""" - cuNumeric.eye([T,] rows::Int) + cuNumeric.eye([T=Float32,] rows::Int) -Create a 2D identity `NDArray` of size `rows x rows` with element type `T` -(defaults to `DEFAULT_FLOAT`). +Create a 2D identity `NDArray` of size `rows × rows` with element type `T`. +The default type is Float32 if not specified. """ function eye(::Type{T}, rows::Int) where {T} return nda_eye(Int32(rows), T) @@ -142,30 +142,61 @@ end # conversion from NDArray to Base Julia array # get_ptr is a blocking call that grabs the physical store # we have not tested across multiple processes or devices yet + +function (::Type{<:Array{A}})(arr::NDArray{B,0}) where {A,B} + out = Array{A}(undef) + allowscalar() do + out[] = convert(A, arr[]) + end + return out +end + +function (::Type{<:Array{A}})(arr::NDArray{B,1}) where {A,B} + return make_array(A, Ptr{A}(get_ptr(arr)), size(arr)) +end + +# TODO: query the store's DimOrdering instead of assuming row-major (C-order). +# Julia Arrays are column-major, so for C-ordered buffers we transpose, wrap, +# then transpose back so `Array(nda)` matches `nda[i,j]`. function (::Type{<:Array{A}})(arr::NDArray{B}) where {A,B} - dims = Base.size(arr) - ptr = Ptr{A}(get_ptr(arr)) - return make_array(A, ptr, dims) + t = transpose(arr) + w = make_array(B, Ptr{B}(get_ptr(t)), size(t)) + out = collect(permutedims(w, reverse(1:ndims(w)))) + return A === B ? out : copyto!(Array{A}(undef, size(arr)), out) end function (::Type{<:Array})(arr::NDArray{B}) where {B} - dims = Base.size(arr) - ptr = Ptr{B}(get_ptr(arr)) - return make_array(B, ptr, dims) + return Array{B}(arr) end # conversion from Base Julia array to NDArray -function (::Type{<:NDArray{T}})(arr::Array{T,N}) where {T,N} +# Julia Arrays are column-major; Legate stores are row-major. For N>=2 we +# materialize a C-ordered buffer via permutedims, attach it with the original +# shape, and keep that buffer as `parent` for lifetime. +function _nda_from_julia_array(arr::Array{T,0}) where {T} + return cuNumeric.nda_attach_external(arr) +end + +function _nda_from_julia_array(arr::Array{T,1}) where {T} return cuNumeric.nda_attach_external(arr) end +function _nda_from_julia_array(arr::Array{T,N}) where {T,N} + tmp = collect(permutedims(arr, reverse(ntuple(identity, Val(N))))) + return cuNumeric.nda_attach_external(tmp; shape=size(arr)) +end + +function (::Type{<:NDArray{T}})(arr::Array{T,N}) where {T,N} + return _nda_from_julia_array(arr) +end + function (::Type{<:NDArray{A}})(arr::Array{B,N}) where {A,B,N} - # If types differ, we cast in Julia first (creating a temp) then attach - return cuNumeric.nda_attach_external(A.(arr)) + # If types differ, cast in Julia first (creating a temp) then attach + return _nda_from_julia_array(convert(Array{A}, arr)) end function (::Type{<:NDArray})(arr::Array{T,N}) where {T,N} - return cuNumeric.nda_attach_external(arr) + return _nda_from_julia_array(arr) end # Base.convert(::Type{<:NDArray{T}}, a::A) where {T, A} = NDArray(T(a))::NDArray{T} @@ -220,6 +251,7 @@ size(arr, 2) """ Base.size(arr::NDArray{<:Any,N}) where {N} = cuNumeric.shape(arr) Base.size(arr::NDArray, dim::Int) = Base.size(arr)[dim] +Base.isempty(arr::NDArray) = any(==(0), size(arr)) @doc""" Base.firstindex(arr::NDArray, dim::Int) @@ -246,9 +278,8 @@ Base.view(arr::NDArray, inds...) = arr[inds...] # NDArray slices are views by de Base.IndexStyle(::NDArray) = IndexCartesian() function Base.show(io::IO, arr::NDArray{T,0}) where {T} - println(io, "0-dimensional NDArray{$(T),0}") allowscalar() do - print(io, arr[]) + print(io, "NDArray{$(T),0}(", repr(arr[]), ")") end end @@ -259,13 +290,13 @@ function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T,0}) where {T} end end -function Base.show(io::IO, arr::NDArray{T,D}) where {T,D} - println(io, "NDArray{$(T),$(D)}") - Base.print_array(io, Array(arr)) +function Base.show(io::IO, arr::NDArray{T,N}) where {T,N} + print(io, "NDArray{$(T),$(N)} with size ", size(arr)) end -function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T}) where {T} - Base.show(io, arr) +function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T,N}) where {T,N} + println(io, "NDArray{$(T),$(N)} with size ", size(arr)) + Base.print_array(io, Array(arr)) end function Base.print(arr::NDArray{T}) where {T} @@ -375,39 +406,48 @@ function _setindex!(::Val{N}, arr::NDArray{Bool,N}, value::Bool, idxs::Vararg{In end #### START OF SLICING #### +# LHS slices from `nda_get_slice` are invisible to `@analyze_lifetimes`; 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) + s = nda_get_slice(lhs, slices) + copyto!(s, rhs) + destroy!(s) + return nothing +end + function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Colon, j::Int64) - s = nda_get_slice(lhs, slice_array((0, Base.size(lhs, 1)), (j-1, j))) - copyto!(s, rhs); + _setindex_slice!(lhs, rhs, slice_array((0, Base.size(lhs, 1)), (j-1, j))) end function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Int64, j::Colon) - s = nda_get_slice(lhs, slice_array((i-1, i))) - copyto!(s, rhs); + _setindex_slice!(lhs, rhs, slice_array((i-1, i))) end function Base.setindex!(lhs::NDArray, rhs::NDArray, i::UnitRange, j::Colon) - s = nda_get_slice(lhs, slice_array((first(i) - 1, last(i)), (0, Base.size(lhs, 2)))) - copyto!(s, rhs) + _setindex_slice!( + lhs, rhs, slice_array((first(i) - 1, last(i)), (0, Base.size(lhs, 2))) + ) end function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Colon, j::UnitRange) - s = nda_get_slice(lhs, slice_array((0, Base.size(lhs, 1)), (first(j) - 1, last(j)))) - copyto!(s, rhs) + _setindex_slice!( + lhs, rhs, slice_array((0, Base.size(lhs, 1)), (first(j) - 1, last(j))) + ) end function Base.setindex!(lhs::NDArray, rhs::NDArray, i::UnitRange, j::Int64) - s = nda_get_slice(lhs, slice_array((first(i) - 1, last(i)), (j-1, j))) - copyto!(s, rhs) + _setindex_slice!(lhs, rhs, slice_array((first(i) - 1, last(i)), (j-1, j))) end function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Int64, j::UnitRange) - s = nda_get_slice(lhs, slice_array((i-1, i), (first(j) - 1, last(j)))) - copyto!(s, rhs) + _setindex_slice!(lhs, rhs, slice_array((i-1, i), (first(j) - 1, last(j)))) end function Base.setindex!(lhs::NDArray, rhs::NDArray, i::UnitRange, j::UnitRange) - s = nda_get_slice(lhs, slice_array((first(i) - 1, last(i)), (first(j) - 1, last(j)))) - copyto!(s, rhs) + _setindex_slice!( + lhs, rhs, slice_array((first(i) - 1, last(i)), (first(j) - 1, last(j))) + ) end function Base.getindex(arr::NDArray, i::Colon, j::Int64) @@ -456,13 +496,15 @@ function Base.setindex!(arr::NDArray{T}, rhs::NDArray{T}, c::Vararg{Colon,N}) wh end function Base.setindex!(arr::NDArray{T,2}, val::T, i::Colon, j::Int64) where {T} - s = nda_get_slice(arr, to_cpp_init_slice(slice(0, Base.size(arr, 1)), slice(j-1, j))) + s = nda_get_slice(arr, slice_array((0, Base.size(arr, 1)), (j-1, j))) nda_fill_array(s, val) + destroy!(s) end function Base.setindex!(arr::NDArray{T,2}, val::T, i::Int64, j::Colon) where {T} - s = nda_get_slice(arr, to_cpp_init_slice(slice(i-1, i))) + s = nda_get_slice(arr, slice_array((i-1, i))) nda_fill_array(s, val) + destroy!(s) end Base.fill!(arr::NDArray{T}, val::T) where {T} = nda_fill_array(arr, val) @@ -605,34 +647,40 @@ function ones() end @doc""" - cuNumeric.rand!(arr::NDArray) + cuNumeric.rand!(arr::NDArray{Float64}) -Fills `arr` with AbstractFloats uniformly at random. +Fill `arr` in-place with uniform random `Float64` values. +""" +Random.rand!(arr::NDArray{Float64}) = cuNumeric.nda_random(arr, 0) +function Random.rand!(arr::NDArray{T}) where {T} + error("rand! only supports NDArray{Float64} for now. Cast with cuNumeric.as_type.") +end - cuNumeric.rand(NDArray, dims::Int...) - cuNumeric.rand(NDArray, dims::Tuple) +# Backend only generates Float64. Same-type path needs no cast; other floats +# convert then eagerly drop the Float64 source so it cannot leak until GC. +@doc""" + cuNumeric.rand([T=Float32,] dims::Int...) + cuNumeric.rand([T=Float32,] dims::Tuple) -Create a new `NDArray` of element type Float64, filled with uniform random values. +Create a new `NDArray` filled with uniform random values. -The backend currently supports only `Float64` with uniform distribution. -In order to support other Floats, we type convert for the user automatically. -This can create extra allocations. +The backend currently supports only `Float64` draws. Other floating types are +converted automatically. # Examples ```@repl cuNumeric.rand(2, 2) cuNumeric.rand((4, 1)) -A = cuNumeric.zeros(2, 2); cuNumeric.rand!(A) +A = cuNumeric.zeros(Float64, 2, 2); cuNumeric.rand!(A) ``` """ -Random.rand!(arr::NDArray{Float64}) = cuNumeric.nda_random(arr, 0) -function Random.rand!(arr::NDArray{T}) where {T} - error("rand! only supports NDArray{Float64} for now. Cast with cuNumeric.as_type.") -end +rand(::Type{Float64}, dims::Dims) = cuNumeric.nda_random_array(dims) function rand(::Type{T}, dims::Dims) where {T<:AbstractFloat} arrfp64 = cuNumeric.nda_random_array(dims) - return cuNumeric.as_type(arrfp64, T) + arr = cuNumeric.as_type(arrfp64, T) + destroy!(arrfp64) + return arr end rand(::Type{T}, dims::Int...) where {T<:AbstractFloat} = cuNumeric.rand(T, dims) @@ -641,28 +689,39 @@ rand(dims::Int...) = cuNumeric.rand(DEFAULT_FLOAT, dims) #### OPERATIONS #### @doc""" - reshape(arr::NDArray, dims::Dims{N}; copy::Bool = false) where {N} - reshape(arr::NDArray, dim::Int64; copy::Bool = false) + reshape(arr::NDArray, dims::Dims{N}; copy::Val{C}=Val(false)) where {N,C} + reshape(arr::NDArray, dims::Int...; copy::Val{C}=Val(false)) where {C} Return a new `NDArray` reshaped to the specified dimensions. +By default (`copy=Val(false)`) the result shares data with `arr`. +Pass `copy=Val(true)` to allocate a deep copy; the intermediate reshape +view is then destroyed eagerly. Use `Val` (not a runtime `Bool`) so the +return type stays concrete — a `Bool` branch widens inference. + # Examples ```@repl arr = cuNumeric.ones(4, 3) reshape(arr, (3, 4)) reshape(arr, 12) +reshape(arr, (3, 4); copy=Val(true)) ``` """ -#*USNTABLE USE Val{false} IF WE REALLY WANT THIS FLAG -function reshape(arr::NDArray, i::Dims{N}; copy::Bool=false) where {N} +# `copy` is a type parameter via Val{C}, so the default path constant-folds +# and stays type-stable (needed by solve's 1D-rhs reshape). +function reshape(arr::NDArray, i::Dims{N}; copy::Val{C}=Val(false)) where {N,C} reshaped = nda_reshape_array(arr, i) - return copy ? copy(reshaped) : reshaped + if C + copied = Base.copy(reshaped) + destroy!(reshaped) + return copied + end + return reshaped end -#*USNTABLE USE Val{false} IF WE REALLY WANT THIS FLAG -function reshape(arr::NDArray, i::Int...; copy::Bool=false) - return reshape(arr, i; copy=copy) +function reshape(arr::NDArray, i::Int...; copy::Val{C}=Val(false)) where {C} + return reshape(arr, i; copy=Val{C}()) end # Ignore the scalar indexing here... diff --git a/src/ndarray/promotion.jl b/src/ndarray/promotion.jl index 77d53749a..b9dcbac0e 100644 --- a/src/ndarray/promotion.jl +++ b/src/ndarray/promotion.jl @@ -11,6 +11,12 @@ end unchecked_promote_arr(arr::NDArray{T}, ::Type{T}) where {T} = arr unchecked_promote_arr(arr::NDArray{T}, ::Type{S}) where {T,S} = as_type(arr, S) +# Unfused broadcast leaves Numbers as scalars until promote; always a fresh 0-d array. +unchecked_promote_arr(x::Number, ::Type{T}) where {T} = NDArray(T(x)) + +# Fusion keeps Numbers as scalars in the PTX arg buffer (no 0-d NDArray). +unchecked_promote_scalar(x::Number, ::Type{T}) where {T} = T(x) +unchecked_promote_scalar(x, ::Type) = x # kinda hacky, but lets us support weird cases like broadcasting literal_pow unchecked_promote_arr(::Base.RefValue{typeof(^)}, ::Type{T}) where {T} = typeof(Base.:(^)) @@ -90,3 +96,8 @@ __my_promote_type(::Type{A}, ::Type{A}) where {A} = A is_wider_type(T, S) && assertpromotion(promote_type, S, T) return T end + +# Flattened fusion (and any n-ary leaf list) folds pairwise. +@inline function __my_promote_type(::Type{A}, ::Type{B}, ::Type{C}, rest::Type...) where {A,B,C} + return __my_promote_type(__my_promote_type(A, B), C, rest...) +end diff --git a/src/ndarray/unary.jl b/src/ndarray/unary.jl index 96e69f981..e9363a2b5 100644 --- a/src/ndarray/unary.jl +++ b/src/ndarray/unary.jl @@ -99,11 +99,14 @@ end end function Base.:(-)(input::NDArray{Bool}) - return -(checked_promote_arr(input, DEFAULT_INT)) + promoted = checked_promote_arr(input, DEFAULT_INT) # always new (Bool → Int) + out = -(promoted) + destroy!(promoted) + return out end function Base.sqrt(input::NDArray{T,2}) where {T} - error("cuNumeric.jl does not support matrix square root.") + return error("cuNumeric.jl does not support matrix square root.") end @inline function __broadcast( @@ -112,18 +115,32 @@ end return nda_unary_op!(out, cuNumeric.SQUARE, input) end +@inline function __broadcast( + ::typeof(Base.literal_pow), out::NDArray{O}, _, input::NDArray{O}, ::Type{Val{-1}} +) where {O} + nda_move(out, O(1) ./ input) #! REPLACE WITH RECIP ONCE FIXED + return out +end + @inline function __broadcast( ::typeof(Base.literal_pow), out::NDArray{O}, _, input::NDArray, ::Type{Val{-1}} ) where {O} - nda_move(out, O(1) ./ checked_promote_arr(input, O)) #! REPLACE WITH RECIP ONCE FIXED + promoted = checked_promote_arr(input, O) # always a new array when eltype ≠ O + nda_move(out, O(1) ./ promoted) #! REPLACE WITH RECIP ONCE FIXED + destroy!(promoted) + return out +end + +@inline function __broadcast(::typeof(Base.inv), out::NDArray{O}, input::NDArray{O}) where {O} + nda_move(out, O(1) ./ input) #! REPLACE WITH RECIP ONCE FIXED return out - # return nda_unary_op!(out, cuNumeric.RECIPROCAL, input) end @inline function __broadcast(::typeof(Base.inv), out::NDArray{O}, input::NDArray) where {O} - nda_move(out, O(1) ./ checked_promote_arr(input, O)) #! REPLACE WITH RECIP ONCE FIXED + promoted = checked_promote_arr(input, O) # always a new array when eltype ≠ O + nda_move(out, O(1) ./ promoted) #! REPLACE WITH RECIP ONCE FIXED + destroy!(promoted) return out - # return nda_unary_op!(out, cuNumeric.RECIPROCAL, checked_promote_arr(input,O)) end #! NEEDS TO SUPPORT inv and ^ -1 @@ -163,11 +180,14 @@ for (julia_fn, op_code) in floaty_unary_ops_no_args return nda_unary_op!(out, $(op_code), input) end - # If input is not already float, promote to that + # If input is not already float, promote to that (temp always new → always destroy) @inline function __broadcast( f::typeof($julia_fn), out::NDArray{A}, input::NDArray{B} ) where {A<:SUPPORTED_FLOAT_TYPES,B<:Union{SUPPORTED_INT_TYPES,Bool}} - return __broadcast(f, out, checked_promote_arr(input, A)) + promoted = checked_promote_arr(input, A) + result = __broadcast(f, out, promoted) + destroy!(promoted) + return result end end end @@ -253,29 +273,53 @@ global const unary_reduction_map = Dict{Function,UnaryRedCode}( #! IT WOULD BE NICE IF THESE JUST RETURNED SCALARS WHEN APPROPRIATE # #*TODO HOW TO GET THESE ACTING ON CERTAIN DIMS +function _unary_reduction_apply(out, op_code, input::NDArray{T}, ::Type{T}) where {T} + return nda_unary_reduction(out, op_code, input) +end + +function _unary_reduction_apply(out, op_code, input::NDArray, ::Type{U}) where {U} + promoted = unchecked_promote_arr(input, U) # always a new array when U ≠ eltype + result = nda_unary_reduction(out, op_code, promoted) + destroy!(promoted) + return result +end + +function _unary_reduction_axes_apply(op_code, input::NDArray{T}, ::Type{T}, axes) where {T} + return nda_unary_reduction_axes(op_code, input, axes, true) +end + +function _unary_reduction_axes_apply(op_code, input::NDArray, ::Type{U}, axes) where {U} + promoted = unchecked_promote_arr(input, U) # always a new array when U ≠ eltype + result = nda_unary_reduction_axes(op_code, promoted, axes, true) + destroy!(promoted) + return result +end + function _unary_reduction_impl(base_func, op_code, input::NDArray{T}, ::Colon) where {T} T_OUT = Base.promote_op(base_func, Vector{T}) is_wider_type(T_OUT, T) && assertpromotion(base_func, T, T_OUT) out = cuNumeric.zeros(T_OUT) - return nda_unary_reduction(out, op_code, unchecked_promote_arr(input, T_OUT)) + return _unary_reduction_apply(out, op_code, input, T_OUT) end function _unary_reduction_impl(base_func, op_code, input::NDArray{T,N}, dims::Integer) where {T,N} T_OUT = Base.promote_op(base_func, Vector{T}) is_wider_type(T_OUT, T) && assertpromotion(base_func, T, T_OUT) axes = Int32[dims - 1] - return nda_unary_reduction_axes(op_code, unchecked_promote_arr(input, T_OUT), axes, true) + return _unary_reduction_axes_apply(op_code, input, T_OUT, axes) end function _unary_reduction_impl(base_func, op_code, input::NDArray{T,N}, dims::Tuple) where {T,N} if length(dims) > 1 - error("$(base_func): reducing over multiple dimensions is not yet supported. Got dims=$dims") + error( + "$(base_func): reducing over multiple dimensions is not yet supported. Got dims=$dims" + ) end # single element tuple T_OUT = Base.promote_op(base_func, Vector{T}) is_wider_type(T_OUT, T) && assertpromotion(base_func, T, T_OUT) axes = Int32[dims[1] - 1] - return nda_unary_reduction_axes(op_code, unchecked_promote_arr(input, T_OUT), axes, true) + return _unary_reduction_axes_apply(op_code, input, T_OUT, axes) end # Generate code for all unary reductions. diff --git a/src/scoping.jl b/src/scoping.jl index 6297d6866..47ae01cea 100644 --- a/src/scoping.jl +++ b/src/scoping.jl @@ -1,4 +1,4 @@ -export @analyze_lifetimes +export @analyze_lifetimes, @show_lifetimes @doc""" @analyze_lifetimes expr @@ -7,16 +7,21 @@ 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. + +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) - esc(process_ndarray_scope(block)) + return esc(process_ndarray_scope(block)) end const counter = Ref(0) function maybe_insert_delete(var::NDArray) - cuNumeric.nda_destroy_array(var.ptr) - var.ptr = Ptr{Cvoid}(0) + return cuNumeric.destroy!(var) end maybe_insert_delete(x) = x @@ -98,11 +103,50 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) last_use[v] = maximum(idxs) end + # `F_u = tmp1` aliases one NDArray under two names; resolve to a canonical rep + # so it's freed once (double-free is masked only by the ptr=0 null-out). + function canon(v) + seen = Set{Symbol}() + while haskey(alias_map, v) && !(v in seen) + push!(seen, v) + v = alias_map[v] + end + return v + end + + is_indexed_assign(s) = s isa Expr && s.head == :(=) && !(s.args[1] isa Symbol) + result_symbol(s) = + if s isa Symbol + s + else + (s isa Expr && s.head == :(=) && s.args[1] isa Symbol ? s.args[1] : nothing) + end + # Pass 2: insert finalizers out = Any[] n = length(stmts) + freed = Set{Symbol}() + + # The block's value escapes to the caller, except for `A[...] = rhs`: Julia + # returns `rhs` there, but that's a dead temp nobody consumes — free it and + # return `nothing` rather than leak it or hand back a dangling handle. + terminal_indexed = n > 0 && is_indexed_assign(stmts[n]) + + protected = Set{Symbol}() + if n > 0 && !terminal_indexed + rs = result_symbol(stmts[n]) + rs isa Symbol && push!(protected, canon(rs)) + end + + function emit_delete!(v) + c = canon(v) + (c in freed || c in protected) && return nothing + push!(freed, c) + return push!(out, :(cuNumeric.maybe_insert_delete($v))) + end + for (i, stmt) in enumerate(stmts) - # detect aliasing: v = w means don't finalize w + # `v = w` aliases w, so don't finalize w at this statement. skip_finalize = Set{Symbol}() if stmt isa Expr && stmt.head == :(=) lhs, rhs = stmt.args @@ -111,36 +155,21 @@ function insert_finalizers(exprs::Vector, assigned_vars::Set{Symbol}) end end - if i == n - # Capture result of the last statement + if i == n && !terminal_indexed res_var = Symbol(:res, counter[]) counter[] += 1 push!(out, :($res_var = $stmt)) - - # Insert finalizers for the last statement - for (v, lasti) in last_use - if lasti == i && v ∈ assigned_vars && !(v ∈ skip_finalize) - # Do not delete if the result of the block is exactly this variable - # or if it's an assignment to this variable. - is_result = (stmt === v) - if stmt isa Expr && stmt.head == :(=) && stmt.args[1] === v - is_result = true - end - if !is_result - push!(out, :(cuNumeric.maybe_insert_delete($v))) - end - end - end - # Return the captured result - push!(out, res_var) else push!(out, stmt) - for (v, lasti) in last_use - if lasti == i && v ∈ assigned_vars && !(v ∈ skip_finalize) - push!(out, :(cuNumeric.maybe_insert_delete($v))) - end + end + + for (v, lasti) in last_use + if lasti == i && v ∈ assigned_vars && !(v ∈ skip_finalize) + emit_delete!(v) end end + + i == n && push!(out, terminal_indexed ? :nothing : res_var) end return out @@ -162,6 +191,11 @@ function insert_finalizers(block::Expr, assigned_vars::Set{Symbol}) end function process_ndarray_scope(block) + # Broadcast trees need the fusion-aware hoisting; otherwise the + # plain analysis treats every call (dotted or not) as a real allocation. + @static if FUSE_BROADCAST_EXPRS + return process_broadcast_scope(block) + end assigned_vars = Set{Symbol}() # Process the entire block at once so lifetimes are tracked across statements rewritten = find_ndarray_assignments(block, assigned_vars) @@ -270,3 +304,166 @@ function find_ndarray_assignments(ex, assigned_vars::Set{Symbol}) return Expr(:block, temps..., new_ex) end end + +# Broadcast-fusion-aware lifetime analysis. Under fusion, dotted operators +# (.+, .*, …) form a lazy Broadcasted tree lowered to one PTX kernel, so their +# intermediate nodes are not real NDArrays — only slices and the tree root are. +# Non-broadcast sub-expressions break the tree and are hoisted like any call. + +is_broadcast_op(op) = op isa Symbol && startswith(string(op), ".") + +function find_broadcast_assignments(ex, assigned_vars::Set{Symbol}) + local_assigned = Set{Symbol}() + + function fresh_tmp(expr) + counter[] += 1 + tmp = Symbol(:tmp, counter[]) + push!(local_assigned, tmp) + return tmp, [:($tmp = $expr)] + end + + # Rewrite each arg with `f`, collecting the temps each one hoists. + function maphoist(f, args) + new_args, hoisted = Any[], Expr[] + for a in args + na, ts = f(a) + push!(new_args, na) + append!(hoisted, ts) + end + return new_args, hoisted + end + + # Inside a broadcast tree: hoist slices, keep dotted ops/f.(…) lazy, and + # delegate anything else to rewrite() (it breaks the tree → real NDArray). + function rewrite_bcast(e)::Tuple{Any,Vector{Expr}} + e isa Expr || return e, Expr[] + e.head == :ref && return fresh_tmp(e) + if e.head == :call && is_broadcast_op(e.args[1]) + args, hoisted = maphoist(rewrite_bcast, e.args[2:end]) + return Expr(:call, e.args[1], args...), hoisted + end + if e.head == :. && length(e.args) == 2 && + e.args[2] isa Expr && e.args[2].head == :tuple + args, hoisted = maphoist(rewrite_bcast, e.args[2].args) + return Expr(:., e.args[1], Expr(:tuple, args...)), hoisted + end + return rewrite(e) + end + + function rewrite(e)::Tuple{Any,Vector{Expr}} + e isa Expr || return e, Expr[] + + if e.head == :(=) + lhs, rhs = e.args + lhs isa Symbol && push!(local_assigned, lhs) + new_rhs, temps = rewrite(rhs) + return :($lhs = $new_rhs), temps + end + + # .= RHS is a broadcast tree: only the slices inside it are hoisted. + if e.head == :(.=) + lhs, rhs = e.args + new_lhs, lts = rewrite(lhs) + new_rhs, rts = rewrite_bcast(rhs) + return Expr(:(.=), new_lhs, new_rhs), vcat(lts, rts) + end + + e.head == :ref && return fresh_tmp(e) + + # Broadcast root: hoist the fused result as one temp. + if e.head == :call && is_broadcast_op(e.args[1]) + inner, hoisted = rewrite_bcast(e) + tmp, bind = fresh_tmp(inner) + return tmp, vcat(hoisted, bind) + end + + # Regular call: hoist it and its args. + if e.head == :call + args, hoisted = maphoist(rewrite, e.args[2:end]) + tmp, bind = fresh_tmp(Expr(:call, e.args[1], args...)) + return tmp, vcat(hoisted, bind) + end + + # Other exprs: recurse, splicing hoisted temps inline within blocks. + new_args, hoisted = Any[], Expr[] + is_block = e.head == :block || e.head == :begin + for a in e.args + na, ts = rewrite(a) + if is_block && !(a isa LineNumberNode) + append!(new_args, ts) + push!(new_args, na) + else + push!(new_args, na) + append!(hoisted, ts) + end + end + return Expr(e.head, new_args...), hoisted + end + + new_ex, temps = rewrite(ex) + union!(assigned_vars, local_assigned) + if new_ex isa Expr && new_ex.head == :block + return Expr(:block, temps..., new_ex.args...) + else + return Expr(:block, temps..., new_ex) + end +end + +function process_broadcast_scope(block) + assigned_vars = Set{Symbol}() + rewritten = find_broadcast_assignments(block, assigned_vars) + result = insert_finalizers(rewritten, assigned_vars) + counter[] = 0 + return result +end + +# Pretty-print the @analyze_lifetimes rewrite (see @show_lifetimes below). +_is_delete_call(s) = Meta.isexpr(s, :call) && s.args[1] == :(cuNumeric.maybe_insert_delete) + +# Flatten nested begin/blocks into a linear statement list, dropping line nodes. +function _flatten_stmts(x) + stmts = Any[] + function walk(e) + if Meta.isexpr(e, (:block, :begin)) + foreach(walk, e.args) + elseif !(e isa LineNumberNode) + push!(stmts, e) + end + end + walk(x) + return stmts +end + +function print_lifetime_analysis(block; io::IO=stdout) + rule = "-"^60 + stmts = _flatten_stmts(process_ndarray_scope(block)) + mode = FUSE_BROADCAST_EXPRS ? "fusion-aware" : "plain" + + println(io, "@analyze_lifetimes expansion ($mode analysis)\n", rule) + + n = 0 + for s in stmts + if _is_delete_call(s) + printstyled(io, lpad("✗ free ", 11), s.args[2], "\n"; color=:red) + else + n += 1 + println(io, lpad(n, 4), " ", s) + end + 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) + return :(print_lifetime_analysis($(QuoteNode(block)))) +end diff --git a/src/util.jl b/src/util.jl index 96856ab78..5ca776267 100644 --- a/src/util.jl +++ b/src/util.jl @@ -1,4 +1,5 @@ export get_time_microseconds, get_time_nanoseconds +export issue_mapping_fence, issue_execution_fence @doc""" Returns the timestamp in microseconds. Blocks on all Legate operations @@ -16,6 +17,21 @@ function get_time_nanoseconds() return Legate.time_nanoseconds() end +""" + issue_mapping_fence() + +Insert a Legate mapping fence (DAG-only; does not block the Julia caller). +""" +issue_mapping_fence() = Legate.issue_mapping_fence() + +""" + issue_execution_fence(block::Bool) + +Insert a Legate execution fence. `block=true` waits until prior ops finish; +`block=false` only inserts a DAG node (Julia can keep submitting). +""" +issue_execution_fence(; block::Bool=false) = Legate.issue_execution_fence(block) + function Experimental(setting::Bool) task_local_storage(:Experimental, setting) end diff --git a/src/utilities/cuda_stubs.jl b/src/utilities/cuda_stubs.jl deleted file mode 100644 index caead31a4..000000000 --- a/src/utilities/cuda_stubs.jl +++ /dev/null @@ -1,92 +0,0 @@ -## This file contains stubs for methods implemented in -## the CUDA package extensions not implemented -## elsewhere in the package. - -export @cuda_task, @launch -export map_ndarray_cuda_types, launch, CUDATask - -function ptx_task end -function map_ndarray_cuda_types end -function launch end - -struct CUDATask - func::String - argtypes::NTuple{N,Type} where {N} -end - -""" - @cuda_task(f(args...)) - -Compile a Julia GPU kernel to PTX, register it with the Legate runtime, -and return a `CUDATask` object for later launch. - -# Arguments -- `f` — The name of the Julia CUDA.jl GPU kernel function to compile. -- `args...` — Example arguments to the kernel, used to determine the - argument type signature when generating PTX. - -# Description -This macro automates the process of: -1. Inferring the CUDA argument types for the given `args` using - `map_ndarray_cuda_types`. -2. Using `CUDA.code_ptx` to compile the specified GPU kernel - (`f`) into raw PTX text for the inferred types. -3. Extracting the kernel's function symbol name from the PTX using - `extract_kernel_name`. -4. Registering the compiled PTX and kernel name with the Legate runtime - via `ptx_task`, making it available for GPU execution. -5. Returning a `CUDATask` struct that stores the kernel name and type signature, - which can be used to configure and launch the kernel later. - -# Notes -- The `args...` are not executed; they are used solely for type inference. -- This macro is intended for use with the Legate runtime and - assumes a CUDA context is available. -- Make sure your kernel code is GPU-compatible and does not rely on - unsupported Julia features. - -# Example -```julia -mytask = @cuda_task my_kernel(A, B, C) -``` -""" -macro cuda_task end - -""" - @launch(; task, blocks=(1,), threads=(256,), inputs=(), outputs=(), scalars=()) - -Launch a GPU kernel (previously registered via [`@cuda_task`](@ref)) through the Legate runtime. - -# Keywords -- `task` — A `CUDATask` object, typically returned by [`@cuda_task`](@ref). -- `blocks` — Tuple or single element specifying the CUDA grid dimensions. Defaults to `(1,)`. -- `threads` — Tuple or single element specifying the CUDA block dimensions. Defaults to `(256,)`. -- `inputs` — Tuple or single element of input NDArray objects. -- `outputs` — Tuple or single element of output NDArray objects. -- `scalars` — Tuple or single element of scalar values. - -# Description -The `@launch` macro validates the provided keywords, ensuring only -the allowed set (`:task`, `:blocks`, `:threads`, `:inputs`, `:outputs`, `:scalars`) -are present. It then expands to a call to `cuNumeric.launch`, -passing the given arguments to the Legate runtime for execution. - -This macro is meant to provide a concise, declarative syntax for -launching GPU kernels, separating kernel compilation (via `@cuda_task`) -from execution configuration. - -# Notes -- `task` **must** be a kernel registered with the runtime, usually from `@cuda_task`. -- All keyword arguments must be specified as assignments, e.g. `blocks=(2,2)` not positional arguments. -- Defaults are chosen for single-block, 256-thread 1D launches. -- The macro escapes its body so that the values of inputs/outputs/scalars are captured - from the surrounding scope at macro expansion time. - -# Example -```julia -mytask = @cuda_task my_kernel(A, B, C) - -@launch task=mytask blocks=(8,8) threads=(32,32) inputs=(A, B) outputs=(C) -``` -""" -macro launch end diff --git a/test/runtests.jl b/test/runtests.jl index 55bb4af1d..03caf4e68 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,6 +41,8 @@ end using cuNumeric VERBOSE && cuNumeric.versioninfo() +@info "Broadcast fusion: FUSE_BROADCAST_EXPRS=$(cuNumeric.FUSE_BROADCAST_EXPRS) FUSE_BROADCAST_MIN_OPS=$(cuNumeric.FUSE_BROADCAST_MIN_OPS)" + # TODO # After loading cuNumeric, we should verify that the Legate config has set a GPU device # Right now, if you have a gpu device, but your LEGATE_CONFIG is cpu only, @@ -57,6 +59,7 @@ include("tests/unary_tests.jl") include("tests/binary_tests.jl") include("tests/scoping.jl") include("tests/scoping-advanced.jl") +include("tests/broadcast_fusion_tests.jl") @testset verbose = true "AXPY" begin N = 100 @@ -88,11 +91,11 @@ end @testset for T in Base.uniontypes(cuNumeric.SUPPORTED_ARRAY_TYPES) allowpromotion(true) do - test_unary_function_set(cuNumeric.floaty_unary_ops_no_args, T, N) + return test_unary_function_set(cuNumeric.floaty_unary_ops_no_args, T, N) end allowpromotion(T == Bool) do - test_unary_function_set(cuNumeric.unary_op_map_no_args, T, N) + return test_unary_function_set(cuNumeric.unary_op_map_no_args, T, N) end # Special cases for unary ops that dont use . syntax @testset "- (Negation)" begin @@ -204,7 +207,7 @@ end @testset for T in Base.uniontypes(cuNumeric.SUPPORTED_ARRAY_TYPES) allowpromotion(true) do test_binary_function_set(cuNumeric.floaty_binary_op_map, T, N) - test_binary_function_set(cuNumeric.binary_op_map, T, N) + return test_binary_function_set(cuNumeric.binary_op_map, T, N) end arr_jl = my_rand(T, N) @@ -461,11 +464,17 @@ end end if run_gpu_tests + @testset verbose = true "Broadcast Fusion" begin + test_broadcast_fusion() + test_broadcast_fusion_edge_cases() + test_broadcast_fusion_ptx_cache() + end + # @testset verbose = true "CUDA Tests" begin # cuda_unaryop(rtol(Float32)) # cuda_binaryop(rtol(Float32)) # end - @warn "CUDA tests are turned off inside Pkg.test for now. --check-bounds=yes causes issues." + @warn "CUDA @cuda_task tests are turned off inside Pkg.test for now. --check-bounds=yes causes issues." else @warn "The CUDA tests will not be run as a CUDA-enabled device is not available" end diff --git a/test/tests/axpy_advanced.jl b/test/tests/axpy_advanced.jl index e090e3a37..18950f7c1 100644 --- a/test/tests/axpy_advanced.jl +++ b/test/tests/axpy_advanced.jl @@ -1,4 +1,4 @@ -#= Copyright 2026 Northwestern University, +#= Copyright 2026 Northwestern University, * Carnegie Mellon University University * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -20,8 +20,6 @@ #= Purpose of test: daxpy_advanced -- add overloading support for [double/float scalar] * NDArray -- equavalence operator between a cuNumeric and Julia array without looping - -- result == (α_cpu * x_cpu + y_cpu) - -- (α_cpu * x_cpu + y_cpu) == -- NDArray copy method allocates a new NDArray and copies all elements -- NDArray assign method assigns the contents from one NDArray to another NDArray -- x[:] colon notation for reading entire 1D NDArray to a Julia array @@ -110,9 +108,8 @@ function axpy_advanced(T, N) @test is_same(y_cpu_1D, y_1d) result = α .* x .+ y + result_cpu = α * x_cpu + y_cpu - # check results - @test is_same(result, (α * x_cpu + y_cpu)) - @test is_same(α * x_cpu + y_cpu, result) # LHS and RHS switched + @test safe_compare(result_cpu, result, atol(T), rtol(T)) end end diff --git a/test/tests/broadcast_fusion_tests.jl b/test/tests/broadcast_fusion_tests.jl new file mode 100644 index 000000000..ceaa84593 --- /dev/null +++ b/test/tests/broadcast_fusion_tests.jl @@ -0,0 +1,507 @@ +#= 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 +=# + +#= Broadcast fusion arg buffer tests. + * Tests every combination of array and scalar argument positions + * to verify the arg_map protocol correctly reconstructs the CUDA + * kernel arg buffer. + * + * Pattern naming convention: + * A = NDArray input + * S = scalar input + * Position = left-to-right in the broadcast expression + * "same" = same array appears multiple times (deduplication test) +=# + +_broadcast_fusion_user_add(x, y) = x + y + +function test_broadcast_fusion(; T=Float32, N=100, atol=1e-5, rtol=1e-5) + # Create test arrays with known non-zero values + julia_a = rand(T, N) + julia_b = rand(T, N) + julia_c = rand(T, N) + + a = @allowscalar NDArray(julia_a) + b = @allowscalar NDArray(julia_b) + c = @allowscalar NDArray(julia_c) + + s1 = T(2.5) + s2 = T(1.0) + s3 = T(0.5) + + # two different arrays + @testset "A + B (two different arrays)" begin + expected = julia_a .+ julia_b + result = a .+ b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array, deduplication + @testset "A + A (same array twice)" begin + expected = julia_a .+ julia_a + result = a .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # array then scalar + @testset "A + scalar (array first)" begin + expected = julia_a .+ s1 + result = a .+ s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar then array + @testset "scalar + A (scalar first)" begin + expected = s1 .+ julia_a + result = s1 .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # array, two scalars, fused + @testset "A * scalar - scalar (fused)" begin + expected = julia_a .* s1 .- s2 + result = a .* s1 .- s2 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar-array-scalar + @testset "scalar * A + scalar" begin + expected = s1 .* julia_a .+ s2 + result = s1 .* a .+ s2 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # two arrays then scalar + @testset "A + B + scalar" begin + expected = julia_a .+ julia_b .+ s1 + result = a .+ b .+ s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar then two arrays + @testset "scalar + A + B" begin + expected = s1 .+ julia_a .+ julia_b + result = s1 .+ a .+ b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array twice + scalar (dedup) + @testset "A + A + scalar (dedup + scalar)" begin + expected = julia_a .+ julia_a .+ s1 + result = a .+ a .+ s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # three different arrays + @testset "A + B + C (three arrays)" begin + expected = julia_a .+ julia_b .+ julia_c + result = a .+ b .+ c + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array three times (triple dedup) + @testset "A + A + A (triple dedup)" begin + expected = julia_a .+ julia_a .+ julia_a + result = a .+ a .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # two scalars then array + @testset "scalar * scalar + A" begin + expected = s1 .* s2 .+ julia_a + result = s1 .* s2 .+ a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # multiply, two arrays (different PTX kernel name collision test) + @testset "A * B (multiply)" begin + expected = julia_a .* julia_b + result = a .* b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array squared, dedup + @testset "A * A (self multiply, dedup)" begin + expected = julia_a .* julia_a + result = a .* a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # subtraction, order matters + @testset "A - B (subtraction)" begin + expected = julia_a .- julia_b + result = a .- b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # scalar minus array + @testset "scalar - A" begin + expected = s1 .- julia_a + result = s1 .- a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # three arrays, mixed ops + @testset "A * B + C (three arrays, mixed ops)" begin + expected = julia_a .* julia_b .+ julia_c + result = a .* b .+ c + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # two arrays fused, then scaled + @testset "(A + B) * scalar" begin + expected = (julia_a .+ julia_b) .* s1 + result = (a .+ b) .* s1 + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + if cuNumeric.FUSE_BROADCAST_EXPRS + @testset "z .= scalar * f.(A, B)" begin + expected = T(2.0) .* (julia_a .+ julia_b) + z = cuNumeric.zeros(T, (N,)) + @analyze_lifetimes begin + z .= T(2.0) .* _broadcast_fusion_user_add.(a, b) + end + @allowscalar @test cuNumeric.compare(expected, z, atol, rtol) + end + end + + # scalar-array pairs + @testset "scalar * A + scalar * B" begin + expected = s1 .* julia_a .+ s2 .* julia_b + result = s1 .* a .+ s2 .* b + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end + + # same array subtracted, should be all zeros + @testset "A - A (same array, expect zeros)" begin + expected = julia_a .- julia_a + result = a .- a + @allowscalar @test cuNumeric.compare(expected, result, atol, rtol) + end +end + +#= Edge cases for linear-only broadcast fusion. + * Complements `test_broadcast_fusion` with size extremes, 2D same-shape, + * fusion gating for shape mismatch, 0-d fallback, and dest/input aliasing. +=# +function test_broadcast_fusion_edge_cases(; T=Float32, atol=1e-5, rtol=1e-5) + s1 = T(2.5) + s2 = T(1.0) + + @testset "very small 1D (N=1)" begin + ja = T[1.5] + jb = T[2.25] + a = @allowscalar NDArray(ja) + b = @allowscalar NDArray(jb) + result = a .+ b .* s1 .- s2 + @allowscalar @test cuNumeric.compare(ja .+ jb .* s1 .- s2, result, atol, rtol) + end + + @testset "very small 1D (N=2)" begin + ja = T[1.5, 3.0] + jb = T[2.5, 4.0] + a = @allowscalar NDArray(ja) + b = @allowscalar NDArray(jb) + result = a .+ b + @allowscalar @test cuNumeric.compare(ja .+ jb, result, atol, rtol) + result = s1 .* a .- b + @allowscalar @test cuNumeric.compare(s1 .* ja .- jb, result, atol, rtol) + end + + @testset "empty / zero-size 1D" begin + # NDArray supports size (0,). `_copyto!` short-circuits on isempty + # before fusion, so this only checks the empty path does not crash. + e = cuNumeric.zeros(T, 0) + result = e .+ e + @test isempty(result) + @test size(result) == (0,) + result = e .+ s1 + @test isempty(result) + @test size(result) == (0,) + end + + @testset "empty / zero-size 2D" begin + e = cuNumeric.zeros(T, (0, 3)) + result = e .+ e + @test isempty(result) + @test size(result) == (0, 3) + end + + @testset "large-ish 1D same-shape fused" begin + N = 10_000 + ja = rand(T, N) + jb = rand(T, N) + a = @allowscalar NDArray(ja) + b = @allowscalar NDArray(jb) + result = a .+ b .* s1 .- s2 + @allowscalar @test cuNumeric.compare(ja .+ jb .* s1 .- s2, result, atol, rtol) + # Gate: same-shape leaves should be fusible. + dest = cuNumeric.zeros(T, N) + bc = Base.Broadcast.instantiate(Base.broadcasted(+, a, b)) + @test cuNumeric.can_fuse_linear_broadcast(dest, bc) + end + + @testset "large-ish 2D same-shape fused" begin + M, N = 128, 256 + ja = rand(T, M, N) + jb = rand(T, M, N) + a = @allowscalar NDArray(ja) + b = @allowscalar NDArray(jb) + result = a .+ b .* s1 + @allowscalar @test cuNumeric.compare(ja .+ jb .* s1, result, atol, rtol) + result = a .+ a .* b .- s2 + @allowscalar @test cuNumeric.compare(ja .+ ja .* jb .- s2, result, atol, rtol) + dest = cuNumeric.zeros(T, M, N) + bc = Base.Broadcast.instantiate(Base.broadcasted(+, a, b)) + @test cuNumeric.can_fuse_linear_broadcast(dest, bc) + end + + @testset "scalar gaps: A .^ 2 and scalar*A*scalar" begin + N = 64 + ja = rand(T, N) + a = @allowscalar NDArray(ja) + # Literal power uses RefValue{Val} / static-arg lowering. + result = a .^ 2 + @allowscalar @test cuNumeric.compare(ja .^ 2, result, atol, rtol) + result = s1 .* a .* s2 + @allowscalar @test cuNumeric.compare(s1 .* ja .* s2, result, atol, rtol) + end + + # Gray-Scott-style slice stencils: strided views must fuse correctly via + # CuStridedDeviceArray (packed Legate element strides). + @testset "fused slice stencils (X/Y and slice dest)" begin + N = 32 + ja = rand(T, N, N) + u = @allowscalar NDArray(ja) + two = T(2) + + # X-shifted (vary first index) + result_x = + u[3:end, 2:(end - 1)] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ + u[1:(end - 2), 2:(end - 1)] + expected_x = + ja[3:end, 2:(end - 1)] .- two .* ja[2:(end - 1), 2:(end - 1)] .+ + ja[1:(end - 2), 2:(end - 1)] + @allowscalar @test cuNumeric.compare(expected_x, result_x, atol, rtol) + + dest_x = similar(result_x) + bc_x = Base.Broadcast.instantiate( + Base.broadcasted( + +, + Base.broadcasted( + -, + u[3:end, 2:(end - 1)], + Base.broadcasted(*, 2, u[2:(end - 1), 2:(end - 1)]), + ), + u[1:(end - 2), 2:(end - 1)], + ), + ) + @test cuNumeric.can_fuse_linear_broadcast(dest_x, bc_x) + + # Y-shifted (vary second index) — previously failed under dense packing + result_y = + u[2:(end - 1), 3:end] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+ + u[2:(end - 1), 1:(end - 2)] + expected_y = + ja[2:(end - 1), 3:end] .- two .* ja[2:(end - 1), 2:(end - 1)] .+ + ja[2:(end - 1), 1:(end - 2)] + @allowscalar @test cuNumeric.compare(expected_y, result_y, atol, rtol) + + # Assign into a slice destination + out = @allowscalar NDArray(zeros(T, N, N)) + out[2:(end - 1), 2:(end - 1)] = result_x .+ result_y + expected_out = zeros(T, N, N) + expected_out[2:(end - 1), 2:(end - 1)] = expected_x .+ expected_y + @allowscalar @test cuNumeric.compare(expected_out, out, atol, rtol) + end + + @testset "fused/unfused scalar promotion parity" begin + N = 32 + ja = rand(T, N) + a = @allowscalar NDArray(ja) + dest = cuNumeric.zeros(T, N) + + # Bare Int64: fusible, host-promoted to T, same result as unfused. + bc_i64 = Base.Broadcast.instantiate(Base.broadcasted(*, 2, a)) + @test cuNumeric.can_fuse_linear_broadcast(dest, bc_i64) + result = 2 .* a + @allowscalar @test cuNumeric.compare(T(2) .* ja, result, atol, rtol) + @allowscalar @test cuNumeric.compare( + result, + cuNumeric.unravel_broadcast_tree(bc_i64), + atol, + rtol, + ) + + # Wider Float64 literal: fused and unfused throw the same promotion error. + if T === Float32 + bc_f64 = Base.Broadcast.instantiate(Base.broadcasted(+, a, 1.0)) + err_fused = try + a .+ 1.0 + nothing + catch e + sprint(showerror, e) + end + err_unfused = try + cuNumeric.unravel_broadcast_tree(bc_f64) + nothing + catch e + sprint(showerror, e) + end + @test err_fused !== nothing + @test err_unfused !== nothing + @test err_fused == err_unfused + @test occursin("Implicit promotion", err_fused) + @test_throws "Implicit promotion" a .+ 1.0 + end + end + + @testset "shape-mismatched broadcast refuses fusion" begin + # Linear fusion requires every NDArray leaf to match dest shape. + # Matrix .+ vector must fall back to the unfused path. + # + # NOTE: the unfused matrix.+vector path currently disagrees with Julia + # broadcasting semantics; do not assert equality with `ja .+ jv` here. + M, N = 64, 32 + ja = rand(T, M, N) + jv = rand(T, M) + a = @allowscalar NDArray(ja) + v = @allowscalar NDArray(jv) + dest = cuNumeric.zeros(T, M, N) + bc = Base.Broadcast.instantiate(Base.broadcasted(+, a, v)) + @test !cuNumeric.can_fuse_linear_broadcast(dest, bc) + + # Unfused fallback should not crash (correctness vs Julia is known-wrong). + result = a .+ v + @test size(result) == (M, N) + end + + @testset "0-d scalar NDArray (fusion refused, unfused ok)" begin + # RunPTXBroadcastTask only supports dims in [1, 6]; can_fuse refuses + # 0-d so `_copyto!` falls back to unfused. + # + # NOTE: `z1 .+ z2` still errors after a successful `copyto!` because + # `Broadcast.copy` for NDArrayStyle{0} unwraps with + # `dest[CartesianIndex()]`, which is not implemented. Test via + # `copyto!` into an explicit 0-d dest instead. + z1 = @allowscalar NDArray(T(2)) + z2 = @allowscalar NDArray(T(3)) + @test ndims(z1) == 0 + dest = cuNumeric.zeros(T) + bc = Base.Broadcast.instantiate(Base.broadcasted(+, z1, z2)) + @test !cuNumeric.can_fuse_linear_broadcast(dest, bc) + copyto!(dest, bc) + @allowscalar @test dest[] == T(5) + + dest2 = cuNumeric.zeros(T) + bc2 = Base.Broadcast.instantiate(Base.broadcasted(+, Base.broadcasted(*, z1, s1), z2)) + @test !cuNumeric.can_fuse_linear_broadcast(dest2, bc2) + copyto!(dest2, bc2) + @allowscalar @test dest2[] ≈ T(2) * s1 + T(3) atol = atol rtol = rtol + end + + @testset "dest aliases an input" begin + N = 128 + ja = rand(T, N) + jb = rand(T, N) + a = @allowscalar NDArray(copy(ja)) + b = @allowscalar NDArray(jb) + a .+= b + @allowscalar @test cuNumeric.compare(ja .+ jb, a, atol, rtol) + + a2 = @allowscalar NDArray(copy(ja)) + a2 .= a2 .* s1 .+ b + @allowscalar @test cuNumeric.compare(ja .* s1 .+ jb, a2, atol, rtol) + end +end + +#= Broadcast fusion PTX compilation cache. + * Verifies `_BCAST_PTX_CACHE` grows on first fused launch of a signature and + * is reused (no new entry) on a second launch of the same signature. + * Gated on `FUSE_BROADCAST_EXPRS` + `HAS_CUDA`; skips otherwise. + * With `FUSE_BROADCAST_MIN_OPS > 1`, single-op exprs are unfused — tests + * should set min ops to 1 (LocalPreferences / ENV) to exercise the cache. +=# +function test_broadcast_fusion_ptx_cache(; T=Float32, N=64) + if !(cuNumeric.FUSE_BROADCAST_EXPRS && cuNumeric.HAS_CUDA) + @info "Skipping PTX cache tests (need FUSE_BROADCAST_EXPRS && HAS_CUDA)" + return nothing + end + if cuNumeric.FUSE_BROADCAST_MIN_OPS > 1 + @info "Skipping PTX cache tests (need FUSE_BROADCAST_MIN_OPS <= 1 to fuse single-op exprs)" + return nothing + end + + cache = cuNumeric._BCAST_PTX_CACHE + cache_lock = cuNumeric._BCAST_PTX_CACHE_LOCK + cache_len() = + lock(cache_lock) do + return length(cache) + end + clear_cache!() = + lock(cache_lock) do + empty!(cache) + return nothing + end + + @testset "PTX cache hit / miss" begin + clear_cache!() + @test cache_len() == 0 + + a = @allowscalar NDArray(rand(T, N)) + b = @allowscalar NDArray(rand(T, N)) + + # First fused launch of a signature should compile and cache. + _ = a .+ b + n1 = cache_len() + @test n1 >= 1 + + # Same signature again should hit the cache (no new entry). + _ = a .+ b + @test cache_len() == n1 + + # Different op should miss and add a new entry. + _ = a .* b + n2 = cache_len() + @test n2 > n1 + + # Same multiply signature again should hit. + _ = a .* b + @test cache_len() == n2 + + # Different element type should miss and add another entry. + a64 = @allowscalar NDArray(rand(Float64, N)) + b64 = @allowscalar NDArray(rand(Float64, N)) + _ = a64 .+ b64 + n3 = cache_len() + @test n3 > n2 + + _ = a64 .+ b64 + @test cache_len() == n3 + + # Nested fused expression should miss, then hit on re-run. + _ = a .+ b .* a + n4 = cache_len() + @test n4 > n3 + + _ = a .+ b .* a + @test cache_len() == n4 + end +end diff --git a/test/tests/lifetime.jl b/test/tests/lifetime.jl index 9e5f037e9..c6253ae8a 100644 --- a/test/tests/lifetime.jl +++ b/test/tests/lifetime.jl @@ -1,4 +1,4 @@ -#= Copyright 2026 Northwestern University, +#= Copyright 2026 Northwestern University, * Carnegie Mellon University University * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -17,10 +17,23 @@ * Ethan Meitz =# -@testset "Zero-Copy Verification" begin +@testset "Array ↔ NDArray value roundtrip (row-major attach)" begin A = rand(Float64, 4, 4) NA = NDArray(A) + @allowscalar begin + @test all(A .== Array(NA)) + for i in 1:4, j in 1:4 + @test NA[i, j] == A[i, j] + end + end +end + +@testset "Zero-Copy Verification (1D)" begin + # 1D attach remains zero-copy; N>=2 copies into a C-ordered buffer. + A = rand(Float64, 16) + NA = NDArray(A) + @allowscalar begin @test all(A .== Array(NA)) end @@ -28,16 +41,16 @@ @test pointer(A) == cuNumeric.get_ptr(NA) # modify julia array, verify ndarray sees it - A[1, 1] = 99.0 + A[1] = 99.0 @allowscalar begin - @test NA[1, 1] == 99.0 + @test NA[1] == 99.0 end # modify ndarray, verify julia array sees it @allowscalar begin - NA[2, 2] = 88.0 + NA[2] = 88.0 end - @test A[2, 2] == 88.0 + @test A[2] == 88.0 end @testset "Lifetime Protection" begin @@ -76,8 +89,8 @@ end # the temporary Float64 array should be kept alive by NA.parent GC.gc(true) - GC.gc(true) - GC.gc(true) + GC.gc(true) + GC.gc(true) @allowscalar begin @test NA[1] == 1.0 @@ -90,3 +103,21 @@ end @test NA[1] == 42.0 end end + +@testset "2D attach parent lifetime" begin + function create_2d() + local_A = rand(Float32, 8, 8) + local_A[1, 1] = 3.14f0 + return NDArray(local_A), local_A[1, 1] + end + + NA, expected_val = create_2d() + GC.gc(true) + GC.gc(true) + GC.gc(true) + + @allowscalar begin + @test NA[1, 1] == expected_val + @test Array(NA)[1, 1] == expected_val + end +end diff --git a/test/tests/util.jl b/test/tests/util.jl index 0035b7764..a477d0576 100644 --- a/test/tests/util.jl +++ b/test/tests/util.jl @@ -1,4 +1,4 @@ -#= Copyright 2026 Northwestern University, +#= Copyright 2026 Northwestern University, * Carnegie Mellon University University * * Licensed under the Apache License, Version 2.0 (the "License"); @@ -110,6 +110,9 @@ end function safe_compare(x::AbstractArray{T}, y::NDArray{T}, rtol, atol) where {T} for CI in CartesianIndices(x) if !safe_isapprox(x[CI], y[Tuple(CI)...], rtol, atol) + println("Failed at index $(Tuple(CI))") + println("x[$(Tuple(CI))] = $(x[CI])") + println("y[$(Tuple(CI))] = $(y[Tuple(CI)...])") return false end end