Make host Tiled broadcast shape-correct

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

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)
-    Broadcasted{Nothing}(bc.f, new_args, nothing)
+"""
+    Extruded{K}(arr::AbstractTileArray)
+
+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 Extruded{K, A}
+    x::A
+end
+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}
-    dest_ta = _to_tiled_bc(dest)
-    tiled_bc = _to_tiled_bc(bc)
+    # Match Base semantics: size-1 dimensions expand, other mismatches throw.
+    Base.Broadcast.check_broadcast_axes(axes(dest), bc.args...)
+    isempty(dest) && return
+
+    dest_ta = cuTileconvert(dest)
+    tiled_bc = preprocess(bc, Val(N))
 
     ts = _compute_tile_sizes(size(dest))
     grid = ntuple(i -> cld(size(dest, i), ts[i]), N)
@@ -65,31 +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)
+# 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 _bc_func(f) = f
-@inline _bc_func(::Constant{Type{T}, T}) where {T} = T
-
 @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)...)
+@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, convert(Tile{$T}, result))
+        # 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
@@ -122,6 +156,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)))

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

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
 
 

test/host/broadcast.jl

@@ -120,4 +120,116 @@ 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 "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)
+        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 "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 TypeError ct.Tiled(C) .= ct.Tiled(C) .+ (1:16)  # no device storage
+        @test_throws ArgumentError ct.Tiled(A_cpu) .= 0  # CPU destination
+    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