diff --git a/README.md b/README.md
index 08eacd126..cad89042a 100644
--- a/README.md
+++ b/README.md
@@ -61,7 +61,7 @@ See [Kernel Fusion](https://julialegate.github.io/cuNumeric.jl/dev/perf/kernel_f
### The `@accelerate` macro
-`@accelerate` optimizes straight-line array code by fusing eligible CUDA broadcasts and releasing dead temporary `NDArray`s. See [The `@accelerate` Macro](https://julialegate.github.io/cuNumeric.jl/dev/perf/reduce_allocations) for usage guidance.
+`@accelerate` fuses eligible GPU broadcasts within and across statements, then releases materialized temporary `NDArray`s after their last use on CPU or GPU. See [The `@accelerate` Macro](https://julialegate.github.io/cuNumeric.jl/dev/perf/reduce_allocations) for usage guidance.
### Benchmarks
diff --git a/benchmark/benchmarks.toml b/benchmark/benchmarks.toml
index dc36587ca..9bedc6036 100644
--- a/benchmark/benchmarks.toml
+++ b/benchmark/benchmarks.toml
@@ -36,16 +36,40 @@ T = "Float32"
gpus = [1, 2, 4, 8]
cpus = 16
fusion = [true, false]
-N = [2000, 2832, 4000, 5656]
-M = [2000, 2832, 4000, 5656]
+N = [24000, 33944, 48000, 67888]
+M = [24000, 33944, 48000, 67888]
-[[grayscott_accelerated]]
+[[grayscott_function_accelerated]]
T = "Float32"
gpus = [1, 2, 4, 8]
cpus = 16
fusion = [true, false]
-N = [2000, 2832, 4000, 5656]
-M = [2000, 2832, 4000, 5656]
+N = [24000, 33944, 48000, 67888]
+M = [24000, 33944, 48000, 67888]
+
+[[grayscott_begin_accelerated]]
+T = "Float32"
+gpus = [1, 2, 4, 8]
+cpus = 16
+fusion = [true, false]
+N = [24000, 33944, 48000, 67888]
+M = [24000, 33944, 48000, 67888]
+
+[[grayscott_let_accelerated]]
+T = "Float32"
+gpus = [1, 2, 4, 8]
+cpus = 16
+fusion = [true, false]
+N = [24000, 33944, 48000, 67888]
+M = [24000, 33944, 48000, 67888]
+
+[[grayscott_expression_accelerated]]
+T = "Float32"
+gpus = [1, 2, 4, 8]
+cpus = 16
+fusion = [true, false]
+N = [24000, 33944, 48000, 67888]
+M = [24000, 33944, 48000, 67888]
#################################
# Monte-Carlo Integration #
diff --git a/benchmark/src/benchmarks/grayscott.jl b/benchmark/src/benchmarks/grayscott.jl
index c4da0673c..0301b913c 100644
--- a/benchmark/src/benchmarks/grayscott.jl
+++ b/benchmark/src/benchmarks/grayscott.jl
@@ -92,63 +92,64 @@ function check_benchmark_correctness(
return (u_ok && v_ok) ? "pass" : "fail"
end
-# Variant description:
-# baseline: as written
-# accelerated: step wrapped in @accelerate
-let body = quote
- # currently we don't have NDArray^x working yet. every operator is dotted
- # so each rhs fuses into a single broadcast kernel rather than shattering
- # into bare +/-/* binary tasks.
- F_u = (
- (
- .-u[2:(end - 1), 2:(end - 1)] .*
- (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)])
- ) .+ args.f .* (1.0f0 .- u[2:(end - 1), 2:(end - 1)])
- )
- F_v = (
- (
- u[2:(end - 1), 2:(end - 1)] .*
- (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)])
- ) .- (args.f + args.k) .* v[2:(end - 1), 2:(end - 1)]
- )
- # 2-D Laplacian via slicing, excluding boundaries
- u_lap = (
- (
- u[3:end, 2:(end - 1)] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+
- u[1:(end - 2), 2:(end - 1)]
- ) ./ args.dx^2 .+
- (
- u[2:(end - 1), 3:end] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+
- u[2:(end - 1), 1:(end - 2)]
- ) ./ args.dx^2
- )
- v_lap = (
- (
- v[3:end, 2:(end - 1)] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+
- v[1:(end - 2), 2:(end - 1)]
- ) ./ args.dx^2 .+
- (
- v[2:(end - 1), 3:end] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+
- v[2:(end - 1), 1:(end - 2)]
- ) ./ args.dx^2
- )
-
- # Forward-Euler step for all interior points
- u_new[2:(end - 1), 2:(end - 1)] =
- ((args.c_u .* u_lap) .+ F_u) .* args.dt .+ u[2:(end - 1), 2:(end - 1)]
- v_new[2:(end - 1), 2:(end - 1)] =
- ((args.c_v .* v_lap) .+ F_v) .* args.dt .+ v[2:(end - 1), 2:(end - 1)]
-
- # Periodic boundary conditions
- u_new[:, 1] = u[:, end - 1]
- u_new[:, end] = u[:, 2]
- u_new[1, :] = u[end - 1, :]
- u_new[end, :] = u[2, :]
- v_new[:, 1] = v[:, end - 1]
- v_new[:, end] = v[:, 2]
- v_new[1, :] = v[end - 1, :]
- v_new[end, :] = v[2, :]
- end
+# Shared syntax tree keeps every Gray-Scott variant on the exact same workload.
+const GRAYSCOTT_STEP_BODY = quote
+ # currently we don't have NDArray^x working yet. every operator is dotted
+ # so each rhs fuses into a single broadcast kernel rather than shattering
+ # into bare +/-/* binary tasks.
+ F_u = (
+ (
+ .-u[2:(end - 1), 2:(end - 1)] .*
+ (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)])
+ ) .+ args.f .* (1.0f0 .- u[2:(end - 1), 2:(end - 1)])
+ )
+ F_v = (
+ (
+ u[2:(end - 1), 2:(end - 1)] .*
+ (v[2:(end - 1), 2:(end - 1)] .* v[2:(end - 1), 2:(end - 1)])
+ ) .- (args.f + args.k) .* v[2:(end - 1), 2:(end - 1)]
+ )
+ # 2-D Laplacian via slicing, excluding boundaries
+ u_lap = (
+ (
+ u[3:end, 2:(end - 1)] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+
+ u[1:(end - 2), 2:(end - 1)]
+ ) ./ args.dx^2 .+
+ (
+ u[2:(end - 1), 3:end] .- 2 .* u[2:(end - 1), 2:(end - 1)] .+
+ u[2:(end - 1), 1:(end - 2)]
+ ) ./ args.dx^2
+ )
+ v_lap = (
+ (
+ v[3:end, 2:(end - 1)] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+
+ v[1:(end - 2), 2:(end - 1)]
+ ) ./ args.dx^2 .+
+ (
+ v[2:(end - 1), 3:end] .- 2 .* v[2:(end - 1), 2:(end - 1)] .+
+ v[2:(end - 1), 1:(end - 2)]
+ ) ./ args.dx^2
+ )
+
+ # Forward-Euler step for all interior points
+ u_new[2:(end - 1), 2:(end - 1)] =
+ ((args.c_u .* u_lap) .+ F_u) .* args.dt .+ u[2:(end - 1), 2:(end - 1)]
+ v_new[2:(end - 1), 2:(end - 1)] =
+ ((args.c_v .* v_lap) .+ F_v) .* args.dt .+ v[2:(end - 1), 2:(end - 1)]
+
+ # Periodic boundary conditions
+ u_new[:, 1] = u[:, end - 1]
+ u_new[:, end] = u[:, 2]
+ u_new[1, :] = u[end - 1, :]
+ u_new[end, :] = u[2, :]
+ v_new[:, 1] = v[:, end - 1]
+ v_new[:, end] = v[:, 2]
+ v_new[1, :] = v[end - 1, :]
+ v_new[end, :] = v[2, :]
+end
+
+# Original baseline and recommended function-form benchmark.
+let body = deepcopy(GRAYSCOTT_STEP_BODY)
@eval _gs_step!(b::GrayScottBaseline, u, v, u_new, v_new, args::GSParams) = $body
@eval @accelerate function _gs_step!(
b::GrayScottAccelerated, u, v, u_new, v_new, args::GSParams
diff --git a/benchmark/src/benchmarks/grayscott_accelerate_forms.jl b/benchmark/src/benchmarks/grayscott_accelerate_forms.jl
new file mode 100644
index 000000000..2fd181b9c
--- /dev/null
+++ b/benchmark/src/benchmarks/grayscott_accelerate_forms.jl
@@ -0,0 +1,96 @@
+# Compare the four scope contracts of `@accelerate` on one shared Gray-Scott step.
+# Each type has a distinct result name so benchmark runs produce separate CSVs.
+
+abstract type AbstractGrayScottAccelerateForm{T} <: AbstractGrayScott{T} end
+
+Base.@kwdef struct GrayScottFunctionAccelerated{T} <:
+ AbstractGrayScottAccelerateForm{T}
+ N::Int
+ M::Int
+end
+
+Base.@kwdef struct GrayScottBeginAccelerated{T} <: AbstractGrayScottAccelerateForm{T}
+ N::Int
+ M::Int
+end
+
+Base.@kwdef struct GrayScottLetAccelerated{T} <: AbstractGrayScottAccelerateForm{T}
+ N::Int
+ M::Int
+end
+
+Base.@kwdef struct GrayScottExpressionAccelerated{T} <:
+ AbstractGrayScottAccelerateForm{T}
+ N::Int
+ M::Int
+end
+
+name(::GrayScottFunctionAccelerated) = "grayscott_function_accelerated"
+name(::GrayScottBeginAccelerated) = "grayscott_begin_accelerated"
+name(::GrayScottLetAccelerated) = "grayscott_let_accelerated"
+name(::GrayScottExpressionAccelerated) = "grayscott_expression_accelerated"
+
+# Function form is the reusable default: arguments and the return value survive,
+# while non-returned locals may fuse across statements or die after their last use.
+let body = deepcopy(GRAYSCOTT_STEP_BODY)
+ @eval @accelerate function _gs_step!(
+ b::GrayScottFunctionAccelerated, u, v, u_new, v_new, args::GSParams
+ )
+ $body
+ end
+end
+
+# `begin` adds no scope. Every named local remains visible, so it measures the
+# multi-output/materialized path rather than eliminating named intermediates.
+let body = deepcopy(GRAYSCOTT_STEP_BODY)
+ @eval function _gs_step!(
+ b::GrayScottBeginAccelerated, u, v, u_new, v_new, args::GSParams
+ )
+ @accelerate begin
+ $body
+ end
+ end
+end
+
+# `let` is a hard one-off scope. Only its result escapes, allowing aggressive
+# inter-statement fusion and last-use cleanup for all other local temporaries.
+let body = deepcopy(GRAYSCOTT_STEP_BODY)
+ @eval function _gs_step!(
+ b::GrayScottLetAccelerated, u, v, u_new, v_new, args::GSParams
+ )
+ @accelerate let
+ $body
+ end
+ end
+end
+
+# Expression form has no multi-statement scope. Accelerating each RHS preserves
+# fusion inside that expression but deliberately materializes statement results,
+# isolating intra-expression fusion from the inter-statement rewrite cases above.
+function accelerate_grayscott_rhs(body::Expr)
+ statements = Any[]
+ for statement in body.args
+ if statement isa LineNumberNode
+ push!(statements, statement)
+ elseif statement isa Expr && statement.head === :(=)
+ lhs, rhs = statement.args
+ push!(statements, :($lhs = @accelerate $rhs))
+ else
+ error("Gray-Scott expression benchmark expected assignments; got $(repr(statement))")
+ end
+ end
+ return Expr(:block, statements...)
+end
+
+let body = accelerate_grayscott_rhs(deepcopy(GRAYSCOTT_STEP_BODY))
+ @eval function _gs_step!(
+ b::GrayScottExpressionAccelerated, u, v, u_new, v_new, args::GSParams
+ )
+ $body
+ end
+end
+
+register_benchmark("grayscott_function_accelerated", GrayScottFunctionAccelerated)
+register_benchmark("grayscott_begin_accelerated", GrayScottBeginAccelerated)
+register_benchmark("grayscott_let_accelerated", GrayScottLetAccelerated)
+register_benchmark("grayscott_expression_accelerated", GrayScottExpressionAccelerated)
diff --git a/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..bb266637d
--- /dev/null
+++ b/docs/src/index.md
@@ -0,0 +1,96 @@
+```@raw html
+
+ 
+ cuNumeric.jl
+
+```
+
+[](https://julialegate.github.io/cuNumeric.jl/dev) [](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [](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 installation can take a while because it includes several large dependencies, such as the CUDA SDK. To use a local cupynumeric build, see [Build Modes](https://julialegate.github.io/cuNumeric.jl/dev/install).
+
+```julia
+using cuNumeric
+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](https://julialegate.github.io/cuNumeric.jl/dev/configuration/hardware).
+
+### How `NDArray`s work
+
+The semantics of `NDArray` closely mirror Julia's `Array`, and in most cases it is a drop-in replacement. You can use the same constructors (i.e., `zeros`, `ones`, `rand`), broadcasting, slicing, and linear algebra. Under the hood a few details differ from Base, and knowing them can help you write fast code.
+
+**Data may live across many devices.** An `NDArray` is a logical array whose physical buffers can be partitioned over GPUs and CPUs by the Legate runtime. You write ordinary array code and Legate decides where the data lives and how/when it is communicated between devices. As a result, elementwise indexing (i.e. `arr[1]`) is slow (and is prevented by default). Scalar indexing like this forces synchronization and blocks other tasks from executing.
+
+**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 communicate with the Julia runtime. When you need a plain Julia number, call `unwrap`:
+
+```julia
+s = sum(A) # NDArray{T,0}
+x = unwrap(s) # T, e.g. Float32
+```
+
+**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](https://julialegate.github.io/cuNumeric.jl/dev/api_initialization) and [NDArray Reference](https://julialegate.github.io/cuNumeric.jl/dev/api). For common performance pitfalls, see [Patterns to Avoid](https://julialegate.github.io/cuNumeric.jl/dev/perf/patterns_to_avoid).
+
+### Kernel Fusion
+
+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](https://julialegate.github.io/cuNumeric.jl/dev/perf/kernel_fusion) and [Debugging](https://julialegate.github.io/cuNumeric.jl/dev/debugging) for controls and diagnostics.
+
+### The `@accelerate` macro
+
+`@accelerate` fuses eligible GPU broadcasts within and across statements, then releases materialized temporary `NDArray`s after their last use on CPU or GPU. See [The `@accelerate` Macro](https://julialegate.github.io/cuNumeric.jl/dev/perf/reduce_allocations) for usage guidance.
+
+### Benchmarks
+
+Results and reproduction instructions live under [Benchmark Results](https://julialegate.github.io/cuNumeric.jl/dev/benchmarks/results) and [How to Benchmark](https://julialegate.github.io/cuNumeric.jl/dev/benchmarks/howto).
+
+### Try an example
+
+```julia
+using cuNumeric
+
+integrand(x) = @. exp(-x^2)
+
+@accelerate function monte_carlo(N, x_max)
+ Ω = 2 * x_max
+ raw_samples = cuNumeric.rand(N)
+ samples = @. Ω * raw_samples - x_max
+ return (Ω / N) * sum(integrand(samples))
+end
+
+N = 1_000_000
+x_max = 10.0f0
+estimate = monte_carlo(N, x_max)
+
+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/perf/reduce_allocations.md b/docs/src/perf/reduce_allocations.md
index 4ff66f988..b8ebec7c5 100644
--- a/docs/src/perf/reduce_allocations.md
+++ b/docs/src/perf/reduce_allocations.md
@@ -1,40 +1,16 @@
# The `@accelerate` Macro
-`@accelerate` optimizes straight-line array code at macro-expansion time:
+`@accelerate` optimizes straight-line array code by coordinating three related jobs:
-- On CPU and GPU, static last-use analysis releases non-returned temporary `NDArray`s as soon as they become dead.
-- With CUDA broadcast fusion enabled, eligible dotted expressions can be combined into fewer GPU kernels.
+1. **Fusion within a broadcast expression.** On CUDA, an eligible dotted expression such as `@. A + B * C` can run as one kernel. CPU execution uses the normal unfused path.
+2. **Fusion across broadcast statements.** A single-use broadcast result can be substituted into its consumer, producing fewer GPU kernel launches.
+3. **Temporary lifetime analysis.** After rewriting the code, the macro releases materialized, non-returned `NDArray`s after their final use on CPU or GPU.
-Arguments are never freed, mutations remain visible, and returned values remain materialized. The macro has four forms:
+These jobs must happen together: an intermediate that fuses into its consumer is never allocated, while an intermediate that cannot fuse is materialized and then released after its last use.
-```julia
-@accelerate function f(args...)
- # reusable straight-line kernel
-end
-
-@accelerate begin
- # soft scope
-end
-
-@accelerate let
- # hard scope
-end
+## Use the function form by default
-result = @accelerate expr
-```
-
-## Choose a form
-
-| Form | Use it when | What remains available afterward |
-|---|---|---|
-| `@accelerate function ... end` | Defining a reusable update or compute kernel. This is the recommended default. | Caller-owned arguments, their mutations, and returned values. |
-| `@accelerate begin ... end` | Named results must remain in the surrounding scope. | Every named binding created by the block. |
-| `@accelerate let ... end` | Writing a one-off multi-statement calculation whose intermediates should not escape. | Only the block result. |
-| `@accelerate expr` | Accelerating one expression without introducing a new scope. | The materialized expression result. |
-
-## 1. Function form
-
-Use this form by default. Function arguments belong to the caller and are never freed. Non-returned locals may be fused into consumers or released after their final use.
+Annotate a reusable straight-line function:
```julia
@accelerate function update!(C, A, B)
@@ -44,9 +20,20 @@ Use this form by default. Function arguments belong to the caller and are never
end
```
-## 2. `begin` form
+On an eligible GPU path, `combined` can be folded into the second broadcast so the chain runs as one kernel. On CPU, or when fusion is ineligible, `combined` is materialized and released after the update. Function arguments belong to the caller and are never released by `@accelerate`; returned values also remain valid.
-`begin` preserves normal Julia scope. Named bindings remain available after the block, so the macro cannot discard them. Eligible same-shape CUDA chains may still run as one multi-output kernel.
+## Choose a form
+
+The forms differ in which values must remain available, which determines how aggressively the macro may fuse or release intermediates.
+
+| Form | Use it when | Fusion and lifetime behavior |
+|---|---|---|
+| `@accelerate function ... end` | Defining reusable array code. This is the recommended default. | Arguments and returned values are protected. Non-returned locals may fuse into consumers or be released after their last use. |
+| `@accelerate begin ... end` | Named results must remain in the current scope. | `begin` creates no new Julia scope, so every named binding is protected. An eligible same-shape CUDA chain may still use one multi-output kernel, but each named result is materialized. |
+| `@accelerate let ... end` | Writing a one-off multi-statement calculation when only its result is needed. | `let` creates a local scope. Only the result escapes; other locals may fuse away or be released after their last use. |
+| `@accelerate expr` | Evaluating one expression without named intermediates. | The result is materialized and returned. Eligible operations fuse within the expression, and transient temporaries are released. |
+
+For example, choose `begin` when both names are needed afterward:
```julia
@accelerate begin
@@ -54,47 +41,31 @@ end
shifted = @. product + 1
end
-# Both names are still defined here.
consume(product, shifted)
```
-## 3. `let` form
-
-`let` creates a hard scope. Only its result escapes, giving `@accelerate` freedom to fuse or release every non-returned intermediate.
+Choose `let` when only the final result should escape:
```julia
result = @accelerate let
product = @. A * B
@. product + 1
end
-
-# `product` is not defined here.
```
-## 4. Expression form
-
-Use the expression form when there are no named intermediates:
+For a single unnamed expression, use:
```julia
result = @accelerate (@. A + B * C)
```
-The result is materialized before it is returned.
-
-## Writing the body
-
-Apply `@.` to each elementwise right-hand side. Do not wrap the entire `@accelerate` body in `@.`: block-wide dotting changes `x = ...` into `x .= ...` and changes an ordinary call such as `bc!(...)` into `bc!.(...)`.
-
-Prefer ordinary assignment for a single-use intermediate:
-
-```julia
-temporary = @. A * B
-C .= @. temporary + 1
-```
-
-An explicit in-place `.=` mutation is observable and therefore forms a fusion boundary. Ordinary function calls run in program order and also form boundaries; `@accelerate` does not rewrite inside a called function unless that function is separately annotated.
+## Writing an accelerated body
-The body must be straight-line. Control flow (`if`, loops, `try`, short-circuit operators) and nested functions are rejected. Put control flow outside the accelerated region:
+- Apply `@.` to each elementwise right-hand side. Applying it to the entire body would change `x = ...` into `x .= ...` and `f(...)` into `f.(...)`.
+- Use ordinary `=` for a disposable intermediate. This allows a single-use producer to fuse into its consumer.
+- Use `.=` when the mutation must be visible. The destination write is preserved, although an eligible producer may fuse into it.
+- Keep control flow outside the accelerated body. Loops, conditionals, `try`, short-circuit operators, and nested functions are rejected.
+- Ordinary function calls run in program order and form rewrite boundaries. Annotate the called function separately if its body should also be accelerated.
```julia
for _ in 1:nsteps
@@ -102,4 +73,4 @@ for _ in 1:nsteps
end
```
-See [Kernel Fusion](./kernel_fusion.md) for fusion requirements and [`@show_lifetimes`](../debugging.md#inspect-lifetime-rewrites-with-show_lifetimes) to inspect the exact rewrite without executing it.
+See [Kernel Fusion](./kernel_fusion.md) for CUDA fusion requirements and [`@show_lifetimes`](../debugging.md#inspect-lifetime-rewrites-with-show_lifetimes) to inspect the exact rewrite without executing it.
diff --git a/examples/gray-scott.py b/examples/gray-scott.py
index dce4e6cab..5eab0b8c8 100644
--- a/examples/gray-scott.py
+++ b/examples/gray-scott.py
@@ -1,82 +1,54 @@
-# python equivalent of gray-scott.jl to test the GC problem
+"""cuPyNumeric equivalent of examples/gray-scott.jl."""
import cupynumeric as np
-# import matplotlib.animation as animation
-# from IPython.display import HTML
-# import matplotlib.pyplot as plt
-
-def greyScottSys(u, v, dx, dt, c_u, c_v, f, k):
- # u,v are arrays
- # dx,dt are space and time steps
- # c_u, c_v, f, k are constant paramaters
-
- #create new u array
+def step(u, v, dx, dt, c_u, c_v, feed, kill):
u_new = np.zeros_like(u)
v_new = np.zeros_like(v)
- #calculate F_u and F_v functions
- F_u = (-u[1:-1,1:-1]*(v[1:-1,1:-1]**2)) + f*(1-u[1:-1,1:-1])
- F_v = (u[1:-1,1:-1]*(v[1:-1,1:-1]**2)) - (f+k)*v[1:-1,1:-1]
-
- # 2-D Laplacian of f using array slicing, excluding boundaries
- # For an N x N array f, f_lap is the N-1 x N-1 array in the "middle"
- u_lap = (u[2:,1:-1] - 2*u[1:-1,1:-1] + u[:-2,1:-1]) / dx**2\
- + (u[1:-1,2:] - 2*u[1:-1,1:-1] + u[1:-1,:-2]) / dx**2
- v_lap = (v[2:,1:-1] - 2*v[1:-1,1:-1] + v[:-2,1:-1]) / dx**2\
- + (v[1:-1,2:] - 2*v[1:-1,1:-1] + v[1:-1,:-2]) / dx**2
-
- # Forward-Euler time step for all points except the boundaries
- u_new[1:-1,1:-1] = ((c_u * u_lap) + F_u)*dt + u[1:-1,1:-1]
- v_new[1:-1,1:-1] = ((c_v * v_lap) + F_v)*dt + v[1:-1,1:-1]
-
- # Apply periodic boundary conditions
- u_new[:,0] = u[:,-2]
- u_new[:,-1] = u[:,1]
- u_new[0,:] = u[-2,:]
- u_new[-1,:] = u[1,:]
- v_new[:,0] = v[:,-2]
- v_new[:,-1] = v[:,1]
- v_new[0,:] = v[-2,:]
- v_new[-1,:] = v[1,:]
-
+ u_mid = u[1:-1, 1:-1]
+ v_mid = v[1:-1, 1:-1]
+ reaction = u_mid * v_mid**2
+ f_u = -reaction + feed * (1 - u_mid)
+ f_v = reaction - (feed + kill) * v_mid
+
+ u_lap = (
+ u[2:, 1:-1] - 2 * u_mid + u[:-2, 1:-1]
+ + u[1:-1, 2:] - 2 * u_mid + u[1:-1, :-2]
+ ) / dx**2
+ v_lap = (
+ v[2:, 1:-1] - 2 * v_mid + v[:-2, 1:-1]
+ + v[1:-1, 2:] - 2 * v_mid + v[1:-1, :-2]
+ ) / dx**2
+
+ u_new[1:-1, 1:-1] = (c_u * u_lap + f_u) * dt + u_mid
+ v_new[1:-1, 1:-1] = (c_v * v_lap + f_v) * dt + v_mid
+
+ u_new[:, 0] = u[:, -2]
+ u_new[:, -1] = u[:, 1]
+ u_new[0, :] = u[-2, :]
+ u_new[-1, :] = u[1, :]
+ v_new[:, 0] = v[:, -2]
+ v_new[:, -1] = v[:, 1]
+ v_new[0, :] = v[-2, :]
+ v_new[-1, :] = v[1, :]
return u_new, v_new
+def gray_scott(n=4000, n_steps=100):
+ dx = 1.0
+ dt = dx / 5
+ u = np.ones((n, n))
+ v = np.zeros((n, n))
+ seed = min(150, n)
+ u[:seed, :seed] = np.random.rand(seed, seed)
+ v[:seed, :seed] = np.random.rand(seed, seed)
-# initial conditions and discretizaiton
-dx = 1
-dt = dx/5
-u = np.ones((4000,4000))
-v = np.zeros((4000,4000))
-u[:150,:150] = np.random.rand(150,150)
-v[:150,:150] = np.random.rand(150,150)
-
-
-# fig = plt.figure()
-
-c_u = 1
-c_v = 0.3
-f = 0.03
-k = 0.06
-
-# t_final = 1000
-
-# ims = []
-n_steps = 100 # number of steps to take
-frame_interval = 200 # steps to take between making plots
-
-# build a list of images
-for n in range(n_steps) :
-
- ## This may need to be changed.
- u,v = greyScottSys(u, v, dx, dt, c_u, c_v, f, k)
+ for _ in range(n_steps):
+ u, v = step(u, v, dx, dt, 1.0, 0.3, 0.03, 0.06)
+ return u, v
- # ## Store frames when n is a multiple of frame_interval
- # if n%frame_interval == 0:
- # im = plt.imshow(u, vmin=0, vmax=1) # Show a plot of u.
- # ims.append([im]) # append single image to the list of images
-# anim = animation.ArtistAnimation(fig, ims, interval=100, repeat=False)
-# HTML(anim.to_jshtml())
+if __name__ == "__main__":
+ gray_scott()
diff --git a/src/cuda/cuda_ptx_task.jl b/src/cuda/cuda_ptx_task.jl
index 87962b7fa..44213a47a 100644
--- a/src/cuda/cuda_ptx_task.jl
+++ b/src/cuda/cuda_ptx_task.jl
@@ -2,7 +2,11 @@ export @cuda_task, @launch, CUDATask
struct CUDATask
func::String
- argtypes::NTuple{N,Type} where {N} #! THIS IS TYPE UNSTABLE
+ argtypes::Vector{DataType}
+
+ function CUDATask(func, argtypes)
+ return new(convert(String, func), collect(DataType, argtypes))
+ end
end
#! JUST PASS TYPES HERE INSTEAD OF CALLING typeof()
@@ -16,57 +20,59 @@ function to_stdvec(::Type{T}, vec) where {T}
return stdvec
end
-function add_padding(arr::NDArray, dims::Dims{N}; copy=false) where {N}
- old_size = size(arr)
-
- @assert all(dims .>= old_size) "newdims must be ≥ current dims elementwise"
- new = zeros(eltype(arr), dims)
+@inline _launch_shape(arr::NDArray) = _launch_shape(arr, _padding(arr))
+@inline _launch_shape(arr::NDArray, ::Nothing) = size(arr)
+@inline _launch_shape(::NDArray, padding::PaddedStorage) = padding.shape
+
+@inline _physical_array(arr::NDArray, ::Nothing) = arr
+@inline _physical_array(::NDArray, padding::PaddedStorage) = padding.backing
+
+function _ensure_launch_padding!(arr::NDArray{T,N}, target_shape; copy=false) where {T,N}
+ padding = _padding(arr)
+ !isnothing(padding) && padding.shape == target_shape && return arr
+ isnothing(padding) && size(arr) == target_shape && return arr
+ @assert all(target_shape .>= size(arr)) "cannot pad $(size(arr)) to $target_shape"
+
+ padded = zeros(T, target_shape)
+ slices = ntuple(d -> (0, size(arr, d)), N)
+ logical = nda_get_slice(padded, slice_array(slices...))
+ aliases_parent = !isnothing(arr.parent)
+ copy && !aliases_parent && copyto!(logical, arr)
+ storage = PaddedStorage{T,N}(
+ padded,
+ aliases_parent ? logical : nothing,
+ target_shape,
+ )
- if copy # due to being an input. we don't need to copy outputs
- slices = ntuple(d -> (0, Int(old_size[d])), length(old_size))
- s = nda_get_slice(new, slice_array(slices...))
- copyto!(s, arr)
- destroy!(s)
+ if aliases_parent
+ old_padding = _padding(arr)
+ arr.padding = storage
+ !isnothing(old_padding) && _destroy_padded_storage!(old_padding)
+ else
+ destroy!(arr)
+ arr.ptr = logical.ptr
+ arr.nbytes = logical.nbytes
+ arr.padding = storage
+ logical.ptr = Ptr{Cvoid}(0)
+ logical.nbytes = 0
end
-
- nda_destroy_array(arr.ptr)
- register_free!(arr.nbytes)
-
- # update pointer & update metadata
- arr.ptr = new.ptr
- arr.nbytes = new.nbytes
- arr.padding = old_size # remember the prior (before the padding)
-
- # julia GC will call finalizer, but we manually cleaned it
- new.ptr = Ptr{Cvoid}(0)
- new.nbytes = 0
- return new.padding = nothing
-end
-
-function add_padding(arr::NDArray, i::Int64; copy=false)
- return add_padding(arr, (i,); copy=copy)
+ return arr
end
-function check_sz!(arr, maxshape; copy=false)
- sz = cuNumeric.size(arr)
- if maxshape != nothing
- # currently require all ndarray inputs to be equal
- alligned_equal_size = sz == maxshape
- if !alligned_equal_size
- cuNumeric.add_padding(arr, maxshape; copy=copy)
- new_size = padded_shape(arr)
- @warn "[Padding Added] $sz output is now $new_size"
- end
+function _sync_to_launch_padding!(arr::NDArray)
+ padding = _padding(arr)
+ if !isnothing(padding) && !isnothing(padding.staging)
+ copyto!(padding.staging, arr)
end
+ return nothing
end
-function check_sz(arr, maxshape)
- sz = cuNumeric.size(arr)
- if maxshape != nothing
- # currently require all ndarray inputs to be equal
- alligned_equal_size = sz == maxshape
- @assert alligned_equal_size
+function _sync_from_launch_padding!(arr::NDArray)
+ padding = _padding(arr)
+ if !isnothing(padding) && !isnothing(padding.staging)
+ copyto!(arr, padding.staging)
end
+ return nothing
end
# `get_store` returns a Julia-owned `LogicalArrayImplAllocated` that shares the
@@ -74,42 +80,32 @@ end
# array into the task; if we leave the temporary alive until GC, store refcounts
# stay elevated and framebuffer reclaim stalls (fusion 1-GPU OOM under load).
# Finalize the temporary immediately after the copy into the task.
-function _add_task_array!(add_to, task, arr::NDArray)
- st = cuNumeric.get_store(arr)
- var = add_to(task, st)
- finalize(st)
- return var
+function _add_task_array!(add_to, task, arr::NDArray; physical=false)
+ task_arr = physical ? _physical_array(arr, _padding(arr)) : arr
+ st = cuNumeric.get_store(task_arr)
+ try
+ return add_to(task, st)
+ finally
+ finalize(st)
+ end
end
function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}},
outputs::Tuple{Vararg{NDArray}}, scalars::Tuple{Vararg{Any}};
- blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing, validate_shapes=true)
- max_shape = if validate_shapes
- # Generic PTX tasks retain the existing padding/shape behavior.
- ndarrays = vcat(inputs..., outputs...) # returns (nbytes, position)
- mx = findmax(arr -> arr.nbytes, ndarrays) # first elem nbytes
- shape = size(ndarrays[mx[2]]) # second elem max position
- @assert !isnothing(shape)
- shape
- else
- # Fused linear broadcast verifies shapes match
- nothing
- end
-
+ blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing)
rt = Legate.get_runtime()
lib = cuNumeric.get_lib()
task = Legate.create_auto_task(rt, lib, taskid)
+ physical = taskid == cuNumeric.RUN_PTX
input_vars = Vector{Legate.Variable}()
for arr in inputs
- validate_shapes && check_sz!(arr, max_shape; copy=true)
- push!(input_vars, _add_task_array!(Legate.add_input, task, arr))
+ push!(input_vars, _add_task_array!(Legate.add_input, task, arr; physical))
end
output_vars = Vector{Legate.Variable}()
for arr in outputs
- validate_shapes && check_sz!(arr, max_shape; copy=false)
- push!(output_vars, _add_task_array!(Legate.add_output, task, arr))
+ push!(output_vars, _add_task_array!(Legate.add_output, task, arr; physical))
end
# Reserved scalars: kernel_name (0), blocks (1,2,3), threads (4,5,6)
@@ -136,15 +132,34 @@ function Launch(kernel::CUDATask, inputs::Tuple{Vararg{NDArray}},
end
function launch(kernel::CUDATask, inputs, outputs, scalars;
- blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing, validate_shapes=true)
- return Launch(kernel,
- isa(inputs, Tuple) ? inputs : (inputs,),
- isa(outputs, Tuple) ? outputs : (outputs,),
+ blocks, threads, taskid=cuNumeric.RUN_PTX, ctx=nothing)
+ input_tuple = isa(inputs, Tuple) ? inputs : (inputs,)
+ output_tuple = isa(outputs, Tuple) ? outputs : (outputs,)
+
+ # Custom tasks require equal physical shapes. Keep the padded backing so
+ # repeated launches do not allocate or copy again.
+ if taskid == cuNumeric.RUN_PTX
+ arrays = (input_tuple..., output_tuple...)
+ if !isempty(arrays)
+ rank = ndims(first(arrays))
+ @assert all(ndims(arr) == rank for arr in arrays) "custom task arrays must have equal ranks"
+ max_shape = ntuple(d -> maximum(_launch_shape(arr)[d] for arr in arrays), rank)
+ foreach(arr -> _ensure_launch_padding!(arr, max_shape; copy=true), input_tuple)
+ foreach(arr -> _ensure_launch_padding!(arr, max_shape), output_tuple)
+ foreach(_sync_to_launch_padding!, input_tuple)
+ end
+ end
+
+ result = Launch(kernel,
+ input_tuple,
+ output_tuple,
isa(scalars, Tuple) ? scalars : (scalars,);
blocks=isa(blocks, Tuple) ? blocks : (blocks,),
threads=isa(threads, Tuple) ? threads : (threads,),
- taskid=taskid, ctx=ctx, validate_shapes=validate_shapes,
+ taskid=taskid, ctx=ctx,
)
+ taskid == cuNumeric.RUN_PTX && foreach(_sync_from_launch_padding!, output_tuple)
+ return result
end
function ptx_task(ptx::String, kernel_name)
diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl
index ff80a9dd0..64149a9f7 100644
--- a/src/ndarray/broadcast_fusion.jl
+++ b/src/ndarray/broadcast_fusion.jl
@@ -752,7 +752,6 @@ function fuse_broadcast_tree!(dest::D, bc::B) where {D<:NDArray,B<:Base.Broadcas
threads=fkm.threads,
taskid=cuNumeric.RUN_PTX_BROADCAST,
ctx=fkm.ctx,
- validate_shapes=false,
)
end
@@ -1045,7 +1044,6 @@ function _fused_multi_launch!(out_arrs::Tuple, seg_bcs::Tuple)
task, tuple(input_ndarrays...), out_arrs,
(Int32(length(argmap)), argmap..., actual_scalars...);
blocks=1, threads=threads, taskid=cuNumeric.RUN_PTX_BROADCAST, ctx=ctx,
- validate_shapes=false,
)
return out_arrs
end
diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl
index 8b7e67671..a8b6f8118 100644
--- a/src/ndarray/detail/ndarray.jl
+++ b/src/ndarray/detail/ndarray.jl
@@ -43,21 +43,24 @@ get_n_dim(ptr::NDArray_t) = Int(ccall((:nda_array_dim, libnda), Int32, (NDArray_
abstract type AbstractNDArray{T<:SUPPORTED_TYPES,N} <: AbstractArray{T,N} end
+# Runtime padding uses an abstract field to break the recursive storage definition.
+abstract type AbstractPaddedStorage{T,N} end
+
@doc"""
The NDArray type represents a multi-dimensional array in cuNumeric.
It is a wrapper around a Legate array and provides various methods for array manipulation and operations.
Finalizer calls `nda_destroy_array` to clean up the underlying Legate array when the NDArray is garbage collected.
"""
-mutable struct NDArray{T,N,PADDED,P} <: AbstractNDArray{T,N}
+mutable struct NDArray{T,N,P} <: AbstractNDArray{T,N}
ptr::NDArray_t
nbytes::Int64
- padding::Union{Nothing,NTuple{N,Int}}
+ padding::Union{Nothing,AbstractPaddedStorage{T,N}}
parent::P
function NDArray(ptr::NDArray_t, ::Type{T}, ::Val{N}) where {T,N}
nbytes = cuNumeric.nda_nbytes(ptr)
cuNumeric.register_alloc!(nbytes)
- handle = new{T,N,false,Nothing}(ptr, nbytes, nothing, nothing)
+ handle = new{T,N,Nothing}(ptr, nbytes, nothing, nothing)
finalizer(_finalize_ndarray!, handle)
return handle
end
@@ -66,26 +69,53 @@ mutable struct NDArray{T,N,PADDED,P} <: AbstractNDArray{T,N}
function NDArray(ptr::NDArray_t, ::Type{T}, ::Val{N}, parent::P) where {T,N,P}
nbytes = cuNumeric.nda_nbytes(ptr)
cuNumeric.register_alloc!(nbytes)
- handle = new{T,N,false,P}(ptr, nbytes, nothing, parent)
+ handle = new{T,N,P}(ptr, nbytes, nothing, parent)
finalizer(_finalize_ndarray!, handle)
return handle
end
end
+struct PaddedStorage{T,N} <: AbstractPaddedStorage{T,N}
+ backing::NDArray{T,N,Nothing}
+ staging::Union{Nothing,NDArray{T,N,NDArray{T,N,Nothing}}}
+ shape::NTuple{N,Int}
+end
+
+# Narrow the abstract field to its concrete storage type.
+@inline _padding(arr::NDArray{T,N}) where {T,N} =
+ arr.padding::Union{Nothing,PaddedStorage{T,N}}
+
+function _finalize_padded_storage!(storage::PaddedStorage)
+ !isnothing(storage.staging) && finalize(storage.staging)
+ finalize(storage.backing)
+ return nothing
+end
+
+function _destroy_padded_storage!(storage::PaddedStorage)
+ !isnothing(storage.staging) && destroy!(storage.staging)
+ destroy!(storage.backing)
+ return nothing
+end
+
# May run off the launch thread, so defer the Legate free to drain_pending_frees!.
# Accounting is atomic and safe to do here immediately.
function _finalize_ndarray!(arr::NDArray)
ptr = arr.ptr
- ptr == C_NULL && return nothing
arr.ptr = Ptr{Cvoid}(0)
nbytes = arr.nbytes
arr.nbytes = 0
- nbytes > 0 && register_free!(nbytes)
- _enqueue_free!(ptr)
+ padding = _padding(arr)
+ arr.padding = nothing
+
+ if ptr != C_NULL
+ nbytes > 0 && register_free!(nbytes)
+ _enqueue_free!(ptr)
+ end
+ !isnothing(padding) && _finalize_padded_storage!(padding)
return nothing
end
-@inline _is_ndarray_slice(arr::NDArray) = arr.parent isa NDArray
+@inline _is_ndarray_slice(arr::NDArray) = arr.parent isa NDArray || !isnothing(_padding(arr))
"""
destroy!(arr::NDArray)
@@ -102,6 +132,9 @@ function destroy!(arr::NDArray)
arr.nbytes = 0
nbytes > 0 && register_free!(nbytes)
end
+ padding = _padding(arr)
+ arr.padding = nothing
+ !isnothing(padding) && _destroy_padded_storage!(padding)
return arr
end
@@ -591,9 +624,7 @@ end
Return the size of the given `NDArray`.
"""
-shape(arr::NDArray{<:Any,N,true}) where {N} = arr.padding
-
-function shape(arr::NDArray{<:Any,N,false}) where {N}
+function shape(arr::NDArray{<:Any,N}) where {N}
shp = cuNumeric.nda_array_shape(arr)
return ntuple(i -> Int(shp[i]), Val(N))
end
diff --git a/src/ndarray/promotion.jl b/src/ndarray/promotion.jl
index e905c1d5e..cb1034098 100644
--- a/src/ndarray/promotion.jl
+++ b/src/ndarray/promotion.jl
@@ -35,6 +35,14 @@ unchecked_promote_arr(::Base.RefValue{Val{V}}, ::Type{T}) where {T,V} = Val{V}
__checked_promote_op(op, ::Type{Tuple{A}}) where {A} = __checked_promote_op(op, A)
__checked_promote_op(op, ::Type{Tuple{A,B}}) where {A,B} = __checked_promote_op(op, A, B)
+# Julia flattens dotted `+` and `*` chains into n-ary Broadcasted nodes. Fold
+# their input types pairwise, matching both the binary C API and fused path.
+@inline function __checked_promote_op(
+ op::Union{typeof(+),typeof(*)}, ::Type{Args}
+) where {Args<:Tuple{Any,Any,Any,Vararg{Any}}}
+ return _checked_promote_associative(op, Args.parameters...)
+end
+
# Path for literal powers
@inline function __checked_promote_op(
f::typeof(Base.literal_pow), a::Type{Tuple{_,ARR_TYPE,Val{POWER}}}
diff --git a/src/scoping/accelerate.jl b/src/scoping/accelerate.jl
index bcbbd3806..3b0268c69 100644
--- a/src/scoping/accelerate.jl
+++ b/src/scoping/accelerate.jl
@@ -212,22 +212,22 @@ end
@accelerate let ... end
@accelerate expr
-Fuse straight-line array code into fewer kernel launches and free temporaries.
-The body must be straight-line — control flow and nested/anonymous functions are
-rejected. Four forms, by scope:
+Optimize straight-line array code by coordinating CUDA broadcast fusion within
+expressions, fusion across broadcast statements, and scope-aware cleanup of
+materialized temporaries. Control flow and nested/anonymous functions are
+rejected. Four forms determine which values must remain valid:
- * **function** (preferred): args are caller-owned; only the returned value stays
- materialized, so non-returned intermediates fuse away and are freed.
- * **`begin`**: 1:1 Julia scope — every named binding stays live; on GPU,
- same-shape chains may fuse into one multi-output launch; anonymous temps
- (slices) are freed.
- * **`let`**: hard scope — combines single-use producers and frees every
- non-returned temporary; only the returned value(s) escape. Maximum reuse.
- * **expression**: materializes and returns one expression without introducing
- a new scope.
+ * **function** (preferred): arguments and returned values are protected;
+ non-returned locals may fuse into consumers or be freed after their last use.
+ * **`begin`**: creates no new Julia scope, so every named binding stays live;
+ eligible GPU chains may use one multi-output kernel that materializes them.
+ * **`let`**: creates a local scope; only the result escapes, so other locals may
+ fuse away or be freed after their last use.
+ * **expression**: materializes and returns one expression; eligible operations
+ fuse within it and transient temporaries are released.
```julia
-@accelerate function step(u, v) # c freed; w returned
+@accelerate function step(u, v) # c may fuse away; the result is returned
c = u .* v
return c .^ 2
end
diff --git a/test/Project.toml b/test/Project.toml
index bd7ab116b..6e9f7cd1a 100644
--- a/test/Project.toml
+++ b/test/Project.toml
@@ -1,5 +1,6 @@
[deps]
CNPreferences = "3e078157-ea10-49d5-bf32-908f777cd46f"
+CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
CUDACore = "bd0ed864-bdfe-4181-a5ed-ce625a5fdea2"
InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
diff --git a/test/analysis/promotion.jl b/test/analysis/promotion.jl
index 7b1fb39a7..422e5766a 100644
--- a/test/analysis/promotion.jl
+++ b/test/analysis/promotion.jl
@@ -19,3 +19,8 @@
@test safe_compare(r1, r2, atol(Float64), rtol(Float64))
end
end
+
+@testset "Flattened associative broadcast promotion" begin
+ @test @inferred(cuNumeric.__checked_promote_op(+, NTuple{5,Float64})) === Float64
+ @test @inferred(cuNumeric.__checked_promote_op(*, NTuple{4,Int32})) === Int32
+end
diff --git a/test/analysis/type_stability.jl b/test/analysis/type_stability.jl
index 208cee677..ce81c1cb0 100644
--- a/test/analysis/type_stability.jl
+++ b/test/analysis/type_stability.jl
@@ -41,6 +41,8 @@ function _type_stable_accelerate_expr(a, b)
return @accelerate (@. (a + b) * 2.0f0)
end
+_type_stable_cuda_argtypes(task::cuNumeric.CUDATask) = task.argtypes
+
@testset verbose = true "core" begin
a = cuNumeric.zeros(5)
b = cuNumeric.zeros(Float64, 3, 4)
@@ -90,6 +92,18 @@ end
@test @inferred(cuNumeric.NDArray(rand(Float32, 3, 3))) !== nothing
end
+@testset verbose = true "custom CUDA metadata" begin
+ task = cuNumeric.CUDATask("kernel", (Float32, Int32))
+ @test isconcretetype(typeof(task))
+ @test all(isconcretetype, fieldtypes(typeof(task)))
+ @test @inferred(_type_stable_cuda_argtypes(task)) == DataType[Float32, Int32]
+
+ storage_type = cuNumeric.PaddedStorage{Float32,1}
+ @test all(isconcretetype, Base.uniontypes(fieldtype(storage_type, :backing)))
+ @test all(isconcretetype, Base.uniontypes(fieldtype(storage_type, :staging)))
+ @test all(isconcretetype, Base.uniontypes(fieldtype(storage_type, :shape)))
+end
+
@testset verbose = true "conversion" begin
# cast to array, as_type
a = cuNumeric.zeros(Float64, 5, 5)
diff --git a/test/cuda.jl/fusion_compare.jl b/test/cuda.jl/fusion_compare.jl
index dd4934fca..be0d5090e 100644
--- a/test/cuda.jl/fusion_compare.jl
+++ b/test/cuda.jl/fusion_compare.jl
@@ -1,3 +1,9 @@
+using CUDA: CUDA, @cuda
+using CUDACore: blockDim, blockIdx, threadIdx
+import CUDACore: i32
+
+cuNumeric.Experimental(true)
+
function unfused_cunumeric(u, v, f, k)
F_u = (
(
@@ -101,8 +107,8 @@ function run_unfused_baseline(N, u, v)
end
function fusion_test(; N=1024, atol=1.0f-6, rtol=1.0f-6)
- u = cuNumeric.as_type(cuNumeric.random(Float32, (N, N)), Float32)
- v = cuNumeric.as_type(cuNumeric.random(Float32, (N, N)), Float32)
+ u = cuNumeric.rand(Float32, (N, N))
+ v = cuNumeric.rand(Float32, (N, N))
# using CUDA
u_base = CUDA.rand(Float32, (N, N))
@@ -117,8 +123,14 @@ function fusion_test(; N=1024, atol=1.0f-6, rtol=1.0f-6)
Fu_fused, Fv_fused = run_fused_cunumeric(N, u, v)
Fu_unfused, Fv_unfused = run_unfused_cunumeric(N, u, v)
- @test isapprox(Fu_fused, Fu_unfused; atol=atol, rtol=rtol)
- @test isapprox(Fv_fused, Fv_unfused; atol=atol, rtol=rtol)
+ @test isapprox(Array(Fu_fused), Array(Fu_unfused); atol=atol, rtol=rtol)
+ @test isapprox(Array(Fv_fused), Array(Fv_unfused); atol=atol, rtol=rtol)
end
-fusion_test()
+try
+ @testset "2D fusion comparison" begin
+ fusion_test()
+ end
+finally
+ cuNumeric.Experimental(false)
+end
diff --git a/test/cuda.jl/fusion_compare_1d.jl b/test/cuda.jl/fusion_compare_1d.jl
index 8163a5db1..8efab2989 100644
--- a/test/cuda.jl/fusion_compare_1d.jl
+++ b/test/cuda.jl/fusion_compare_1d.jl
@@ -1,4 +1,10 @@
+using CUDA: CUDA, @cuda
+using CUDACore: blockDim, blockIdx, threadIdx
+import CUDACore: i32
+
+cuNumeric.Experimental(true)
+
function unfused_cunumeric(u, v, f, k)
F_u = (
(
@@ -102,8 +108,8 @@ function run_unfused_baseline(N, u, v)
end
function fusion_test(; N=1024*1024, atol=1.0f-6, rtol=1.0f-6)
- u = cuNumeric.as_type(cuNumeric.rand(NDArray, N), Float32)
- v = cuNumeric.as_type(cuNumeric.rand(NDArray, N), Float32)
+ u = cuNumeric.rand(Float32, N)
+ v = cuNumeric.rand(Float32, N)
# using CUDA
u_base = CUDA.rand(Float32, N)
@@ -116,8 +122,14 @@ function fusion_test(; N=1024*1024, atol=1.0f-6, rtol=1.0f-6)
# using cuNumeric
Fu_fused, Fv_fused = run_fused_cunumeric(N, u, v)
Fu_unfused, Fv_unfused = run_unfused_cunumeric(N, u, v)
- @test isapprox(Fu_fused, Fu_unfused; atol=atol, rtol=rtol)
- @test isapprox(Fv_fused, Fv_unfused; atol=atol, rtol=rtol)
+ @test isapprox(Array(Fu_fused), Array(Fu_unfused); atol=atol, rtol=rtol)
+ @test isapprox(Array(Fv_fused), Array(Fv_unfused); atol=atol, rtol=rtol)
end
-fusion_test()
+try
+ @testset "1D fusion comparison" begin
+ fusion_test()
+ end
+finally
+ cuNumeric.Experimental(false)
+end
diff --git a/test/cuda.jl/padding.jl b/test/cuda.jl/padding.jl
new file mode 100644
index 000000000..f38e875a8
--- /dev/null
+++ b/test/cuda.jl/padding.jl
@@ -0,0 +1,224 @@
+#= Copyright 2025 Northwestern University,
+ * Carnegie Mellon University University
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Author(s): David Krasowska
+ * Ethan Meitz
+=#
+
+#= Purpose of test: cuda
+ -- Validate custom-kernel padding, synchronization, and lifetime management
+=#
+
+using CUDACore: blockDim, blockIdx, threadIdx
+import CUDACore: i32
+
+cuNumeric.Experimental(true)
+
+function padding_add(a, b, c, N)
+ i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
+ if i <= N
+ @inbounds c[i] = a[i] + b[i]
+ end
+ return nothing
+end
+
+function padding_mul(a, c, b, N)
+ i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
+ if i <= N
+ @inbounds b[i] = a[i] * c[i]
+ end
+ return nothing
+end
+
+function cuda_padding_lifetime()
+ N = 1_000_000
+ M = N - 2
+ threads = 256
+ blocks = cld(M, threads)
+ initial_bytes = cuNumeric.current_device_bytes[]
+
+ a = cuNumeric.ones(Float32, N)
+ b = cuNumeric.ones(Float32, N)
+ c = cuNumeric.zeros(Float32, M)
+ task = cuNumeric.@cuda_task padding_add(a, b, c, UInt32(M))
+ unpadded_bytes = cuNumeric.current_device_bytes[]
+
+ try
+ @test @inferred(
+ cuNumeric.launch(
+ task, (a, b), c, UInt32(M); threads=threads, blocks=blocks
+ )
+ ) === nothing
+ padded_bytes = cuNumeric.current_device_bytes[]
+
+ @test @inferred(cuNumeric._launch_shape(c)) == (N,)
+ @test @inferred(cuNumeric._sync_to_launch_padding!(c)) === nothing
+ @test @inferred(cuNumeric._sync_from_launch_padding!(c)) === nothing
+
+ accounting_ok = true
+ for _ in 1:15
+ cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, b) outputs=c scalars=UInt32(
+ M
+ )
+ accounting_ok &= cuNumeric.current_device_bytes[] == padded_bytes
+ end
+
+ @test padded_bytes > unpadded_bytes
+ @test accounting_ok
+ @test size(c) == (M,)
+ @test all(Array(c) .== 2.0f0)
+ finally
+ cuNumeric.destroy!(a)
+ cuNumeric.destroy!(b)
+ cuNumeric.destroy!(c)
+ end
+ @test cuNumeric.current_device_bytes[] == initial_bytes
+end
+
+function cuda_padding_api_interop()
+ N = 4096
+ M = N - 2
+ threads = 256
+ blocks = cld(M, threads)
+ initial_bytes = cuNumeric.current_device_bytes[]
+
+ a = cuNumeric.ones(Float32, N)
+ b = cuNumeric.ones(Float32, N)
+ c = cuNumeric.zeros(Float32, M)
+ library_result = nothing
+
+ try
+ task = cuNumeric.@cuda_task padding_add(a, b, c, UInt32(M))
+ cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, b) outputs=c scalars=UInt32(
+ M
+ )
+
+ # The broadcast writes through c's logical view into its padded backing.
+ c .= c .* 2.0f0 .+ 0.0f0
+ @test all(Array(c) .== 4.0f0)
+
+ # Non-broadcasted operators also consume the logical shape.
+ library_result = c + c
+ @test size(library_result) == (M,)
+ @test all(Array(library_result) .== 8.0f0)
+
+ # A later custom launch sees the values written by the regular API.
+ task = cuNumeric.@cuda_task padding_mul(a, c, b, UInt32(M))
+ cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, c) outputs=b scalars=UInt32(
+ M
+ )
+ result = Array(b)
+ @test all(result[1:M] .== 4.0f0)
+ @test all(result[(M + 1):N] .== 1.0f0)
+ finally
+ !isnothing(library_result) && cuNumeric.destroy!(library_result)
+ cuNumeric.destroy!(a)
+ cuNumeric.destroy!(b)
+ cuNumeric.destroy!(c)
+ end
+ @test cuNumeric.current_device_bytes[] == initial_bytes
+end
+
+function cuda_padding_slice_output()
+ N = 4096
+ M = N - 2
+ threads = 256
+ blocks = cld(M, threads)
+ initial_bytes = cuNumeric.current_device_bytes[]
+
+ a = cuNumeric.ones(Float32, N)
+ b = cuNumeric.ones(Float32, N)
+ parent = cuNumeric.zeros(Float32, N)
+ output = parent[1:M]
+ task = cuNumeric.@cuda_task padding_add(a, b, output, UInt32(M))
+ unpadded_bytes = cuNumeric.current_device_bytes[]
+
+ try
+ @test @inferred(
+ cuNumeric.launch(
+ task, (a, b), output, UInt32(M); threads=threads, blocks=blocks
+ )
+ ) === nothing
+ padded_bytes = cuNumeric.current_device_bytes[]
+ @test @inferred(cuNumeric._launch_shape(output)) == (N,)
+ @test @inferred(cuNumeric._sync_from_launch_padding!(output)) === nothing
+
+ # Mutate the logical parent view before reusing it as a custom-kernel input.
+ output .= output .* 1.0f0 .+ 1.0f0
+ library_result = output + output
+ @test all(Array(library_result) .== 6.0f0)
+ cuNumeric.destroy!(library_result)
+
+ task = cuNumeric.@cuda_task padding_mul(a, output, b, UInt32(M))
+ @test @inferred(
+ cuNumeric.launch(
+ task, (a, output), b, UInt32(M); threads=threads, blocks=blocks
+ )
+ ) === nothing
+ values = Array(parent)
+ product = Array(b)
+
+ @test padded_bytes > unpadded_bytes
+ @test cuNumeric.current_device_bytes[] == padded_bytes
+ @test all(values[1:M] .== 3.0f0)
+ @test all(product[1:M] .== 3.0f0)
+ @test values[end] == 0.0f0
+ finally
+ cuNumeric.destroy!(output)
+ cuNumeric.destroy!(parent)
+ cuNumeric.destroy!(a)
+ cuNumeric.destroy!(b)
+ end
+ @test cuNumeric.current_device_bytes[] == initial_bytes
+end
+
+Base.@noinline function drop_padded_arrays()
+ N = 4096
+ M = N - 2
+ a = cuNumeric.ones(Float32, N)
+ b = cuNumeric.ones(Float32, N)
+ c = cuNumeric.zeros(Float32, M)
+ task = cuNumeric.@cuda_task padding_add(a, b, c, UInt32(M))
+ cuNumeric.@launch task=task threads=256 blocks=cld(M, 256) inputs=(a, b) outputs=c scalars=UInt32(
+ M
+ )
+ return nothing
+end
+
+function cuda_padding_finalizer()
+ GC.gc(true)
+ cuNumeric.drain_pending_frees!()
+ baseline = cuNumeric.current_device_bytes[]
+
+ drop_padded_arrays()
+ allocated = cuNumeric.current_device_bytes[]
+ GC.gc(true)
+ GC.gc(true)
+ cuNumeric.drain_pending_frees!()
+
+ @test allocated > baseline
+ @test cuNumeric.current_device_bytes[] == baseline
+end
+
+try
+ @testset "Custom CUDA padding" begin
+ cuda_padding_lifetime()
+ cuda_padding_api_interop()
+ cuda_padding_slice_output()
+ cuda_padding_finalizer()
+ end
+finally
+ cuNumeric.Experimental(false)
+end
diff --git a/test/cuda.jl/vecadd.jl b/test/cuda.jl/vecadd.jl
index ea01e7ab3..e2b5de552 100644
--- a/test/cuda.jl/vecadd.jl
+++ b/test/cuda.jl/vecadd.jl
@@ -21,6 +21,11 @@
-- Register various custom kernels using CUDA.jl
=#
+using CUDACore: blockDim, blockIdx, threadIdx
+import CUDACore: i32
+
+cuNumeric.Experimental(true)
+
function kernel_add(a, b, c, N)
i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
if i <= N
@@ -29,9 +34,8 @@ function kernel_add(a, b, c, N)
return nothing
end
-# testing a second kernel
-# on purpose switching inputs and outputs
-function kernel_mul(a, b, c, N)
+# Test a second kernel with `c` as an input and `b` as the output.
+function kernel_mul(a, c, b, N)
i = (blockIdx().x - 1i32) * blockDim().x + threadIdx().x
if i <= N
@inbounds b[i] = a[i] * c[i]
@@ -70,18 +74,16 @@ function cuda_binaryop(max_diff)
N
)
- @test @allowscalar cuNumeric.compare(c, c_cpu, atol(Float32), rtol(Float32))
+ @test @allowscalar cuNumeric.compare(c, c_cpu, max_diff, max_diff)
- for i in 1:N
- @allowscalar b[i] = a[i] * c[i]
- end
+ b_cpu .= a_cpu .* c_cpu
- task = cuNumeric.@cuda_task kernel_mul(a, b, c, UInt32(1))
+ task = cuNumeric.@cuda_task kernel_mul(a, c, b, UInt32(1))
cuNumeric.@launch task=task threads=threads blocks=blocks inputs=(a, c) outputs=b scalars=UInt32(
N
)
- @test @allowscalar cuNumeric.compare(b, b_cpu, atol(Float32), rtol(Float32))
+ @test @allowscalar cuNumeric.compare(b, b_cpu, max_diff, max_diff)
end
function kernel_sin(a, b, N)
@@ -119,5 +121,14 @@ function cuda_unaryop(max_diff)
# TODO explore getting inplace ops working.
cuNumeric.@launch task=task threads=threads blocks=blocks inputs=a outputs=b scalars=UInt32(N)
- @test @allowscalar cuNumeric.compare(b, b_cpu, atol(Float32), rtol(Float32))
+ @test @allowscalar cuNumeric.compare(b, b_cpu, max_diff, max_diff)
+end
+
+try
+ @testset "Custom CUDA kernels" begin
+ cuda_binaryop(1.0f-5)
+ cuda_unaryop(1.0f-5)
+ end
+finally
+ cuNumeric.Experimental(false)
end
diff --git a/test/runtests.jl b/test/runtests.jl
index 9efce84a1..a7ca7bbb7 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -32,18 +32,23 @@ delete!(testsuite, "util")
delete!(testsuite, "array/unary/tests")
delete!(testsuite, "array/binary/tests")
-if !run_gpu_tests
- @warn "CUDA GPU not available, skipping GPU-only tests"
- filter!(test -> !startswith(first(test), "gpu_only/"), testsuite)
+test_args = parse_args(ARGS)
+if filter_tests!(testsuite, test_args)
+ if !run_gpu_tests
+ @warn "CUDA GPU not available, skipping GPU-only tests"
+ filter!(
+ test ->
+ !startswith(first(test), "gpu_only/") &&
+ !startswith(first(test), "cuda.jl/"),
+ testsuite,
+ )
+ end
+
+ if !run_gpu_tests || !cuNumeric.FUSE_BROADCAST_EXPRS
+ @warn "Broadcast fusion is disabled, skipping fusion tests"
+ filter!(test -> !startswith(first(test), "gpu_only/broadcast_fusion"), testsuite)
+ end
end
-if !run_gpu_tests || !cuNumeric.FUSE_BROADCAST_EXPRS
- @warn "Broadcast fusion is disabled, skipping fusion tests"
- filter!(test -> !startswith(first(test), "gpu_only/broadcast_fusion"), testsuite)
-end
-
-# TODO
-# filter out tests for now, but the custom kernel registry should be tested
-filter!(test -> !startswith(first(test), "cuda.jl/"), testsuite)
-
-runtests(cuNumeric, ARGS; testsuite, init_code)
+cuda_tests = filter(test -> startswith(test, "cuda.jl/"), collect(keys(testsuite)))
+runtests(cuNumeric, test_args; testsuite, init_code, serial=cuda_tests)