Support primitive types with unused bits (Julia 1.14)

README.md

@@ -12,6 +12,8 @@ BitPacking.jl provides small bit-packed building blocks for Julia array code:
 - `NarrowTuple` and `@NarrowTuple` pack heterogeneous isbits values, with optional padding fields.
 
 Packages can extend `BitPacking.bitwidth(::Type)` for their own narrow scalar types.
+On Julia 1.14 and later, a primitive type declared with its real width, such as
+`primitive type Float4 4 end`, needs no overload.
 
 ## Examples
 

docs/src/index.md

@@ -11,6 +11,8 @@ Documentation for [BitPacking](https://github.com/MurrellGroup/BitPacking.jl).
 
 ```@docs
 BitPacking.bitwidth
+BitPacking.raw_bits
+BitPacking.from_raw_bits
 ```
 
 ## Narrow-element arrays

src/BitPacking.jl

@@ -3,27 +3,69 @@ module BitPacking
 using Republic
 import Adapt
 
-@public bitwidth
+@public bitwidth, raw_bits, from_raw_bits
 export NArray, NVector, NMatrix
 export NarrowArray, NarrowVector, NarrowMatrix
 export Narrow
 export NarrowTuple, @NarrowTuple
 @public Pad, ZeroPad, OnePad
 
+@static if isdefined(Core, :bitsizeof)
+    # Julia 1.14 lets a primitive type declare fewer bits than it stores, as in
+    # `primitive type Float4 4 end`, which occupies a byte.
+    bitsizeof(::Type{T}) where T = Core.bitsizeof(T)
+else
+    bitsizeof(::Type{T}) where T = 8 * sizeof(T)
+end
+
+has_unused_bits(::Type{T}) where T = isprimitivetype(T) && bitsizeof(T) < 8 * sizeof(T)
+
 """
     bitwidth(T)::Int
     bitwidth(x)::Int
 
 Return the number of value bits used by `T` or by the type of `x`.
 
-The default for bitstypes is `8 * sizeof(T)`, with `Bool` specialized to
-one bit. Packages and users can overload this for narrow primitive types,
-packed containers, and layout markers such as [`ZeroPad`](@ref).
+The default for bitstypes is `Core.bitsizeof(T)` on Julia 1.14 and later, which
+is the declared width of a primitive type such as `primitive type Float4 4 end`,
+and `8 * sizeof(T)` before that. `Bool` is specialized to one bit. Packages and
+users can overload this for narrow types declared with whole bytes, packed
+containers, and layout markers such as [`ZeroPad`](@ref).
 """
-bitwidth(::Type{T}) where T = isbitstype(T) ? sizeof(T) * 8 : error("$T is not a bitstype")
+bitwidth(::Type{T}) where T = isbitstype(T) ? bitsizeof(T) : error("$T is not a bitstype")
 bitwidth(::T) where T = bitwidth(T)
 bitwidth(::Type{Bool}) = 1
 
+"""
+    raw_bits(x)::Unsigned
+
+Return the bits of `x` as the unsigned integer type of the same size, like
+`reinterpret(UInt8, x)` for a one-byte `x`.
+
+Unlike `reinterpret`, this also accepts primitive types that declare fewer bits
+than they store, such as `primitive type Float4 4 end` on Julia 1.14 and later.
+Their unused high bits read as zero.
+
+See also [`from_raw_bits`](@ref).
+"""
+@inline function raw_bits(x::T) where T
+    U = unsigned_type(sizeof(T))
+    return has_unused_bits(T) ? Core.Intrinsics.zext_int(U, x) : reinterpret(U, x)
+end
+
+"""
+    from_raw_bits(T, bits::Unsigned)::T
+
+Return the `T` whose bits are `bits`, an unsigned integer of `sizeof(T)` bytes.
+The inverse of [`raw_bits`](@ref).
+
+For a primitive type that declares fewer bits than it stores, bits of `bits`
+beyond its width are dropped.
+"""
+@inline function from_raw_bits(::Type{T}, bits::Unsigned) where T
+    return has_unused_bits(T) ? Core.Intrinsics.trunc_int(T, bits) : reinterpret(T, bits)
+end
+
 include("packing.jl")
 include("NArray.jl")
 include("NarrowTuple.jl")

src/NArray.jl

@@ -68,7 +68,7 @@ NArray{T,N,S}(data::D) where {T,N,S<:NTuple{N,Any},D<:Unsigned} =
 bitwidth(::Type{T}) where T<:NArray = bitwidth(eltype(T)) * length(T)
 
 function NArray{T,N,S}(xs::NTuple{L,T}) where {T,N,S<:NTuple{N,Any},L}
-    unpacked_bytes = reinterpret.(unsigned_type(sizeof(T)), xs)
+    unpacked_bytes = raw_bits.(xs)
     data = pack(Val(bitwidth(T)), unpacked_bytes)
     arr = NArray{T,N,S,typeof(data)}(data)
     length(arr) == L || throw(ArgumentError("Type expects $(length(arr)) element(s), but got $L."))
@@ -93,7 +93,7 @@ NVector(xs::NTuple{L,T}) where {L,T} = NVector{T,L}(xs)
 
 function Base.Tuple(v::NArray)
     unpacked_bytes = unpack(Val(bitwidth(eltype(v))), Val(length(v)), v.data)
-    xs = reinterpret.(eltype(v), unpacked_bytes)::NTuple{length(v)}
+    xs = from_raw_bits.(eltype(v), unpacked_bytes)::NTuple{length(v)}
     return xs
 end
 
@@ -114,7 +114,7 @@ function Base.getindex(arr::NArray{T}, i::Int) where T
     W = bitwidth(T)
     S = length(arr)
     u = packed_getindex(Val(W), Val(S), arr.data, i)
-    return reinterpret(T, u)
+    return from_raw_bits(T, u)
 end
 
 Base.broadcastable(v::NArray{<:Any,<:Any,S}) where S = SArray{S}(Tuple(v))

src/NarrowArray.jl

@@ -68,13 +68,36 @@ function NarrowArray{T,N,L}(arr::AbstractArray{S,N}) where {T,N,L,S}
         throw(ArgumentError("the first dimension of a NarrowArray{$T,$N,$L} input must be divisible by $chunk_length, got $(size(arr, 1))"))
 
     chunks = if S === T
-        NVector{T,chunk_length}.(reinterpret(NTuple{chunk_length,T}, arr))
+        NVector{T,chunk_length}.(dense_chunks(Val(chunk_length), arr))
     else
-        NVector{T,chunk_length}.(SVector{chunk_length,S}.(reinterpret(NTuple{chunk_length,S}, arr)))
+        NVector{T,chunk_length}.(SVector{chunk_length,S}.(dense_chunks(Val(chunk_length), arr)))
     end
     return NarrowArray(chunks)
 end
 
+# `reinterpret` refuses arrays of a primitive type with unused bits (Julia
+# 1.14+), so for those the conversions between dense values and chunks go
+# through raw bits instead, at the cost of a temporary. They map closures over
+# `T` because `from_raw_bits.(T, x)` does not infer when `T` is a static
+# parameter, and GPU broadcasts need it to.
+
+# The first dimension of `dense` as runs of `L` values, each an `NTuple`.
+function dense_chunks(::Val{L}, dense::AbstractArray{T}) where {L,T}
+    has_unused_bits(T) || return reinterpret(NTuple{L,T}, dense)
+    raw = reinterpret(NTuple{L,unsigned_type(sizeof(T))}, raw_bits.(dense))
+    return Broadcast.broadcasted(chunk -> map(bits -> from_raw_bits(T, bits), chunk), raw)
+end
+
+# The inverse of `dense_chunks`: runs of values, flattened into the first
+# dimension.
+function dense_values(::Type{T}, chunks::AbstractArray) where T
+    has_unused_bits(T) || return reinterpret(T, chunks)
+    raw = reinterpret(unsigned_type(sizeof(T)), raw_chunk.(chunks))
+    return map(bits -> from_raw_bits(T, bits), raw)
+end
+
+raw_chunk(xs) = raw_bits.(Tuple(xs))
+
 NarrowArray{T}(arr::NarrowArray{T}) where T = arr
 NarrowArray{T,N}(arr::NarrowArray{T,N}) where {T,N} = arr
 NarrowArray{T,N,L}(arr::NarrowArray{T,N,L}) where {T,N,L} = arr
@@ -142,10 +165,10 @@ function Broadcast.broadcasted(::Type{T}, arr::NarrowArray) where T
 
     L = length(eltype(parent(arr)))
     chunks = SVector{L,T}.(parent(arr))
-    return reinterpret(T, chunks)
+    return dense_values(T, chunks)
 end
 
-Base.copy(arr::NarrowArray) = reinterpret(eltype(arr), map(SArray, parent(arr)))
+Base.copy(arr::NarrowArray{T,<:Any,L}) where {T,L} = dense_values(T, SVector{L,T}.(parent(arr)))
 
 """
     Narrow(T)
@@ -184,9 +207,9 @@ Narrow(::Type{T}) where T = Narrow{T}()
 # straight into the existing `chunks`, so no packed temporary is allocated.
 function _pack_into!(chunks, ::Type{T}, ::Val{L}, dense::AbstractArray{S}) where {T,L,S}
     if S === T
-        chunks .= NVector{T,L}.(reinterpret(NTuple{L,T}, dense))
+        chunks .= NVector{T,L}.(dense_chunks(Val(L), dense))
     else
-        chunks .= NVector{T,L}.(SVector{L,S}.(reinterpret(NTuple{L,S}, dense)))
+        chunks .= NVector{T,L}.(SVector{L,S}.(dense_chunks(Val(L), dense)))
     end
     return chunks
 end

src/NarrowTuple.jl

@@ -157,10 +157,7 @@ function narrow_mask(::Type{U}, width::Int) where U<:Unsigned
     return U((big(1) << width) - 1)
 end
 
-function field_raw_expr(name, index, T)
-    U = unsigned_type(sizeof(T))
-    return :($U(reinterpret($U, Core.getfield($name, $index))))
-end
+field_raw_expr(name, index) = :(raw_bits(Core.getfield($name, $index)))
 
 function _typed_or_expr(::Type{U}, terms) where U<:Unsigned
     isempty(terms) && return :($U(0))
@@ -259,7 +256,7 @@ bitwidth(::Type{<:NarrowTuple{Ts}}) where Ts = sum(bitwidth, Ts.parameters; init
                 expr = :($D($(narrow_mask(D, W))))
             else
                 U = unsigned_type(sizeof(T))
-                expr = field_raw_expr(:xs, index, T)
+                expr = field_raw_expr(:xs, index)
                 W < 8 * sizeof(U) && (expr = :($expr & $(narrow_mask(U, W))))
                 expr = :($D($expr))
             end
@@ -290,7 +287,7 @@ bitwidth(::Type{<:NarrowTuple{Ts}}) where Ts = sum(bitwidth, Ts.parameters; init
                 T <: OnePad || continue
                 expr = :(UInt8($(_low_mask(width))))
             else
-                expr = field_raw_expr(:xs, field_index, T)
+                expr = field_raw_expr(:xs, field_index)
                 source_shift > 0 && (expr = :($expr >>> $source_shift))
                 width < 8 && (expr = :($expr & $(_low_mask(width))))
                 expr = :(UInt8($expr))
@@ -366,7 +363,7 @@ NarrowTuple(xs::Ts) where Ts<:Tuple = NarrowTuple{layout_tuple_type(Ts)}(xs)
             start > 0 && (expr = :($expr >>> $start))
             W < 8 * sizeof(D) && (expr = :($expr & $(narrow_mask(D, W))))
             expr = :($U($expr))
-            push!(values, :(reinterpret($T, $expr)))
+            push!(values, :(from_raw_bits($T, $expr)))
         end
         return Expr(:tuple, values...)
     end
@@ -402,7 +399,7 @@ NarrowTuple(xs::Ts) where Ts<:Tuple = NarrowTuple{layout_tuple_type(Ts)}(xs)
         end
 
         expr = _typed_or_expr(U, terms)
-        push!(values, :(reinterpret($T, $expr)))
+        push!(values, :(from_raw_bits($T, $expr)))
     end
 
     return Expr(:tuple, values...)

test/runtests.jl

@@ -436,6 +436,8 @@ bits(x) = x
         @test contains(sprint(show, MIME("text/plain"), arr3), "NarrowArray")
     end
 
+    isdefined(Core, :bitsizeof) && include("unused_bits.jl")
+
     TEST_CUTILE && include("cutile.jl")
 
 end

test/unused_bits.jl

@@ -0,0 +1,96 @@
+# Julia 1.14 lets a primitive type declare fewer bits than it stores. These
+# need no `bitwidth` overload, and `reinterpret` refuses them.
+primitive type NativeFloat4 4 end
+primitive type NativeFloat6 6 end
+primitive type NativeUInt24 24 end
+
+native4(x::UInt8) = Core.Intrinsics.trunc_int(NativeFloat4, x)
+native6(x::UInt8) = Core.Intrinsics.trunc_int(NativeFloat6, x)
+Base.convert(::Type{NativeFloat4}, x::UInt8) = native4(x)
+Base.convert(::Type{UInt8}, x::NativeFloat4) = BitPacking.raw_bits(x)
+
+@testset "Primitive types with unused bits" begin
+    @testset "bitwidth and raw bits" begin
+        @test BitPacking.bitwidth(NativeFloat4) == 4
+        @test BitPacking.bitwidth(NativeFloat6) == 6
+        @test BitPacking.bitwidth(NativeUInt24) == 24
+        @test BitPacking.bitwidth(native4(0x01)) == 4
+
+        x = native4(0xfa)
+        @test BitPacking.raw_bits(x) === 0x0a
+        @test BitPacking.from_raw_bits(NativeFloat4, 0x0a) === x
+        @test BitPacking.from_raw_bits(NativeFloat4, 0xfa) === x
+        @test BitPacking.raw_bits(Core.Intrinsics.trunc_int(NativeUInt24, 0xffbcdef0)) === 0x00bcdef0
+
+        @test BitPacking.raw_bits(0x0a) === 0x0a
+        @test BitPacking.raw_bits(true) === 0x01
+        @test BitPacking.from_raw_bits(Bool, 0x01) === true
+        @test BitPacking.from_raw_bits(Int16, 0xffff) === Int16(-1)
+    end
+
+    @testset "NArray" begin
+        v = NVector((native4(0x0a), native4(0x03)))
+        @test v isa NVector{NativeFloat4,2}
+        @test v.data == 0x3a
+        @test Tuple(v) === (native4(0x0a), native4(0x03))
+        @test v[1] === native4(0x0a)
+        @test v[2] === native4(0x03)
+
+        v6 = NVector{NativeFloat6,4}(native6.((0x01, 0x02, 0x03, 0x04)))
+        @test v6.data == (0x81, 0x30, 0x10)
+        @test Tuple(v6) === native6.((0x01, 0x02, 0x03, 0x04))
+    end
+
+    @testset "NarrowTuple" begin
+        x24 = Core.Intrinsics.trunc_int(NativeUInt24, 0x00bcdef0)
+        T = @NarrowTuple{NativeFloat4,NativeUInt24,NativeFloat4}
+        @test T === NarrowTuple{Tuple{NativeFloat4,NativeUInt24,NativeFloat4},UInt32}
+        nt = T(native4(0x0a), x24, native4(0x03))
+        @test nt.data == 0x3bcdef0a
+        @test Tuple(nt) === (native4(0x0a), x24, native4(0x03))
+
+        nt6 = NarrowTuple((native6(0x2a), 0xbc))
+        @test nt6.data == (0x2a, 0x2f)
+        @test Tuple(nt6) === (native6(0x2a), 0xbc)
+    end
+
+    @testset "NarrowArray" begin
+        values = native4.(UInt8[0x01, 0x02, 0x03, 0x04])
+        narrow = NarrowArray{NativeFloat4}(values)
+        @test narrow isa NarrowVector{NativeFloat4}
+        @test collect(reinterpret(UInt8, narrow)) == UInt8[0x21, 0x43]
+        @test narrow[3] === native4(0x03)
+        @test copy(narrow) isa Vector{NativeFloat4}
+        @test copy(narrow) == values
+        @test NativeFloat4.(narrow) == values
+
+        # GPU broadcasts need these to infer
+        @test @inferred(NarrowArray{NativeFloat4}(values)) == narrow
+        @test @inferred(copy(narrow)) == values
+        @test @inferred(broadcast(NativeFloat4, narrow)) == values
+
+        # conversions on either side of the packing
+        @test NarrowArray{NativeFloat4}(UInt8[0x01, 0x02, 0x03, 0x04]) == narrow
+        @test copy(NarrowArray{UInt8}(values)) == UInt8[0x01, 0x02, 0x03, 0x04]
+        @test copy(NarrowArray{UInt8}(narrow)) == UInt8[0x01, 0x02, 0x03, 0x04]
+
+        matrix = NarrowArray{NativeFloat4}(reshape(values, 2, 2))
+        @test collect(reinterpret(UInt8, matrix)) == UInt8[0x21 0x43]
+        @test copy(matrix) == reshape(values, 2, 2)
+
+        values6 = native6.(UInt8[0x01, 0x02, 0x03, 0x04])
+        narrow6 = NarrowArray{NativeFloat6}(values6)
+        @test collect(reinterpret(UInt8, narrow6)) == UInt8[0x81, 0x30, 0x10]
+        @test copy(narrow6) == values6
+        @test copy(reinterpret(NativeFloat6, reinterpret(Bool, narrow6))) == values6
+
+        @test Narrow(NativeFloat4).(values) == narrow
+        @test copy(reinterpret(Narrow(NativeFloat4), UInt8[0x21, 0x43])) == values
+
+        dest = similar(narrow)
+        dest .= reverse(values)
+        @test copy(dest) == reverse(values)
+        dest .= UInt8[0x01, 0x02, 0x03, 0x04]
+        @test copy(dest) == values
+    end
+end