Skip to content

Make host Tiled broadcast shape-correct#264

Merged
maleadt merged 3 commits into
mainfrom
ao/broadcast-shape-fix
Jul 11, 2026
Merged

Make host Tiled broadcast shape-correct#264
maleadt merged 3 commits into
mainfrom
ao/broadcast-shape-fix

Conversation

@AntonOresten

@AntonOresten AntonOresten commented Jul 4, 2026

Copy link
Copy Markdown
Collaborator

Host-level Tiled broadcasts loaded every array leaf with the destination's block index and tile shape, and never validated axes — so shape-mismatched inputs (both Base-legal (1,N) vs (M,N) expansion and outright wrong sizes) silently produced garbage.

Changes

  • Validation: _tiled_broadcast! now calls Broadcast.check_broadcast_axes, so Base-illegal shapes throw DimensionMismatch instead of computing garbage.
  • Size-1 / trailing-dim expansion works: leaves whose shape differs from the destination are wrapped in BroadcastLeaf{P} (P = static mask of expanded dims), load a 1-wide slice at index 1 in those dims, and are expanded by the existing device-side tile broadcast (broadcast_shape/broadcast_to). Exact-shape leaves keep the previous fast path unchanged.
  • Scalar-only RHS: Tiled(C) .= 0 used to MethodError on convert(Tile{T}, ::Number); the kernel result is now routed through _bc_result, which splats scalars via broadcast_to(Tile(T(x)), tile_size) (and also expands sub-shaped result tiles when every leaf is size-1 in a dim).
  • Empty destinations: no-op instead of launching a grid with a 0 dimension.
  • ct.@. return value: in-place assignment (ct.@. C = A + B) now evaluates to the original destination array C rather than the Tiled wrapper.

Tests

9 new testsets in test/host/broadcast.jl: mismatch throws (1D/2D, nested), row/column size-1 expansion, expansion-only RHS, rank expansion (vector + matrix), allocating form with expansion, scalar fill, empty no-op, and macro return identity. Follow-up commit: 0-dim broadcasts now work (delegated to a 1-element 1-D launch), and host-side leaves (CPU Arrays — previously an illegal memory access that poisoned the CUDA context — and ranges) are rejected with a clear ArgumentError before launch. Full host broadcast suite passes on GPU: 37/37.

(Supersedes #261/#263, closed by GitHub during head-branch renames.)

🤖 Generated with Claude Code

@AntonOresten AntonOresten marked this pull request as ready for review July 5, 2026 13:22
@maleadt maleadt force-pushed the ao/broadcast-shape-fix branch from 5daeebd to 0bd80d1 Compare July 10, 2026 19:10
AntonOresten and others added 2 commits July 10, 2026 22:29
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 <noreply@anthropic.com>
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 <noreply@anthropic.com>
@maleadt maleadt force-pushed the ao/broadcast-shape-fix branch from 0bd80d1 to c336349 Compare July 10, 2026 20:52
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 <noreply@anthropic.com>
@maleadt maleadt force-pushed the ao/broadcast-shape-fix branch from c336349 to 976e8b9 Compare July 10, 2026 20:58
@maleadt maleadt merged commit 4faa3b1 into main Jul 11, 2026
1 check passed
@maleadt maleadt deleted the ao/broadcast-shape-fix branch July 11, 2026 07:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants