Skip to content

Commit 26261ac

Browse files
ChrisRackauckasChrisRackauckas
authored
Merge pull request #42 from ChrisRackauckas-Claude/perf-improvements-20260102-050710
Performance: Zero-allocation colsupport using DisjointRange
2 parents 4b7a9d1 + 28a2e25 commit 26261ac

4 files changed

Lines changed: 216 additions & 1 deletion

File tree

‎src/FastAlmostBandedMatrices.jl‎

Lines changed: 75 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,79 @@ import MatrixFactorizations: QR, QRPackedQ, getQ, getR, QRPackedQLayout, AdjQRPa
1414

1515
@reexport using BandedMatrices
1616

17+
# ------------------
18+
# DisjointRange - for zero-allocation colsupport
19+
# ------------------
20+
21+
"""
22+
DisjointRange{T}
23+
24+
A lazy representation of the union of two ranges, supporting iteration and indexing
25+
without heap allocation.
26+
"""
27+
struct DisjointRange{T<:Integer, R1<:AbstractUnitRange{T}, R2<:AbstractUnitRange{T}} <:
28+
AbstractVector{T}
29+
r1::R1
30+
r2::R2
31+
end
32+
33+
Base.size(d::DisjointRange) = (length(d.r1) + length(d.r2),)
34+
Base.length(d::DisjointRange) = length(d.r1) + length(d.r2)
35+
36+
@inline function Base.getindex(d::DisjointRange, i::Integer)
37+
@boundscheck checkbounds(d, i)
38+
n1 = length(d.r1)
39+
if i <= n1
40+
return @inbounds d.r1[i]
41+
else
42+
return @inbounds d.r2[i - n1]
43+
end
44+
end
45+
46+
Base.IndexStyle(::Type{<:DisjointRange}) = IndexLinear()
47+
48+
@inline function Base.iterate(d::DisjointRange)
49+
if !isempty(d.r1)
50+
val, state = iterate(d.r1)
51+
return val, (1, state)
52+
elseif !isempty(d.r2)
53+
val, state = iterate(d.r2)
54+
return val, (2, state)
55+
else
56+
return nothing
57+
end
58+
end
59+
60+
@inline function Base.iterate(d::DisjointRange, state)
61+
which, inner_state = state
62+
if which == 1
63+
next = iterate(d.r1, inner_state)
64+
if next !== nothing
65+
return next[1], (1, next[2])
66+
else
67+
# Switch to r2
68+
if !isempty(d.r2)
69+
val, new_state = iterate(d.r2)
70+
return val, (2, new_state)
71+
else
72+
return nothing
73+
end
74+
end
75+
else
76+
next = iterate(d.r2, inner_state)
77+
if next !== nothing
78+
return next[1], (2, next[2])
79+
else
80+
return nothing
81+
end
82+
end
83+
end
84+
85+
Base.first(d::DisjointRange) = isempty(d.r1) ? first(d.r2) : first(d.r1)
86+
Base.last(d::DisjointRange) = isempty(d.r2) ? last(d.r1) : last(d.r2)
87+
Base.minimum(d::DisjointRange) = min(minimum(d.r1), minimum(d.r2))
88+
Base.maximum(d::DisjointRange) = max(maximum(d.r1), maximum(d.r2))
89+
1790
# ------------------
1891
# AlmostBandedMatrix
1992
# ------------------
@@ -185,7 +258,8 @@ end
185258
if isempty(sup)
186259
return Base.OneTo(r)
187260
else
188-
return vcat(Base.OneTo(min(r, minimum(sup) - 1)), sup)
261+
# Use DisjointRange to avoid heap allocation from vcat
262+
return DisjointRange(Base.OneTo(min(r, minimum(sup) - 1)), sup)
189263
end
190264
end
191265
end

‎test/Project.toml‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
[deps]
2+
AllocCheck = "9b6a8646-10ed-4001-bbdc-1d2f46dfbb1a"
23
Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
34
ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"
5+
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
46
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
57
MatrixFactorizations = "a3b82374-2e81-5b9e-98ce-41277c0e4c87"
68
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"

‎test/alloc_tests.jl‎

