Fix SLEB128 zigzag encoding for 64-bit and small integer types

arhik · arhik · commit be4265f01cf5 · 2026-01-16T14:03:17.000Z
- writer.jl: Changed IntegerIdentityOp.value from Int64 to UInt128 to store full
  64-bit unsigned values. Added mask_to_width() to mask values to correct bit
  width before zigzag encoding.

- core.jl: Added to_uint128() helper to convert signed/unsigned values to
  UInt128 via bit reinterpretation for proper identity value storage.

- examples/reducekernel.jl: Added comprehensive tests for all reduce operations
  (min, max, sum, xor, or, and) on UInt16/32/64, Int16/32/64, and Float16/32/64.

Fixes:
- UInt64 reduce_min and reduce_and now work correctly
- Int16 reduce_max and reduce_and now work correctly
- All small integer types (Int8, Int16, Int32) now encode properly with SLEB128
diff --git a/examples/reducekernel.jl b/examples/reducekernel.jl
@@ -3,18 +3,151 @@ using CUDA
 using cuTile
 import cuTile as ct
 
-elType = UInt16
-function reduceKernel(a::ct.TileArray{elType,1}, b::ct.TileArray{elType,1}, tileSz::ct.Constant{Int})
-    bid = ct.bid(1)
-    tile = ct.load(a, bid, (tileSz[],))
-    result = ct.reduce_min(tile, Val(1))
-    ct.store(b, bid, result)
-    return nothing
+# Kernel factory to properly capture element type and operation
+function makeReduceKernel(::Type{T}, op::Symbol) where {T}
+    reduceFunc = if op == :reduce_min
+        ct.reduce_min
+    elseif op == :reduce_max
+        ct.reduce_max
+    elseif op == :reduce_sum
+        ct.reduce_sum
+    elseif op == :reduce_xor
+        ct.reduce_xor
+    elseif op == :reduce_or
+        ct.reduce_or
+    elseif op == :reduce_and
+        ct.reduce_and
+    end
+
+    @inline function kernel(a::ct.TileArray{T,1}, b::ct.TileArray{T,1}, tileSz::ct.Constant{Int})
+        ct.store(b, ct.bid(1), reduceFunc(ct.load(a, ct.bid(1), (tileSz[],)), Val(1)))
+        return nothing
+    end
+    return kernel
+end
+
+# Test with UInt types
+@testset for elType in [UInt16, UInt32, UInt64]
+    @testset for op in [:reduce_min, :reduce_max, :reduce_sum, :reduce_xor, :reduce_or, :reduce_and]
+        sz = 32
+        N = 2^15
+
+        # Create kernel using factory
+        reduceKernel = try
+            makeReduceKernel(elType, op)
+        catch e
+            @test_broken false
+            rethrow()
+        end
+
+        # Create data and run kernel
+        a_gpu = CUDA.rand(elType, N)
+        b_gpu = CUDA.zeros(elType, cld(N, sz))
+        try
+            CUDA.@sync ct.launch(reduceKernel, cld(length(a_gpu), sz), a_gpu, b_gpu, ct.Constant(sz))
+        catch e
+            @test_broken false
+            rethrow()
+        end
+        res = Array(b_gpu)
+
+        # CPU computation
+        a_cpu = Array(a_gpu)
+        a_reshaped = reshape(a_cpu, sz, :)
+
+        if op == :reduce_min
+            cpu_result = minimum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_max
+            cpu_result = maximum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_sum
+            raw_sum = sum(a_reshaped, dims=1)[:]
+            cpu_result = raw_sum .& typemax(elType)
+        elseif op == :reduce_xor
+            cpu_result = mapslices(x -> reduce(⊻, x), a_reshaped, dims=1)[:]
+        elseif op == :reduce_or
+            cpu_result = mapslices(x -> reduce(|, x), a_reshaped, dims=1)[:]
+        elseif op == :reduce_and
+            cpu_result = mapslices(x -> reduce(&, x), a_reshaped, dims=1)[:]
+        end
+
+        @test cpu_result == res
+    end
+end
+
+# Test with signed Int types
+@testset for elType in [Int16, Int32, Int64]
+    @testset for op in [:reduce_min, :reduce_max, :reduce_sum, :reduce_xor, :reduce_or, :reduce_and]
+        sz = 32
+        N = 2^15
+
+        # Create kernel using factory
+        reduceKernel = try
+            makeReduceKernel(elType, op)
+        catch e
+            @test_broken false
+            rethrow()
+        end
+
+        # Create data and run kernel - use range to get negative values too
+        a_gpu = CuArray{elType}(rand(-1000:1000, N))
+        b_gpu = CUDA.zeros(elType, cld(N, sz))
+        try
+            CUDA.@sync ct.launch(reduceKernel, cld(length(a_gpu), sz), a_gpu, b_gpu, ct.Constant(sz))
+        catch e
+            @test_broken false
+            rethrow()
+        end
+        res = Array(b_gpu)
+
+        # CPU computation
+        a_cpu = Array(a_gpu)
+        a_reshaped = reshape(a_cpu, sz, :)
+
+        if op == :reduce_min
+            cpu_result = minimum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_max
+            cpu_result = maximum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_sum
+            cpu_result = sum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_xor
+            cpu_result = mapslices(x -> reduce(⊻, x), a_reshaped, dims=1)[:]
+        elseif op == :reduce_or
+            cpu_result = mapslices(x -> reduce(|, x), a_reshaped, dims=1)[:]
+        elseif op == :reduce_and
+            cpu_result = mapslices(x -> reduce(&, x), a_reshaped, dims=1)[:]
+        end
+
+        @test cpu_result == res
+    end
 end
 
