From 79906c34700db5181072492fe55a9dc8186ec154 Mon Sep 17 00:00:00 2001 From: AntonOresten Date: Sat, 4 Jul 2026 18:38:36 +0200 Subject: [PATCH] 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 | 74 +++++++++++++++++++++++++++++++----- test/host/broadcast.jl | 85 ++++++++++++++++++++++++++++++++++++++++++ 3 files changed, 156 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 4894dcd1..a050dbfe 100644 --- a/src/broadcast.jl +++ b/src/broadcast.jl @@ -43,18 +43,41 @@ end ## kernel wrapper -_to_tiled_bc(t::Tiled) = TileArray(parent(t)) -_to_tiled_bc(arr::AbstractArray) = TileArray(arr) -_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 = TileArray(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} + # 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 = TileArray(dest) - tiled_bc = _to_tiled_bc(bc) + tiled_bc = _to_tiled_bc(bc, size(dest)) ts = _compute_tile_sizes(size(dest)) grid = ntuple(i -> cld(size(dest, i), ts[i]), N) @@ -70,6 +93,18 @@ end @inline _eval_bc(arr::TileArray, bid, tile_size) = cuTile.load(arr, bid, tile_size) @inline _eval_bc(x::Number, bid, tile_size) = x +# 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 @@ -81,11 +116,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::TileArray{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 @@ -118,7 +161,20 @@ 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))) + wrapped = _wrap_tiled(Base.Broadcast.__dot__(ex)) + if Meta.isexpr(wrapped, :.=) && Meta.isexpr(wrapped.args[1], :call) && + wrapped.args[1].args[1] === Tiled + dest = wrapped.args[1].args[2] + tmp = gensym(:dest) + wrapped = Expr(:block, + :($tmp = $dest), + Expr(:.=, :($Tiled($tmp)), wrapped.args[2]), + tmp) + end + esc(wrapped) end 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