Merge pull request #175 from Cthonios/sparse-gpu-rework

reworking sparse implementation on GPU.

Author: cmhamel

ext/FiniteElementContainersAMDGPUExt.jl

@@ -21,7 +21,12 @@ function AMDGPU.rocSPARSE.ROCSparseMatrixCSC(asm::SparseMatrixAssembler)
   #
 
   # n_dofs = FiniteElementContainers.num_unknowns(asm.dof)
-  n_dofs = length(asm.dof.unknown_dofs)
+  if FiniteElementContainers._is_condensed(asm.dof)
+    n_dofs = length(asm.dof)
+  else
+    n_dofs = length(asm.dof.unknown_dofs)
+  end
+
   return AMDGPU.rocSPARSE.ROCSparseMatrixCSC(
     asm.pattern.csccolptr,
     asm.pattern.cscrowval,

ext/FiniteElementContainersCUDAExt.jl

@@ -16,7 +16,12 @@ function CUDA.CUSPARSE.CuSparseMatrixCSC(asm::SparseMatrixAssembler)
   # they are doing the right thing
   # @assert length(asm.pattern.cscnzval) > 0 "Need to assemble the assembler once with SparseArrays.sparse!(assembler)"
   # @assert all(x -> x != zero(eltype(asm.pattern.cscnzval)), asm.pattern.cscnzval) "Need to assemble the assembler once with SparseArrays.sparse!(assembler)"
-  n_dofs = FiniteElementContainers.num_unknowns(asm.dof)
+  if FiniteElementContainers._is_condensed(asm.dof)
+    n_dofs = length(asm.dof)
+  else
+    n_dofs = FiniteElementContainers.num_unknowns(asm.dof)
+  end
+
   return CUDA.CUSPARSE.CuSparseMatrixCSC(
     asm.pattern.csccolptr,
     asm.pattern.cscrowval,

src/DofManagers.jl

@@ -86,23 +86,12 @@ KA.@kernel function _extract_field_unknowns_kernel!(
 end
 # COV_EXCL_STOP
 
-# COV_EXCL_START
-KA.@kernel function _extract_field_unknowns_kernel!(
-    Uu::V, 
-    dof::DofManager{true, IT, IDs, Var}, 
-    U::AbstractField
-) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
-    N = KA.@index(Global)
-    @inbounds Uu[dof.unknown_dofs[N]] = U[dof.unknown_dofs[N]]
-end
-# COV_EXCL_STOP
-
 function _extract_field_unknowns!(
     Uu::V, 
-    dof::DofManager, 
+    dof::DofManager{false, IT, IDs, Var}, 
     U::AbstractField, 
     backend::KA.Backend
-) where V <: AbstractVector{<:Number}
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     kernel! = _extract_field_unknowns_kernel!(backend)
     kernel!(Uu, dof, U, ndrange = length(Uu))
     return nothing
@@ -118,21 +107,11 @@ function _extract_field_unknowns!(
     return nothing
 end
 
-function _extract_field_unknowns!(
-    Uu::V, 
-    dof::DofManager{true, IT, IDs, Var}, 
-    U::AbstractField, 
-    ::KA.CPU
-) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
-    @views Uu[dof.unknown_dofs] .= U[dof.unknown_dofs]
-    return nothing
-end
-
 function extract_field_unknowns!(
     Uu::V,
-    dof::DofManager,
+    dof::DofManager{false, IT, IDs, Var},
     U::AbstractField
-) where V <: AbstractVector{<:Number}
+) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     backend = KA.get_backend(dof)
     @assert KA.get_backend(U) == backend
     @assert KA.get_backend(Uu) == backend
@@ -168,7 +147,7 @@ KA.@kernel function _update_field_unknowns_kernel!(
     Uu::V
 ) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     N = KA.@index(Global)
-    @inbounds U[dof.unknown_dofs[N]] = Uu[N]
+    @inbounds U.data[dof.unknown_dofs[N]] = Uu[N]
 end
 # COV_EXCL_STOP
 
@@ -179,7 +158,7 @@ KA.@kernel function _update_field_unknowns_kernel!(
     Uu::V
 ) where {V <: AbstractVector{<:Number}, IT, IDs, Var}
     N = KA.@index(Global)
-    @inbounds U[dof.unknown_dofs[N]] = Uu[dof.unknown_dofs[N]]
+    @inbounds U.data[dof.unknown_dofs[N]] = Uu[dof.unknown_dofs[N]]
 end
 # COV_EXCL_STOP
 
@@ -190,7 +169,7 @@ function _update_field_unknowns!(
     backend::KA.Backend
 ) where T <: AbstractVector{<:Number}
     kernel! = _update_field_unknowns_kernel!(backend)
-    kernel!(U, dof, Uu, ndrange = length(Uu))
+    kernel!(U, dof, Uu, ndrange = length(dof.unknown_dofs))
     return nothing
 end
 
@@ -223,7 +202,13 @@ function update_field_unknowns!(
     backend = KA.get_backend(dof)
     @assert KA.get_backend(U) == backend
     @assert KA.get_backend(Uu) == backend
-    # @assert length(dof.unknown_dofs) == length(Uu)
+
+    if _is_condensed(dof)
+        @assert length(Uu) == length(U)
+    else
+        @assert length(Uu) == length(dof.unknown_dofs)
+    end
+
     _update_field_unknowns!(U, dof, Uu, backend)
     return nothing
 end

src/assemblers/Assemblers.jl

@@ -8,7 +8,7 @@ $(TYPEDSIGNATURES)
 """
 KA.get_backend(asm::AbstractAssembler) = KA.get_backend(asm.dof)
 
-function _adjust_matrix_action_entries_for_condensed!(
+function _adjust_matrix_action_entries_for_constraints!(
   Av, constraint_storage, v, ::KA.CPU
   # TODO do we need a penalty scale here as well?
 )
@@ -22,7 +22,25 @@ function _adjust_matrix_action_entries_for_condensed!(
   return nothing
 end
 
-function _adjust_vector_entries_for_condensed!(b, constraint_storage, ::KA.CPU)
+KA.@kernel function _adjust_matrix_action_entries_for_constraints_kernel!(
+  Av, constraint_storage, v
+)
+  I = KA.@index(Global)
+  # modify Av => (I - G) * Av + Gv
+  @inbounds Av[I] = (1. - constraint_storage[I]) * Av[I] + constraint_storage[I] * v[I]
+end
+
+function _adjust_matrix_action_entries_for_constraints!(
+  Av, constraint_storage, v, backend::KA.Backend
+)
+  @assert length(Av) == length(constraint_storage)
+  @assert length(v) == length(constraint_storage)
+  kernel! = _adjust_matrix_action_entries_for_constraints_kernel!(backend)
+  kernel!(Av, constraint_storage, v, ndrange = length(Av))
+  return nothing
+end
+
+function _adjust_vector_entries_for_constraints!(b, constraint_storage, ::KA.CPU)
   @assert length(b) == length(constraint_storage)
   # modify b => (I - G) * b + (Gu - g)
   # but Gu = g, so we don't need that here
@@ -34,6 +52,19 @@ function _adjust_vector_entries_for_condensed!(b, constraint_storage, ::KA.CPU)
   return nothing
 end
 
+KA.@kernel function _adjust_vector_entries_for_constraints_kernel(b, constraint_storage)
+  I = KA.@index(Global)
+  # modify b => (I - G) * b + (Gu - g)
+  @inbounds b[I] = (1. - constraint_storage[I]) * b[I]
+end
+
+function _adjust_vector_entries_for_constraints!(b, constraint_storage, backend::KA.Backend)
+  @assert length(b) == length(constraint_storage)
+  kernel! = _adjust_vector_entries_for_constraints_kernel(backend)
+  kernel!(b, constraint_storage, ndrange = length(b))
+  return nothing
+end
+
 """
 $(TYPEDSIGNATURES)
 Assembly method for an H1Field, e.g. internal force
@@ -162,30 +193,31 @@ function _quadrature_level_state(state::L2QuadratureField, q::Int, e::Int)
   return state_q
 end
 
+
+# function hvp(asm::AbstractAssembler)
+#   if _is_condensed(asm.dof)
+#     _adjust_matrix_action_entries_for_constraints!(
+#       asm.stiffness_action_storage, asm.constraint_storage, 
+#       KA.get_backend(asm)
+#     )
+#     return asm.stiffness_action_storage.data
+#   else
+#     extract_field_unknowns!(
+#       asm.stiffness_action_unknowns,
+#       asm.dof,
+#       asm.stiffness_action_storage
+#     )
+#     return asm.stiffness_action_unknowns
+#   end
+# end
+
+# new approach requiring access to the v that makes Hv
 """
 $(TYPEDSIGNATURES)
 """
-function hvp(asm::AbstractAssembler)
-  if _is_condensed(asm.dof)
-    _adjust_matrix_action_entries_for_condensed!(
-      asm.stiffness_action_storage, asm.constraint_storage, 
-      KA.get_backend(asm)
-    )
-    return asm.stiffness_action_storage.data
-  else
-    extract_field_unknowns!(
-      asm.stiffness_action_unknowns,
-      asm.dof,
-      asm.stiffness_action_storage
-    )
-    return asm.stiffness_action_unknowns
-  end
-end
-
-# new approach requiring access to the v that makes Hv
 function hvp(asm::AbstractAssembler, v)
   if _is_condensed(asm.dof)
-    _adjust_matrix_action_entries_for_condensed!(
+    _adjust_matrix_action_entries_for_constraints!(
       asm.stiffness_action_storage, asm.constraint_storage, v,
       KA.get_backend(asm)
     )
@@ -213,7 +245,7 @@ assumes assemble_vector! has already been called
 """
 function residual(asm::AbstractAssembler)
   if _is_condensed(asm.dof)
-    _adjust_vector_entries_for_condensed!(
+    _adjust_vector_entries_for_constraints!(
       asm.residual_storage, asm.constraint_storage,
       KA.get_backend(asm)
     )

src/assemblers/SparseMatrixAssembler.jl

@@ -111,7 +111,7 @@ function _assemble_element!(
 end
 
 # TODO this only work on CPU right now
-function _adjust_matrix_entries_for_condensed!(
+function _adjust_matrix_entries_for_constraints!(
   A::SparseMatrixCSC, constraint_storage, ::KA.CPU;
   penalty_scale = 1.e6
 )
@@ -145,6 +145,46 @@ function _adjust_matrix_entries_for_condensed!(
   return nothing
 end
 
+KA.@kernel function _adjust_matrix_entries_for_constraints_kernel!(
+  A, constraint_storage, trA;
+  penalty_scale = 1.e6
+)
+  J = KA.@index(Global)
+
+  penalty = penalty_scale * trA / size(A, 2)
+
+  # now modify A => (I - G) * A + G
+  nz = nonzeros(A)
+  rowval = rowvals(A)
+
+  col_start = A.colptr[J]
+  col_end   = A.colptr[J + 1] - 1
+  for k in col_start:col_end
+    # for (I - G) * A term
+    nz[k] = (1. - constraint_storage[J]) * nz[k]
+
+    # for + G term
+    if rowval[k] == J
+      @inbounds nz[k] = nz[k] + penalty * constraint_storage[J]
+    end
+  end
+end
+
+function _adjust_matrix_entries_for_constraints(
+  A, constraint_storage, backend::KA.Backend
+)
+  # first ensure things are the right size
+  @assert size(A, 1) == size(A, 2)
+  @assert length(constraint_storage) == size(A, 2)
+
+  # get trA ahead of time to save some allocations at kernel level
+  trA = tr(A)
+  
+  kernel! = _adjust_matrix_entries_for_constraints_kernel!(backend)
+  kernel!(A, constraint_storage, trA, ndrange = size(A, 2))
+  return nothing
+end
+
 function constraint_matrix(assembler::SparseMatrixAssembler)
   return Diagonal(assembler.constraint_storage)
 end
@@ -153,7 +193,7 @@ function _mass(assembler::SparseMatrixAssembler, ::KA.CPU)
   M = SparseArrays.sparse!(assembler.pattern, assembler.mass_storage)
 
   if _is_condensed(assembler.dof)
-    _adjust_matrix_entries_for_condensed!(M, assembler.constraint_storage, KA.get_getbackend(assembler))
+    _adjust_matrix_entries_for_constraints!(M, assembler.constraint_storage, KA.get_getbackend(assembler))
   end
 
   return M
@@ -163,7 +203,7 @@ function _stiffness(assembler::SparseMatrixAssembler, ::KA.CPU)
   K = SparseArrays.sparse!(assembler.pattern, assembler.stiffness_storage)
 
   if _is_condensed(assembler.dof)
-    _adjust_matrix_entries_for_condensed!(K, assembler.constraint_storage, KA.get_backend(assembler))
+    _adjust_matrix_entries_for_constraints!(K, assembler.constraint_storage, KA.get_backend(assembler))
   end
 
   return K

src/assemblers/Vector.jl

@@ -26,6 +26,8 @@ function assemble_vector!(
       backend
     )
   end
+  
+  return nothing
 end
 
 # CPU Implementation

test/TestAssemblers.jl

@@ -33,5 +33,5 @@ include("poisson/TestPoissonCommon.jl")
   K = stiffness(asm)
   M = mass(asm)
   R = residual(asm)
-  Mv = hvp(asm)
+  Mv = hvp(asm, Vu)
 end

test/poisson/TestPoissonAMDGPU.jl

@@ -17,16 +17,15 @@ bc_func(_, _) = 0.
 
 include("TestPoissonCommon.jl")
 
-function poisson_amdgpu()
+function poisson_amdgpu(use_condensed)
   # do all setup on CPU
   # the mesh for instance is not gpu compatable
   mesh = UnstructuredMesh(mesh_file)
   V = FunctionSpace(mesh, H1Field, Lagrange)
   physics = Poisson()
   props = create_properties(physics)
   u = ScalarFunction(V, :u)
-  dof = DofManager(u)
-  asm = SparseMatrixAssembler(u)
+  asm = SparseMatrixAssembler(u; use_condensed=use_condensed)
 
   dbcs = DirichletBC[
     DirichletBC(:u, :sset_1, bc_func),
@@ -63,7 +62,9 @@ function poisson_amdgpu()
   display(solver.timer)
 end
 
-@time poisson_amdgpu()
-@time poisson_amdgpu()
+@time poisson_amdgpu(false)
+@time poisson_amdgpu(false)
+@time poisson_amdgpu(true)
+@time poisson_amdgpu(true)
 
 # @benchmark poisson_cuda()