Speed up multi-dimensional stencil assembly by 20-40x

Project.toml

@@ -26,6 +26,7 @@ KrylovKit = "0.10"
 LinearAlgebra = "1.10"
 PhysicalConstants = "0.2"
 Printf = "1.10"
+Random = "1.10"
 SparseArrays = "1.10"
 Statistics = "1.10"
 StatsBase = "0.34"
@@ -39,9 +40,10 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "LinearAlgebra", "Statistics", "Suppressor", "Test"]
+test = ["Aqua", "LinearAlgebra", "Random", "Statistics", "Suppressor", "Test"]

src/buildStencil/buildStencilMatrices.jl

@@ -44,6 +44,83 @@ function build_1d_stencil(
 	return matrix
 end
 
+"""
+	accumulate_stencil_blocks!(Is, Js, Vs, block_size, n_outer, total_points,
+	                            periodic_outer, sub_matrix_indices, block_triplets)
+
+Shared core of [`build_2d_stencil`](@ref) and [`build_3d_stencil`](@ref):
+given the (I, J, V) nonzero triplets of each stencil offset's sub-block, plus
+the row/col block bookkeeping the caller works out, append the final
+triplets for one "outer axis" level of block placement to `Is`/`Js`/`Vs`.
+
+Building the result via COO triplet accumulation followed by a single
+`sparse(...)` call (instead of the equivalent-but-much-slower repeated
+`matrix[range_i, range_j] = block` assignment into an already-existing large
+sparse matrix - each such assignment can touch and reallocate a large
+fraction of the matrix's internal CSC storage) is what makes this fast: it
+turns an operation that scaled worse than quadratically in the grid size
+into one that is close to linear.
+
+For each row-block `j`, a small per-`j` `Dict` (at most `stencil_size`
+entries, discarded every iteration) resolves which stencil offset wins the
+placement for a given target column-block *before* any triplets are
+appended, exactly reproducing the original's `matrix[...] = block`
+last-write-wins overwrite semantics - this matters only when a periodic
+wraparound collides with a non-wrapped placement, which can only happen for
+grids smaller than the stencil width, but is reproduced correctly regardless.
+"""
+function accumulate_stencil_blocks!(
+	Is,
+	Js,
+	Vs,
+	block_size,
+	n_outer,
+	total_points,
+	periodic_outer,
+	sub_matrix_indices,
+	block_triplets,
+)
+
+	stencil_size = length(sub_matrix_indices)
+	winners = Dict{Int, Int}()
+
+	for j in 1:n_outer
+
+		empty!(winners)
+		row_offset = (j - 1) * block_size
+
+		for i in 1:stencil_size
+
+			cell_index = j + sub_matrix_indices[i]
+			col_start = (cell_index - 1) * block_size + 1
+
+			if cell_index < 1 || cell_index > n_outer
+				if periodic_outer
+					col_start -= sign(cell_index - 1) * total_points
+				else
+					continue
+				end
+			end
+
+			winners[col_start] = i
+
+		end
+
+		for (col_start, i) in winners
+
+			I1, J1, V1 = block_triplets[i]
+			col_offset = col_start - 1
+
+			append!(Is, I1 .+ row_offset)
+			append!(Js, J1 .+ col_offset)
+			append!(Vs, V1)
+
+		end
+
+	end
+
+	return nothing
+end
 
 """
 	build_2d_stencil(system, n_datapoints, stencil, stencil_size)
@@ -70,48 +147,31 @@ function build_2d_stencil(
 
 	total_points = prod(n_datapoints)
 
-	matrix = spzeros(total_points, total_points)
-
-	for i in 1:stencil_size
-
-		sub_matrix_index = i - stencil_size ÷ 2 - 1
-
-		matrix_1d = build_1d_stencil(
-			system,
-			stencil[i, :],
-			stencil_size,
-		)
-
-		for j in 1:n_datapoints[1]
-
-			cell_index = j + sub_matrix_index
-
-			range_i =
-				((j-1)*n_datapoints[2]+1):j*n_datapoints[2]
-			range_j =
-				((cell_index-1)*n_datapoints[2]+1):cell_index*n_datapoints[2]
-
-			if cell_index < 1 || cell_index > n_datapoints[1]
-				if system.periodic[end-1]
-					range_j =
-						(range_j[1]-sign(
-							cell_index - 1,
-						)*total_points):(range_j[end]-sign(
-							cell_index - 1,
-						)*total_points)
-				else
-					continue
-				end
-			end
-
-			matrix[range_i, range_j] = matrix_1d
-		end
-	end
-
-	return matrix
+	sub_matrix_indices = [i - stencil_size ÷ 2 - 1 for i in 1:stencil_size]
+	block_triplets = [
+		findnz(build_1d_stencil(system, stencil[i, :], stencil_size))
+		for i in 1:stencil_size
+	]
+
+	Is = Int[]
+	Js = Int[]
+	Vs = Float64[]
+
+	accumulate_stencil_blocks!(
+		Is,
+		Js,
+		Vs,
+		n_datapoints[2],
+		n_datapoints[1],
+		total_points,
+		system.periodic[end-1],
+		sub_matrix_indices,
+		block_triplets,
+	)
+
+	return sparse(Is, Js, Vs, total_points, total_points)
 end
 
-
 """
 	build_3d_stencil(system, n_datapoints, stencil, stencil_size)
 
@@ -137,48 +197,27 @@ function build_3d_stencil(
 
 	total_points = prod(n_datapoints)
 
-	matrix = spzeros(total_points, total_points)
-
-	for i in 1:stencil_size
-
-		sub_matrix_index = i - stencil_size ÷ 2 - 1
-
-		matrix_2d = build_2d_stencil(
-			system,
-			n_datapoints[2:3],
-			stencil[:, :, i],
-			stencil_size,
-		)
-
-		for j in 1:n_datapoints[1]
-
-			cell_index = j + sub_matrix_index
-
-			range_i =
-				((j-1)*prod(n_datapoints[2:3])+1):j*prod(
-					n_datapoints[2:3],
-				)
-			range_j =
-				((cell_index-1)*prod(n_datapoints[2:3])+1):cell_index*prod(
-					n_datapoints[2:3],
-				)
-
-			if cell_index < 1 || cell_index > n_datapoints[1]
-				if system.periodic[end-2]
-					range_j =
-						(range_j[1]-sign(
-							cell_index - 1,
-						)*total_points):(range_j[end]-sign(
-							cell_index - 1,
-						)*total_points)
-				else
-					continue
-				end
-			end
-
-			matrix[range_i, range_j] = matrix_2d
-		end
-	end
-
-	return matrix
+	sub_matrix_indices = [i - stencil_size ÷ 2 - 1 for i in 1:stencil_size]
+	block_triplets = [
+		findnz(build_2d_stencil(system, n_datapoints[2:3], stencil[:, :, i], stencil_size))
+		for i in 1:stencil_size
+	]
+
+	Is = Int[]
+	Js = Int[]
+	Vs = Float64[]
+
+	accumulate_stencil_blocks!(
+		Is,
+		Js,
+		Vs,
+		prod(n_datapoints[2:3]),
+		n_datapoints[1],
+		total_points,
+		system.periodic[end-2],
+		sub_matrix_indices,
+		block_triplets,
+	)
+
+	return sparse(Is, Js, Vs, total_points, total_points)
 end

test/runtests.jl

@@ -6,6 +6,7 @@ using Suppressor
 using Statistics
 using SparseArrays
 using LinearAlgebra
+using Random
 using Unitful
 using UnitfulAtomic
 

test/unittests.jl

@@ -4,6 +4,7 @@ include("unittests/test_1Dlaplace.jl")
 include("unittests/test_1Dnabla.jl")
 include("unittests/test_2Dlaplace.jl")
 include("unittests/test_2Dnabla.jl")
+include("unittests/test_buildStencilMatrices.jl")
 
 include("unittests/test_k_paths.jl")
 
@@ -18,6 +19,7 @@ function unittests()
     @testset "test 1D nabla   operator" test_1D∇()
     @testset "test 2D laplace operator" test_2DΔ()
     @testset "test 2D nabla   operator" test_2D∇()
+    @testset "test build_2d/3d_stencil vs reference" test_buildStencilMatrices()
 
     k_intervalls = [1.0, 1.0, 1.0]
     n_kpoints    = 3