Extend tile reduction operations with additional operators

[arhik]

This PR extends the tile reduction operations to support a broader range of reduction operators for both floating-point and integer types.

Changes

1. Core Intrinsics (src/compiler/intrinsics/core.jl)

• Added reduce_mul, reduce_min operations (work with any element type)
• Added reduce_and, reduce_or, reduce_xor bitwise operations (integer-only)
• Removed AbstractFloat constraint from reduce_sum and reduce_max
• Added corresponding emit_intrinsic! handlers for new operations
• Added integer encode methods for and, or, xor operations

2. Operations (src/language/operations.jl)

• Added axis(i::Integer) helper function for 1-based to 0-based axis conversion
• Provides self-documenting API: ct.cumsum(tile, ct.axis(1)) instead of raw Val

3. Documentation (README.md)

• Documented new reduce operations in the operations table
• Added axis(i) helper to the reductions section

Supported Reductions

Operation	Types	Description
reduce_sum	Any	Sum along axis
reduce_mul	Any	Product along axis
reduce_max	Any	Maximum along axis
reduce_min	Any	Minimum along axis
reduce_and	Integer	Bitwise AND along axis
reduce_or	Integer	Bitwise OR along axis
reduce_xor	Integer	Bitwise XOR along axis

[AntonOresten]

What is the point of having an axis function? I don't think it needs to be corrected to 0-based indexing that early, and they don't need to be Val since they can be constant-propagated anyway.

[arhik]

@AntonOresten I agree.

I noted this in the commit related to it.

reduce ops update and axis convenience
axis convenience is a bit helper function for Val.

But I see reduce is already one-based. Not sure if we should go with it. It doesn't harm anything. its just a convenience.

[arhik]

I am still working on it and scan operations. Axis convenience came up when working on scan ops. I just included it in this pull request. Either way it's going to be consensus decision as we go forward.

[AntonOresten]

Maybe the reduction methods can be exposed through Base.reduce, dispatching on the binary operator (+, max, * etc). The ct.reduce_* methods currently drop the reduced dimension, unlike Base.reduce. This might lead to a slippery slope, turning reshape, permute, transpose, etc. into methods of their Base counterparts, but then there's no nice counterpart for extract.

[arhik]

Ideally yes.

using Test
using CUDA
using cuTile
import cuTile as ct

function reduceKernel(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.reduce_min(tile, Val(1))
    ct.store(b, bid, result)
    return nothing
end

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

julia> b[1:5]
5-element CuArray{Float32, 1, CUDA.DeviceMemory}:
0.005346642
0.0034002797
0.029549925
0.14612293
0.0142624555

julia> minimum(a[1:32])
0.005346642f0

julia> minimum(a[33:64])
0.0034002797f0

[arhik]

Yay. Something. I don't have blackwell gpu. I rented vast.ai 5080 gpu. worth it i guess.

[AntonOresten]

Ah, I think there is a valid use of the axis helper, but it is not user-facing. I was rather confused by the emphasis and documentation of this one helper.

[maleadt]

I'm not a fan of the reduce_max etc functions; those were just placeholders. We should probably expose these using more idiomatic constructs, like max or cumsum etc. If possible by actually implementing something underlying like mapreduce! and not Base.sum directly, but I'm not sure how realistic that is (initially it could error out on unsupported configurations, but longer term it could be extended to actually allow some of the more interesting combinations like summing over a function, or only reducing some dimensions).

[arhik]

@maleadt You spoke my mind verbatim. I have a package that is very heavy on very complicated scan, reduce, top-k like computations with event triggering state machine transitions. And none of the above simple reduce_{ops} actually help me with my goal. I would love something like you suggested and that's my hope on this package. I am just pushing reduce_{ops} while exploring the package itself and understanding what it takes. I actually went in the path you suggested first, to check if its possible to generate code for custom julia function.

Also, the package I have, with CUDA.jl based custom kernels, is unique on its own and I don't have a baseline to compare with a reference implementation like other DL and AI package and pipelines enjoy. I am hoping this library will also help me prepare a baseline too. Otherwise I will have to go with c++ path and deal with permutation and combinations of configuration, kernel splitting, fusion etc, which is daunting than actual research work.

Appreciate the work you have put into this.

[AntonOresten]

If this is replaced by #37, could you close this one?

[arhik]

#37 is partial simpler PR. Closing this and will open multiple PR meeting this.