Lines changed: 134 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,134 @@
1+
using AllocCheck
2+
using BenchmarkTools
3+
using FastAlmostBandedMatrices
4+
using FastAlmostBandedMatrices: DisjointRange
5+
using ArrayLayouts: colsupport, rowsupport
6+
using Test
7+
8+
@testset "Allocation Tests" begin
9+
@testset "DisjointRange - Zero Allocations" begin
10+
# Test that DisjointRange operations don't allocate
11+
r1 = Base.OneTo(5)
12+
r2 = 10:15
13+
dr = DisjointRange(r1, r2)
14+
15+
# Test length
16+
allocs = @allocated length(dr)
17+
@test allocs == 0
18+
19+
# Test getindex
20+
allocs = @allocated dr[3]
21+
@test allocs == 0
22+
23+
allocs = @allocated dr[8]
24+
@test allocs == 0
25+
26+
# Test first/last
27+
allocs = @allocated first(dr)
28+
@test allocs == 0
29+
30+
allocs = @allocated last(dr)
31+
@test allocs == 0
32+
33+
# Test iteration (after warmup)
34+
sum_test = 0
35+
for x in dr
36+
sum_test += x
37+
end
38+
allocs = @allocated begin
39+
s = 0
40+
for x in dr
41+
s += x
42+
end
43+
s
44+
end
45+
@test allocs == 0
46+
end
47+
48+
@testset "colsupport - Zero Allocations" begin
49+
n = 100
50+
m = 2
51+
B = brand(Float64, n, n, m + 1, m)
52+
F = rand(Float64, m, n)
53+
A = AlmostBandedMatrix(B, F)
54+
55+
# Warmup
56+
colsupport(A, 5)
57+
colsupport(A, 50)
58+
59+
# Test colsupport for j <= l+u (should return OneTo, no allocation)
60+
allocs = @allocated colsupport(A, 5)
61+
@test allocs == 0
62+
63+
# Test colsupport for j > l+u (now returns DisjointRange instead of vcat)
64+
allocs = @allocated colsupport(A, 50)
65+
@test allocs == 0
66+
end
67+
68+
@testset "rowsupport - Zero Allocations" begin
69+
n = 100
70+
m = 2
71+
B = brand(Float64, n, n, m + 1, m)
72+
F = rand(Float64, m, n)
73+
A = AlmostBandedMatrix(B, F)
74+
75+
# Warmup
76+
rowsupport(A, 1)
77+
rowsupport(A, 50)
78+
79+
# Test rowsupport (always returns UnitRange, no allocation)
80+
allocs = @allocated rowsupport(A, 1)
81+
@test allocs == 0
82+
83+
allocs = @allocated rowsupport(A, 50)
84+
@test allocs == 0
85+
end
86+
87+
@testset "getindex/setindex! - Zero Allocations" begin
88+
n = 100
89+
m = 2
90+
B = brand(Float64, n, n, m + 1, m)
91+
F = rand(Float64, m, n)
92+
A = AlmostBandedMatrix(B, F)
93+
94+
# Warmup
95+
_ = A[50, 50]
96+
A[50, 50] = 1.0
97+
98+
# Test getindex
99+
allocs = @allocated A[50, 50]
100+
@test allocs == 0
101+
102+
# Test setindex! in band part
103+
allocs = @allocated A[50, 50] = 2.0
104+
@test allocs == 0
105+
106+
# Test setindex! in fill part
107+
allocs = @allocated A[1, 50] = 3.0
108+
@test allocs == 0
109+
110+
# Test setindex! in overlapping part
111+
allocs = @allocated A[1, 1] = 4.0
112+
@test allocs == 0
113+
end
114+
115+
@testset "bandpart/fillpart - Zero Allocations" begin
116+
n = 100
117+
m = 2
118+
B = brand(Float64, n, n, m + 1, m)
119+
F = rand(Float64, m, n)
120+
A = AlmostBandedMatrix(B, F)
121+
122+
# Warmup
123+
bandpart(A)
124+
fillpart(A)
125+
126+
# Test bandpart
127+
allocs = @allocated bandpart(A)
128+
@test allocs == 0
129+
130+
# Test fillpart
131+
allocs = @allocated fillpart(A)
132+
@test allocs == 0
133+
end
134+
end

‎test/runtests.jl‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,3 +107,8 @@ using SafeTestsets, Test
107107
@test length(A1.fill.nzval) == 2
108108
end
109109
end
110+
111+
# Allocation tests run separately to avoid precompilation interference
112+
if get(ENV, "GROUP", "all") == "all" || get(ENV, "GROUP", "all") == "nopre"
113+
include("alloc_tests.jl")
114+
end

0 commit comments

Comments
 (0)