-sz = 32
-N = 2^15
-a = CUDA.rand(elType, N)
-b = CUDA.zeros(elType, cld(N, sz))
-CUDA.@sync ct.launch(reduceKernel, cld(length(a), sz), a, b, ct.Constant(sz))
-res = Array(b)
+# Test with Float types
+@testset for elType in [Float16, Float32, Float64]
+    @testset for op in [:reduce_min, :reduce_max, :reduce_sum]
+        sz = 32
+        N = 2^15
+
+        # Create kernel using factory
+        reduceKernel = makeReduceKernel(elType, op)
+
+        # Create data and run kernel
+        a_gpu = CUDA.rand(elType, N)
+        b_gpu = CUDA.zeros(elType, cld(N, sz))
+        CUDA.@sync ct.launch(reduceKernel, cld(length(a_gpu), sz), a_gpu, b_gpu, ct.Constant(sz))
+        res = Array(b_gpu)
+
+        # CPU computation
+        a_cpu = Array(a_gpu)
+        a_reshaped = reshape(a_cpu, sz, :)
+
+        if op == :reduce_min
+            cpu_result = minimum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_max
+            cpu_result = maximum(a_reshaped, dims=1)[:]
+        elseif op == :reduce_sum
+            cpu_result = sum(a_reshaped, dims=1)[:]
+        end
+
+        @test isapprox(cpu_result, res)
+    end
+end
diff --git a/src/bytecode/writer.jl b/src/bytecode/writer.jl
@@ -257,7 +257,7 @@ end
 Integer identity value for binary operations.
 """
 struct IntegerIdentityOp <: IdentityOp
-    value::Int64  # Store as signed Int64, will be reinterpreted as unsigned
+    value::UInt128  # Store as UInt128 to handle all unsigned values up to 64 bits
     type_id::TypeId
     dtype::Type   # Int8, Int16, Int32, Int64, UInt8, etc.
     signed::Bool  # true for signed, false for unsigned
@@ -291,13 +291,47 @@ function encode_tagged_int!(cb::CodeBuilder, identity::IntegerIdentityOp)
     # Type ID
     encode_typeid!(cb.buf, identity.type_id)
     # Value: signed uses zigzag varint, unsigned uses plain varint
+    # Mask value to correct bit width and apply zigzag for signed types
+    masked_value = mask_to_width(identity.value, identity.dtype, identity.signed)
     if identity.signed
-        encode_signed_varint!(cb.buf, identity.value)
+        encode_signed_varint!(cb.buf, masked_value)
     else
-        encode_varint!(cb.buf, UInt64(identity.value))
+        encode_varint!(cb.buf, masked_value)
     end
 end
 
+"""
+    mask_to_width(value, dtype, signed)
+
+Mask a UInt128 value to the correct bit width for the given type and apply zigzag if signed.
+For signed types, this masks first, then applies zigzag encoding.
+"""
+# Signed Int64: mask to 64 bits first, then zigzag encode
+mask_to_width(value::UInt128, ::Type{Int64}, signed::Bool) = 
+    let masked = UInt64(value & 0xFFFFFFFFFFFFFFFF)
+        UInt64((masked << 1) ⊻ (masked >>> 63))
+    end
+# Signed Int32: mask to 32 bits first, then zigzag encode
+mask_to_width(value::UInt128, ::Type{Int32}, signed::Bool) = 
+    let masked = UInt32(value & 0xFFFFFFFF)
+        UInt32((masked << 1) ⊻ (masked >>> 31))
+    end
+# Signed Int16: mask to 16 bits first, then zigzag encode
+mask_to_width(value::UInt128, ::Type{Int16}, signed::Bool) = 
+    let masked = UInt16(value & 0xFFFF)
+        UInt16((masked << 1) ⊻ (masked >>> 15))
+    end
+# Signed Int8: mask to 8 bits first, then zigzag encode
+mask_to_width(value::UInt128, ::Type{Int8}, signed::Bool) = 
+    let masked = UInt8(value & 0xFF)
+        UInt8((masked << 1) ⊻ (masked >>> 7))
+    end
+# Unsigned types: just mask to bit width, no zigzag
+mask_to_width(value::UInt128, ::Type{UInt64}, signed::Bool) = UInt64(value & 0xFFFFFFFFFFFFFFFF)
+mask_to_width(value::UInt128, ::Type{UInt32}, signed::Bool) = UInt32(value & 0xFFFFFFFF)
+mask_to_width(value::UInt128, ::Type{UInt16}, signed::Bool) = UInt16(value & 0xFFFF)
+mask_to_width(value::UInt128, ::Type{UInt8}, signed::Bool) = UInt8(value & 0xFF)
+
 """
     float_to_bits(value, dtype)
 
