Improve unpack_bits (#165)

svilupp · Jun 18, 2024 · 4c6a4d0 · 4c6a4d0 · svilupp · Jun 18, 2024
1 parent eff6682
commit 4c6a4d0
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 3 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -8,7 +8,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Added
 
-### Fixed
+## [0.31.1]
+
+### Updated
+- Improved the implementation of `RAGTools.unpack_bits` to be faster with fewer allocations.
 
 ## [0.31.0]
 

diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "PromptingTools"
 uuid = "670122d1-24a8-4d70-bfce-740807c42192"
 authors = ["J S @svilupp and contributors"]
-version = "0.31.0"
+version = "0.31.1"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"

diff --git a/src/Experimental/RAGTools/utils.jl b/src/Experimental/RAGTools/utils.jl
@@ -537,8 +537,20 @@ binx = unpack_bits(binint)
 @assert bin == binx
 ```
 """
+# function unpack_bits(packed_vector::AbstractVector{UInt64})
+#     return Bool[((x >> i) & 1) == 1 for x in packed_vector for i in 0:63]
+# end
 function unpack_bits(packed_vector::AbstractVector{UInt64})
-    return Bool[((x >> i) & 1) == 1 for x in packed_vector for i in 0:63]
+    n = length(packed_vector)
+    result = Vector{Bool}(undef, n * 64)
+    @inbounds @simd for i in 1:n
+        x = packed_vector[i]
+        for j in 1:64
+            result[(i - 1) * 64 + j] = (x & 1) == 1
+            x >>= 1
+        end
+    end
+    return result
 end
 function unpack_bits(packed_matrix::AbstractMatrix{UInt64})
     num_rows, num_cols = size(packed_matrix)