From 2b5ee25abfb188b11a33e4e97051378201225c08 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Sat, 4 Jul 2026 18:38:36 +0200 Subject: [PATCH 1/3] Make host Tiled broadcast shape-correct Host-level Tiled broadcasts previously loaded every array leaf with the destination's block index and tile shape, silently producing garbage for any shape-mismatched input, and never validated axes at all. Fix: - Validate via Broadcast.check_broadcast_axes: Base-illegal shapes now throw DimensionMismatch instead of computing garbage. - Base-legal size-1 (and missing trailing) dims now work: such leaves are wrapped in BroadcastLeaf{P}, load a 1-wide slice at index 1 in flagged dims, and are expanded by the device-side tile broadcast. - Scalar-only RHS (Tiled(C) .= 0) now fills the destination instead of MethodError-ing on convert(Tile{T}, ::Number). - Empty destinations are a no-op instead of launching a 0-block grid. - ct.@. in-place assignment returns the original destination array. Co-Authored-By: Claude Fable 5 --- README.md | 6 +++ src/broadcast.jl | 64 ++++++++++++++++++++++++++----- test/host/broadcast.jl | 85 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 146 insertions(+), 9 deletions(-) diff --git a/README.md b/README.md index 9bd17c4a..2aa6d24c 100644 --- a/README.md +++ b/README.md @@ -692,6 +692,12 @@ The entire broadcast expression is fused into a single cuTile kernel. Tile sizes are automatically chosen based on array dimensions (power-of-2, budget-based). Works with 1D through N-dimensional arrays. +Standard broadcast shape semantics apply: size-1 dimensions are expanded to the +destination's size (`Tiled(C) .= Tiled(row) .+ Tiled(B)` with a `(1, N)` row), +scalars fill the destination (`Tiled(C) .= 0`), and incompatible shapes throw a +`DimensionMismatch`. For in-place assignment, `ct.@.` returns the original +destination array. + ### Random Number Generation `cuTile.RNG` fills `CuArray`s on the device using the same Philox2x32-7 diff --git a/src/broadcast.jl b/src/broadcast.jl index 61df207d..15f1ee53 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -44,18 +44,41 @@ end ## kernel wrapper -_to_tiled_bc(arr::AbstractArray) = cuTileconvert(arr) -_to_tiled_bc(t::Tiled) = _to_tiled_bc(parent(t)) -_to_tiled_bc(x::Number) = x -_to_tiled_bc(x) = x # fallback for other types -function _to_tiled_bc(bc::Broadcasted) - new_args = map(_to_tiled_bc, bc.args) +""" + BroadcastLeaf{P}(arr::TileArray) + +Kernel-side wrapper for an array leaf whose shape differs from the broadcast +destination's. `P` is an `NTuple{N,Bool}` over the destination dims, flagging +dims where the leaf is size 1 (or absent) and must be loaded 1-wide at index 1, +to be expanded to the full tile shape by the device-side tile broadcast. +""" +struct BroadcastLeaf{P, A} + arr::A +end +BroadcastLeaf{P}(arr::A) where {P, A} = BroadcastLeaf{P, A}(arr) + +_to_tiled_bc(t::Tiled, dsize) = _to_tiled_bc(parent(t), dsize) +function _to_tiled_bc(arr::AbstractArray, dsize::Dims{N}) where N + ta = cuTileconvert(arr) + ndims(arr) == N && size(arr) == dsize && return ta + P = ntuple(d -> size(arr, d) == 1 && dsize[d] > 1, N) + return BroadcastLeaf{P}(ta) +end +_to_tiled_bc(x::Number, dsize) = x +_to_tiled_bc(x, dsize) = x # fallback for other types +function _to_tiled_bc(bc::Broadcasted, dsize) + new_args = map(arg -> _to_tiled_bc(arg, dsize), bc.args) Broadcasted{Nothing}(bc.f, new_args, nothing) end function _tiled_broadcast!(dest::AbstractArray{T,N}, bc::Broadcasted) where {T, N} - dest_ta = _to_tiled_bc(dest) - tiled_bc = _to_tiled_bc(bc) + # Reject shapes Base broadcasting would reject (throws DimensionMismatch); + # size-1 dims are legal and expanded per-leaf in the kernel. + Base.Broadcast.check_broadcast_axes(axes(dest), bc.args...) + isempty(dest) && return + + dest_ta = cuTileconvert(dest) + tiled_bc = _to_tiled_bc(bc, size(dest)) ts = _compute_tile_sizes(size(dest)) grid = ntuple(i -> cld(size(dest, i), ts[i]), N) @@ -74,6 +97,18 @@ end @inline _bc_func(f) = f @inline _bc_func(::Constant{Type{T}, T}) where {T} = T +# Shape-mismatched leaf: dims flagged in P load a 1-wide slice at index 1; +# the tile-level broadcast in `broadcast(bc.f, ...)` expands them to the +# common tile shape. The leaf's rank M may differ from the destination rank +# (trailing dims); `load` pads/squeezes trailing singleton dims to match. +@generated function _eval_bc(leaf::BroadcastLeaf{P}, bid, tile_size) where P + A = fieldtype(leaf, :arr) + M, N = ndims(A), length(P) + shape = [P[d] ? 1 : :(tile_size[$d]) for d in 1:N] + index = [(d > N || P[d]) ? :(Int32(1)) : :(bid[$d]) for d in 1:M] + return :(cuTile.load(leaf.arr, ($(index...),), ($(shape...),))) +end + @inline function _eval_bc(bc::Broadcasted, bid, tile_size) args = _eval_bc_args(bc.args, bid, tile_size) # Use broadcast to get element-wise semantics (not direct call, which @@ -85,11 +120,19 @@ end @inline _eval_bc_args(args::Tuple, bid, tile_size) = (_eval_bc(args[1], bid, tile_size), _eval_bc_args(Base.tail(args), bid, tile_size)...) +# Convert the broadcast result to a dest-eltype, dest-shaped tile. A `Number` +# result comes from a scalar-only RHS (e.g. `Tiled(C) .= 0`); a smaller tile +# arises when every leaf is size-1 in some dim. +@inline _bc_result(x::Number, ::Type{T}, tile_size) where {T} = + broadcast_to(Tile(T(x)), tile_size) +@inline _bc_result(t::Tile, ::Type{T}, tile_size) where {T} = + broadcast_to(convert(Tile{T}, t), tile_size) + @generated function broadcast_kernel(dest::AbstractTileArray{T, N}, bc, tile_size, overflow_grids) where {T, N} quote bids = _unflatten_bids(Val{$N}(), overflow_grids) result = _eval_bc(bc, bids, tile_size) - store(dest, bids, convert(Tile{$T}, result)) + store(dest, bids, _bc_result(result, $T, tile_size)) return end end @@ -122,6 +165,9 @@ the broadcast through cuTile kernels. using cuTile; const ct = cuTile ct.@. C = A + sin(B) # equivalent to: Tiled(C) .= Tiled(A) .+ sin.(Tiled(B)) + +For in-place assignment, the expression evaluates to the original destination +array (`C` above), not its `Tiled` wrapper. """ macro __dot__(ex) esc(_wrap_tiled(Base.Broadcast.__dot__(ex))) diff --git a/test/host/broadcast.jl b/test/host/broadcast.jl index 0797fad5..a8f3e3c9 100644 --- a/test/host/broadcast.jl +++ b/test/host/broadcast.jl @@ -120,4 +120,89 @@ using CUDA ct.Tiled(C) .= ct.Tiled(A) .+ ct.Tiled(B) @test Array(C) ≈ Array(A) .+ Array(B) end + + @testset "shape mismatch throws" begin + A = CUDA.rand(Float32, 11) + C = CUDA.zeros(Float32, 10) + @test_throws DimensionMismatch ct.Tiled(C) .= ct.Tiled(A) + @test_throws DimensionMismatch ct.Tiled(C) .= ct.Tiled(C) .+ ct.Tiled(A) + + A2 = CUDA.rand(Float32, 8, 16) + C2 = CUDA.zeros(Float32, 8, 8) + @test_throws DimensionMismatch ct.Tiled(C2) .= ct.Tiled(A2) .* 2.0f0 + end + + @testset "size-1 dim expansion (row)" begin + m, n = 64, 128 + A = CUDA.rand(Float32, 1, n) + B = CUDA.rand(Float32, m, n) + C = CUDA.zeros(Float32, m, n) + ct.Tiled(C) .= ct.Tiled(A) .+ ct.Tiled(B) + @test Array(C) ≈ Array(A) .+ Array(B) + end + + @testset "size-1 dim expansion (column)" begin + m, n = 64, 128 + A = CUDA.rand(Float32, m, 1) + B = CUDA.rand(Float32, m, n) + C = CUDA.zeros(Float32, m, n) + ct.Tiled(C) .= ct.Tiled(A) .* ct.Tiled(B) + @test Array(C) ≈ Array(A) .* Array(B) + end + + @testset "size-1 expansion only" begin + m, n = 32, 512 + A = CUDA.rand(Float32, 1, n) + C = CUDA.zeros(Float32, m, n) + ct.Tiled(C) .= ct.Tiled(A) .+ 1.0f0 + @test Array(C) ≈ repeat(Array(A) .+ 1.0f0, m, 1) + end + + @testset "rank expansion (vector + matrix)" begin + m, n = 64, 128 + A = CUDA.rand(Float32, m) + B = CUDA.rand(Float32, m, n) + C = CUDA.zeros(Float32, m, n) + ct.Tiled(C) .= ct.Tiled(A) .+ ct.Tiled(B) + @test Array(C) ≈ Array(A) .+ Array(B) + end + + @testset "allocating with expansion" begin + m, n = 64, 128 + A = CUDA.rand(Float32, 1, n) + B = CUDA.rand(Float32, m, 1) + C = ct.Tiled(A) .+ ct.Tiled(B) + @test size(C) == (m, n) + @test Array(C) ≈ Array(A) .+ Array(B) + end + + @testset "scalar RHS fill" begin + C = CUDA.rand(Float32, 1000) + ct.Tiled(C) .= 0 + @test all(iszero, Array(C)) + + D = CUDA.zeros(Float32, 33, 65) + ct.Tiled(D) .= 2.5f0 + @test all(==(2.5f0), Array(D)) + end + + @testset "empty array no-op" begin + C = CUDA.zeros(Float32, 0) + ct.Tiled(C) .= ct.Tiled(C) .+ 1.0f0 # must not launch a 0-block grid + @test isempty(Array(C)) + + D = CUDA.zeros(Float32, 4, 0) + ct.Tiled(D) .= 1.0f0 + @test isempty(Array(D)) + end + + @testset "ct.@. returns destination array" begin + n = 256 + A = CUDA.rand(Float32, n) + B = CUDA.rand(Float32, n) + C = CUDA.zeros(Float32, n) + ret = ct.@. C = A + B + @test ret === C + @test Array(C) ≈ Array(A) .+ Array(B) + end end From 8008c83a9bb1a31af98de1caa76797d442467b95 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Sat, 4 Jul 2026 20:18:48 +0200 Subject: [PATCH 2/3] Handle 0-dim broadcasts, reject host-side leaves 0-dim destinations previously died with an IR error; they now run as a 1-element 1-D broadcast (0-dim leaves are reshaped to 1-element vectors, since the kernel path assumes rank >= 1). CPU Array leaves previously compiled and crashed with an illegal memory access on the device (poisoning the CUDA context), and ranges failed with a cryptic pointer MethodError; both now throw a clear ArgumentError on the host before launch. Co-Authored-By: Claude Fable 5 --- src/broadcast.jl | 24 ++++++++++++++++++++++++ test/host/broadcast.jl | 18 ++++++++++++++++++ 2 files changed, 42 insertions(+) diff --git a/src/broadcast.jl b/src/broadcast.jl index 15f1ee53..101150c2 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -57,8 +57,24 @@ struct BroadcastLeaf{P, A} end BroadcastLeaf{P}(arr::A) where {P, A} = BroadcastLeaf{P, A}(arr) +# Reject leaves the kernel cannot address: a host Array would compile fine but +# dereference a CPU pointer on device (illegal memory access, poisoning the +# CUDA context), and ranges have no pointer at all. +function _check_device_leaf(arr::AbstractArray) + root = arr + while (p = parent(root)) !== root + root = p + end + if root isa Array || root isa AbstractRange + throw(ArgumentError("cuTile broadcast requires device arrays, got $(typeof(arr))")) + end +end + _to_tiled_bc(t::Tiled, dsize) = _to_tiled_bc(parent(t), dsize) +# 0-dim leaves become 1-element vectors; the kernel path assumes rank >= 1. +_to_tiled_bc(arr::AbstractArray{<:Any,0}, dsize::Dims) = _to_tiled_bc(reshape(arr, 1), dsize) function _to_tiled_bc(arr::AbstractArray, dsize::Dims{N}) where N + _check_device_leaf(arr) ta = cuTileconvert(arr) ndims(arr) == N && size(arr) == dsize && return ta P = ntuple(d -> size(arr, d) == 1 && dsize[d] > 1, N) @@ -71,12 +87,20 @@ function _to_tiled_bc(bc::Broadcasted, dsize) Broadcasted{Nothing}(bc.f, new_args, nothing) end +# 0-dim destinations: the kernel path needs at least one dimension, so run as +# a 1-element 1-D broadcast (0-dim leaves load rank-0 tiles and expand). +function _tiled_broadcast!(dest::AbstractArray{T,0}, bc::Broadcasted) where T + Base.Broadcast.check_broadcast_axes(axes(dest), bc.args...) + _tiled_broadcast!(reshape(dest, 1), bc) +end + function _tiled_broadcast!(dest::AbstractArray{T,N}, bc::Broadcasted) where {T, N} # Reject shapes Base broadcasting would reject (throws DimensionMismatch); # size-1 dims are legal and expanded per-leaf in the kernel. Base.Broadcast.check_broadcast_axes(axes(dest), bc.args...) isempty(dest) && return + _check_device_leaf(dest) dest_ta = cuTileconvert(dest) tiled_bc = _to_tiled_bc(bc, size(dest)) diff --git a/test/host/broadcast.jl b/test/host/broadcast.jl index a8f3e3c9..d0a02f52 100644 --- a/test/host/broadcast.jl +++ b/test/host/broadcast.jl @@ -196,6 +196,24 @@ using CUDA @test isempty(Array(D)) end + @testset "0-dim arrays" begin + C = CUDA.zeros(Float32) + ct.Tiled(C) .= 1.0f0 + @test Array(C)[] == 1.0f0 + + A = CUDA.fill(2.0f0) + ct.Tiled(C) .= ct.Tiled(A) .+ 1.0f0 + @test Array(C)[] == 3.0f0 + end + + @testset "host-side leaves rejected" begin + C = CUDA.zeros(Float32, 16) + A_cpu = rand(Float32, 16) + @test_throws ArgumentError ct.Tiled(C) .= ct.Tiled(A_cpu) .* 2.0f0 + @test_throws ArgumentError ct.Tiled(C) .= ct.Tiled(C) .+ (1:16) + @test_throws ArgumentError ct.Tiled(A_cpu) .= 0 # CPU destination + end + @testset "ct.@. returns destination array" begin n = 256 A = CUDA.rand(Float32, n) From 976e8b965d6e7fe3f2aed07131accfbd41b046ae Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Fri, 10 Jul 2026 22:47:42 +0200 Subject: [PATCH 3/3] Simplify broadcast leaf handling Align the kernel-wrapper machinery with Base.Broadcast's vocabulary and trim the helper zoo introduced by the shape fixes: - Replace BroadcastLeaf{P} with Extruded{K}, mirroring Base.Broadcast.Extruded but with the keep-mask in the type domain (it determines the static tile shape). Every array leaf is wrapped, so the separate exact-shape load path disappears. - Replace _to_tiled_bc with preprocess, the counterpart of Base.Broadcast.preprocess. - Drop _bc_result: scalars now convert through a new convert(Tile{T}, ::Number) method, and the store site uniformly expands with broadcast_to (a codegen no-op for full-shaped tiles). - Drop _eval_bc_args in favor of a tuple map. - Move host-memory rejection from a parent-walking blocklist into TileArray construction itself: pointer(arr) returning a plain Ptr means host memory. This also covers launch() arguments and mapreduce, not just broadcast. - Drop the macro return-value rewrite: materialize! already returns the parent array since #258. Co-Authored-By: Claude Fable 5 --- src/broadcast.jl | 123 ++++++++++++++----------------------- src/language/operations.jl | 1 + src/language/types.jl | 19 ++++-- test/host/broadcast.jl | 11 +++- 4 files changed, 69 insertions(+), 85 deletions(-) diff --git a/src/broadcast.jl b/src/broadcast.jl index 101150c2..a5206515 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -45,64 +45,40 @@ end ## kernel wrapper """ - BroadcastLeaf{P}(arr::TileArray) + Extruded{K}(arr::AbstractTileArray) -Kernel-side wrapper for an array leaf whose shape differs from the broadcast -destination's. `P` is an `NTuple{N,Bool}` over the destination dims, flagging -dims where the leaf is size 1 (or absent) and must be loaded 1-wide at index 1, -to be expanded to the full tile shape by the device-side tile broadcast. +Array leaf of a `Broadcasted` tree. Like `Base.Broadcast.Extruded`, `K` flags +the dimensions along which the leaf has extent, except that the mask lives in +the type domain because it determines the static shape of the tile to load: +flagged dimensions cover the tile, the others are loaded 1-wide at index 1 and +expanded by the device-side tile broadcast. """ -struct BroadcastLeaf{P, A} - arr::A +struct Extruded{K, A} + x::A end -BroadcastLeaf{P}(arr::A) where {P, A} = BroadcastLeaf{P, A}(arr) - -# Reject leaves the kernel cannot address: a host Array would compile fine but -# dereference a CPU pointer on device (illegal memory access, poisoning the -# CUDA context), and ranges have no pointer at all. -function _check_device_leaf(arr::AbstractArray) - root = arr - while (p = parent(root)) !== root - root = p - end - if root isa Array || root isa AbstractRange - throw(ArgumentError("cuTile broadcast requires device arrays, got $(typeof(arr))")) - end -end - -_to_tiled_bc(t::Tiled, dsize) = _to_tiled_bc(parent(t), dsize) -# 0-dim leaves become 1-element vectors; the kernel path assumes rank >= 1. -_to_tiled_bc(arr::AbstractArray{<:Any,0}, dsize::Dims) = _to_tiled_bc(reshape(arr, 1), dsize) -function _to_tiled_bc(arr::AbstractArray, dsize::Dims{N}) where N - _check_device_leaf(arr) - ta = cuTileconvert(arr) - ndims(arr) == N && size(arr) == dsize && return ta - P = ntuple(d -> size(arr, d) == 1 && dsize[d] > 1, N) - return BroadcastLeaf{P}(ta) -end -_to_tiled_bc(x::Number, dsize) = x -_to_tiled_bc(x, dsize) = x # fallback for other types -function _to_tiled_bc(bc::Broadcasted, dsize) - new_args = map(arg -> _to_tiled_bc(arg, dsize), bc.args) - Broadcasted{Nothing}(bc.f, new_args, nothing) -end - -# 0-dim destinations: the kernel path needs at least one dimension, so run as -# a 1-element 1-D broadcast (0-dim leaves load rank-0 tiles and expand). -function _tiled_broadcast!(dest::AbstractArray{T,0}, bc::Broadcasted) where T - Base.Broadcast.check_broadcast_axes(axes(dest), bc.args...) - _tiled_broadcast!(reshape(dest, 1), bc) +Extruded{K}(x::A) where {K, A} = Extruded{K, A}(x) + +# Counterpart of `Base.Broadcast.preprocess`: convert the array leaves of a +# Broadcasted tree into Extruded TileArrays, given the destination rank. +preprocess(t::Tiled, rank::Val) = preprocess(parent(t), rank) +preprocess(arr::AbstractArray{<:Any, 0}, rank::Val) = preprocess(reshape(arr, 1), rank) +function preprocess(arr::AbstractArray, ::Val{N}) where N + keeps = ntuple(d -> size(arr, d) != 1, Val(N)) + # The typeassert rejects leaves without device storage, like ranges, + # which `cuTileconvert` passes through unchanged. + Extruded{keeps}(cuTileconvert(arr)::AbstractTileArray) end +preprocess(bc::Broadcasted, rank::Val) = + Broadcasted{Nothing}(bc.f, map(arg -> preprocess(arg, rank), bc.args), nothing) +preprocess(x, rank::Val) = x function _tiled_broadcast!(dest::AbstractArray{T,N}, bc::Broadcasted) where {T, N} - # Reject shapes Base broadcasting would reject (throws DimensionMismatch); - # size-1 dims are legal and expanded per-leaf in the kernel. + # Match Base semantics: size-1 dimensions expand, other mismatches throw. Base.Broadcast.check_broadcast_axes(axes(dest), bc.args...) isempty(dest) && return - _check_device_leaf(dest) dest_ta = cuTileconvert(dest) - tiled_bc = _to_tiled_bc(bc, size(dest)) + tiled_bc = preprocess(bc, Val(N)) ts = _compute_tile_sizes(size(dest)) grid = ntuple(i -> cld(size(dest, i), ts[i]), N) @@ -112,51 +88,42 @@ function _tiled_broadcast!(dest::AbstractArray{T,N}, bc::Broadcasted) where {T, Constant(ts), Constant(overflow)) end +# The kernel operates on tiles of rank >= 1, so run 0-dim broadcasts (both +# destinations, above, and leaves, in `preprocess`) as 1-element vectors. +_tiled_broadcast!(dest::AbstractArray{<:Any,0}, bc::Broadcasted) = + _tiled_broadcast!(reshape(dest, 1), bc) -## kernel - -@inline _eval_bc(arr::AbstractTileArray, bid, tile_size) = cuTile.load(arr, bid, tile_size) -@inline _eval_bc(x::Number, bid, tile_size) = x -@inline _bc_func(f) = f -@inline _bc_func(::Constant{Type{T}, T}) where {T} = T +## kernel -# Shape-mismatched leaf: dims flagged in P load a 1-wide slice at index 1; -# the tile-level broadcast in `broadcast(bc.f, ...)` expands them to the -# common tile shape. The leaf's rank M may differ from the destination rank -# (trailing dims); `load` pads/squeezes trailing singleton dims to match. -@generated function _eval_bc(leaf::BroadcastLeaf{P}, bid, tile_size) where P - A = fieldtype(leaf, :arr) - M, N = ndims(A), length(P) - shape = [P[d] ? 1 : :(tile_size[$d]) for d in 1:N] - index = [(d > N || P[d]) ? :(Int32(1)) : :(bid[$d]) for d in 1:M] - return :(cuTile.load(leaf.arr, ($(index...),), ($(shape...),))) +# Evaluate one node of the Broadcasted tree for the current block, like +# `Base.Broadcast._broadcast_getindex` does per element. A leaf's rank may be +# lower than the destination's (trailing dimensions); `load` completes the +# tile shape with trailing singleton dimensions. +@generated function _eval_bc(leaf::Extruded{K}, bid, tile_size) where K + M = ndims(fieldtype(leaf, :x)) + shape = [K[d] ? :(tile_size[$d]) : 1 for d in 1:length(K)] + index = [K[d] ? :(bid[$d]) : :(Int32(1)) for d in 1:M] + return :(cuTile.load(leaf.x, ($(index...),), ($(shape...),))) end - +@inline _eval_bc(x::Number, bid, tile_size) = x @inline function _eval_bc(bc::Broadcasted, bid, tile_size) - args = _eval_bc_args(bc.args, bid, tile_size) + args = map(arg -> _eval_bc(arg, bid, tile_size), bc.args) # Use broadcast to get element-wise semantics (not direct call, which # would dispatch to e.g. matmul for * on tiles) broadcast(_bc_func(bc.f), args...) end -@inline _eval_bc_args(::Tuple{}, bid, tile_size) = () -@inline _eval_bc_args(args::Tuple, bid, tile_size) = - (_eval_bc(args[1], bid, tile_size), _eval_bc_args(Base.tail(args), bid, tile_size)...) - -# Convert the broadcast result to a dest-eltype, dest-shaped tile. A `Number` -# result comes from a scalar-only RHS (e.g. `Tiled(C) .= 0`); a smaller tile -# arises when every leaf is size-1 in some dim. -@inline _bc_result(x::Number, ::Type{T}, tile_size) where {T} = - broadcast_to(Tile(T(x)), tile_size) -@inline _bc_result(t::Tile, ::Type{T}, tile_size) where {T} = - broadcast_to(convert(Tile{T}, t), tile_size) +@inline _bc_func(f) = f +@inline _bc_func(::Constant{Type{T}, T}) where {T} = T @generated function broadcast_kernel(dest::AbstractTileArray{T, N}, bc, tile_size, overflow_grids) where {T, N} quote bids = _unflatten_bids(Val{$N}(), overflow_grids) result = _eval_bc(bc, bids, tile_size) - store(dest, bids, _bc_result(result, $T, tile_size)) + # The result can be smaller than the tile (all leaves size-1 in some + # dimension) or even a scalar (scalar-only right-hand side). + store(dest, bids, broadcast_to(convert(Tile{$T}, result), tile_size)) return end end diff --git a/src/language/operations.jl b/src/language/operations.jl index a164d438..d2d9cbbb 100644 --- a/src/language/operations.jl +++ b/src/language/operations.jl @@ -1015,6 +1015,7 @@ end @inline Base.convert(::Type{Tile{T}}, tile::Tile{T}) where {T} = tile @inline Base.convert(::Type{Tile{T2}}, tile::Tile{T1, Shape}) where {T1, T2, Shape} = map(T2, tile) +@inline Base.convert(::Type{Tile{T}}, x::Number) where {T} = Tile(convert(T, x)) #============================================================================= Reduction diff --git a/src/language/types.jl b/src/language/types.jl index 2e58e6c3..b1fd962f 100644 --- a/src/language/types.jl +++ b/src/language/types.jl @@ -215,7 +215,7 @@ end """ TileArray(arr) -Create a TileArray from a CUDA array (CuArray or similar). +Create a TileArray from a device array (CuArray or similar). Automatically extracts pointer, sizes, strides, and computes ArraySpec. This method works with any array type that supports: @@ -224,18 +224,25 @@ This method works with any array type that supports: - `strides(arr)` - returns array strides """ function TileArray(arr::AbstractArray{T, N}) where {T, N} - # Use reinterpret to handle both Ptr and CuPtr (device pointers) - ptr = reinterpret(Ptr{T}, pointer(arr)) sizes = NTuple{N, Int32}(Int32.(size(arr))) strides_val = NTuple{N, Int32}(Int32.(strides(arr))) - TileArray(ptr, sizes, strides_val) + TileArray(device_pointer(arr), sizes, strides_val) end function TileArray(arr::PermutedDimsArray{T, N}) where {T, N} - ptr = reinterpret(Ptr{T}, pointer(parent(arr))) sizes = NTuple{N, Int32}(Int32.(size(arr))) strides_val = NTuple{N, Int32}(Int32.(strides(arr))) - TileArray(ptr, sizes, strides_val) + TileArray(device_pointer(parent(arr)), sizes, strides_val) +end + +# Device arrays hand out device pointer types (e.g. `CuPtr`), which TileArray +# stores reinterpreted as `Ptr`. A `pointer` that already is a `Ptr` means host +# memory, which would compile fine but fault when dereferenced on device. +function device_pointer(arr::AbstractArray{T}) where T + ptr = pointer(arr) + ptr isa Ptr && + throw(ArgumentError("cannot create a TileArray from host memory ($(typeof(arr)))")) + reinterpret(Ptr{T}, ptr) end diff --git a/test/host/broadcast.jl b/test/host/broadcast.jl index d0a02f52..7d86bfa5 100644 --- a/test/host/broadcast.jl +++ b/test/host/broadcast.jl @@ -158,6 +158,15 @@ using CUDA @test Array(C) ≈ repeat(Array(A) .+ 1.0f0, m, 1) end + @testset "size-1 expansion in nested broadcast" begin + m, n = 64, 128 + A = CUDA.rand(Float32, 1, n) + B = CUDA.rand(Float32, m, n) + C = CUDA.zeros(Float32, m, n) + ct.Tiled(C) .= ct.Tiled(B) .+ sin.(ct.Tiled(A)) + @test Array(C) ≈ Array(B) .+ sin.(Array(A)) + end + @testset "rank expansion (vector + matrix)" begin m, n = 64, 128 A = CUDA.rand(Float32, m) @@ -210,7 +219,7 @@ using CUDA C = CUDA.zeros(Float32, 16) A_cpu = rand(Float32, 16) @test_throws ArgumentError ct.Tiled(C) .= ct.Tiled(A_cpu) .* 2.0f0 - @test_throws ArgumentError ct.Tiled(C) .= ct.Tiled(C) .+ (1:16) + @test_throws TypeError ct.Tiled(C) .= ct.Tiled(C) .+ (1:16) # no device storage @test_throws ArgumentError ct.Tiled(A_cpu) .= 0 # CPU destination end