tile scan ops related changes

[arhik]

Scan Operations - Pull Request Summary

Overview

This PR adds comprehensive scan (parallel prefix sum) operation support to the cuTile compiler, building on the existing reduce operations framework with shared abstractions.

Key Changes

1. Commit ca48190: Consolidate Common Identity Operations

• Unified IdentityOp abstraction between reduce and scan operations

• Reused the same encoding patterns since mathematically both ops are identical

• Left room for future dispatch paths if scan needs custom handling

  Files modified:

  • src/bytecode/encodings.jl - Simplified encoding interface
  • src/bytecode/writer.jl - Removed 96 lines of duplicate code
  • src/compiler/intrinsics/core.jl - Unified identity handling

2. Commit 1840d16: Add Min/Max Support for Scan

• Added encode_scan_body methods for :min and :max operations

• Float types: MinFOp/MaxFOp (no signedness needed)

• Integer types: MinIOp/MaxIOp with signedness parameter

• Aligns scan implementation with existing reduce pattern

  Files modified:

  • src/compiler/intrinsics/core.jl - 8 new lines

3. Commit 2769c85: Simple Scan Kernel Example

• Added examples/scanKernel.jl demonstrating basic usage

• Shows how to construct and execute scan operations

  Files added:

  • examples/scanKernel.jl - 62 lines

  Supported Operations (as of this PR)

Operation	Float Types	Integer Types
:add	✅ AddFOp	✅ AddIOp
:mul	✅ MulFOp	✅ MulIOp
:min	✅ MinFOp	✅ MinIOp (with signedness)
:max	✅ MaxFOp	✅ MaxIOp (with signedness)

Architectural Decisions

1. Shared Identity Abstraction: Both reduce and scan now use OperationIdentity for dispatch
2. Signedness Handling: Integer min/max respect is_signed(T) parameter
3. Minimal Code Duplication: 96 lines removed by consolidating writer logic
4. Extensible Design: Easy to add new operations following the encode_scan_body pattern

[arhik]

using Test
using CUDA
using cuTile
import cuTile as ct

# intra tile scan
function cumsum_1d_kernel(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1},
                          tile_size::ct.Constant{Int})
    bid = ct.bid(1)
    tile = ct.load(a, bid, (tile_size[],))
    result = ct.cumsum(tile, ct.axis(1))
    ct.store(b, bid, result)
    return nothing
end

sz = 32
N = 2^15
a = CUDA.rand(Float32, N)
b = CUDA.zeros(Float32, N)
CUDA.@sync ct.launch(cumsum_1d_kernel, cld(length(a), sz), a, b, ct.Constant(sz))

# This is supposed to be a single pass kernel but its simpler version than memory ordering version.
# The idea is to show scan operation.

# CSDL phase 1: Intra-tile scan + store tile sums
function cumsum_csdl_phase1(a::ct.TileArray{Float32,1}, b::ct.TileArray{Float32,1},
                            tile_sums::ct.TileArray{Float32,1},
                            tile_size::ct.Constant{Int})
    bid = ct.bid(1)
    tile = ct.load(a, bid, (tile_size[],))
    result = ct.cumsum(tile, ct.axis(1))
    ct.store(b, bid, result)
    tile_sum = ct.extract(result, (tile_size[],), (1,))
    ct.store(tile_sums, bid, tile_sum)
    return
end

# CSDL phase 2: Decoupled lookback to accumulate previous tile sums
function cumsum_csdl_phase2(b::ct.TileArray{Float32,1},
                            tile_sums::ct.TileArray{Float32,1},
                            tile_size::ct.Constant{Int})
    bid = ct.bid(1)
    prev_sum = ct.zeros((tile_size[],), Float32)
    k = Int32(bid)
    while k > 1
        tile_sum_k = ct.load(tile_sums, (k,), (1,))
        prev_sum = prev_sum .+ tile_sum_k
        k -= Int32(1)
    end
    tile = ct.load(b, bid, (tile_size[],))
    result = tile .+ prev_sum
    ct.store(b, bid, result)
    return nothing
end

n = length(a)
num_tiles = cld(n, sz)
tile_sums = CUDA.zeros(Float32, num_tiles)
CUDA.@sync ct.launch(cumsum_csdl_phase1, num_tiles, a, b, tile_sums, ct.Constant(sz))
CUDA.@sync ct.launch(cumsum_csdl_phase2, num_tiles, b, tile_sums, ct.Constant(sz))

b_cpu = cumsum(a |> collect, dims=1)
@test isapprox(b |> collect, b_cpu)

using BenchmarkTools

@benchmark CUDA.@sync begin
    CUDA.@sync ct.launch(cumsum_csdl_phase1, num_tiles, a, b, tile_sums, ct.Constant(sz))
    CUDA.@sync ct.launch(cumsum_csdl_phase2, num_tiles, b, tile_sums, ct.Constant(sz))
end

[arhik]

Wanted to add modular sum and modular prod. Next PR maybe.

[maleadt]

This seems to conflict weirdly with #21? How did you develop this?

[arhik]

I branched out from #21 's last commit ff8b77b and made changes for scan. What conflicts do you see ? I made a silly forced push in tilereduce commit. That must be problematic. I can't think of any other reasons.

Edit: Actually I think branched much early and merged the reduce branch into the tile scan branch and continued on that merge. I was working parallelly on those branches. It was not strictly linear work. I am not sure what you mean by "How did you develop this?". So I am trying to answer all possibilities I can think of. In case you are asking about my setup; I don't have a compatible GPU so I was launching 5060-5090 GPU's on avast.ai with CUDA 13.1 filter applied for the development.

[maleadt]

I branched out from #21 's last commit ff8b77b and made changes for scan.

Ah ok. You can change the target branch of this PR to make that clear; right now the diff is massive with significant overlap with that other PR. Some of the conflicting changes I saw was encode_tagged_integer! here vs encode_tagged_int! in the other PR, making them incompatible (unless one is based off of another and that change is part of the follow-up PR).

[arhik]

#39 will be simpler PR. Closing this