20
20
struct 🦋workspace{T}
21
21
A:: Matrix{T}
22
22
b:: Vector{T}
23
- B:: Matrix{T}
24
23
ws:: Vector{T}
25
24
U:: Matrix{T}
26
25
V:: Matrix{T}
@@ -33,18 +32,17 @@ struct 🦋workspace{T}
33
32
xn = 4 - M % 4
34
33
b = [b; rand (xn)]
35
34
end
36
- B = similar (A)
37
35
U, V = (similar (A), similar (A))
38
36
ws = 🦋generate_random! (A)
39
- new {eltype(A)} (A, b, B, ws, U, V, out)
37
+ materializeUV (U, V, ws)
38
+ new {eltype(A)} (A, b, ws, U, V, out)
40
39
end
41
40
end
42
41
43
42
function 🦋lu! (workspace: :🦋workspace, M, thread)
44
- (;A, b, B, ws, U, V, out) = workspace
45
- 🦋mul! (copyto! (B, A), ws)
46
- materializeUV (U, V, ws)
47
- F = RecursiveFactorization. lu! (B, Val (false ), thread)
43
+ (;A, b, ws, U, V, out) = workspace
44
+ 🦋mul! (A, ws)
45
+ F = RecursiveFactorization. lu! (A, Val (false ), thread)
48
46
sol = V * (F \ (U' * b))
49
47
out .= @view sol[1 : M]
50
48
out
@@ -55,14 +53,14 @@ const butterfly_workspace = 🦋workspace;
55
53
function 🦋mul_level! (A, u, v)
56
54
M, N = size (A)
57
55
@assert M == length (u) && N == length (v)
58
- Mh = M >>> 1
59
- Nh = N >>> 1
60
- @turbo for n in 1 : Nh
61
- for m in 1 : Mh
56
+ M_half = M >>> 1
57
+ N_half = N >>> 1
58
+ @turbo for n in 1 : N_half
59
+ for m in 1 : M_half
62
60
A11 = A[m, n]
63
- A21 = A[m + Mh , n]
64
- A12 = A[m, n + Nh ]
65
- A22 = A[m + Mh , n + Nh ]
61
+ A21 = A[m + M_half , n]
62
+ A12 = A[m, n + N_half ]
63
+ A22 = A[m + M_half , n + N_half ]
66
64
67
65
T1 = A11 + A12
68
66
T2 = A21 + A22
@@ -74,32 +72,32 @@ function 🦋mul_level!(A, u, v)
74
72
C22 = T3 - T4
75
73
76
74
u1 = u[m]
77
- u2 = u[m + Mh ]
75
+ u2 = u[m + M_half ]
78
76
v1 = v[n]
79
- v2 = v[n + Nh ]
77
+ v2 = v[n + N_half ]
80
78
81
79
A[m, n] = u1 * C11 * v1
82
- A[m + Mh , n] = u2 * C21 * v1
83
- A[m, n + Nh ] = u1 * C12 * v2
84
- A[m + Mh , n + Nh ] = u2 * C22 * v2
80
+ A[m + M_half , n] = u2 * C21 * v1
81
+ A[m, n + N_half ] = u1 * C12 * v2
82
+ A[m + M_half , n + N_half ] = u2 * C22 * v2
85
83
end
86
84
end
87
85
end
88
86
89
87
function 🦋mul! (A, uv)
90
88
M, N = size (A)
91
89
@assert M == N
92
- Mh = M >>> 1
90
+ M_half = M >>> 1
93
91
94
- U₁ = @view (uv[1 : Mh ])
95
- V₁ = @view (uv[(Mh + 1 ): (M)])
96
- U₂ = @view (uv[(1 + M): (M + Mh )])
97
- V₂ = @view (uv[(1 + M + Mh ): (2 * M)])
92
+ U₁ = @view (uv[1 : M_half ])
93
+ V₁ = @view (uv[(M_half + 1 ): (M)])
94
+ U₂ = @view (uv[(1 + M): (M + M_half )])
95
+ V₂ = @view (uv[(1 + M + M_half ): (2 * M)])
98
96
99
- 🦋mul_level! (@view (A[1 : Mh , 1 : Mh ]), U₁, V₁)
100
- 🦋mul_level! (@view (A[Mh + 1 : M, 1 : Mh ]), U₂, V₁)
101
- 🦋mul_level! (@view (A[1 : Mh, Mh + 1 : M]), U₁, V₂)
102
- 🦋mul_level! (@view (A[Mh + 1 : M, Mh + 1 : M]), U₂, V₂)
97
+ 🦋mul_level! (@view (A[1 : M_half , 1 : M_half ]), U₁, V₁)
98
+ 🦋mul_level! (@view (A[M_half + 1 : M, 1 : M_half ]), U₂, V₁)
99
+ 🦋mul_level! (@view (A[1 : M_half, M_half + 1 : M]), U₁, V₂)
100
+ 🦋mul_level! (@view (A[M_half + 1 : M, M_half + 1 : M]), U₂, V₂)
103
101
104
102
U = @view (uv[(1 + 2 * M): (3 * M)])
105
103
V = @view (uv[(1 + 3 * M): (4 * M)])
@@ -128,7 +126,7 @@ function 🦋!(C::SparseBandedMatrix, A::Diagonal, B::Diagonal)
128
126
C
129
127
end
130
128
131
- function 🦋2 !(C, A:: Diagonal , B:: Diagonal )
129
+ function 🦋! (C, A:: Diagonal , B:: Diagonal )
132
130
@assert size (A) == size (B)
133
131
A1 = size (A, 1 )
134
132
@@ -144,27 +142,27 @@ end
144
142
145
143
function materializeUV (U, V, uv)
146
144
M = size (U, 1 )
147
- Mh = M >>> 1
145
+ M_half = M >>> 1
148
146
149
- U₁u, U₁l = diagnegbottom (@view (uv[1 : Mh ])) # Mh
150
- U₂u, U₂l = diagnegbottom (@view (uv[(1 + 2 * Mh ): (M + Mh )])) # Mh
151
- V₁u, V₁l = diagnegbottom (@view (uv[(Mh + 1 ): (2 * Mh )])) # Mh
152
- V₂u, V₂l = diagnegbottom (@view (uv[(1 + 3 * Mh ): (2 * Mh + M)])) # Mh
147
+ U₁u, U₁l = diagnegbottom (@view (uv[1 : M_half ])) # M_half
148
+ U₂u, U₂l = diagnegbottom (@view (uv[(1 + 2 * M_half ): (M + M_half )])) # M_half
149
+ V₁u, V₁l = diagnegbottom (@view (uv[(M_half + 1 ): (2 * M_half )])) # M_half
150
+ V₂u, V₂l = diagnegbottom (@view (uv[(1 + 3 * M_half ): (2 * M_half + M)])) # M_half
153
151
Uu, Ul = diagnegbottom (@view (uv[(1 + 2 * M): (3 * M)])) # M
154
152
Vu, Vl = diagnegbottom (@view (uv[(1 + 3 * M): (4 * M)])) # M
155
153
156
154
Bu2 = SparseBandedMatrix {typeof(uv[1])} (undef, M, M)
157
155
158
- 🦋2 !(view (Bu2, 1 : Mh , 1 : Mh ), U₁u, U₁l)
159
- 🦋2 !(view (Bu2, Mh + 1 : M, Mh + 1 : M), U₂u, U₂l)
156
+ 🦋! (view (Bu2, 1 : M_half , 1 : M_half ), U₁u, U₁l)
157
+ 🦋! (view (Bu2, M_half + 1 : M, M_half + 1 : M), U₂u, U₂l)
160
158
161
159
Bu1 = SparseBandedMatrix {typeof(uv[1])} (undef, M, M)
162
160
🦋! (Bu1, Uu, Ul)
163
161
164
162
Bv2 = SparseBandedMatrix {typeof(uv[1])} (undef, M, M)
165
163
166
- 🦋2 !(view (Bv2, 1 : Mh , 1 : Mh ), V₁u, V₁l)
167
- 🦋2 !(view (Bv2, Mh + 1 : M, Mh + 1 : M), V₂u, V₂l)
164
+ 🦋! (view (Bv2, 1 : M_half , 1 : M_half ), V₁u, V₁l)
165
+ 🦋! (view (Bv2, M_half + 1 : M, M_half + 1 : M), V₂u, V₂l)
168
166
169
167
Bv1 = SparseBandedMatrix {typeof(uv[1])} (undef, M, M)
170
168
🦋! (Bv1, Vu, Vl)
0 commit comments