Overlay fill, zeros, and ones from Base (#123)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: AntonOresten

README.md

@@ -140,9 +140,10 @@ uses standard Julia syntax and is overlaid on `Base`.
 ### Construction
 | Operation | Description |
 |-----------|-------------|
-| `ct.zeros(shape, T)` | Zero-filled tile |
-| `ct.full(shape, value, T)` | Constant-filled tile |
-| `ct.arange(shape, T)` | Sequence `[1, 2, 3, ...]` |
+| `zeros(T, dims...)` | Zero-filled tile (Base overlay) |
+| `ones(T, dims...)` | One-filled tile (Base overlay) |
+| `fill(value, dims...)` | Constant-filled tile (Base overlay) |
+| `ct.arange(shape, T)` / `ct.arange(n, T)` | Sequence `[1, 2, 3, ..., n]` |
 
 ### Shape
 | Operation | Description |

examples/batchmatmul.jl

@@ -23,7 +23,7 @@ function batch_matmul_kernel(A::ct.TileArray{T,3}, B::ct.TileArray{T,3}, C::ct.T
     num_k = cld(K, Int32(tk))
 
     # Initialize accumulator with Float32 for precision
-    acc = ct.full((tm, tn), zero(Float32), Float32)
+    acc = zeros(Float32, tm, tn)
 
     # K reduction loop
     k = Int32(1)

examples/layernorm.jl

@@ -29,7 +29,7 @@ function layer_norm_fwd(X::ct.TileArray{Float32, 2}, W::ct.TileArray{Float32, 1}
     N = size(X, 2)
 
     # Compute mean
-    mean = ct.full((1, TILE_N), 0.0f0, Float32)
+    mean = zeros(Float32, (1, TILE_N))
     j = Int32(1)
     while j <= num_tiles
         tx = ct.load(X, (bid_m, j), (1, TILE_N); padding_mode=ct.PaddingMode.Zero)
@@ -40,12 +40,12 @@ function layer_norm_fwd(X::ct.TileArray{Float32, 2}, W::ct.TileArray{Float32, 1}
     ct.store(Mean, bid_m, mean)
 
     # Compute variance
-    var = ct.full((1, TILE_N), 0.0f0, Float32)
+    var = zeros(Float32, (1, TILE_N))
     j = Int32(1)
     while j <= num_tiles
         tx = ct.load(X, (bid_m, j), (1, TILE_N); padding_mode=ct.PaddingMode.Zero)
         # Mask for valid elements
-        mask = reshape(((j - Int32(1)) * Int32(TILE_N) .+ ct.arange((TILE_N,), Int32)) .<= N, (1, TILE_N))
+        mask = reshape(((j - Int32(1)) * Int32(TILE_N) .+ ct.arange(TILE_N, Int32)) .<= N, (1, TILE_N))
         centered_tx = ifelse.(mask, tx .- mean, 0.0f0)
         var = var .+ (centered_tx .^ 2.0f0)
         j += Int32(1)
@@ -93,7 +93,7 @@ bid_m and j are 1-indexed (block ID and tile index).
     wdy = tw .* tdy
 
     # Mask for valid elements
-    indices = ct.arange((TILE_N,), Int32)
+    indices = ct.arange(TILE_N, Int32)
     offset = (j - Int32(1)) * Int32(TILE_N)
     global_indices = offset .+ indices
     mask = reshape(global_indices .<= N, (1, TILE_N))
@@ -131,8 +131,8 @@ function layer_norm_bwd_dx(DX::ct.TileArray{Float32, 2}, DY::ct.TileArray{Float3
     rstd = ct.load(Rstd, bid_m, (1,); padding_mode=ct.PaddingMode.Zero)
 
     # First pass: compute c1 and c2 reduction terms
-    c1 = ct.full((1, TILE_N), 0.0f0, Float32)
-    c2 = ct.full((1, TILE_N), 0.0f0, Float32)
+    c1 = zeros(Float32, (1, TILE_N))
+    c2 = zeros(Float32, (1, TILE_N))
     j = Int32(1)
     while j <= num_tiles
         _, xhat, wdy = bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N)
@@ -190,8 +190,8 @@ function layer_norm_bwd_dx_partial_dwdb(DX::ct.TileArray{Float32, 2}, DY::ct.Til
     rstd = ct.load(Rstd, bid_m, (1,); padding_mode=ct.PaddingMode.Zero)
 
     # First pass: compute c1 and c2 reduction terms
-    c1 = ct.full((1, TILE_N), 0.0f0, Float32)
-    c2 = ct.full((1, TILE_N), 0.0f0, Float32)
+    c1 = zeros(Float32, (1, TILE_N))
+    c2 = zeros(Float32, (1, TILE_N))
     j = Int32(1)
     while j <= num_tiles
         _, xhat, wdy = bwd_helper(X, W, DY, bid_m, j, mean, rstd, TILE_N, N)
@@ -253,8 +253,8 @@ function layer_norm_bwd_dwdb(DW::ct.TileArray{Float32, 2}, DB::ct.TileArray{Floa
     bid_n = ct.bid(1)
     num_tiles = ct.num_tiles(DW, 2, (TILE_N, TILE_M))
 
-    dw = ct.zeros((TILE_N, TILE_M), Float32)
-    db = ct.zeros((TILE_N, TILE_M), Float32)
+    dw = zeros(Float32, (TILE_N, TILE_M))
+    db = zeros(Float32, (TILE_N, TILE_M))
     i = Int32(1)
     while i <= num_tiles
         dw = dw .+ ct.load(DW, (bid_n, i), (TILE_N, TILE_M); padding_mode=ct.PaddingMode.Zero)

examples/matmul.jl

@@ -39,7 +39,7 @@ function matmul_kernel(A::ct.TileArray{T,2}, B::ct.TileArray{T,2}, C::ct.TileArr
     num_k = ct.num_tiles(A, 2, (tm, tk))
 
     # Initialize accumulator with Float32 for precision
-    acc = ct.full((tm, tn), zero(Float32), Float32)
+    acc = zeros(Float32, tm, tn)
 
     # K reduction loop - accumulate partial products
     # NOTE: Uses while-loop pattern. Native `for k in 0:n` syntax generates complex

examples/vadd.jl

@@ -31,7 +31,7 @@ function vec_add_kernel_1d_gather(a::ct.TileArray{T,1}, b::ct.TileArray{T,1}, c:
                                    tile::Int) where {T}
     bid = ct.bid(1)
     # Create index tile for this block's elements
-    offsets = ct.arange((tile,), Int32)
+    offsets = ct.arange(tile, Int32)
     base = ct.Tile((bid - Int32(1)) * Int32(tile))
     indices = ct.broadcast_to(base, (tile,)) .+ offsets
 

src/compiler/intrinsics/core.jl

@@ -177,9 +177,9 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.constant), args)
 
     # Extract shape
     shape = get_constant(ctx, args[1])
-    shape isa Tuple || throw(IRError("full() shape must be a compile-time constant tuple"))
+    shape isa Tuple || throw(IRError("fill() shape must be a compile-time constant tuple"))
     tile_shape = collect(Int, shape)
-    validate_tile_shape(tile_shape, "full")
+    validate_tile_shape(tile_shape, "fill")
 
     # Extract dtype from Type{T} argument
     elem_type = @something get_constant(ctx, args[3]) throw(IRError("constant() requires a compile-time element type"))
@@ -188,7 +188,7 @@ function emit_intrinsic!(ctx::CGCtx, ::typeof(Intrinsics.constant), args)
     tile_type = tile_type!(tt, dtype, tile_shape)
 
     tv = emit_value!(ctx, args[2])
-    tv === nothing && throw(IRError("full() value must be a constant or a runtime scalar"))
+    tv === nothing && throw(IRError("fill() value must be a constant or a runtime scalar"))
     if tv.constant !== nothing
         # Compile-time constant: use ConstantOp directly
         value_bytes = constant_to_bytes(something(tv.constant), elem_type)

src/language/operations.jl

@@ -419,48 +419,30 @@ end
  Factory
 =============================================================================#
 
-public arange, full, zeros
+public arange
 
 """
     arange(shape::NTuple{1, Int}, dtype::Type{T}) -> Tile{T, shape}
+    arange(n::Int, dtype::Type{T}) -> Tile{T, (n,)}
 
-Create a 1D tile with values [1, 2, 3, ..., shape[1]] (1-indexed).
+Create a 1D tile with values [1, 2, 3, ..., n] (1-indexed).
 
 # Example
 ```julia
 indices = ct.arange((16,), Int32)  # Creates Tile with [1, 2, 3, ..., 16]
+indices = ct.arange(16, Int32)     # Same thing, scalar form
 ```
 """
 @inline arange(shape::NTuple{1, Int}, ::Type{T}) where {T} =
     Intrinsics.iota(shape, T) .+ one(T)
+@inline arange(n::Int, ::Type{T}) where {T} = arange((n,), T)
 
-"""
-    full(shape::NTuple{N, Int}, value, dtype::Type{T}) -> Tile{T, shape}
-
-Create a tile filled with a constant value.
-
-# Example
-```julia
-ones_tile = ct.full((32, 32), 1.0f0, Float32)
-```
-"""
-@inline full(shape::NTuple{N, Int}, value::Tile, ::Type{T}) where {N, T} =
-    Intrinsics.constant(shape, convert(Tile{T}, value), T)
-@inline full(shape::NTuple{N, Int}, value, ::Type{T}) where {N, T} =
+# Internal: create a tile filled with a constant value.
+# Used by Base.fill/zeros/ones overlays (see overlays.jl).
+@inline _full(value, ::Type{T}, shape::NTuple{N, Int}) where {N, T} =
     Intrinsics.constant(shape, Tile(T(value)), T)
-
-"""
-    zeros(shape::NTuple{N, Int}, dtype::Type{T}) -> Tile{T, shape}
-
-Create a tile filled with zeros.
-
-# Example
-```julia
-zeros_tile = ct.zeros((32, 32), Float32)
-```
-"""
-@inline zeros(shape::NTuple{N, Int}, ::Type{T}) where {N, T} =
-    full(shape, zero(T), T)
+@inline _full(value::Tile, ::Type, shape::NTuple{N, Int}) where {N} =
+    Intrinsics.constant(shape, value, eltype(value))
 
 #=============================================================================
  Shape & DType
@@ -522,6 +504,7 @@ reshaped = reshape(tile, (2, 16))  # Shape (2, 16), still 32 elements
     size(tile) === shape && return tile
     Intrinsics.reshape(tile, shape)
 end
+@inline Base.reshape(tile::Tile{T}, dims::Int...) where {T} = reshape(tile, dims)
 
 """
     permutedims(tile::Tile{T, S}, perm) -> Tile{T, permuted_shape}
@@ -917,7 +900,7 @@ end
 @inline function _matmul(a::Tile{T1}, b::Tile, ::Val{2}) where {T1}
     M = size(a, 1)
     N = size(b, 2)
-    acc = zeros((M, N), T1)
+    acc = zeros(T1, (M, N))
     muladd(a, b, acc)
 end
 
@@ -926,7 +909,7 @@ end
     B = max(size(a, 1), size(b, 1))  # Broadcast batch dimension
     M = size(a, 2)
     N = size(b, 3)
-    acc = zeros((B, M, N), T1)
+    acc = zeros(T1, (B, M, N))
     muladd(a, b, acc)
 end
 

src/language/overlays.jl

@@ -109,3 +109,17 @@ for F in Floats
         end
     end
 end
+
+
+#=============================================================================
+ Tile Constructors
+=============================================================================#
+
+# Base.fill/zeros/ones return Tiles in kernel context, matching Julia's standard API.
+# Marked non-foldable because they return differently-typed objects.
+Base.Experimental.@consistent_overlay cuTileMethodTable @inline Base.fill(v, dims::NTuple{N, Int}) where {N} =
+    _full(v, typeof(v), dims)
+Base.Experimental.@consistent_overlay cuTileMethodTable @inline Base.zeros(::Type{T}, dims::NTuple{N, Int}) where {T, N} =
+    _full(zero(T), T, dims)
+Base.Experimental.@consistent_overlay cuTileMethodTable @inline Base.ones(::Type{T}, dims::NTuple{N, Int}) where {T, N} =
+    _full(one(T), T, dims)