diff --git a/README.md b/README.md
index b51c551b..39dd0896 100644
--- a/README.md
+++ b/README.md
@@ -2,26 +2,41 @@
cuNumeric.jl
+
+[](https://julialegate.github.io/cuNumeric.jl/dev/) [](https://app.codecov.io/github/JuliaLegate/cuNumeric.jl) [](https://opensource.org/licenses/MIT)
[](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 wraps and extends the [cuPyNumeric](https://github.com/nv-legate/cupynumeric) library from NVIDIA to bring distributed array computing on GPUs and CPUs to Julia. The central type is `NDArray`, which behaves like Julia's `Array` or the `CuArray` from [CUDA.jl](https://github.com/juliagpu/cuda.jl), but executes across multiple GPUs/CPUs. We implement array-level operations on `NDArray` which can be composed into larger programs without the need for explicit MPI calls or writing CUDA kernels.
+cuNumeric.jl requires x86 Linux, an NVIDIA GPU, and Julia >= 1.10. If ARM support is of interest open an issue.
cuNumeric.jl requires x86 Linux, an NVIDIA GPU, and Julia >= 1.10. If ARM support is of interest open an issue.
### Quick Start
+cuNumeric.jl can be installed with the Julia package manager. Activate your preferred environment and then from the Julia REPL run:
+
+
cuNumeric.jl can be installed with the Julia package manager. Activate your preferred environment and then from the Julia REPL run:
```julia
using Pkg
Pkg.add(url = "https://github.com/JuliaLegate/cuNumeric.jl", rev = "main")
+using Pkg
+Pkg.add(url = "https://github.com/JuliaLegate/cuNumeric.jl", rev = "main")
```
The first time might take awhile as it has to install multiple large dependencies such as the CUDA SDK (if you have an NVIDIA GPU). To use a local build of cupynumeric.so, see [Build Modes](./install.md).
+The first time might take awhile as it has to install multiple large dependencies such as the CUDA SDK (if you have an NVIDIA GPU). To use a local build of cupynumeric.so, see [Build Modes](./install.md).
+
```julia
using cuNumeric
+using cuNumeric
cuNumeric.versioninfo()
```
@@ -34,41 +49,57 @@ For more details, see [Hardware](./configuration/hardware.md).
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.
+**Data may live across many devices.** An `NDArray` is a logical array whose physical buffers can be partitioned over GPUs and CPUs by the Legate runtime. You write ordinary array code and Legate decides where the data lives and how/when it is communicated between devices. As a result, elementwise indexing (i.e. `arr[1]`) is slow (and is prevented by default). Scalar indexing like this forces synchronization and blocks other tasks from executing. Functions like `println` result in data being copied to the host and can also be slow.
**Slices are views.** Indexing an `NDArray` with ranges returns a view onto the same store, not a copy. That differs from Base Julia, where `A[1:n]` allocates a new `Array`. Mutations through an `NDArray` slice are visible through other aliases of the same data.
-**Reductions return arrays, not Julia scalars.** Reductions such as `sum(A)` produce a **0D or 1D** `NDArray` (axis reductions produce a lower-rank `NDArray`), rather than a bare `Float64` / `Float32`. That keeps the Legate task graph asynchronous instead of forcing synchronization to communite with the Julia runtime. When you need a plain Julia number, call `unwrap`:
+**Reductions return arrays, not Julia scalars.** Reductions such as `sum(A)` produce a **0D or 1D** `NDArray` (axis reductions produce a lower-rank `NDArray`), rather than a bare `Float64` / `Float32`. That keeps the Legate task graph asynchronous instead of forcing synchronization to communite with the Julia runtime. When you need a plain Julia number, call `unwrap` or `only`:
```julia
s = sum(A) # NDArray{T,0}
x = unwrap(s) # T, e.g. Float32
+x2 = only(s)
+```
+s = sum(A) # NDArray{T,0}
+x = unwrap(s) # T, e.g. Float32
```
+**The Legate runtime builds a DAG asynchronously.** Calling `cuNumeric.zeros` or `A .+ B` records work into Legate's task graph rather than blocking until every GPU kernel finishes. Results are materialized when you need them (for example `println`, `unwrap`, or converting with `Array(A)`). Hiding latency enables performant code.
**The Legate runtime builds a DAG asynchronously.** Calling `cuNumeric.zeros` or `A .+ B` records work into Legate's task graph rather than blocking until every GPU kernel finishes. Results are materialized when you need them (for example `println`, `unwrap`, or converting with `Array(A)`). Hiding latency enables performant code.
+For API details see [Initialization](./api_initialization.md) and [NDArray Reference](./api.md). For anti-patterns that kill performance, see [Patterns to Avoid](./perf/patterns_to_avoid.md).
For API details see [Initialization](./api_initialization.md) and [NDArray Reference](./api.md). For anti-patterns that kill performance, see [Patterns to Avoid](./perf/patterns_to_avoid.md).
+### Kernel Fusion
### Kernel Fusion
+Nested broadcast expressions fuse into a single kernel by default when on GPU. Prefer `@.` for multi-op elementwise code so every operator is dotted and the expression stays completely fused. Even just forgetting the `.` on unary negation (i.e., `y .= -a .+ b .* c`) will result in unfused code. Use the following pattern instead.
Nested broadcast expressions fuse into a single kernel by default when on GPU. Prefer `@.` for multi-op elementwise code so every operator is dotted and the expression stays completely fused. Even just forgetting the `.` on unary negation (i.e., `y .= -a .+ b .* c`) will result in unfused code. Use the following pattern instead.
+```julia
+y .= @. -a + b * c
```julia
y .= @. -a + b * c
```
See [Kernel Fusion](./perf/kernel_fusion.md) and [Debugging](./debugging.md) for controls and pretty printers.
+See [Kernel Fusion](./perf/kernel_fusion.md) and [Debugging](./debugging.md) for controls and pretty printers.
+
### Helping the Garbage Collector
+Many calls such as array slicing and un-fused broadcasts allocate a new `NDArray`. The Legate runtime keeps track of all references to the underlying data and will not free the memory until Julia's GC frees the `NDArray` handles. Because Julia's GC runs on memory pressure and an `NDArray` only stores a pointer (i.e., Julia's GC does not know the true size), many dead buffers accumulate and can cause out-of-memory errors.
Many calls such as array slicing and un-fused broadcasts allocate a new `NDArray`. The Legate runtime keeps track of all references to the underlying data and will not free the memory until Julia's GC frees the `NDArray` handles. Because Julia's GC runs on memory pressure and an `NDArray` only stores a pointer (i.e., Julia's GC does not know the true size), many dead buffers accumulate and can cause out-of-memory errors.
+`@analyze_lifetimes` performs a **static last-use analysis** at macro-expansion time and inserts eager calls to immediately free unused `NDArrays`. These buffers can then be reused by legate later for same-sized allocations.
`@analyze_lifetimes` performs a **static last-use analysis** at macro-expansion time and inserts eager calls to immediately free unused `NDArrays`. These buffers can then be reused by legate later for same-sized allocations.
```julia
@analyze_lifetimes begin
result = @. A[1:end, :] + B[1:end, :]
C .= @. result * 2.0f0
+ result = @. A[1:end, :] + B[1:end, :]
+ C .= @. result * 2.0f0
end
```
@@ -100,3 +131,11 @@ More worked examples (initialization, Gray-Scott, …) are in the documentation
### Known Limitations
- There is no support for `Float16` or `ComplexF16`
+- Arrays with 4 or more dimensions might have worse performance
+- Maximum array dimension is 6
+
+
+### Known Deviations from Base Julia
+- Reductions return 0D stores instead of scalars
+- Slices return views
+- `inv` does not throw `SingularException` for singular matrices
diff --git a/TODO.md b/TODO.md
index 6637e87f..b1bbe0b0 100644
--- a/TODO.md
+++ b/TODO.md
@@ -5,4 +5,58 @@
- Support Ints on methods that takes floats
- Programatic manipulation of Legate hardware config (not currently possible)
- Add Aqua.jl to CI to ensure we didn't pirate any types
-- Fix CodeCov reports
+
+## Base
+
+Easy `Base` / `AbstractArray` gaps for `NDArray`. Module helpers
+(`cuNumeric.reshape` / `transpose` / `unique` / …) often exist, but the
+corresponding `Base` methods are missing, so calls fall through to
+`AbstractArray` and may scalar-index. Wire up `Base.*` when convenient.
+
+**P0**
+- `Base.reshape`, `Base.vec`
+- `fill!` (convertible eltypes)
+- `collect`
+
+**P0 done**
+- `iszero` / `isone` (on-device reductions → `Bool` via scalar sync; `isone` square 2D only)
+
+**P1**
+- `transpose` / `adjoint`
+- `unique`
+- `ones_like`
+- `dropdims`
+
+**P2**
+- `extrema`, `mean`
+- `count` / nonzero
+- `abs2`
+- numeric `all` / `any`
+- `sum(f, A)` specials
+
+**P3**
+- `copyto!(NDArray, AbstractArray)`
+- `floor` / `ceil` / `clamp`
+- 2D `permutedims`
+- `diff`
+
+## LinearAlgebra
+
+Starter list of easy/medium LA gaps. Prefer wiring `LinearAlgebra` entry
+points so they do not fall through to scalar-indexing Base paths.
+
+**BLAS-1 style (dense `NDArray`)**
+- `axpy!` / `axpby!` — common scale-and-add; examples use broadcast today
+- `scal!` and related in-place scale
+- `LinearAlgebra.dot` if not already Base-wired to `nda_dot`
+
+**Reductions / traces**
+- `LinearAlgebra.tr` for dense 2D `NDArray` — `cuNumeric.trace` exists; `tr` is
+ already wired for `Diagonal{<:NDArray}`
+
+**Diagonal vs fallthrough (context)**
+- Already on-device for `Diagonal`: `mul!` / `lmul!` / `rmul!`, `\` / `/`,
+ `tr`, `norm` / `opnorm`, many predicates — see `docs/src/linalg.md`
+- Still fallthrough / unsupported on `Diagonal` (e.g. `svd`, `pinv`,
+ `cholesky`, host `AbstractArray` RHS): leave alone unless fixing is cheap;
+ densify intentionally when needed
diff --git a/deps/build.jl b/deps/build.jl
index efcc2368..68ac071c 100644
--- a/deps/build.jl
+++ b/deps/build.jl
@@ -42,7 +42,7 @@ function build_cpp_wrapper(
)
@info "libcunumeric_jl_wrapper: Building C++ Wrapper Library"
isdir(install_root) && (rm(install_root; recursive=true); mkdir(install_root))
- bld_command = `$(joinpath(repo_root, "scripts/build_cpp_wrapper.sh")) $repo_root $cupynumeric_loc $legate_loc $blas_loc $install_root 8`
+ bld_command = `$(joinpath(repo_root, "scripts/build_cpp_wrapper.sh")) $repo_root $cupynumeric_loc $legate_loc $blas_loc $install_root $(Threads.nthreads())`
return BuildTools.run_build_wrapper_script(
repo_root, bld_command; cuda_root, cuda_enabled, log_dir=@__DIR__
)
diff --git a/docs/src/api_initialization.md b/docs/src/api_initialization.md
index bd8ddacc..fdca3584 100644
--- a/docs/src/api_initialization.md
+++ b/docs/src/api_initialization.md
@@ -2,49 +2,65 @@
Constructors for new `NDArray`s. Default floating-point type is `Float32`.
-## zeros
+## Basic Initialization
+
+### zeros
```@docs
cuNumeric.zeros
```
-## ones
+### ones
```@docs
cuNumeric.ones
```
-## fill
+### fill
```@docs
cuNumeric.fill
```
-## trues
+### trues
```@docs
cuNumeric.trues
```
-## falses
+### falses
```@docs
cuNumeric.falses
```
-## eye
+## Special Matrices
-```@docs
-cuNumeric.eye
+### Diagonal
+
+Construct a `Diagonal` matrix whose elements (diagonal only) are stored in an `NDArray`.
+`LinearAlgebra.I` can be used to construct dense identity matrices as well.
+
+```julia
+using LinearAlgebra
+using cuNumeric
+
+D = Diagonal(cuNumeric.ones(Float32, 5)) # preferred for diagonal work
+I32 = NDArray{Float32}(I, 5, 5) # dense Float32 identity
+Ib = NDArray(I, 5, 5) # Bool identity
```
-## rand
+See [Linear Algebra](./linalg.md#diagonal-and-identity) for preferred patterns.
+
+## Random Numbers
+
+### rand
```@docs
cuNumeric.rand
```
-## rand!
+### rand!
```@docs
Random.rand!(::NDArray{<:cuNumeric.SUPPORTED_FLOAT_TYPES})
diff --git a/docs/src/examples/initialization.md b/docs/src/examples/initialization.md
index afd45363..d187828d 100644
--- a/docs/src/examples/initialization.md
+++ b/docs/src/examples/initialization.md
@@ -3,6 +3,7 @@
Create `NDArray`s with the usual Julia-style constructors. The default element type is `Float32` unless you pass one.
```julia
+using LinearAlgebra
using cuNumeric
# Zeros / ones / fill
@@ -15,9 +16,10 @@ F = cuNumeric.fill(7.5f0, (2, 3))
T = cuNumeric.trues(2, 3)
Fbool = cuNumeric.falses(2, 3)
-# Identity
-I = cuNumeric.eye(5)
-I16 = cuNumeric.eye(Float32, 5)
+# Identity: prefer Diagonal / I; densify only when needed (no eye)
+D = Diagonal(cuNumeric.ones(Float32, 5))
+I32 = NDArray{Float32}(I, 5, 5) # dense identity
+Ib = NDArray(I, 5, 5) # Bool identity
# Uniform / normal random values (native Float32 and Float64)
R = cuNumeric.rand(4, 4)
diff --git a/docs/src/examples/special_mat.md b/docs/src/examples/special_mat.md
new file mode 100644
index 00000000..56e64ce6
--- /dev/null
+++ b/docs/src/examples/special_mat.md
@@ -0,0 +1,23 @@
+# Special Matrices
+
+Currently cuNumeric only supports `LinearAlgebra.Diagonal`. Other special matrix types like `Tridiagonal` and `Symmetric` will follow.
+
+Diagonal matrices are common, require only storage of the diagonal elements and are often simple to compute operations on (i.e. `LinearAlgebra.inv`). `LinearAlgebra.Diagonal` matrices can be constructed from 1D or 2D `NDArray`s and have certain operations implemented (i.e., `eigen` and `inv`).
+
+
+```julia
+using cuNumeric
+using LinearAlgebra
+
+one_dim = cuNumeric.NDArray([1,2,3,4,5])
+two_dim = cuNumeric.rand(5,5)
+
+D1 = Diagonal(one_dim)
+D2 = Diagonal(two_dim)
+
+evals, evecs = eigen(D1)
+D1_inv = inv(D1)
+
+D1 ./= 2 # stays diagonal
+arr = D2 .+ two_dim # densifies because `two_dim` is not guranteed to be diagonal
+```
diff --git a/docs/src/linalg.md b/docs/src/linalg.md
index 103c2c40..1052ca4a 100644
--- a/docs/src/linalg.md
+++ b/docs/src/linalg.md
@@ -81,11 +81,113 @@ supported.
These helpers live on `NDArray` and are also listed in the Public API:
- `cuNumeric.transpose`
-- `cuNumeric.eye`
- `cuNumeric.diag` (2D to 1D)
- `cuNumeric.trace`
+## Diagonal and identity
+
+Prefer structured `LinearAlgebra` types over materializing a full matrix.
+
+Wrap a 1D `NDArray` in `Diagonal` for scale / solve / inverse along a diagonal.
+Use `LinearAlgebra.I` (`UniformScaling`) for `A + I`, `D + I`, and `A * I`.
+`D + I` stays a `Diagonal`; `A + I` returns a dense `NDArray`.
+
+```julia
+using LinearAlgebra
+using cuNumeric
+
+d = cuNumeric.ones(Float32, 64)
+D = Diagonal(d) # preferred: keep diagonal structure
+A = cuNumeric.rand(Float32, 64, 64)
+v = cuNumeric.rand(Float32, 64)
+
+y = D * v # scale a vector
+B = D * A # scale rows
+X = D \ A # scale columns by 1 ./ d
+Di = D + I # still Diagonal
+C = A + I # dense NDArray
+```
+
+### Supported `Diagonal{<:NDArray}` APIs
+
+These paths stay on-device (no host densify for the math). RHS / other operands
+must be `NDArray` unless noted.
+
+**Construction / display**
+
+- `Diagonal(v::NDArray{<:Any,1})` — wrap without copying
+- `Diagonal(A::NDArray{<:Any,2})` — `Diagonal(diag(A))`
+- `Matrix(D)` / `Matrix{T}(D)` — densify to a host `Matrix` (conversion only)
+- `show` — densifies `.diag` for printing only
+
+**Multiply / divide / inverse**
+
+- `D * A`, `A * D`, `D * v` for 2D / 1D `NDArray`
+- `mul!`, `lmul!`, `rmul!` with `NDArray`
+- `D \ B`, `A / D`, `ldiv!`, `rdiv!` with `NDArray`
+- `inv(D)` — reciprocal on-device; zeros become Inf (no `SingularException`)
+- `det(D)` — 0-dimensional `NDArray` product of the diagonal
+
+**`NDArray` ± `Diagonal`**
+
+- `A + D`, `D + A`, `A - D`, `D - A` for square 2D `NDArray`
+
+**UniformScaling (`I`)**
+
+- `NDArray{T}(I, m, n)` / `NDArray(I, …)`, `copyto!(A, I)`, `one(A)`, `oneunit(A)`
+- `A ± I`, `A * I`, `I * A`
+- `D ± I`, `D * I`, `I * D`, `copyto!(D, I)` — `D ± I` stays `Diagonal`
+
+**Broadcast**
+
+- Structure- or zero-preserving broadcasts on `Diagonal` (e.g. `D .* c`,
+ `D .*= c`, `D .+ D`) lower to 1D broadcast on `.diag`
+- Densifying out-of-place broadcasts (e.g. `D .+ 1`, `D .+ A`) materialize a
+ dense `NDArray`, matching Base’s densify-to-`Matrix` behavior
+- In-place densifying writes into `Diagonal` (e.g. `D .+= 1`, `D .+= A`) still
+ throw `ArgumentError` (off-diagonal / densify), matching Base
+
+**Eigen / reductions / predicates / norms**
+
+- `eigvals(D)`, `eigen(D)`, `eigvecs(D)` — unsorted; values are a copy of the
+ diagonal (`NDArray`), vectors are `NDArray` identity. Keyword `sortby` is not
+ supported on this method.
+- `tr`, `sum`, `prod`, `maximum`, `minimum` — 0-dimensional `NDArray` (not a Julia scalar)
+- `iszero`, `isone`, `istriu`, `istril`, `ishermitian`, `issymmetric`, `isposdef` — 0-dimensional `NDArray{Bool}`
+- `opnorm(D)` / `opnorm(D, p)` for `p ∈ {1, 2, Inf}` — 0-dimensional `NDArray`
+- `norm(D)` / `norm(D, p)` for finite `p` (including `±Inf`); off-diagonals are zero — 0-dimensional `NDArray`
+- `cond(D)` / `cond(D, p)` for `p ∈ {1, 2, Inf}` — 0-dimensional `NDArray`
+- `logdet(D)` for real `Diagonal` only — 0-dimensional `NDArray`
+
+**Helpers on dense `NDArray`**
+
+- `cuNumeric.diag` / `LinearAlgebra.diag` (2D → 1D), `cuNumeric.trace` / `LinearAlgebra.tr` (2D square → 0D)
+- `iszero(A)` — all elements `== zero(T)` → 0-dimensional `NDArray{Bool}`
+- `isone(A)` — square 2D vs `_eye(T, n)` → 0-dimensional `NDArray{Bool}` (non-square → `false`)
+
+### Unsupported / fallthrough
+
+Other `LinearAlgebra` operations on `Diagonal{<:NDArray}` (for example `svd`,
+`svdvals`, `pinv`, `logabsdet`, complex `logdet`, `kron`, `cholesky`, host
+`AbstractArray` RHS for `\` / `/` / `ldiv!` / `rdiv!`, or `eigen(...; sortby)`)
+are **not** specially implemented. They fall through to Base and typically fail
+with the package’s scalar-indexing error (NDArray does not support scalar
+indexing without `@allowscalar`). There are no “not implemented” `ArgumentError`
+stubs for these.
+
+Only densify when you truly need a full identity matrix:
+
+```julia
+E = NDArray{Float32}(I, 64, 64) # dense identity
+copyto!(A, I) # fill an existing array
+E2 = one(A) # same shape / eltype as A
+```
+
+Avoid building a dense identity (or densifying `D` with `Matrix(D)`) just to
+scale or shift; prefer `Diagonal` and `I` instead. There is no public `eye`.
+
## Not available yet
-There is no public `cholesky`, `eig`, `lu`, matrix `inv`, or `ldiv!` yet.
+There is no public dense-matrix `cholesky`, `eig`, `lu`, matrix `inv`, or
+`ldiv!` yet (beyond the `Diagonal` / `NDArray` paths listed above).
Elementwise `inv` / `^-1` are unary operations, not matrix inverse.
diff --git a/src/cuNumeric.jl b/src/cuNumeric.jl
index d86f36fd..e05a6e19 100644
--- a/src/cuNumeric.jl
+++ b/src/cuNumeric.jl
@@ -168,6 +168,7 @@ const FUSE_BROADCAST_EXPRS = CNPreferences.FUSE_BROADCAST
const FUSE_BROADCAST_MIN_OPS = CNPreferences.FUSE_BROADCAST_MIN_OPS
# Functionality
+include("ndarray/diagonal.jl")
include("ndarray/promotion.jl")
include("cuda/cuda_ptx_task.jl")
include("ndarray/broadcast_fusion.jl")
diff --git a/src/ndarray/binary.jl b/src/ndarray/binary.jl
index 07745f48..5631cfa5 100644
--- a/src/ndarray/binary.jl
+++ b/src/ndarray/binary.jl
@@ -136,8 +136,10 @@ function Base.:(+)(rhs1::NDArray{A,N}, rhs2::NDArray{B,N}) where {A,B,N}
return _nda_binary_op_promoted!(out, cuNumeric.ADD, rhs1, rhs2)
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.:(*)(val::V, arr::NDArray{A}) where {A,V<:Number}
+ return _mul_scalar(__my_promote_type(A, V), val, arr)
+end
+Base.:(*)(arr::NDArray{A}, val::V) where {A,V<:Number} = 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}
@@ -156,14 +158,14 @@ function Base.:(*)(rhs1::NDArray{A,2}, rhs2::NDArray{B,2}) where {A,B}
end
function Base.:(*)(rhs1::NDArray{Bool,2}, rhs2::NDArray{Bool,2})
- throw(
+ return throw(
ArgumentError("cuNumeric.jl does not support matrix multiplication of two Boolean arrays")
)
end
function Base.:(*)(rhs1::NDArray{<:Integer,2}, rhs2::NDArray{<:Integer,2})
#* this is a stupid.....
- throw(
+ return throw(
ArgumentError("cuNumeric.jl does not support matrix multiplication of two Integer arrays")
)
end
@@ -221,14 +223,14 @@ end
function LinearAlgebra.mul!(out::NDArray, rhs1::NDArray{Bool,2}, rhs2::NDArray{Bool,2})
#* Could just promote both inputs to Int32
- throw(
+ return throw(
ArgumentError("cuNumeric.jl does not support matrix multiplication of two Boolean arrays")
)
end
function LinearAlgebra.mul!(out::NDArray, rhs1::NDArray{<:Integer,2}, rhs2::NDArray{<:Integer,2})
#* this is a stupid.....
- throw(
+ return throw(
ArgumentError("cuNumeric.jl does not support matrix multiplication of two Integer arrays")
)
end
diff --git a/src/ndarray/broadcast.jl b/src/ndarray/broadcast.jl
index 2d64c65f..8df0b0ee 100644
--- a/src/ndarray/broadcast.jl
+++ b/src/ndarray/broadcast.jl
@@ -10,7 +10,7 @@ function map_cuda_type(::Type{cuNumeric.NDArrayStyle{N}}) where {N}
end # Also can be HostMemory or UnifiedMemory
function _nd_forbid_mix()
- throw(
+ return throw(
ArgumentError(
"Broadcast between NDArray and other array types is not supported. " *
"Convert explicitly to a single array type before broadcasting.",
@@ -26,6 +26,20 @@ Base.BroadcastStyle(::DefaultArrayStyle{0}, a::NDArrayStyle) = a
Base.BroadcastStyle(::NDArrayStyle, ::DefaultArrayStyle) = _nd_forbid_mix()
Base.BroadcastStyle(::DefaultArrayStyle, ::NDArrayStyle) = _nd_forbid_mix()
+# Like Base Diagonal vs Array: structured Diagonal style wins over dense NDArray
+# so D.+A uses StructuredMatrixStyle{Diagonal} (densify to NDArray in diagonal.jl)
+# instead of ArrayConflict → host Matrix + scalar indexing.
+function Base.BroadcastStyle(
+ ::LinearAlgebra.StructuredMatrixStyle{<:Diagonal}, ::NDArrayStyle
+)
+ return LinearAlgebra.StructuredMatrixStyle{Diagonal}()
+end
+function Base.BroadcastStyle(
+ ::NDArrayStyle, ::LinearAlgebra.StructuredMatrixStyle{<:Diagonal}
+)
+ return LinearAlgebra.StructuredMatrixStyle{Diagonal}()
+end
+
Base.broadcastable(A::NDArray) = A
#* IS THERE A BETTER WAY TO ALLOCATE THE NEW ARRAY???
@@ -37,6 +51,16 @@ Base.similar(arr::NDArray{T}, dims::Base.DimOrInd...) where {T} = similar(arr, T
Base.similar(arr::NDArray, ::Type{T}) where {T} = similar(arr, T, size(arr))
#* IS THERE A BETTER WAY TO ALLOCATE THE NEW ARRAY???
+# Prefer Dims over the axes catch-all: with StaticArrays loaded (GPU CI via CUDA),
+# `similar(::Type{<:AbstractArray}, ::Tuple{})` is otherwise ambiguous between
+# Base, StaticArrays, and our catch-all (0-d broadcast uses axes `()`).
+Base.similar(::Type{NDArray{T}}, dims::Dims{N}) where {T,N} = cuNumeric.zeros(T, dims)
+function Base.similar(
+ ::Type{NDArray{T}},
+ shape::Tuple{Union{Integer,Base.OneTo},Vararg{Union{Integer,Base.OneTo}}},
+) where {T}
+ return cuNumeric.zeros(T, map(Int, Base.to_shape.(shape)))
+end
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}
return similar(NDArray{ElType}, axes(bc))
@@ -55,7 +79,7 @@ end
# Get depth of Broadcast tree recursively
# Need to call instantiate first
-bcast_depth(bc::Base.Broadcast.Broadcasted) = maximum(bcast_depth, bc.args, init=0) + 1;
+bcast_depth(bc::Base.Broadcast.Broadcasted) = maximum(bcast_depth, bc.args; init=0) + 1;
bcast_depth(::Any) = 0
struct BrokenBroadcast{T} end
@@ -63,18 +87,22 @@ 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
+# Use cuNumeric promotion (`__recip_type` for inv, etc.), not Base.combine_eltypes
+# — e.g. inv.(Int32) must allocate Float32, not Float64.
+@inline function _broadcast_copy_eltype(bc::Broadcasted)
+ return __checked_promote_op(bc.f, Base.Broadcast.eltypes(bc.args))
+end
+
function Broadcast.copy(bc::Broadcasted{<:NDArrayStyle{0}})
- ElType = Broadcast.combine_eltypes(bc.f, bc.args)
+ ElType = _broadcast_copy_eltype(bc)
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()]
+ return copyto!(similar(bc, ElType), bc)
end
@inline function Broadcast.copy(bc::Broadcasted{<:NDArrayStyle})
- ElType = Broadcast.combine_eltypes(bc.f, bc.args)
+ ElType = _broadcast_copy_eltype(bc)
if ElType == Union{} || !Base.allocatedinline(ElType)
ElType = BrokenBroadcast{ElType}
end
@@ -162,7 +190,9 @@ end
end
@inline _copyto_unfused!(dest::NDArray{T}, temp_result::NDArray{T}) where {T} =
- _store_broadcast_result!(dest, temp_result)
+ _store_broadcast_result!(
+ dest, temp_result
+ )
@inline function _copyto_unfused!(dest::NDArray{T}, temp_result::NDArray) where {T}
promoted = checked_promote_arr(temp_result, T)
diff --git a/src/ndarray/broadcast_fusion.jl b/src/ndarray/broadcast_fusion.jl
index 6611edf1..728deec1 100644
--- a/src/ndarray/broadcast_fusion.jl
+++ b/src/ndarray/broadcast_fusion.jl
@@ -270,8 +270,22 @@ Also refuses 0-d destinations: `RunPTXBroadcastTask` only supports dims in
end
# Same-shaped operands do not need Broadcast's dynamic index projection.
+#
+# Size-1 dimensions are special: Broadcast marks them `keeps=false` even when the
+# leaf shape matches `dest` (e.g. length-1 vectors). Linear `I` is still valid in
+# that case because the dimension only has index 1.
+@inline function _extruded_ok_for_fusion(x::Base.Broadcast.Extruded)
+ keeps = x.keeps
+ for i in eachindex(keeps)
+ if !keeps[i] && size(x.x, i) != 1
+ return false
+ end
+ end
+ return true
+end
+
@inline function _unwrap_fusion_arg(x::Base.Broadcast.Extruded)
- if all(x.keeps)
+ if _extruded_ok_for_fusion(x)
return x.x
end
throw(
diff --git a/src/ndarray/detail/ndarray.jl b/src/ndarray/detail/ndarray.jl
index 386850a1..8b7e6767 100644
--- a/src/ndarray/detail/ndarray.jl
+++ b/src/ndarray/detail/ndarray.jl
@@ -41,7 +41,7 @@ end
get_n_dim(ptr::NDArray_t) = Int(ccall((:nda_array_dim, libnda), Int32, (NDArray_t,), ptr))
-abstract type AbstractNDArray{T<:SUPPORTED_TYPES,N} end
+abstract type AbstractNDArray{T<:SUPPORTED_TYPES,N} <: AbstractArray{T,N} end
@doc"""
The NDArray type represents a multi-dimensional array in cuNumeric.
@@ -475,7 +475,7 @@ function nda_trace(
(NDArray_t, Int32, Int32, Int32, Legate.LegateTypeAllocated),
arr.ptr, offset, a1, a2, legate_type)
end
- return NDArray(ptr, T, Val(1))
+ return NDArray(ptr, T, Val(0))
end
# transpose reverses the axes: element type and rank are preserved
diff --git a/src/ndarray/diagonal.jl b/src/ndarray/diagonal.jl
new file mode 100644
index 00000000..ae1d306a
--- /dev/null
+++ b/src/ndarray/diagonal.jl
@@ -0,0 +1,557 @@
+###### diag / _eye / trace ######
+
+@doc"""
+ cuNumeric.diag(arr::NDArray; k=0)
+
+Extract the k-th diagonal from a 2D `NDArray`.
+"""
+function diag(arr::NDArray; k::Int=0)
+ return nda_diag(arr, Int32(k))
+end
+
+LinearAlgebra.diag(arr::NDArray{<:Any,2}, k::Integer=0) = nda_diag(arr, Int32(k))
+
+# Internal dense identity used by UniformScaling / Diagonal densify helpers.
+# Prefer `LinearAlgebra.I` / `NDArray{T}(I, n, n)` / `one(A)` in user code.
+function _eye(::Type{T}, rows::Int) where {T}
+ return nda_eye(Int32(rows), T)
+end
+_eye(rows::Int) = _eye(DEFAULT_FLOAT, rows)
+
+@doc"""
+ cuNumeric.trace(arr::NDArray; offset=0, a1=0, a2=1)
+
+Compute the trace (sum of a diagonal) of the `NDArray`.
+Returns a 0-dimensional `NDArray`. The accumulator type follows promotions of
+other reductions like `sum`.
+"""
+function trace(arr::NDArray{T,2}; offset::Int=0, a1::Int=0, a2::Int=1) where {T}
+ LinearAlgebra.checksquare(arr)
+ T_OUT = Base.promote_op(Base.sum, Vector{T})
+ return nda_trace(arr, Int32(offset), Int32(a1), Int32(a2), T_OUT)
+end
+
+function LinearAlgebra.tr(arr::NDArray{<:Any,2})
+ return cuNumeric.trace(arr)
+end
+
+###### Diagonal constructors ######
+
+const DiagonalNDArray{T} = Diagonal{T,<:NDArray{T,1}}
+
+function LinearAlgebra.Diagonal(arr::NDArray{T,1}) where {T}
+ return Diagonal{T,typeof(arr)}(arr)
+end
+
+function LinearAlgebra.Diagonal(arr::NDArray{T,2}) where {T}
+ return Diagonal(diag(arr))
+end
+
+# Note: Matrix{T} === Array{T,2}, so do not also define Array{T,2}(...).
+# Use Base.zeros — bare `zeros` resolves to cuNumeric.zeros inside this module.
+function Base.Matrix{T}(D::DiagonalNDArray) where {T}
+ dv = Array(D.diag)
+ n = length(dv)
+ B = Base.zeros(T, n, n)
+ @inbounds for i in 1:n
+ B[i, i] = dv[i]
+ end
+ return B
+end
+Base.Matrix(D::DiagonalNDArray{T}) where {T} = Matrix{T}(D)
+
+function Base.show(io::IO, D::DiagonalNDArray)
+ return show(io, Diagonal(Array(D.diag)))
+end
+
+function Base.show(io::IO, ::MIME"text/plain", D::DiagonalNDArray)
+ # Keep the real Diagonal{T,<:NDArray} in the summary; only densify the
+ # diagonal vector for Base's ⋅-style body formatting.
+ summary(io, D)
+ isempty(D) && return nothing
+ println(io, ":")
+ Base.print_array(io, Diagonal(Array(D.diag)))
+ return nothing
+end
+
+###### Diagonal operators ######
+
+@inline _diag_vec(D::DiagonalNDArray) = D.diag
+@inline _row_scale(d::NDArray{<:Any,1}) = reshape(d, (length(d), 1))
+@inline _col_scale(d::NDArray{<:Any,1}) = reshape(d, (1, length(d)))
+
+function Base.:*(D::DiagonalNDArray, A::NDArray{<:Any,2})
+ size(A, 1) == size(D, 1) || throw(
+ DimensionMismatch(
+ "matrix is $(size(A,1))×$(size(A,2)), but diagonal is $(size(D,1))×$(size(D,2))"
+ ),
+ )
+ return _row_scale(_diag_vec(D)) .* A
+end
+
+function Base.:*(A::NDArray{<:Any,2}, D::DiagonalNDArray)
+ size(A, 2) == size(D, 1) || throw(
+ DimensionMismatch(
+ "matrix is $(size(A,1))×$(size(A,2)), but diagonal is $(size(D,1))×$(size(D,2))"
+ ),
+ )
+ return A .* _col_scale(_diag_vec(D))
+end
+
+function Base.:*(D::DiagonalNDArray, v::NDArray{<:Any,1})
+ length(v) == size(D, 1) || throw(
+ DimensionMismatch("vector length $(length(v)) does not match diagonal $(size(D,1))")
+ )
+ return _diag_vec(D) .* v
+end
+
+function LinearAlgebra.lmul!(D::DiagonalNDArray, B::NDArray)
+ return copyto!(B, D * B)
+end
+
+function LinearAlgebra.rmul!(A::NDArray, D::DiagonalNDArray)
+ return copyto!(A, A * D)
+end
+
+function LinearAlgebra.mul!(C::NDArray, D::DiagonalNDArray, A::NDArray)
+ return copyto!(C, D * A)
+end
+
+function LinearAlgebra.mul!(C::NDArray, A::NDArray, D::DiagonalNDArray)
+ return copyto!(C, A * D)
+end
+
+function Base.:\(D::DiagonalNDArray, B::NDArray{<:Any,1})
+ length(B) == size(D, 1) || throw(
+ DimensionMismatch("vector length $(length(B)) does not match diagonal $(size(D,1))")
+ )
+ return B ./ _diag_vec(D)
+end
+
+function Base.:\(D::DiagonalNDArray, B::NDArray{<:Any,2})
+ size(B, 1) == size(D, 1) || throw(
+ DimensionMismatch(
+ "matrix is $(size(B,1))×$(size(B,2)), but diagonal is $(size(D,1))×$(size(D,2))"
+ ),
+ )
+ return B ./ _row_scale(_diag_vec(D))
+end
+
+function Base.:/(A::NDArray{<:Any,2}, D::DiagonalNDArray)
+ size(A, 2) == size(D, 1) || throw(
+ DimensionMismatch(
+ "matrix is $(size(A,1))×$(size(A,2)), but diagonal is $(size(D,1))×$(size(D,2))"
+ ),
+ )
+ return A * inv(D)
+end
+
+function LinearAlgebra.ldiv!(D::DiagonalNDArray, B::NDArray)
+ return copyto!(B, D \ B)
+end
+
+function LinearAlgebra.rdiv!(A::NDArray, D::DiagonalNDArray)
+ return copyto!(A, A / D)
+end
+
+function Base.inv(D::DiagonalNDArray{T}) where {T}
+ # Base Julia checks and throws a SingularException. We cannot do
+ # that without unwrapping the NDArray to a Julia scalar.
+ return Diagonal(inv.(_diag_vec(D)))
+end
+
+LinearAlgebra.det(D::DiagonalNDArray) = prod(_diag_vec(D))
+LinearAlgebra.tr(D::DiagonalNDArray{<:Number}) = sum(_diag_vec(D))
+Base.sum(D::DiagonalNDArray) = sum(_diag_vec(D))
+
+# Generic `prod(::AbstractMatrix)` walks every entry (scalar-indexing). For n>1 a
+# Diagonal has off-diagonal zeros, so the product is zero — match Base, as 0D.
+function Base.prod(D::DiagonalNDArray{T}) where {T<:Number}
+ n = size(D, 1)
+ n == 0 && return NDArray(one(T))
+ n == 1 && return prod(_diag_vec(D))
+ return NDArray(zero(T))
+end
+
+function Base.maximum(D::DiagonalNDArray{T}) where {T<:Number}
+ maxdiag = maximum(_diag_vec(D))
+ size(D, 1) > 1 && return max.(zero(T), maxdiag)
+ return maxdiag
+end
+
+function Base.minimum(D::DiagonalNDArray{T}) where {T<:Number}
+ mindiag = minimum(_diag_vec(D))
+ size(D, 1) > 1 && return min.(zero(T), mindiag)
+ return mindiag
+end
+
+Base.iszero(D::DiagonalNDArray) = iszero(_diag_vec(D))
+function Base.isone(D::DiagonalNDArray{T}) where {T}
+ return all(_diag_vec(D) .== one(T))
+end
+
+# Base walks `iszero(D.diag)` by scalar iteration; keep on-device via `iszero(D)`.
+function LinearAlgebra.istriu(D::DiagonalNDArray, k::Integer=0)
+ return k <= 0 ? NDArray(true) : iszero(D)
+end
+function LinearAlgebra.istril(D::DiagonalNDArray, k::Integer=0)
+ return k >= 0 ? NDArray(true) : iszero(D)
+end
+
+# Real Diagonal is always Hermitian/symmetric in Base; Complex Hermitian needs isreal(diag).
+LinearAlgebra.ishermitian(D::DiagonalNDArray{<:Real}) = NDArray(true)
+function LinearAlgebra.ishermitian(D::DiagonalNDArray{<:Complex})
+ return all(imag(_diag_vec(D)) .== zero(real(eltype(D))))
+end
+LinearAlgebra.issymmetric(D::DiagonalNDArray{<:Number}) = NDArray(true)
+
+# Base `isposdef(D) = all(isposdef, D.diag)` scalar-iterates.
+function LinearAlgebra.isposdef(D::DiagonalNDArray{T}) where {T<:Real}
+ isempty(D) && return NDArray(true)
+ return all(_diag_vec(D) .> zero(T))
+end
+function LinearAlgebra.isposdef(D::DiagonalNDArray{T}) where {T<:Complex}
+ # isposdef(z) = isreal(z) && real(z) > 0 — keep on-device, no host densify.
+ d = _diag_vec(D)
+ return all((imag(d) .== zero(real(T))) .& (real(d) .> zero(real(T))))
+end
+
+###### Eigen / related ######
+
+# Base: eigvals(D::Diagonal{<:Number}) = copy(D.diag). Keep NDArray (package style).
+function LinearAlgebra.eigvals(D::DiagonalNDArray{<:Number}; permute::Bool=true, scale::Bool=true)
+ return copy(_diag_vec(D))
+end
+
+# Unsorted eigen: values are a copy of the diagonal (NDArray); vectors are NDArray I.
+# Keyword `sortby` is not accepted on this override.
+function LinearAlgebra.eigen(
+ D::DiagonalNDArray;
+ permute::Bool=true,
+ scale::Bool=true,
+)
+ Td = Base.promote_op(/, eltype(D), eltype(D))
+ return Eigen(copy(_diag_vec(D)), _eye(Td, size(D, 1)))
+end
+
+function LinearAlgebra.eigvecs(
+ D::DiagonalNDArray;
+ permute::Bool=true,
+ scale::Bool=true,
+)
+ return eigen(D; permute=permute, scale=scale).vectors
+end
+
+# Real logdet is sum(log.(diag)) on-device. Complex logdet / other Base LinearAlgebra
+# ops without overrides fall through and may scalar-index `.diag`.
+LinearAlgebra.logdet(D::DiagonalNDArray{<:Real}) = sum(log.(_diag_vec(D)))
+
+# Operator / entrywise norms from the diagonal only (no host densify).
+function LinearAlgebra.opnorm(D::DiagonalNDArray, p::Real=2)
+ if !(p == 1 || p == 2 || p == Inf)
+ throw(ArgumentError(lazy"invalid p-norm p=$p. Valid: 1, 2, Inf"))
+ end
+ isempty(D) && return NDArray(float(real(zero(eltype(D)))))
+ return maximum(abs.(_diag_vec(D)))
+end
+
+function LinearAlgebra.norm(D::DiagonalNDArray, p::Real=2)
+ # Off-diagonals are zero, so the matrix vec-norm equals the diag vec-norm.
+ d = abs.(_diag_vec(D))
+ if p == 2
+ return sqrt.(sum(d .^ 2))
+ elseif p == 1
+ return sum(d)
+ elseif p == Inf
+ return isempty(D) ? NDArray(float(real(zero(eltype(D))))) : maximum(d)
+ elseif p == -Inf
+ return isempty(D) ? NDArray(float(real(zero(eltype(D))))) : minimum(d)
+ else
+ return sum(d .^ p) .^ (one(p) / p)
+ end
+end
+
+function LinearAlgebra.cond(D::DiagonalNDArray, p::Real=2)
+ if !(p == 1 || p == 2 || p == Inf)
+ throw(ArgumentError(lazy"invalid p-norm p=$p. Valid: 1, 2, Inf"))
+ end
+ isempty(D) && return NDArray(float(one(real(eltype(D)))))
+ dabs = abs.(_diag_vec(D))
+ return maximum(dabs) ./ minimum(dabs)
+end
+
+function Base.:+(A::NDArray{T,2}, D::DiagonalNDArray) where {T}
+ size(A, 1) == size(A, 2) == size(D, 1) || throw(
+ DimensionMismatch(
+ "matrix is $(size(A,1))×$(size(A,2)), but diagonal is $(size(D,1))×$(size(D,2))"
+ ),
+ )
+ return A + (_row_scale(_diag_vec(D)) .* _eye(eltype(D), size(D, 1)))
+end
+Base.:+(D::DiagonalNDArray, A::NDArray{<:Any,2}) = A + D
+
+Base.:-(A::NDArray{<:Any,2}, D::DiagonalNDArray) = A + (-D)
+Base.:-(D::DiagonalNDArray, A::NDArray{<:Any,2}) = D + (-A)
+
+###### UniformScaling (LinearAlgebra.I) ######
+
+@inline function _uniformscaling_eye(::Type{R}, n::Integer, λ) where {R}
+ E = _eye(R, Int(n))
+ return isone(λ) ? E : nda_multiply_scalar(E, R(λ))
+end
+
+function NDArray{T}(J::LinearAlgebra.UniformScaling, dims::Dims{2}) where {T}
+ A = zeros(T, dims)
+ copyto!(A, J)
+ return A
+end
+function NDArray{T}(J::LinearAlgebra.UniformScaling, m::Integer, n::Integer) where {T}
+ return NDArray{T}(J, Dims((Int(m), Int(n))))
+end
+NDArray(J::LinearAlgebra.UniformScaling{T}, dims::Dims{2}) where {T} = NDArray{T}(J, dims)
+function NDArray(J::LinearAlgebra.UniformScaling{T}, m::Integer, n::Integer) where {T}
+ return NDArray{T}(J, Dims((Int(m), Int(n))))
+end
+
+function Base.copyto!(A::NDArray{T,2}, J::LinearAlgebra.UniformScaling) where {T}
+ m, n = size(A)
+ if iszero(J.λ)
+ return fill!(A, zero(T))
+ elseif m == n
+ return copyto!(A, _uniformscaling_eye(T, m, J.λ))
+ else
+ fill!(A, zero(T))
+ k = min(m, n)
+ A[1:k, 1:k] = _uniformscaling_eye(T, k, J.λ)
+ return A
+ end
+end
+
+function Base.:+(A::NDArray{T,2}, J::LinearAlgebra.UniformScaling) where {T}
+ LinearAlgebra.checksquare(A)
+ R = Base.promote_op(+, T, typeof(J.λ))
+ return A + _uniformscaling_eye(R, size(A, 1), J.λ)
+end
+Base.:+(J::LinearAlgebra.UniformScaling, A::NDArray{<:Any,2}) = A + J
+
+Base.:-(A::NDArray{<:Any,2}, J::LinearAlgebra.UniformScaling) = A + (-J)
+function Base.:-(J::LinearAlgebra.UniformScaling, A::NDArray{<:Any,2})
+ return (-A) + J
+end
+
+# Scale by λ without promoting the array (A * I must not Bool→Float32 promote).
+function Base.:*(A::NDArray{T}, J::LinearAlgebra.UniformScaling) where {T}
+ return _mul_scalar(T, J.λ, A)
+end
+function Base.:*(J::LinearAlgebra.UniformScaling, A::NDArray{T}) where {T}
+ return _mul_scalar(T, J.λ, A)
+end
+
+function Base.one(A::NDArray{T,2}) where {T}
+ LinearAlgebra.checksquare(A)
+ return _eye(T, size(A, 1))
+end
+function Base.oneunit(A::NDArray{T,2}) where {T}
+ LinearAlgebra.checksquare(A)
+ return _eye(T, size(A, 1))
+end
+
+###### Diagonal ↔ UniformScaling ######
+
+# Keep Diagonal structure: D + λI == Diagonal(d .+ λ), not a dense matrix.
+function Base.:+(D::DiagonalNDArray{T}, J::LinearAlgebra.UniformScaling) where {T}
+ R = Base.promote_op(+, T, typeof(J.λ))
+ return Diagonal(_diag_vec(D) .+ convert(R, J.λ))
+end
+Base.:+(J::LinearAlgebra.UniformScaling, D::DiagonalNDArray) = D + J
+
+Base.:-(D::DiagonalNDArray, J::LinearAlgebra.UniformScaling) = D + (-J)
+function Base.:-(J::LinearAlgebra.UniformScaling, D::DiagonalNDArray{T}) where {T}
+ R = Base.promote_op(-, typeof(J.λ), T)
+ return Diagonal(convert(R, J.λ) .- _diag_vec(D))
+end
+
+function Base.:*(D::DiagonalNDArray{T}, J::LinearAlgebra.UniformScaling) where {T}
+ return Diagonal(_mul_scalar(T, J.λ, _diag_vec(D)))
+end
+Base.:*(J::LinearAlgebra.UniformScaling, D::DiagonalNDArray) = D * J
+
+function Base.copyto!(D::DiagonalNDArray{T}, J::LinearAlgebra.UniformScaling) where {T}
+ fill!(_diag_vec(D), convert(T, J.λ))
+ return D
+end
+
+###### Diagonal broadcast ######
+
+# LinearAlgebra's StructuredMatrixStyle{Diagonal} path (structuredbroadcast.jl)
+# fills `dest.diag[i]` via `Broadcast._broadcast_getindex(bc, (i,i))`, which
+# scalar-indexes the NDArray. CUDA/GPUArrays hit the same Base path and do not
+# special-case Diagonal broadcast either.
+#
+# For structure-preserving / zero-preserving broadcasts we lower to a 1D
+# broadcast on `.diag` (NDArrayStyle). Fusion then applies to that vector
+# broadcast as usual; there is no separate Diagonal-matrix fusion.
+#
+# Densifying out-of-place broadcasts (e.g. `D .+ 1`, `D .+ A`) allocate a dense
+# NDArray and expand Diagonal leaves to `diag .* I` before the usual NDArray
+# `copyto!` path — matching Base, which densifies to Matrix. In-place writes
+# into Diagonal that would fill off-diagonals still throw ArgumentError.
+
+@inline _has_diagonal_ndarray(@nospecialize(_)) = false
+@inline _has_diagonal_ndarray(::DiagonalNDArray) = true
+@inline function _has_diagonal_ndarray(bc::Broadcast.Broadcasted)
+ return _has_diagonal_ndarray_args(bc.args)
+end
+@inline _has_diagonal_ndarray_args(::Tuple{}) = false
+@inline function _has_diagonal_ndarray_args(args::Tuple)
+ return _has_diagonal_ndarray(getfield(args, 1)) ||
+ _has_diagonal_ndarray_args(Base.tail(args))
+end
+
+# Dense NDArray leaves (not wrapped in Diagonal). Used to avoid scalar-indexing
+# when building the in-place densifying off-diagonal ArgumentError.
+@inline _has_plain_ndarray(@nospecialize(_)) = false
+@inline _has_plain_ndarray(::NDArray) = true
+@inline function _has_plain_ndarray(bc::Broadcast.Broadcasted)
+ return _has_plain_ndarray_args(bc.args)
+end
+@inline _has_plain_ndarray_args(::Tuple{}) = false
+@inline function _has_plain_ndarray_args(args::Tuple)
+ return _has_plain_ndarray(getfield(args, 1)) ||
+ _has_plain_ndarray_args(Base.tail(args))
+end
+
+@inline _first_diagonal_ndarray_diag(D::DiagonalNDArray) = D.diag
+@inline function _first_diagonal_ndarray_diag(bc::Broadcast.Broadcasted)
+ return _first_diagonal_ndarray_diag_args(bc.args)
+end
+@inline function _first_diagonal_ndarray_diag_args(args::Tuple)
+ a = getfield(args, 1)
+ return if _has_diagonal_ndarray(a)
+ _first_diagonal_ndarray_diag(a)
+ else
+ _first_diagonal_ndarray_diag_args(Base.tail(args))
+ end
+end
+
+# Replace Diagonal leaves with their `.diag` vectors; keep scalars / Refs / etc.
+@inline _diag_bc_arg(D::Diagonal) = D.diag
+@inline function _diag_bc_arg(bc::Broadcast.Broadcasted)
+ return Broadcast.broadcasted(bc.f, map(_diag_bc_arg, bc.args)...)
+end
+@inline _diag_bc_arg(x) = x
+
+@inline function _diagonal_broadcast_preserves_structure(bc::Broadcast.Broadcasted)
+ # `+` / `-` on Diagonal+Diagonal are zero-preserving (Base's fzeropreserving),
+ # not isstructurepreserving. Prefer either so D.+D lowers to `.diag` broadcast.
+ return LinearAlgebra.isstructurepreserving(bc) || LinearAlgebra.fzeropreserving(bc)
+end
+
+# Materialize Diagonal{NDArray} as a dense matrix (same pattern as A ± D).
+@inline function _densify_diagonal_ndarray(D::DiagonalNDArray)
+ d = _diag_vec(D)
+ n = length(d)
+ # nda_eye(1) currently yields an unreadable store; build 1×1 from the diag.
+ n <= 1 && return reshape(copy(d), (n, n))
+ return _row_scale(d) .* _eye(eltype(D), n)
+end
+
+# For densifying Structured → dense NDArray copyto!: expand Diagonal leaves.
+@inline _expand_diagonal_ndarray_bc_arg(D::DiagonalNDArray) = _densify_diagonal_ndarray(D)
+@inline function _expand_diagonal_ndarray_bc_arg(bc::Broadcast.Broadcasted)
+ return Broadcast.broadcasted(bc.f, map(_expand_diagonal_ndarray_bc_arg, bc.args)...)
+end
+@inline _expand_diagonal_ndarray_bc_arg(x) = x
+
+# Off-diagonal Diagonal getindex uses `diagzero` (no `.diag` read), so evaluating
+# an off-diagonal broadcast index is safe without `@allowscalar` when every
+# non-Diagonal leaf is a scalar. Dense NDArray leaves must not be indexed.
+# Used for in-place densifying rejection only (out-of-place densifies).
+function _throw_densifying_diagonal_broadcast(bc::Broadcast.Broadcasted)
+ axs = axes(bc)
+ if length(axs) >= 2 && length(axs[1]) >= 2 && length(axs[2]) >= 2
+ if !_has_plain_ndarray(bc)
+ v = @inbounds Broadcast._broadcast_getindex(bc, CartesianIndex(2, 1))
+ throw(
+ ArgumentError(
+ "cannot set off-diagonal entry (2, 1) to a nonzero value ($v)"
+ ),
+ )
+ end
+ throw(
+ ArgumentError(
+ "cannot set off-diagonal entry (2, 1) to a nonzero value; " *
+ "broadcast over Diagonal with NDArray diagonal would densify",
+ ),
+ )
+ end
+ return throw(
+ ArgumentError(
+ "broadcast over Diagonal with NDArray diagonal is not structure-preserving " *
+ "and would densify; in-place densifying broadcast is not supported",
+ ),
+ )
+end
+
+function Base.similar(
+ bc::Broadcast.Broadcasted{LinearAlgebra.StructuredMatrixStyle{Diagonal}},
+ ::Type{ElType},
+) where {ElType}
+ inds = axes(bc)
+ n = length(inds[1])
+ if _has_diagonal_ndarray(bc)
+ if _diagonal_broadcast_preserves_structure(bc)
+ d = _first_diagonal_ndarray_diag(bc)
+ return Diagonal(similar(d, ElType, (n,)))
+ end
+ # Match Base: densify out-of-place to a dense array (NDArray, not Matrix).
+ return similar(NDArray{ElType}, inds)
+ elseif _diagonal_broadcast_preserves_structure(bc)
+ return LinearAlgebra.structured_broadcast_alloc(bc, Diagonal, ElType, n)
+ else
+ return similar(
+ convert(Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{ndims(bc)}}, bc),
+ ElType,
+ )
+ end
+end
+
+@inline function _copyto_diagonal_ndarray!(
+ dest::DiagonalNDArray, bc::Broadcast.Broadcasted
+)
+ axes(bc) == axes(dest) || Broadcast.throwdm(axes(bc), axes(dest))
+ # Lower to NDArrayStyle vector broadcast so `_copyto!` / fusion apply.
+ copyto!(_diag_vec(dest), Broadcast.instantiate(_diag_bc_arg(bc)))
+ return dest
+end
+
+# Out-of-place densify: destination is dense NDArray from `similar` above.
+function Base.copyto!(
+ dest::NDArray,
+ bc::Broadcast.Broadcasted{LinearAlgebra.StructuredMatrixStyle{Diagonal}},
+)
+ axes(dest) == axes(bc) || Broadcast.throwdm(axes(dest), axes(bc))
+ isempty(dest) && return dest
+ expanded = Broadcast.instantiate(_expand_diagonal_ndarray_bc_arg(bc))
+ return _copyto!(dest, expanded)
+end
+
+function Base.copyto!(
+ dest::DiagonalNDArray,
+ bc::Broadcast.Broadcasted{<:LinearAlgebra.StructuredMatrixStyle},
+)
+ if !LinearAlgebra.isvalidstructbc(dest, bc)
+ # 1×1: Base's generic path only writes the diagonal; no off-diagonals to reject.
+ size(dest, 1) <= 1 || return _throw_densifying_diagonal_broadcast(bc)
+ end
+ return _copyto_diagonal_ndarray!(dest, bc)
+end
+
+# Safety net if a densifying Structured broadcast is converted to Nothing
+# (Base's `isvalidstructbc` fallback) before reaching the method above.
+function Base.copyto!(dest::DiagonalNDArray, bc::Broadcast.Broadcasted{Nothing})
+ if !_diagonal_broadcast_preserves_structure(bc)
+ size(dest, 1) <= 1 || return _throw_densifying_diagonal_broadcast(bc)
+ end
+ return _copyto_diagonal_ndarray!(dest, bc)
+end
diff --git a/src/ndarray/ndarray.jl b/src/ndarray/ndarray.jl
index 8a7f9afa..3f531b40 100644
--- a/src/ndarray/ndarray.jl
+++ b/src/ndarray/ndarray.jl
@@ -20,6 +20,8 @@
export unwrap
+# See TODO.md (Base / LinearAlgebra sections) for AbstractArray and LA gaps.
+
@doc"""
cuNumeric.transpose(arr::NDArray)
@@ -29,39 +31,6 @@ function transpose(arr::NDArray)
return nda_transpose(arr)
end
-@doc"""
- cuNumeric.eye([T=Float32,] rows::Int)
-
-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)
-end
-function eye(rows::Int)
- return eye(DEFAULT_FLOAT, rows)
-end
-
-@doc"""
- cuNumeric.trace(arr::NDArray; offset=0, a1=0, a2=1)
-
-Compute the trace (sum of a diagonal) of the `NDArray`.
-The accumulator type follows promotions of other reductions like 'sum'.
-"""
-function trace(arr::NDArray{T}; offset::Int=0, a1::Int=0, a2::Int=1) where {T}
- T_OUT = Base.promote_op(Base.sum, Vector{T})
- return nda_trace(arr, Int32(offset), Int32(a1), Int32(a2), T_OUT)
-end
-
-@doc"""
- cuNumeric.diag(arr::NDArray; k=0)
-
-Extract the k-th diagonal from a 2D `NDArray`.
-"""
-function diag(arr::NDArray; k::Int=0)
- return nda_diag(arr, Int32(k))
-end
-
@doc"""
cuNumeric.ravel(arr::NDArray)
@@ -111,7 +80,10 @@ copyto!(a, b);
a[1,1]
```
"""
-Base.copyto!(arr::NDArray{T,N}, other::NDArray{T,N}) where {T,N} = nda_assign(arr, other)
+@inline function Base.copyto!(arr::NDArray{T,N}, other::NDArray{T,N}) where {T,N}
+ nda_assign(arr, other)
+ return arr
+end
@doc"""
as_type(arr::NDArray, t::Type{T}) where {T}
@@ -143,16 +115,24 @@ end
# 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)
+# NDArray-specific overrides of Core's AbstractArray constructors (NDArray <:
+# AbstractArray): exact `Array{T}` / `Array{T,N}` / `Array` signatures so we win
+# over `Array{T,N}(::AbstractArray)` (which would scalar-index). Bulk path uses
+# `_copy_to_julia_array`; 1-d dispatches same-type (zero-copy) vs convert.
+function (::Type{Array{T}})(arr::NDArray{S,0}) where {T,S}
+ out = Array{T,0}(undef)
allowscalar() do
- return out[] = convert(A, arr[])
+ return out[] = convert(T, 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))
+function (::Type{Array{T}})(arr::NDArray{T,1}) where {T}
+ return make_array(T, Ptr{T}(get_ptr(arr)), size(arr))
+end
+
+function (::Type{Array{T}})(arr::NDArray{S,1}) where {T,S}
+ return T.(make_array(S, Ptr{S}(get_ptr(arr)), size(arr)))
end
# Copy logically into Julia's column-major storage.
@@ -168,14 +148,14 @@ function _copy_to_julia_array(arr::NDArray{T,N}) where {T,N}
return out
end
-function (::Type{<:Array{A}})(arr::NDArray{B}) where {A,B}
+function (::Type{Array{T}})(arr::NDArray{S,N}) where {T,S,N}
out = _copy_to_julia_array(arr)
- return A === B ? out : copyto!(Array{A}(undef, size(arr)), out)
+ return T === S ? out : copyto!(Array{T}(undef, size(arr)), out)
end
-function (::Type{<:Array})(arr::NDArray{B}) where {B}
- return Array{B}(arr)
-end
+(::Type{Array{T,N}})(arr::NDArray{S,N}) where {T,S,N} = Array{T}(arr)
+
+(::Type{Array})(arr::NDArray{T,N}) where {T,N} = Array{T}(arr)
# conversion from Base Julia array to NDArray
# Julia Arrays are column-major; Legate stores are row-major. For N>=2 we
@@ -243,7 +223,7 @@ dim(::NDArray{T,N}) where {T,N} = N::Int
Base.ndims(::NDArray{T,N}) where {T,N} = N::Int
@doc"""
Base.size(arr::NDArray)
- Base.size(arr::NDArray, dim::Int)
+ Base.size(arr::NDArray, dim::Integer)
Return the size of the given `NDArray`.
@@ -258,12 +238,13 @@ 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.size(arr::NDArray, dim::Integer) = dim <= ndims(arr) ? size(arr)[dim] : 1
Base.isempty(arr::NDArray) = any(==(0), size(arr))
+Base.length(arr::NDArray) = prod(size(arr))
@doc"""
- Base.firstindex(arr::NDArray, dim::Int)
- Base.lastindex(arr::NDArray, dim::Int)
+ Base.firstindex(arr::NDArray, dim::Integer)
+ Base.lastindex(arr::NDArray, dim::Integer)
Base.lastindex(arr::NDArray)
Provide the first and last valid indices along a given dimension `dim` for `NDArray`.
@@ -276,44 +257,33 @@ lastindex(arr, 2)
lastindex(arr)
```
"""
-Base.firstindex(arr::NDArray, dim::Int) = 1
-Base.lastindex(arr::NDArray, dim::Int) = Base.size(arr, dim)
-Base.lastindex(arr::NDArray) = Base.size(arr, 1)
+Base.firstindex(arr::NDArray, dim::Integer) = 1
+Base.lastindex(arr::NDArray, dim::Integer) = size(arr, dim)
+Base.lastindex(arr::NDArray) = length(arr)
+Base.IndexStyle(::Type{<:NDArray}) = IndexCartesian()
Base.axes(arr::NDArray) = Base.OneTo.(size(arr))
Base.view(arr::NDArray, inds...) = arr[inds...] # NDArray slices are views by default.
-Base.IndexStyle(::NDArray) = IndexCartesian()
-
function Base.show(io::IO, arr::NDArray{T,0}) where {T}
- allowscalar() do
- return print(io, "NDArray{$(T),0}(", repr(arr[]), ")")
- end
+ print(io, summary(arr), "(")
+ @allowscalar show(io, arr[])
+ return print(io, ")")
end
-function Base.show(io::IO, ::MIME"text/plain", arr::NDArray{T,0}) where {T}
- println(io, "0-dimensional NDArray{$(T),0}")
- allowscalar() do
- return print(io, arr[])
- end
+# Used by print(arr), println(arr), and nested displays
+function Base.show(io::IO, arr::NDArray)
+ return show(io, Array(arr))
end
-function Base.show(io::IO, arr::NDArray{T,N}) where {T,N}
- return print(io, "NDArray{$(T),$(N)} with size ", size(arr))
-end
+# Used for full REPL display
+function Base.show(io::IO, ::MIME"text/plain", arr::NDArray)
+ summary(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))
- return Base.print_array(io, Array(arr))
-end
-
-function Base.print(arr::NDArray{T}) where {T}
- return Base.show(stdout, arr)
-end
+ isempty(arr) && return nothing
-function Base.println(arr::NDArray{T}) where {T}
- Base.show(stdout, arr)
- return print("\n")
+ println(io, ":")
+ return Base.print_array(io, Array(arr))
end
#### ARRAY INDEXING AND SLICES ####
@@ -332,7 +302,7 @@ end
Overloads `Base.getindex` and `Base.setindex!` to support multidimensional indexing and slicing on `cuNumeric.NDArray`s.
-Slicing supports combinations of `Int`, `UnitRange`, and `Colon()` for selecting ranges of rows and columns.
+Slicing supports combinations of `Integer`, `UnitRange`, and `Colon()` for selecting ranges of rows and columns.
The use of all colons (`arr[:]`, `arr[:, :]`, etc.) returns a new Julia `Array` containing a copy of the data.
Assignment also supports:
@@ -349,10 +319,13 @@ Array(A)
```
"""
##### REGULAR ARRAY INDEXING ####
-function Base.getindex(arr::NDArray{T,N}, idxs::Vararg{Int,N}) where {T<:SUPPORTED_NUMERIC_TYPES,N}
+@inline function Base.getindex(
+ arr::NDArray{T,N}, idxs::Vararg{Integer,N}
+) where {T<:SUPPORTED_NUMERIC_TYPES,N}
+ @boundscheck checkbounds(arr, idxs...)
assertscalar("getindex")
acc = NDArrayAccessor{T,N}()
- return read(acc, arr.ptr, to_cpp_index(idxs))
+ return read(acc, arr.ptr, to_cpp_index(Int.(idxs)))
end
function Base.getindex(arr::NDArray{T,0}) where {T<:SUPPORTED_NUMERIC_TYPES}
@@ -362,10 +335,11 @@ function Base.getindex(arr::NDArray{T,0}) where {T<:SUPPORTED_NUMERIC_TYPES}
return read(acc, arr.ptr, zero_index)
end
-function Base.getindex(arr::NDArray{Bool,N}, idxs::Vararg{Int,N}) where {N}
+@inline function Base.getindex(arr::NDArray{Bool,N}, idxs::Vararg{Integer,N}) where {N}
+ @boundscheck checkbounds(arr, idxs...)
assertscalar("getindex")
acc = NDArrayAccessor{CxxWrap.CxxBool,N}()
- return read(acc, arr.ptr, to_cpp_index(idxs))
+ return read(acc, arr.ptr, to_cpp_index(Int.(idxs)))
end
function Base.getindex(arr::NDArray{Bool,0})
@@ -376,17 +350,26 @@ function Base.getindex(arr::NDArray{Bool,0})
end
#! TODO SUPPORT CONVERSION OF VALUES
-function Base.setindex!(arr::NDArray{T,N}, value::T, idxs::Vararg{Int,N}) where {T,N}
+@inline function Base.setindex!(
+ arr::NDArray{T,N}, value::T, idxs::Vararg{Integer,N}
+) where {T,N}
+ @boundscheck checkbounds(arr, idxs...)
assertscalar("setindex!")
return _setindex!(Val{N}(), arr, value, idxs...)
end
-function Base.setindex!(arr::NDArray{Complex{T},N}, value::T, idxs::Vararg{Int,N}) where {T,N}
+@inline function Base.setindex!(
+ arr::NDArray{Complex{T},N}, value::T, idxs::Vararg{Integer,N}
+) where {T,N}
+ @boundscheck checkbounds(arr, idxs...)
assertscalar("setindex!")
return _setindex!(Val{N}(), arr, Complex{T}(value), idxs...)
end
-function Base.setindex!(arr::NDArray{T,N}, value, idxs::Vararg{Int,N}) where {T,N}
+@inline function Base.setindex!(
+ arr::NDArray{T,N}, value, idxs::Vararg{Integer,N}
+) where {T,N}
+ @boundscheck checkbounds(arr, idxs...)
assertscalar("setindex!")
return _setindex!(Val{N}(), arr, convert(T, value), idxs...)
end
@@ -402,15 +385,17 @@ function _setindex!(::Val{0}, arr::NDArray{Bool,0}, value::Bool)
end
function _setindex!(
- ::Val{N}, arr::NDArray{T,N}, value::T, idxs::Vararg{Int,N}
+ ::Val{N}, arr::NDArray{T,N}, value::T, idxs::Vararg{Integer,N}
) where {T<:SUPPORTED_NUMERIC_TYPES,N}
acc = NDArrayAccessor{T,N}()
- return write(acc, arr.ptr, to_cpp_index(idxs), value)
+ return write(acc, arr.ptr, to_cpp_index(Int.(idxs)), value)
end
-function _setindex!(::Val{N}, arr::NDArray{Bool,N}, value::Bool, idxs::Vararg{Int,N}) where {N}
+function _setindex!(
+ ::Val{N}, arr::NDArray{Bool,N}, value::Bool, idxs::Vararg{Integer,N}
+) where {N}
acc = NDArrayAccessor{CxxWrap.CxxBool,N}()
- return write(acc, arr.ptr, to_cpp_index(idxs), value)
+ return write(acc, arr.ptr, to_cpp_index(Int.(idxs)), value)
end
#### START OF SLICING ####
@@ -424,98 +409,183 @@ function _setindex_slice!(lhs::NDArray, rhs::NDArray, slices)
return nothing
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Colon, j::Int64)
- return _setindex_slice!(lhs, rhs, slice_array((0, Base.size(lhs, 1)), (j-1, j)))
+@inline _zero_based_index(i::Integer) = (Int(i) - 1, Int(i))
+@inline _zero_based_range(i::AbstractUnitRange{<:Integer}) = (Int(first(i)) - 1, Int(last(i)))
+
+@inline function Base.setindex!(
+ lhs::NDArray{T,2}, rhs::NDArray, ::Colon, j::Integer
+) where {T}
+ @boundscheck checkbounds(lhs, :, j)
+ return _setindex_slice!(
+ lhs, rhs, slice_array((0, size(lhs, 1)), _zero_based_index(j))
+ )
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Int64, j::Colon)
- return _setindex_slice!(lhs, rhs, slice_array((i-1, i)))
+@inline function Base.setindex!(
+ lhs::NDArray{T,2}, rhs::NDArray, i::Integer, ::Colon
+) where {T}
+ @boundscheck checkbounds(lhs, i, :)
+ return _setindex_slice!(lhs, rhs, slice_array(_zero_based_index(i)))
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::UnitRange, j::Colon)
+@inline function Base.setindex!(
+ lhs::NDArray{T,2}, rhs::NDArray, i::AbstractUnitRange{<:Integer}, ::Colon
+) where {T}
+ @boundscheck checkbounds(lhs, i, :)
return _setindex_slice!(
- lhs, rhs, slice_array((first(i) - 1, last(i)), (0, Base.size(lhs, 2)))
+ lhs, rhs, slice_array(_zero_based_range(i), (0, size(lhs, 2)))
)
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Colon, j::UnitRange)
+@inline function Base.setindex!(
+ lhs::NDArray{T,2}, rhs::NDArray, ::Colon, j::AbstractUnitRange{<:Integer}
+) where {T}
+ @boundscheck checkbounds(lhs, :, j)
return _setindex_slice!(
- lhs, rhs, slice_array((0, Base.size(lhs, 1)), (first(j) - 1, last(j)))
+ lhs, rhs, slice_array((0, size(lhs, 1)), _zero_based_range(j))
)
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::UnitRange, j::Int64)
- return _setindex_slice!(lhs, rhs, slice_array((first(i) - 1, last(i)), (j-1, j)))
+@inline function Base.setindex!(
+ lhs::NDArray{T,2},
+ rhs::NDArray,
+ i::AbstractUnitRange{<:Integer},
+ j::Integer,
+) where {T}
+ @boundscheck checkbounds(lhs, i, j)
+ return _setindex_slice!(
+ lhs, rhs, slice_array(_zero_based_range(i), _zero_based_index(j))
+ )
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::Int64, j::UnitRange)
- return _setindex_slice!(lhs, rhs, slice_array((i-1, i), (first(j) - 1, last(j))))
+@inline function Base.setindex!(
+ lhs::NDArray{T,2},
+ rhs::NDArray,
+ i::Integer,
+ j::AbstractUnitRange{<:Integer},
+) where {T}
+ @boundscheck checkbounds(lhs, i, j)
+ return _setindex_slice!(
+ lhs, rhs, slice_array(_zero_based_index(i), _zero_based_range(j))
+ )
end
-function Base.setindex!(lhs::NDArray, rhs::NDArray, i::UnitRange, j::UnitRange)
+@inline function Base.setindex!(
+ lhs::NDArray{T,2},
+ rhs::NDArray,
+ i::AbstractUnitRange{<:Integer},
+ j::AbstractUnitRange{<:Integer},
+) where {T}
+ @boundscheck checkbounds(lhs, i, j)
return _setindex_slice!(
- lhs, rhs, slice_array((first(i) - 1, last(i)), (first(j) - 1, last(j)))
+ lhs, rhs, slice_array(_zero_based_range(i), _zero_based_range(j))
)
end
-function Base.getindex(arr::NDArray, i::Colon, j::Int64)
- return nda_get_slice(arr, slice_array((0, Base.size(arr, 1)), (j-1, j)))
+@inline function Base.getindex(arr::NDArray{T,2}, ::Colon, j::Integer) where {T}
+ @boundscheck checkbounds(arr, :, j)
+ return nda_get_slice(
+ arr, slice_array((0, size(arr, 1)), _zero_based_index(j))
+ )
end
-function Base.getindex(arr::NDArray, i::Int64, j::Colon)
- return nda_get_slice(arr, slice_array((i-1, i)))
+@inline function Base.getindex(arr::NDArray{T,2}, i::Integer, ::Colon) where {T}
+ @boundscheck checkbounds(arr, i, :)
+ return nda_get_slice(arr, slice_array(_zero_based_index(i)))
end
-function Base.getindex(arr::NDArray, i::UnitRange, j::Colon)
+@inline function Base.getindex(
+ arr::NDArray{T,2}, i::AbstractUnitRange{<:Integer}, ::Colon
+) where {T}
+ @boundscheck checkbounds(arr, i, :)
return nda_get_slice(
- arr, slice_array((first(i) - 1, last(i)), (0, Base.size(arr, 2)))
+ arr, slice_array(_zero_based_range(i), (0, size(arr, 2)))
)
end
-function Base.getindex(arr::NDArray, i::Colon, j::UnitRange)
+@inline function Base.getindex(
+ arr::NDArray{T,2}, ::Colon, j::AbstractUnitRange{<:Integer}
+) where {T}
+ @boundscheck checkbounds(arr, :, j)
return nda_get_slice(
- arr, slice_array((0, Base.size(arr, 1)), (first(j) - 1, last(j)))
+ arr, slice_array((0, size(arr, 1)), _zero_based_range(j))
)
end
-function Base.getindex(arr::NDArray, i::UnitRange, j::Int64)
- return nda_get_slice(arr, slice_array((first(i) - 1, last(i)), (j-1, j)))
-end
-
-function Base.getindex(arr::NDArray, i::Int64, j::UnitRange)
- return nda_get_slice(arr, slice_array((i-1, i), (first(j) - 1, last(j))))
+@inline function Base.getindex(
+ arr::NDArray{T,2}, i::AbstractUnitRange{<:Integer}, j::Integer
+) where {T}
+ @boundscheck checkbounds(arr, i, j)
+ return nda_get_slice(
+ arr, slice_array(_zero_based_range(i), _zero_based_index(j))
+ )
end
-function Base.getindex(arr::NDArray, i::UnitRange, j::UnitRange)
+@inline function Base.getindex(
+ arr::NDArray{T,2}, i::Integer, j::AbstractUnitRange{<:Integer}
+) where {T}
+ @boundscheck checkbounds(arr, i, j)
return nda_get_slice(
- arr, slice_array((first(i) - 1, last(i)), (first(j) - 1, last(j)))
+ arr, slice_array(_zero_based_index(i), _zero_based_range(j))
)
end
-function Base.getindex(arr::NDArray, i::UnitRange)
+@inline function Base.getindex(
+ arr::NDArray{T,2},
+ i::AbstractUnitRange{<:Integer},
+ j::AbstractUnitRange{<:Integer},
+) where {T}
+ @boundscheck checkbounds(arr, i, j)
return nda_get_slice(
- arr, slice_array((first(i) - 1, last(i)))
+ arr, slice_array(_zero_based_range(i), _zero_based_range(j))
)
end
-Base.getindex(arr::NDArray{T}, c::Vararg{Colon,N}) where {T,N} = Base.copy(arr)
-function Base.setindex!(arr::NDArray{T}, rhs::NDArray{T}, c::Vararg{Colon,N}) where {T,N}
+@inline function Base.getindex(
+ arr::NDArray, i::AbstractUnitRange{<:Integer}
+)
+ @boundscheck checkbounds(arr, i)
+ return nda_get_slice(arr, slice_array(_zero_based_range(i)))
+end
+
+@inline function Base.getindex(
+ arr::NDArray{T}, c::Vararg{Colon,N}
+) where {T,N}
+ @boundscheck checkbounds(arr, c...)
+ return Base.copy(arr)
+end
+
+@inline function Base.setindex!(
+ arr::NDArray{T}, rhs::NDArray{T}, c::Vararg{Colon,N}
+) where {T,N}
+ @boundscheck checkbounds(arr, c...)
return Base.copyto!(arr, rhs)
end
-function Base.setindex!(arr::NDArray{T,2}, val::T, i::Colon, j::Int64) where {T}
- s = nda_get_slice(arr, slice_array((0, Base.size(arr, 1)), (j-1, j)))
+@inline function Base.setindex!(
+ arr::NDArray{T,2}, val::T, ::Colon, j::Integer
+) where {T}
+ @boundscheck checkbounds(arr, :, j)
+ s = nda_get_slice(
+ arr, slice_array((0, size(arr, 1)), _zero_based_index(j))
+ )
nda_fill_array(s, val)
return destroy!(s)
end
-function Base.setindex!(arr::NDArray{T,2}, val::T, i::Int64, j::Colon) where {T}
- s = nda_get_slice(arr, slice_array((i-1, i)))
+@inline function Base.setindex!(
+ arr::NDArray{T,2}, val::T, i::Integer, ::Colon
+) where {T}
+ @boundscheck checkbounds(arr, i, :)
+ s = nda_get_slice(arr, slice_array(_zero_based_index(i)))
nda_fill_array(s, val)
return destroy!(s)
end
-Base.fill!(arr::NDArray{T}, val::T) where {T} = nda_fill_array(arr, val)
+@inline function Base.fill!(arr::NDArray{T}, val::T) where {T}
+ nda_fill_array(arr, val)
+ return arr
+end
#### INITIALIZATION OF NDARRAYS ####
@doc"""
@@ -692,8 +762,16 @@ function reshape(arr::NDArray, i::Int...; copy::Val{C}=Val(false)) where {C}
end
# Ignore the scalar indexing here...
-unwrap(x::NDArray{<:Any,0}) = @allowscalar x[]
-unwrap(x::NDArray{<:Any,1}) = @allowscalar x[][1] # assumes 1 element
+Base.only(x::NDArray{T,0}) where {T} = @allowscalar x[]
+
+function Base.only(x::NDArray{T,N}) where {T,N}
+ length(x) == 1 ||
+ throw(ArgumentError("collection must contain exactly 1 element"))
+
+ return @allowscalar x[firstindex(x)]
+end
+
+unwrap(x::NDArray) = only(x)
@doc"""
==(arr1::NDArray, arr2::NDArray)
diff --git a/src/ndarray/unary.jl b/src/ndarray/unary.jl
index 0c0df9ca..743ccbb6 100644
--- a/src/ndarray/unary.jl
+++ b/src/ndarray/unary.jl
@@ -349,6 +349,17 @@ function Base.any(input::NDArray{Bool}; dims=Colon())
return _bool_reduction_impl(cuNumeric.ANY, input, dims)
end
+# Compare on-device against `zero(T)` / `_eye(T, n)` (identity filled with `one(T)`).
+# Returns a 0D `NDArray{Bool}` — not a Julia `Bool`.
+function Base.iszero(A::NDArray{T}) where {T}
+ return all(A .== zero(T))
+end
+function Base.isone(A::NDArray{T,2}) where {T}
+ m, n = size(A)
+ m != n && return NDArray(false) # LinearAlgebra.isone: only square matrices
+ return all(A .== _eye(T, m))
+end
+
# Boolean multiplication is logical conjunction. cuPyNumeric's PROD reduction
# uses a numeric fill identity, which Legate rejects for a Boolean target.
function Base.prod(input::NDArray{Bool}; dims=Colon())
diff --git a/src/warnings.jl b/src/warnings.jl
index db0ae37f..c2cfc99b 100644
--- a/src/warnings.jl
+++ b/src/warnings.jl
@@ -13,7 +13,7 @@ function repl_frontend_task()
if !isassigned(_repl_frontend_task)
_repl_frontend_task[] = get_repl_frontend_task()
end
- _repl_frontend_task[]
+ return _repl_frontend_task[]
end
@noinline function get_repl_frontend_task()
if isdefined(Base, :active_repl)
@@ -69,7 +69,7 @@ function assertscalar(op::String)
return nothing
end
- _assertscalar(op, behavior)
+ return _assertscalar(op, behavior)
end
"""
@@ -94,7 +94,7 @@ function assertpromotion(op, ::Type{FROM}, ::Type{TO}) where {FROM,TO}
return nothing
end
- _assertpromotion(op, behavior, FROM, TO)
+ return _assertpromotion(op, behavior, FROM, TO)
end
@noinline function _assertscalar(op, behavior)
@@ -119,26 +119,193 @@ end
return nothing
end
+const _CUNUMERIC_MODULE = @__MODULE__
+
+# Sentinel for stack frames with no recoverable module. Prefer this over `nothing`
+# so `_module_of_stackframe` is type-stable as `Module`. Must not match Base /
+# LinearAlgebra / cuNumeric checks below.
+const _UNKNOWN_STACK_MODULE = Module(:__cuNumeric_unknown_stack_module__, false, false)
+
+@inline function _is_cunumeric_module(m::Module)
+ m === _UNKNOWN_STACK_MODULE && return false
+ m === _CUNUMERIC_MODULE && return true
+ pm = parentmodule(m)
+ while pm !== m
+ pm === _CUNUMERIC_MODULE && return true
+ m = pm
+ pm = parentmodule(m)
+ end
+ return false
+end
+
+@inline function _is_linalg_module(m::Module)
+ m === _UNKNOWN_STACK_MODULE && return false
+ m === LinearAlgebra && return true
+ nameof(m) === :LinearAlgebra && return true
+ pm = parentmodule(m)
+ while pm !== m
+ (pm === LinearAlgebra || nameof(pm) === :LinearAlgebra) && return true
+ m = pm
+ pm = parentmodule(m)
+ end
+ return false
+end
+
+# Note: LinearAlgebra (and other stdlibs) often have parentmodule === Base, so callers
+# must check `_is_linalg_module` before treating a frame as Base.
+@inline function _is_base_module(m::Module)
+ m === _UNKNOWN_STACK_MODULE && return false
+ m === Base && return true
+ nameof(m) === :Base && return true
+ pm = parentmodule(m)
+ while pm !== m
+ (pm === Base || nameof(pm) === :Base) && return true
+ m = pm
+ pm = parentmodule(m)
+ end
+ return false
+end
+
+_module_from_def(def::Method) = def.module
+_module_from_def(def::Module) = def
+_module_from_def(_) = _UNKNOWN_STACK_MODULE
+
+_module_of_linfo(linfo::Core.MethodInstance) = _module_from_def(linfo.def)
+_module_of_linfo(linfo::Method) = linfo.module
+# Julia 1.12+ often stores a CodeInstance on stack frames; unwrap to MethodInstance.
+_module_of_linfo(linfo::Core.CodeInstance) = _module_of_linfo(linfo.def)
+_module_of_linfo(_) = _UNKNOWN_STACK_MODULE
+
+# When `frame.linfo` is missing (common for inlined frames), recover module from
+# the source path so LinearAlgebra callers are not skipped and the walk does not
+# fall through to loader frames like `Base.include_string`.
+function _module_from_file(file)
+ (file === nothing || file === :none) && return _UNKNOWN_STACK_MODULE
+ f = string(file)
+ # Match stdlib path segments; avoid false positives on user paths when possible.
+ if occursin(r"(?:^|[/\\])LinearAlgebra(?:[/\\]|$)", f)
+ return LinearAlgebra
+ end
+ if occursin(r"(?:^|[/\\])cuNumeric(?:\.jl)?(?:[/\\]|$)", f)
+ return _CUNUMERIC_MODULE
+ end
+ # Base frames commonly appear as `./abstractarray.jl`, `./set.jl`, etc.
+ if startswith(f, "./") || occursin(r"(?:^|[/\\])[Bb]ase(?:[/\\]|$)", f)
+ return Base
+ end
+ return _UNKNOWN_STACK_MODULE
+end
+
+function _module_of_stackframe(frame::Base.StackTraces.StackFrame)
+ m = _module_of_linfo(frame.linfo)
+ m !== _UNKNOWN_STACK_MODULE && return m
+ return _module_from_file(frame.file)
+end
+
+# Keyword bodies often look like `#cholesky!#272`; surface `cholesky!`.
+function _clean_stack_func_name(fname)
+ fname_sym = ifelse(fname isa Symbol, fname, Symbol(string(fname)))
+ s = string(fname_sym)
+ m = match(r"^#([^#]+)#\d+$", s)
+ return m === nothing ? fname_sym : Symbol(m.captures[1])
+end
+
+# Frames that are never the user-facing "triggering" API for enrichment:
+# - loaders / client entry (`include_string` / Julia 1.12 `IncludeInto` from
+# `include`ing tests, etc.)
+# - keyword-call wrappers (`kwcall`) that would otherwise outrank cholesky/svd
+# - AbstractArray iteration/indexing plumbing between NDArray getindex and the
+# real stdlib caller (e.g. LinearAlgebra.cholesky / Base.unique)
+const _SKIP_STACK_FUNCS = Set{Symbol}((
+ :include_string,
+ :include,
+ :include_relative,
+ :_include,
+ :IncludeInto, # Julia 1.12+ callable include wrapper (Base.IncludeInto)
+ :eval,
+ :exec_options,
+ :_start,
+ :invokelatest,
+ :error,
+ :stacktrace,
+ :kwcall, # Base keyword-call wrapper; do not steal blame from cholesky/svd
+ :iterate,
+ :getindex,
+ :setindex!,
+ :indexed_iterate,
+ Symbol("macro expansion"),
+ Symbol("top-level scope"),
+))
+
+"""
+Best-effort: outermost Base or LinearAlgebra frame above cuNumeric scalar-index
+frames. Walk innermost-first, skip cuNumeric/Core (and frames with unknown
+module), indexing/iteration plumbing, keyword-call wrappers (`kwcall`), and
+Base loader frames (`include_string`, `IncludeInto` on Julia 1.12+, etc.).
+Keep updating the candidate while still in Base/LinearAlgebra (last one wins)
+so attribution names the user-facing API (`LinearAlgebra.svd`) rather than an
+inner helper (`Base.lt`) or a loader (`Base.include_string`). Stop at the first
+user/other-package frame and return that candidate, or `nothing` for the plain
+message (e.g. user `Main`). Check LinearAlgebra before Base — stdlibs often
+parent to Base.
+"""
+@noinline function _scalar_indexing_stdlib_caller()
+ caller = nothing
+ for frame in stacktrace()
+ m = _module_of_stackframe(frame)
+ # Skip frames with unknown module (same as previous `nothing` skip).
+ m === _UNKNOWN_STACK_MODULE && continue
+ (m === Core || _is_cunumeric_module(m)) && continue
+ clean_name = _clean_stack_func_name(frame.func)
+ clean_name in _SKIP_STACK_FUNCS && continue
+ if _is_linalg_module(m)
+ caller = (:LinearAlgebra, clean_name)
+ elseif _is_base_module(m)
+ caller = (:Base, clean_name)
+ else
+ # User / other package code — keep the last stdlib candidate, if any.
+ break
+ end
+ end
+ return caller
+end
+
+# Returns (enriched::Bool, desc::String). Enriched = Base or LinearAlgebra stdlib caller.
function scalardesc(op)
- desc = """Invocation of $op resulted in scalar indexing of an NDArray.
+ caller = _scalar_indexing_stdlib_caller()
+ if caller !== nothing
+ modname, fname = caller
+ # Base/LinearAlgebra AbstractArray fallback — name the outer API first.
+ # No "Scalar indexing is disallowed." header (not part of the enriched template).
+ return true,
+ "`$modname.$fname` fell back to an AbstractArray implementation, which scalar-indexed an `NDArray`. " *
+ "This $modname path is probably not implemented yet for `NDArray`. " *
+ "Using `allowscalar` or `@allowscalar` might allow this function to work slowly, but it has not been tested."
+ end
+
+ # Plain user-level scalar indexing (unchanged).
+ return false, """Invocation of $op resulted in scalar indexing of an `NDArray`.
This is typically caused by calling an iterating implementation of a method.
- This is very slow and should be avoided.
+ This is very slow and should be avoided. This can also happen if an external
+ method (i.e., LinearAlgebra.kron) is not re-implemented in cuNumeric.jl. Because
+ `NDArray`s subtype `AbstractArray`, the method call will dispatch to the
+ `AbstractArray` implementation, which often iterates over the array.
If you want to allow scalar iteration, use `allowscalar` or `@allowscalar`
to enable scalar iteration globally or for the operations in question."""
end
function promotiondesc(op, ::Type{FROM}, ::Type{TO}) where {FROM,TO}
- desc = """Invocation of $op resulted in implicit promotion of an NDArray from $(FROM) to
- wider type: $(TO). This is typically caused by mixing NDArrays or literals
- with different precision. This can cause extra copies of data and is slow.
+ return desc = """Invocation of $op resulted in implicit promotion of an NDArray from $(FROM) to
+ wider type: $(TO). This is typically caused by mixing NDArrays or literals
+ with different precision. This can cause extra copies of data and is slow.
- If you want to allow implicit promotion to wider types, use `allowpromotion` or `@allowpromotion`
- to enable implicit promotion."""
+ If you want to allow implicit promotion to wider types, use `allowpromotion` or `@allowpromotion`
+ to enable implicit promotion."""
end
@noinline function warnscalar(op)
- desc = scalardesc(op)
+ _, desc = scalardesc(op)
@warn("""Performing scalar indexing on task $(current_task()).
$desc""")
end
@@ -150,9 +317,14 @@ end
end
@noinline function errorscalar(op)
- desc = scalardesc(op)
- error("""Scalar indexing is disallowed.
- $desc""")
+ enriched, desc = scalardesc(op)
+ if enriched
+ error(desc)
+ else
+ # Plain path keeps the historical disallow header.
+ error("""Scalar indexing is disallowed.
+ $desc""")
+ end
end
@noinline function errordouble(op, ::Type{FROM}, ::Type{TO}) where {FROM,TO}
@@ -165,7 +337,7 @@ end
# NOTE: This is deprecated and should not be used from user logic. A proper solution to
# this problem will be introduced in https://github.com/JuliaLang/julia/pull/39217
macro __tryfinally(ex, fin)
- Expr(:tryfinally,
+ return Expr(:tryfinally,
:($(esc(ex))),
:($(esc(fin))),
)
@@ -185,7 +357,7 @@ See also: [`@allowscalar`](@ref).
allowscalar
function allowscalar(f::Base.Callable)
- task_local_storage(f, :ScalarIndexing, ScalarAllowed)
+ return task_local_storage(f, :ScalarIndexing, ScalarAllowed)
end
function allowscalar(allow::Bool=true)
diff --git a/test/analysis/type_stability.jl b/test/analysis/type_stability.jl
index db527d19..add8e8a0 100644
--- a/test/analysis/type_stability.jl
+++ b/test/analysis/type_stability.jl
@@ -158,7 +158,7 @@ end
M = cuNumeric.zeros(T, 4, 3)
sq = cuNumeric.zeros(T, 5, 5)
v = cuNumeric.zeros(T, 8)
- @test @inferred(cuNumeric.eye(T, 5)) !== nothing
+ @test @inferred(NDArray{T}(I, 5, 5)) !== nothing
@test @inferred(cuNumeric.transpose(M)) !== nothing
@test @inferred(cuNumeric.trace(sq)) !== nothing
@test @inferred(cuNumeric.diag(sq)) !== nothing
diff --git a/test/array/diagonal.jl b/test/array/diagonal.jl
new file mode 100644
index 00000000..a04e1629
--- /dev/null
+++ b/test/array/diagonal.jl
@@ -0,0 +1,811 @@
+#= Copyright 2026 Northwestern University,
+ * Carnegie Mellon University University
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ * Author(s): David Krasowska
+ * Ethan Meitz
+=#
+
+# Coverage for src/ndarray/diagonal.jl: diag/_eye/trace, Diagonal, UniformScaling.
+
+const DIAGONAL_NUMERIC_TYPES = Base.uniontypes(cuNumeric.SUPPORTED_NUMERIC_TYPES)
+const DIAGONAL_ARRAY_TYPES = Base.uniontypes(cuNumeric.SUPPORTED_ARRAY_TYPES)
+const DIAGONAL_FLOAT_TYPES = Base.uniontypes(cuNumeric.SUPPORTED_FLOAT_TYPES)
+const DIAGONAL_COMPLEX_TYPES = Base.uniontypes(cuNumeric.SUPPORTED_COMPLEX_TYPES)
+
+_nonzero_diag(::Type{T}, n) where {T<:AbstractFloat} = abs.(my_rand(T, n)) .+ one(T)
+function _nonzero_diag(::Type{T}, n) where {T<:Complex}
+ return Complex.(abs.(real(my_rand(T, n))) .+ one(real(T)), zero(real(T)))
+end
+_nonzero_diag(::Type{T}, n) where {T<:Integer} = T.(collect(2:(n + 1)))
+_nonzero_diag(::Type{Bool}, n) = fill(true, n)
+
+function _host_diag_compare(ref, out, ::Type{T}) where {T}
+ allowscalar() do
+ @test cuNumeric.compare(ref, out, atol(T), rtol(T))
+ end
+end
+
+function _host_scalar_compare(ref, out::NDArray{<:Any,0}, ::Type{T}) where {T}
+ allowscalar() do
+ @test ref ≈ out[] atol=atol(T) rtol=rtol(T)
+ end
+end
+
+function _host_bool_compare(ref::Bool, out::NDArray{Bool,0})
+ allowscalar() do
+ @test out[] == ref
+ end
+end
+
+function _host_matrix_compare(ref::AbstractMatrix, out::NDArray, ::Type{T}) where {T}
+ allowscalar() do
+ @test cuNumeric.compare(ref, out, atol(T), rtol(T))
+ end
+end
+
+###### diag / identity / trace ######
+
+@testset "diag" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ A = my_rand(T, 6, 5)
+ nda = NDArray(A)
+ @testset "k=$k" for k in (-2, 0, 2)
+ ref = diag(A, k)
+ _host_diag_compare(ref, cuNumeric.diag(nda; k=k), T)
+ _host_diag_compare(ref, LinearAlgebra.diag(nda, k), T)
+ end
+ end
+end
+
+@testset "identity via I / _eye" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ n = 4
+ ref = Matrix{T}(I, n, n)
+ _host_matrix_compare(ref, NDArray{T}(I, n, n), T)
+ _host_matrix_compare(ref, cuNumeric._eye(T, n), T)
+ end
+ # Untyped NDArray(I, ...) uses Bool; typed / _eye default to Float32 densify
+ ref_bool = Matrix{Bool}(I, 3, 3)
+ _host_matrix_compare(ref_bool, NDArray(I, 3, 3), Bool)
+ ref = Matrix{Float32}(I, 3, 3)
+ _host_matrix_compare(ref, NDArray{Float32}(I, 3, 3), Float32)
+ _host_matrix_compare(ref, cuNumeric._eye(3), Float32)
+end
+
+@testset "trace" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ A = my_rand(T, 5, 5)
+ nda = NDArray(A)
+ @testset "offset=$k" for k in (-2, -1, 0, 1, 2)
+ ref = sum(diag(A, k))
+ out = cuNumeric.trace(nda; offset=k)
+ @test out isa NDArray{<:Any,0}
+ _host_scalar_compare(ref, out, eltype(ref))
+ end
+ end
+end
+
+###### Diagonal constructors / densify / show ######
+
+@testset "Diagonal constructors" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ d = my_rand(T, 4)
+ v = NDArray(d)
+ D = Diagonal(v)
+ @test D isa Diagonal{T,<:NDArray{T,1}}
+ @test D.diag === v
+ allowscalar() do
+ @test cuNumeric.compare(d, D.diag, atol(T), rtol(T))
+ @test Matrix(D) ≈ Matrix(Diagonal(d)) atol=atol(T) rtol=rtol(T)
+ @test Matrix{T}(D) ≈ Matrix(Diagonal(d)) atol=atol(T) rtol=rtol(T)
+ end
+
+ A = my_rand(T, 4, 4)
+ D2 = Diagonal(NDArray(A))
+ allowscalar() do
+ @test cuNumeric.compare(diag(A), D2.diag, atol(T), rtol(T))
+ end
+ end
+end
+
+@testset "Diagonal show" begin
+ D = Diagonal(NDArray(Float32[1, 2, 3]))
+ s = sprint(show, D)
+ @test occursin("1.0", s) && occursin("2.0", s)
+ plain = sprint(show, MIME"text/plain"(), D)
+ # Summary must reflect the real Diagonal{T,<:NDArray} type, not Vector.
+ @test occursin("NDArray", plain)
+ @test occursin(string(typeof(D)), plain)
+ @test occursin("1.0", plain)
+ # Match Base Diagonal formatting (⋅ off-diagonals), not a dense Matrix dump.
+ @test occursin("⋅", plain)
+ dense = sprint(show, MIME"text/plain"(), Matrix(D))
+ @test plain != dense
+end
+
+###### Diagonal operators ######
+
+@testset "Diagonal * Diagonal / ± / scalar" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ d = my_rand(T, 3)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ allowscalar() do
+ @test Matrix(D * D) ≈ Matrix(Dh * Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D + D) ≈ Matrix(Dh + Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D - D) ≈ Matrix(Dh - Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(T(3) * D) ≈ Matrix(T(3) * Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D * T(3)) ≈ Matrix(Dh * T(3)) atol=atol(T) rtol=rtol(T)
+ end
+ end
+end
+
+@testset "Diagonal broadcast on diag" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ d = my_rand(T, 5)
+ c = T(5)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(copy(d)))
+
+ ref = collect((Dh .* c).diag)
+
+ # Out-of-place: stays Diagonal with NDArray diag
+ D2 = D .* c
+ @test D2 isa Diagonal{T,<:NDArray{T,1}}
+ @test D2.diag isa NDArray{T,1}
+ _host_diag_compare(ref, D2.diag, T)
+
+ # In-place fused assign
+ D3 = Diagonal(NDArray(copy(d)))
+ D3 .= D3 .* c
+ @test D3 isa Diagonal{T,<:NDArray{T,1}}
+ _host_diag_compare(ref, D3.diag, T)
+
+ # In-place .*=
+ D4 = Diagonal(NDArray(copy(d)))
+ D4 .*= c
+ @test D4 isa Diagonal{T,<:NDArray{T,1}}
+ _host_diag_compare(ref, D4.diag, T)
+
+ # Zero-preserving Diagonal .+ Diagonal (matches Base structure)
+ ref_add = collect((Dh .+ Dh).diag)
+ D_add = D .+ D
+ @test D_add isa Diagonal{T,<:NDArray{T,1}}
+ _host_diag_compare(ref_add, D_add.diag, T)
+
+ D_add2 = Diagonal(NDArray(copy(d)))
+ D_add2 .+= D_add2
+ @test D_add2 isa Diagonal{T,<:NDArray{T,1}}
+ _host_diag_compare(ref_add, D_add2.diag, T)
+
+ D_add3 = Diagonal(NDArray(copy(d)))
+ D_add3 .= D_add3 .+ D_add3
+ @test D_add3 isa Diagonal{T,<:NDArray{T,1}}
+ _host_diag_compare(ref_add, D_add3.diag, T)
+ end
+
+ # Exact repro from the bug report
+ a = Diagonal(cuNumeric.ones(Int32, 5))
+ a .*= Int32(5)
+ @test a isa Diagonal{Int32,<:NDArray{Int32,1}}
+ @test Array(a.diag) == fill(Int32(5), 5)
+end
+
+@testset "scalar indexing message: LinearAlgebra / Base vs plain" begin
+ # Unsupported LA/Base fallbacks enrich; intentional user scalar indexing stays plain.
+ allowscalar(false)
+
+ function _assert_enriched_fallback(msg, modfunc)
+ @test occursin("`$modfunc` fell back to an AbstractArray implementation", msg)
+ @test occursin("which scalar-indexed an `NDArray`", msg)
+ @test occursin("path is probably not implemented yet for `NDArray`", msg)
+ @test occursin("allowscalar", msg)
+ @test occursin("@allowscalar", msg)
+ @test occursin("might allow this function to work slowly", msg)
+ @test occursin("it has not been tested", msg)
+ # Enriched path replaces the generic iterating-method lead and omits the plain header.
+ @test !occursin("typically caused by calling an iterating implementation", msg)
+ @test !occursin("Scalar indexing is disallowed", msg)
+ @test !occursin("If you want to allow scalar iteration", msg)
+ @test !occursin("triggered via", msg)
+ @test startswith(lstrip(msg), "`$modfunc`")
+ end
+
+ err = @test_throws ErrorException cholesky(Diagonal(NDArray(Float32[2, 3, 4])))
+ msg = sprint(showerror, err.value)
+ _assert_enriched_fallback(msg, "LinearAlgebra.cholesky")
+
+ # sortperm-based LA path must blame svd, not the inner Base.lt helper.
+ err_svd = @test_throws ErrorException svd(Diagonal(NDArray(Float32[2, 3, 4])))
+ msg_svd = sprint(showerror, err_svd.value)
+ _assert_enriched_fallback(msg_svd, "LinearAlgebra.svd")
+ @test !occursin("Base.lt", msg_svd)
+
+ # Base AbstractArray fallback (unique) should enrich with Base..
+ err_base = @test_throws ErrorException unique(NDArray(Float32[1, 2, 1]))
+ msg_base = sprint(showerror, err_base.value)
+ _assert_enriched_fallback(msg_base, "Base.unique")
+
+ # Call through a Main function so the first non-cuNumeric/Core frame is user code,
+ # not Base.include_string / Base.IncludeInto (Julia 1.12+) / client frames from
+ # `include`ing this test file.
+ function _plain_scalar_index_probe()
+ a = NDArray(Float32[1, 2, 3])
+ return a[1]
+ end
+ err_plain = try
+ _plain_scalar_index_probe()
+ nothing
+ catch e
+ e
+ end
+ @test err_plain isa ErrorException
+ msg_plain = sprint(showerror, err_plain)
+ @test occursin("Scalar indexing is disallowed", msg_plain)
+ @test occursin("typically caused by calling an iterating implementation", msg_plain)
+ @test occursin("If you want to allow scalar iteration", msg_plain)
+ @test !occursin("triggered via", msg_plain)
+ @test !occursin("fell back to an AbstractArray", msg_plain)
+ @test !occursin("probably not implemented yet", msg_plain)
+ @test !occursin("might allow this function to work slowly", msg_plain)
+ @test !occursin("it has not been tested", msg_plain)
+end
+
+@testset "NDArray iszero / isone" begin
+ z = cuNumeric.zeros(Float32, 3)
+ o = cuNumeric.ones(Float32, 3)
+ _host_bool_compare(true, iszero(z))
+ _host_bool_compare(false, iszero(o))
+ I = one(cuNumeric.zeros(Float32, 3, 3))
+ _host_bool_compare(true, isone(I))
+ _host_bool_compare(false, isone(cuNumeric.ones(Float32, 3, 3)))
+ _host_bool_compare(false, isone(cuNumeric.zeros(Float32, 3, 3)))
+ _host_bool_compare(false, isone(cuNumeric.ones(Float32, 2, 3)))
+end
+
+@testset "Diagonal densifying broadcast" begin
+ D = Diagonal(NDArray(Float32[1, 2, 3]))
+ expected = Float32[2 1 1; 1 3 1; 1 1 4]
+
+ # Out-of-place densifies to dense NDArray (Base densifies to Matrix).
+ R = D .+ Float32(1)
+ @test R isa NDArray{Float32,2}
+ @test Array(R) == expected
+
+ A = cuNumeric.ones(Float32, 3, 3)
+ R2 = D .+ A
+ @test R2 isa NDArray{Float32,2}
+ @test Array(R2) == expected
+ R3 = A .+ D
+ @test R3 isa NDArray{Float32,2}
+ @test Array(R3) == expected
+
+ # In-place still rejects off-diagonal writes (Base-style ArgumentError).
+ err_inp = @test_throws ArgumentError D .+= Float32(1)
+ @test occursin("off-diagonal", sprint(showerror, err_inp.value))
+
+ D_inp = Diagonal(NDArray(Float32[1, 2, 3]))
+ err_mat_inp = @test_throws ArgumentError D_inp .+= A
+ @test occursin("off-diagonal", sprint(showerror, err_mat_inp.value))
+
+ # Structure-preserving scale still stays Diagonal without @allowscalar.
+ D2 = Diagonal(NDArray(Float32[1, 2, 3]))
+ D2 .*= Float32(4)
+ @test Array(D2.diag) == Float32[4, 8, 12]
+ D3 = D2 .* Float32(2)
+ @test D3 isa Diagonal{Float32,<:NDArray{Float32,1}}
+ @test Array(D3.diag) == Float32[8, 16, 24]
+
+ # 1×1 has no off-diagonals: Base allows in-place densifying-classified ops.
+ D1 = Diagonal(NDArray(Float32[5]))
+ D1 .+= Float32(1)
+ _host_diag_compare(Float32[6], D1.diag, Float32)
+ D1 .*= Float32(2)
+ _host_diag_compare(Float32[12], D1.diag, Float32)
+
+ # 1×1 out-of-place densifies to NDArray (Base returns Matrix).
+ R1 = Diagonal(NDArray(Float32[5])) .+ Float32(1)
+ @test R1 isa NDArray{Float32,2}
+ @test size(R1) == (1, 1)
+ _host_matrix_compare(Float32[6;;], R1, Float32)
+end
+
+@testset "Diagonal * NDArray / mul! / lmul! / rmul!" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ n, m = 3, 4
+ d = my_rand(T, n)
+ dm = my_rand(T, m)
+ Ah = my_rand(T, n, m) # n×m
+ As = my_rand(T, n, n) # n×n
+ vh = my_rand(T, n)
+ Dh = Diagonal(d)
+ Dm = Diagonal(dm)
+ D = Diagonal(NDArray(d))
+ D_m = Diagonal(NDArray(dm))
+ A = NDArray(Ah)
+ A_nd = NDArray(As)
+ v = NDArray(vh)
+
+ _host_diag_compare(Dh * vh, D * v, T)
+ _host_matrix_compare(Dh * Ah, D * A, T) # D (n)× A (n×m)
+ _host_matrix_compare(Ah * Dm, A * D_m, T) # A (n×m) × D (m)
+ _host_matrix_compare(As * Dh, A_nd * D, T) # square A*D
+
+ C = cuNumeric.zeros(T, n, m)
+ mul!(C, D, A)
+ _host_matrix_compare(Dh * Ah, C, T)
+
+ Cs = cuNumeric.zeros(T, n, n)
+ mul!(Cs, A_nd, D)
+ _host_matrix_compare(As * Dh, Cs, T)
+
+ B = copy(A_nd)
+ lmul!(D, B)
+ _host_matrix_compare(Dh * As, B, T)
+
+ B = copy(A_nd)
+ rmul!(B, D)
+ _host_matrix_compare(As * Dh, B, T)
+ end
+end
+
+@testset "Diagonal \\ / / inv / ldiv! / rdiv!" begin
+ # Floats: full path including singular CUDA-style behavior
+ @testset verbose=true for T in DIAGONAL_FLOAT_TYPES
+ n = 3
+ d = _nonzero_diag(T, n)
+ Ah = my_rand(T, n, n)
+ vh = my_rand(T, n)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ A = NDArray(Ah)
+ v = NDArray(vh)
+
+ _host_diag_compare(Dh \ vh, D \ v, T)
+ _host_matrix_compare(Dh \ Ah, D \ A, T)
+ _host_matrix_compare(Ah / Dh, A / D, T)
+ allowscalar() do
+ @test Matrix(inv(D)) ≈ Matrix(inv(Dh)) atol=atol(T) rtol=rtol(T)
+ end
+
+ B = copy(A)
+ ldiv!(D, B)
+ _host_matrix_compare(Dh \ Ah, B, T)
+
+ B = copy(A)
+ rdiv!(B, D)
+ _host_matrix_compare(Ah / Dh, B, T)
+
+ # Singular: zeros become Inf/NaN (no SingularException; that needs a host Bool)
+ d0 = copy(d)
+ d0[2] = zero(T)
+ D0 = Diagonal(NDArray(d0))
+ r = D0 \ v
+ Di = inv(D0)
+ Q = A / D0
+ allowscalar() do
+ @test any(isinf, Array(r)) || any(isnan, Array(r))
+ @test any(isinf, Array(Di.diag)) || any(isnan, Array(Di.diag))
+ @test any(isinf, Array(Q)) || any(isnan, Array(Q))
+ end
+ end
+
+ # Complex: \ works via ./ ; inv/A/D blocked by missing __recip_type
+ @testset verbose=true for T in DIAGONAL_COMPLEX_TYPES
+ n = 3
+ d = _nonzero_diag(T, n)
+ Ah = my_rand(T, n, n)
+ vh = my_rand(T, n)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ A = NDArray(Ah)
+ v = NDArray(vh)
+ _host_diag_compare(Dh \ vh, D \ v, T)
+ _host_matrix_compare(Dh \ Ah, D \ A, T)
+ end
+
+ # Integers with explicit allowpromotion for inv / \
+ # `\` uses ./ → typically float(T) (Float64 for Int32/Int64).
+ # `inv` / `A / D` use cuNumeric.__recip_type (Float32 for Int32, Float64 for Int64),
+ # which is intentional NDArray promotion — not Base's float(Int32)==Float64.
+ @testset verbose=true for T in (Int32, Int64)
+ n = 3
+ d = _nonzero_diag(T, n)
+ Ah = my_rand(T, n, n)
+ vh = my_rand(T, n)
+ FT = float(T)
+ RT = cuNumeric.__recip_type(T)
+ Dh_div = Diagonal(FT.(d))
+ Dh_inv = Diagonal(RT.(d))
+ D = Diagonal(NDArray(d))
+ A = NDArray(Ah)
+ v = NDArray(vh)
+ allowpromotion() do
+ _host_diag_compare(Dh_div \ FT.(vh), D \ v, FT)
+ _host_matrix_compare(Dh_div \ FT.(Ah), D \ A, FT)
+ allowscalar() do
+ @test Matrix(inv(D)) ≈ Matrix(inv(Dh_inv)) atol=atol(RT) rtol=rtol(RT)
+ end
+ _host_matrix_compare(RT.(Ah) / Dh_inv, A / D, RT)
+ end
+ end
+end
+
+@testset "Diagonal det" begin
+ # Floats: det is prod of the diagonal (0D NDArray, not a Julia scalar)
+ @testset verbose=true for T in DIAGONAL_FLOAT_TYPES
+ d = my_rand(T, 3)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ @test det(D) isa NDArray{<:Any,0}
+ _host_scalar_compare(det(Dh), det(D), T)
+ _host_scalar_compare(det(Matrix(D)), det(D), T)
+
+ # 1×1
+ d1 = T[T(5)]
+ D1 = Diagonal(NDArray(d1))
+ _host_scalar_compare(det(Diagonal(d1)), det(D1), T)
+ end
+
+ # Integers: prod may widen (e.g. Int32 → Int64); needs allowpromotion
+ @testset verbose=true for T in (Int32, Int64)
+ d = _nonzero_diag(T, 3)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ allowpromotion() do
+ _host_scalar_compare(det(Dh), det(D), T)
+ _host_scalar_compare(det(Matrix(D)), det(D), T)
+ end
+ # 1×1
+ D1 = Diagonal(NDArray(T[7]))
+ allowpromotion() do
+ _host_scalar_compare(T(7), det(D1), T)
+ end
+ end
+end
+
+@testset "NDArray ± Diagonal" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ n = 3
+ d = my_rand(T, n)
+ Ah = my_rand(T, n, n)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ A = NDArray(Ah)
+ _host_matrix_compare(Ah + Dh, A + D, T)
+ _host_matrix_compare(Dh + Ah, D + A, T)
+ _host_matrix_compare(Ah - Dh, A - D, T)
+ _host_matrix_compare(Dh - Ah, D - A, T)
+ end
+
+ # Bool promotes under +
+ @testset "Bool with allowpromotion" begin
+ Ah = Bool[1 0; 0 1]
+ d = Bool[true, true]
+ A = NDArray(Ah)
+ D = Diagonal(NDArray(d))
+ allowpromotion() do
+ allowscalar() do
+ @test Array(A + D) == Ah + Diagonal(d)
+ @test Array(A - D) == Ah - Diagonal(d)
+ end
+ end
+ end
+end
+
+###### UniformScaling ######
+
+# Prefer typed UniformScaling (one(T)*I) so narrow integers do not promote against
+# Int64 λ from 2I / -I. Plain I (Bool λ) is fine for +/* on numeric arrays.
+
+@testset "UniformScaling constructors / copyto!" begin
+ @testset verbose=true for T in DIAGONAL_ARRAY_TYPES
+ n = 3
+ J1 = one(T) * I
+ ref = Matrix{T}(J1, n, n)
+ E = NDArray{T}(J1, n, n)
+ _host_matrix_compare(ref, E, T)
+ E2 = NDArray{T}(I, (n, n)) # Bool λ → ones on diagonal still
+ _host_matrix_compare(Matrix{T}(I, n, n), E2, T)
+
+ C = cuNumeric.zeros(T, n, n)
+ copyto!(C, J1)
+ _host_matrix_compare(ref, C, T)
+
+ C = cuNumeric.ones(T, n, n)
+ copyto!(C, zero(T) * I)
+ _host_matrix_compare(zeros(T, n, n), C, T)
+
+ # Rectangular scaled identity (skip Bool: Bool(2) is inexact)
+ if T != Bool
+ J2 = T(2) * I
+ R = NDArray{T}(J2, 2, 3)
+ allowscalar() do
+ @test Array(R) == Matrix{T}(J2, 2, 3)
+ end
+ C = cuNumeric.zeros(T, 2, 3)
+ copyto!(C, J2)
+ allowscalar() do
+ @test Array(C) == Matrix{T}(J2, 2, 3)
+ end
+ else
+ R = NDArray{Bool}(I, 2, 3)
+ allowscalar() do
+ @test Array(R) == Matrix{Bool}(I, 2, 3)
+ end
+ end
+ end
+
+ # Untyped NDArray(I, ...) uses λ's type (Bool for I)
+ E = NDArray(I, 2, 2)
+ @test eltype(E) == Bool
+ allowscalar() do
+ @test Array(E) == Bool[1 0; 0 1]
+ end
+ E = NDArray(I, (2, 2))
+ allowscalar() do
+ @test Array(E) == Bool[1 0; 0 1]
+ end
+end
+
+@testset "NDArray ± / * UniformScaling / one / oneunit" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ Ah = my_rand(T, 3, 3)
+ A = NDArray(Ah)
+ J1 = one(T) * I
+ J2 = T(2) * I
+
+ _host_matrix_compare(Ah + J1, A + J1, T)
+ _host_matrix_compare(J1 + Ah, J1 + A, T)
+ _host_matrix_compare(Ah * I, A * I, T)
+ _host_matrix_compare(I * Ah, I * A, T)
+ _host_matrix_compare(Ah * J2, A * J2, T)
+ _host_matrix_compare(J2 * Ah, J2 * A, T)
+ _host_matrix_compare(Matrix{T}(I, 3, 3), one(A), T)
+ _host_matrix_compare(Matrix{T}(I, 3, 3), oneunit(A), T)
+
+ # Subtraction / A+2I: signed & float/complex only (unsigned -one wraps; skip)
+ if T <: Union{AbstractFloat,Complex} || (T <: Signed)
+ _host_matrix_compare(Ah - J1, A - J1, T)
+ _host_matrix_compare(J1 - Ah, J1 - A, T)
+ _host_matrix_compare(Ah + J2, A + J2, T)
+ end
+ end
+
+ # Plain I / 2I (Bool/Int64 λ) — natural API for floats & complex
+ @testset verbose=true for T in (DIAGONAL_FLOAT_TYPES..., DIAGONAL_COMPLEX_TYPES...)
+ Ah = my_rand(T, 3, 3)
+ A = NDArray(Ah)
+ _host_matrix_compare(Ah + I, A + I, T)
+ _host_matrix_compare(I + Ah, I + A, T)
+ _host_matrix_compare(Ah - I, A - I, T)
+ _host_matrix_compare(I - Ah, I - A, T)
+ _host_matrix_compare(Ah + 2I, A + 2I, T)
+ _host_matrix_compare(Ah * (2I), A * (2I), T)
+ _host_matrix_compare((2I) * Ah, (2I) * A, T)
+ end
+
+ # Bool: * I keeps Bool; ± I needs promotion
+ @testset "Bool UniformScaling" begin
+ Ah = Bool[1 0; 0 1]
+ A = NDArray(Ah)
+ allowscalar() do
+ @test Array(A * I) == Ah
+ @test eltype(A * I) == Bool
+ @test Array(one(A)) == Ah
+ @test Array(oneunit(A)) == Ah
+ end
+ allowpromotion() do
+ allowscalar() do
+ @test Array(A + I) == Ah + I
+ @test Array(I - A) == I - Ah
+ end
+ end
+ end
+end
+
+###### Diagonal ↔ UniformScaling ######
+
+@testset "Diagonal ± / * UniformScaling / copyto!" begin
+ @testset verbose=true for T in DIAGONAL_NUMERIC_TYPES
+ d = my_rand(T, 3)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ J1 = one(T) * I
+ J2 = T(2) * I
+
+ Di = D + J1
+ @test Di isa Diagonal
+ allowscalar() do
+ @test Matrix(Di) ≈ Matrix(Dh + J1) atol=atol(eltype(Di)) rtol=rtol(eltype(Di))
+ @test Matrix(J1 + D) ≈ Matrix(J1 + Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D * I) ≈ Matrix(Dh * I) atol=atol(T) rtol=rtol(T)
+ @test Matrix(I * D) ≈ Matrix(I * Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D * J2) ≈ Matrix(Dh * J2) atol=atol(T) rtol=rtol(T)
+ end
+
+ if T <: Union{AbstractFloat,Complex} || (T <: Signed)
+ allowscalar() do
+ @test Matrix(D - J1) ≈ Matrix(Dh - J1) atol=atol(T) rtol=rtol(T)
+ @test Matrix(J1 - D) ≈ Matrix(J1 - Dh) atol=atol(eltype((J1 - D).diag)) rtol=rtol(
+ eltype((J1 - D).diag)
+ )
+ end
+ end
+
+ copyto!(D, J1)
+ allowscalar() do
+ @test Array(D.diag) == ones(T, 3)
+ end
+ copyto!(D, zero(T) * I)
+ allowscalar() do
+ @test Array(D.diag) == zeros(T, 3)
+ end
+ end
+
+ # Plain I on floats
+ @testset verbose=true for T in DIAGONAL_FLOAT_TYPES
+ d = my_rand(T, 3)
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+ allowscalar() do
+ @test Matrix(D + I) ≈ Matrix(Dh + I) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D - I) ≈ Matrix(Dh - I) atol=atol(T) rtol=rtol(T)
+ @test Matrix(I - D) ≈ Matrix(I - Dh) atol=atol(T) rtol=rtol(T)
+ @test Matrix(D * (2I)) ≈ Matrix(Dh * (2I)) atol=atol(T) rtol=rtol(T)
+ end
+ end
+
+ @testset "Bool Diagonal + I with allowpromotion" begin
+ D = Diagonal(NDArray(Bool[true, false]))
+ allowpromotion() do
+ Di = D + I
+ @test Di isa Diagonal
+ allowscalar() do
+ @test Array(Di.diag) == [2, 1]
+ end
+ end
+ end
+end
+
+###### Structured Diagonal vs densified Matrix path ######
+
+# Compare Diagonal{NDArray} ops to the same math on densified host Matrices,
+# and check that D + I stays Diagonal while A + I densifies to NDArray.
+
+@testset "Diagonal vs dense" begin
+ @testset verbose=true for T in DIAGONAL_FLOAT_TYPES
+ n = 3
+ d = _nonzero_diag(T, n)
+ Ah = my_rand(T, n, n)
+ vh = my_rand(T, n)
+ Dh = Diagonal(d)
+ Md = Matrix(Dh) # densified host counterpart
+ D = Diagonal(NDArray(d))
+ A = NDArray(Ah)
+ v = NDArray(vh)
+
+ _host_matrix_compare(Md * Ah, D * A, T)
+ _host_matrix_compare(Ah * Md, A * D, T)
+ _host_diag_compare(Md \ vh, D \ v, T)
+ _host_matrix_compare(Ah / Md, A / D, T)
+ allowscalar() do
+ @test Matrix(inv(D)) ≈ inv(Md) atol=atol(T) rtol=rtol(T)
+ end
+ _host_matrix_compare(Ah + Md, A + D, T)
+
+ Di = D + I
+ @test Di isa Diagonal
+ allowscalar() do
+ @test Matrix(Di) ≈ Md + I atol=atol(T) rtol=rtol(T)
+ end
+
+ Ai = A + I
+ @test Ai isa NDArray
+ @test !(Ai isa Diagonal)
+ _host_matrix_compare(Ah + I, Ai, T)
+ _host_matrix_compare(Ah * I, A * I, T)
+ end
+
+ # Mul / add / structure for remaining numeric types (skip inv / div)
+ @testset verbose=true for T in (DIAGONAL_COMPLEX_TYPES..., Int32, Int64)
+ n = 3
+ d = _nonzero_diag(T, n)
+ Ah = my_rand(T, n, n)
+ Dh = Diagonal(d)
+ Md = Matrix(Dh)
+ D = Diagonal(NDArray(d))
+ A = NDArray(Ah)
+ J1 = one(T) * I
+
+ _host_matrix_compare(Md * Ah, D * A, T)
+ _host_matrix_compare(Ah * Md, A * D, T)
+ _host_matrix_compare(Ah + Md, A + D, T)
+
+ Di = D + J1
+ @test Di isa Diagonal
+ allowscalar() do
+ @test Matrix(Di) ≈ Md + Matrix{T}(J1, n, n) atol=atol(T) rtol=rtol(T)
+ end
+
+ Ai = A + J1
+ @test Ai isa NDArray
+ @test !(Ai isa Diagonal)
+ _host_matrix_compare(Ah + J1, Ai, T)
+ end
+end
+
+###### Native LinearAlgebra API (supported Diagonal paths) ######
+
+@testset "Diagonal eigen / eigvals native" begin
+ @testset verbose=true for T in DIAGONAL_FLOAT_TYPES
+ d = T[T(3), T(1), T(2)]
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+
+ # eigvals is a copy of the diagonal (NDArray)
+ λ = eigvals(D)
+ @test λ isa NDArray{T,1}
+ _host_diag_compare(eigvals(Dh), λ, T)
+ @test λ !== D.diag
+
+ # unsorted eigen: values == diag copy, vectors == I (NDArray)
+ F = eigen(D)
+ @test F.values isa NDArray{T,1}
+ _host_diag_compare(d, F.values, T)
+ @test F.vectors isa NDArray{T,2}
+ _host_matrix_compare(Matrix{T}(I, 3, 3), F.vectors, T)
+
+ @test eigvecs(D) isa NDArray{T,2}
+ _host_matrix_compare(Matrix{T}(I, 3, 3), eigvecs(D), T)
+ end
+end
+
+@testset "Diagonal native reductions / predicates / norms" begin
+ @testset verbose=true for T in DIAGONAL_FLOAT_TYPES
+ d = abs.(my_rand(T, 3)) .+ one(T) # positive → isposdef
+ Dh = Diagonal(d)
+ D = Diagonal(NDArray(d))
+
+ _host_scalar_compare(tr(Dh), tr(D), T)
+ _host_scalar_compare(sum(Dh), sum(D), T)
+ _host_scalar_compare(zero(T), prod(D), T) # n>1 off-diagonals
+ _host_scalar_compare(maximum(Dh), maximum(D), T)
+ _host_scalar_compare(minimum(Dh), minimum(D), T)
+ _host_bool_compare(isposdef(Dh), isposdef(D))
+ _host_bool_compare(true, issymmetric(D))
+ _host_bool_compare(true, ishermitian(D))
+ @test isdiag(D)
+ _host_bool_compare(false, iszero(D))
+ _host_bool_compare(false, isone(D))
+ _host_bool_compare(true, iszero(Diagonal(cuNumeric.zeros(T, 3))))
+ _host_bool_compare(true, isone(Diagonal(cuNumeric.ones(T, 3))))
+ _host_bool_compare(true, istriu(D))
+ _host_bool_compare(true, istril(D))
+ _host_bool_compare(false, istriu(D, 1))
+ _host_bool_compare(false, istril(D, -1))
+
+ _host_scalar_compare(opnorm(Dh), opnorm(D), T)
+ _host_scalar_compare(norm(Dh), norm(D), T)
+ _host_scalar_compare(cond(Dh), cond(D), T)
+ _host_scalar_compare(logdet(Dh), logdet(D), T)
+
+ # matrix functions via f.(diag) broadcast (Base Diagonal methods)
+ allowscalar() do
+ @test Matrix(sqrt(D)) ≈ Matrix(sqrt(Dh)) atol=atol(T) rtol=rtol(T)
+ @test Matrix(exp(D)) ≈ Matrix(exp(Dh)) atol=atol(T) rtol=rtol(T)
+ end
+ end
+end
diff --git a/test/array/linalg.jl b/test/array/linalg.jl
index 836d4724..19a25d00 100644
--- a/test/array/linalg.jl
+++ b/test/array/linalg.jl
@@ -111,7 +111,7 @@ end
@testset verbose=true for T in Base.uniontypes(cuNumeric.SUPPORTED_NUMERIC_TYPES)
n = 5
ref = Matrix{T}(I, n, n)
- out = cuNumeric.eye(T, n)
+ out = NDArray{T}(I, n, n)
allowscalar() do
@test safe_compare(ref, out, atol(T), rtol(T))
end
@@ -126,7 +126,7 @@ end
ref = sum(diag(A)) # widens ints like trace's accumulator
out = cuNumeric.trace(nda)
allowscalar() do
- @test ref ≈ out[1] atol=atol(eltype(ref)) rtol=rtol(eltype(ref))
+ @test ref ≈ out[] atol=atol(eltype(ref)) rtol=rtol(eltype(ref))
end
end
end
@@ -140,7 +140,7 @@ end
ref = sum(diag(A, k))
out = cuNumeric.trace(nda; offset=k)
allowscalar() do
- @test ref ≈ out[1] atol=atol(eltype(ref)) rtol=rtol(eltype(ref))
+ @test ref ≈ out[] atol=atol(eltype(ref)) rtol=rtol(eltype(ref))
end
end
end