@@ -585,7 +619,7 @@ function finalize_function!(func_buf::Vector{UInt8}, cb::CodeBuilder,
 end
 
 #=============================================================================
- Optimization Hints 
+ Optimization Hints
 =============================================================================#
 
 """
diff --git a/src/compiler/intrinsics/core.jl b/src/compiler/intrinsics/core.jl
@@ -667,43 +667,61 @@ is_signed(::Type{T}) where T <: AbstractFloat = false
 Return the identity value for a binary operation (reduce, scan, etc.).
 Identity must satisfy: identity ⊕ x = x for the operation.
 """
+
+"""
+    to_uint128(value, dtype)
+
+Convert an integer value to UInt128 for storage in IntegerIdentityOp.
+For signed types, this returns the two's complement bit representation.
+"""
+# Unsigned types: directly convert
+to_uint128(value::UInt64) = UInt128(value)
+to_uint128(value::UInt32) = UInt128(value)
+to_uint128(value::UInt16) = UInt128(value)
+to_uint128(value::UInt8) = UInt128(value)
+# Signed types: reinterpret as unsigned first, then convert
+to_uint128(value::Int64) = UInt128(reinterpret(UInt64, value))
+to_uint128(value::Int32) = UInt128(reinterpret(UInt32, value))
+to_uint128(value::Int16) = UInt128(reinterpret(UInt16, value))
+to_uint128(value::Int8) = UInt128(reinterpret(UInt8, value))
+
 # Addition identity: 0 + x = x
 operation_identity(::Val{:add}, dtype, ::Type{T}) where T <: AbstractFloat =
     FloatIdentityOp(zero(T), dtype, T)
 operation_identity(::Val{:add}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(zero(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(zero(T)), dtype, T, is_signed(T))
 
 # Maximum identity: max(typemin(T), x) = x
 operation_identity(::Val{:max}, dtype, ::Type{T}) where T <: AbstractFloat =
     FloatIdentityOp(typemin(T), dtype, T)
 operation_identity(::Val{:max}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(typemin(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(typemin(T)), dtype, T, is_signed(T))
 
 # Multiplication identity: 1 * x = x
 operation_identity(::Val{:mul}, dtype, ::Type{T}) where T <: AbstractFloat =
     FloatIdentityOp(one(T), dtype, T)
 operation_identity(::Val{:mul}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(one(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(one(T)), dtype, T, is_signed(T))
 
 # Minimum identity: min(typemax(T), x) = x
 operation_identity(::Val{:min}, dtype, ::Type{T}) where T <: AbstractFloat =
     FloatIdentityOp(typemax(T), dtype, T)
 operation_identity(::Val{:min}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(typemax(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(typemax(T)), dtype, T, is_signed(T))
 
 # AND identity: all bits set (x & identity == x)
 # For signed: -one(T) has all bits set in two's complement
 # For unsigned: typemax(T) has all bits set
 operation_identity(::Val{:and}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(is_signed(T) ? -one(T) : typemax(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(is_signed(T) ? -one(T) : typemax(T)), dtype, T, is_signed(T))
 
 # OR identity: 0 | x = x
 operation_identity(::Val{:or}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(zero(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(zero(T)), dtype, T, is_signed(T))
 
 # XOR identity: 0 ⊕ x = x
 operation_identity(::Val{:xor}, dtype, ::Type{T}) where T <: Integer =
-    IntegerIdentityOp(zero(T), dtype, T, is_signed(T))
+    IntegerIdentityOp(to_uint128(zero(T)), dtype, T, is_signed(T))
 
 #=============================================================================
  Reduce Body Operations - dispatch on Val{fn} and elem_type
