only supporting inplace methods now in enzyme assembly methods. Addin…

Project.toml

@@ -1,6 +1,6 @@
 name = "FiniteElementContainers"
 uuid = "d08262e4-672f-4e7f-a976-f2cea5767631"
-version = "0.14.2"
+version = "0.14.3"
 authors = ["Craig M. Hamel <cmhamel32@gmail.com> and contributors"]
 
 [deps]
@@ -50,7 +50,7 @@ Atomix = "1"
 BlockArrays = "1"
 CUDA = "5, 6"
 DocStringExtensions = "0.9"
-DynamicExpressions = "2.5.2"
+DynamicExpressions = "2.5"
 Exodus = "0.15"
 ForwardDiff = "1"
 GPUArrays = "11"
@@ -74,6 +74,7 @@ TimerOutputs = "0.5"
 julia = "1"
 
 [extras]
+AMDGPU = "21141c5a-9bdb-4563-92ae-f87d6854732e"
 Gmsh = "705231aa-382f-11e9-3f0c-b7cb4346fdeb"
 JuliaC = "acedd4c2-ced6-4a15-accc-2607eb759ba2"
 PartitionedArrays = "5a9dfac6-5c52-46f7-8278-5e2210713be9"
@@ -82,4 +83,4 @@ TestItemRunner = "f8b46487-2199-4994-9208-9a1283c18c0a"
 TestItems = "1c621080-faea-4a02-84b6-bbd5e436b8fe"
 
 [targets]
-test = ["Gmsh", "JuliaC", "PartitionedArrays", "Test", "TestItemRunner", "TestItems"]
+test = ["AMDGPU", "Gmsh", "JuliaC", "PartitionedArrays", "Test", "TestItemRunner", "TestItems"]

examples/parrays-example/script_v4.jl

@@ -6,6 +6,14 @@ using LinearAlgebra
 using PartitionedArrays
 using StaticArrays
 
+struct PLinearSystem{A, B, C, D}
+    residual_vals::A
+    residual_cache::B
+    stiffness_vals::C
+    stiffness_cache::D
+end
+
+
 include("../../test/poisson/TestPoissonCommon.jl")
 f(X, _) = 2. * π^2 * sin(2π * X[1]) * sin(2π * X[2])
 bc_func(_, _) = 0.
@@ -43,21 +51,35 @@ Uu = create_unknowns(asm)
 U = create_field(asm)
 update_field_unknowns!(U, dof, Uu)
 
-assemble_stiffness!(asm, stiffness, U, p)
-K = stiffness(asm)
+# rseidua
+residual_vals = map(asm.local_assemblers) do local_asm
+    local_asm.residual_unknowns
+end
+residual_vals = create_unknowns(dof)
+R, R_cache = pvector(asm.vector_pattern, residual_vals; reuse = true) |> fetch
+
+@time pvector!(R, residual_vals, R_cache) |> wait
+@time pvector!(R, residual_vals, R_cache) |> wait
+
+# system = PLinearSystem()
+
+# assemble_stiffness!(asm, stiffness, U, p)
+# K = stiffness(asm)
+
+# assemble_vector!(asm, residual, U, p)
+# R = residual(asm)
 
-assemble_vector!(asm, residual, U, p)
-R = residual(asm)
+# cg_workspace = CgWorkspace(K, -R)
+# Krylov.cg!(cg_workspace, K, -R)
+# x_new, stats = Krylov.results(cg_workspace)
 
-cg_workspace = CgWorkspace(K, -R)
-Krylov.cg!(cg_workspace, K, -R)
-x_new, stats = Krylov.results(cg_workspace)
+# U = create_field(dof)
+# # update_field_dirichlet_bcs!(U, dof)
+# update_field_unknowns!(U, dof, x_new)
 
-U = create_field(dof)
-# update_field_dirichlet_bcs!(U, dof)
-update_field_unknowns!(U, dof, x_new)
+# pp = PostProcessor(output_file, mesh, true; extra_nodal_names = [names(u)...])
+# write_times(pp, 1, 0.0)
+# write_field(pp, 1, ["u"], U)
+# close(pp)
 
-pp = PostProcessor(output_file, mesh, true; extra_nodal_names = [names(u)...])
-write_times(pp, 1, 0.0)
-write_field(pp, 1, ["u"], U)
-close(pp)
+# epu("output.e.$(num_ranks).")

ext/PartitionedArraysExt.jl

@@ -12,6 +12,7 @@ using PartitionedArrays
 # Helper methods for setup
 #######################################################################
 # NOTE dof_to_unknown returns -1 when the dof 
+# TODO doesn't suport PeriodicBCs
 function _create_field_to_unknown(n_total_fields, dirichlet_dofs)
 	unknown_to_field = Vector{eltype(dirichlet_dofs)}(undef, n_total_fields - length(dirichlet_dofs))
 	ids = 1:n_total_fields
@@ -25,7 +26,7 @@ function _create_field_to_unknown(n_total_fields, dirichlet_dofs)
     field_to_unknown = Dict([(x, y) for (x, y) in zip(unknown_to_field, 1:length(unknown_to_field))])
 
 	for dof in dirichlet_dofs
-		field_to_unknown[dof] = -1
+		field_to_unknown[dof] = FEC.DIRICHLET_DOF
 	end
 
 	new_field_to_unknown = Vector{Int}(undef, length(field_to_unknown))
@@ -260,32 +261,28 @@ function FEC.DofManager(
 		return dof_local
 	end
 
-	# # now need to scatter all to one proc
-	# # out = map(dirichlet_dofs) do ddofs
-	# # 	gather(ddofs, destination = 1)
-	# # end
-	# out = gather(dirichlet_dofs, destination = 1)
-	# dirichlet_dofs = map_main(out) do out_local
-	# 	return reduce(vcat, out_local) |> unique |> sort
-	# end
-
 	# slow approach for now
 	serial_mesh = UnstructuredMesh(mesh_file)
 	serial_V = FunctionSpace(serial_mesh, H1Field, Lagrange)
 	serial_u = eval(typeof(u).name.name){typeof(serial_V)}(serial_V, var_names)
 	serial_dof = DofManager(serial_u)
 	serial_dbcs = DirichletBCs(serial_mesh, serial_dof, dbcs)
-	dirichlet_dofs = FiniteElementContainers.dirichlet_dofs(serial_dbcs)
+	serial_pbcs = PeriodicBCs(serial_mesh, serial_dof, pbcs)
+	dirichlet_dofs = FEC.dirichlet_dofs(serial_dbcs)
+	_, periodic_side_b_dofs = FEC.periodic_dofs(serial_pbcs)
+	constrained_dofs = union(dirichlet_dofs, periodic_side_b_dofs) |> sort |> unique
 
 	# TODO need to cleanup up and finish everthing below
 	unknowns_to_colors = copy(fields_to_colors)
-	deleteat!(unknowns_to_colors, dirichlet_dofs)
+	deleteat!(unknowns_to_colors, constrained_dofs)
 
 	solution_parts = partition_from_color(ranks, unknowns_to_colors)
 	solution_parts = SolutionPartition(unknowns_to_colors, solution_parts)
 
 	# finally create the maps from field to solution and back
-	field_to_unknown, unknown_to_field = _create_field_to_unknown(length(fields_to_colors), dirichlet_dofs)
+	field_to_unknown, unknown_to_field = _create_field_to_unknown(
+		length(fields_to_colors), dirichlet_dofs
+	)
 	return PDofManager(
 		field_parts, solution_parts,
 		field_to_unknown, unknown_to_field,
@@ -425,11 +422,11 @@ function FEC.create_unknowns(pattern::PSparseVectorPattern)
 	return create_unknowns(pattern.dof)
 end
 
-function PartitionedArrays.pvector(pattern::PSparseVectorPattern, vals)
-	vals = map(pattern.unknown_dofs, vals) do dofs, val
-		val[dofs]
-	end
-	return pvector(pattern.Is, vals, pattern.dof.solution_partition.parts)
+function PartitionedArrays.pvector(pattern::PSparseVectorPattern, vals; reuse::Bool = false)
+	# vals = map(pattern.unknown_dofs, vals) do dofs, val
+	# 	val[dofs]
+	# end
+	return pvector(pattern.Is, vals, pattern.dof.solution_partition.parts; reuse = reuse)
 end
 
 function PartitionedArrays.pzeros(pattern::PSparseVectorPattern)
@@ -439,14 +436,34 @@ end
 struct PSparseMatrixAssembler{
 	Assemblers,
 	MatPattern <: PSparseMatrixPattern,
-	VecPattern <: PSparseVectorPattern
+	VecPattern <: PSparseVectorPattern,
+	MapScratch <: AbstractVector{Nothing},
+	RV,
+	RC
 }
 	local_assemblers::Assemblers
 	matrix_pattern::MatPattern
 	vector_pattern::VecPattern
+	map_scratch::MapScratch
+	#
+	residual_vals::RV
+	residual_cache::RC
 end
 
 function FEC.assemble_stiffness!(asm::PSparseMatrixAssembler, func, u, p)
+	# @show @allocated partition(u)
+	# @time map!(asm.map_scratch, asm.local_assemblers) do 
+	# function _assemble_local!(local_asm, func, local_u, local_p)
+	# 	assemble_stiffness!(local_asm, func, local_u, local_p)
+	# end
+	# map!(
+	# 	_assemble_local!, asm.map_scratch, asm.local_assemblers,
+	# 	Iterators.repeated(func),
+	# 	# func,
+	# 	partition(u),
+	# 	p.local_parameters
+	# )
+	# map!(asm.map_scratch,)
 	map(asm.local_assemblers, partition(u), p.local_parameters) do local_asm, local_u, local_p
 		assemble_stiffness!(local_asm, func, local_u, local_p)
 	end
@@ -487,9 +504,22 @@ function FEC.SparseMatrixAssembler(dof::PDofManager)
 	end
 	matrix_pattern = PSparseMatrixPattern(dof)
 	vector_pattern = PSparseVectorPattern(dof)
+	map_scratch = map(dof.local_dof_managers) do local_dof
+		nothing
+	end
+
+	# residual_vals, residual_cache = 
+	residual_temp = map(local_assemblers) do local_asm
+		local_asm.residual_unknowns
+	end
+	# residual_vals, residual_cache = pvector(vector_pattern, residual_temp, true) |> fetch
+	residual_vals, residual_cache = nothing, nothing
+
 	return PSparseMatrixAssembler(
 		local_assemblers,
-		matrix_pattern, vector_pattern
+		matrix_pattern, vector_pattern,
+		map_scratch,
+		residual_vals, residual_cache
 	)
 end