Skip to content

Add various patches for 0.11 #23

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
jorgensd wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Add various patches for 0.11 #23

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion cpp/CMakeLists.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@ cmake_minimum_required(VERSION 3.19)

# ------------------------------------------------------------------------------
# Set project name and version number
project(CUDOLFINX VERSION "0.10.0")
project(CUDOLFINX VERSION "0.11.0")

set(DOXYGEN_CUDOLFINX_VERSION
${CUDOLFINX_VERSION}
Expand Down
249 changes: 123 additions & 126 deletions cpp/cudolfinx/fem/CUDAForm.h
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,14 +1,11 @@
// Copyright (C) 2024 Benjamin Pachev, James D. Trotter
// Copyright (C) 2024-2026 Benjamin Pachev, James D. Trotter, JørgensS. Dokken
//
// This file is part of cuDOLFINX
//
// SPDX-License-Identifier: LGPL-3.0-or-later

#pragma once

#include <dolfinx/fem/Form.h>
#include <dolfinx/fem/DirichletBC.h>
#include <dolfinx/common/IndexMap.h>
#include <cudolfinx/common/CUDA.h>
#include <cudolfinx/common/CUDAStore.h>
#include <cudolfinx/fem/CUDADirichletBC.h>
Expand All @@ -17,231 +14,229 @@
#include <cudolfinx/fem/CUDAFormConstants.h>
#include <cudolfinx/fem/CUDAFormIntegral.h>
#include <cudolfinx/la/CUDAVector.h>
#include <dolfinx/common/IndexMap.h>
#include <dolfinx/common/version.h>
#include <dolfinx/fem/DirichletBC.h>
#include <dolfinx/fem/Form.h>
#include <ranges>
#include <string>
#include <utility>
#include <ufcx.h>
#include <ranges>
#include <utility>

namespace dolfinx {

namespace fem {

/// Consolidates all form classes into one
template <dolfinx::scalar T,
std::floating_point U = dolfinx::scalar_value_t<T>>
class CUDAForm
{
template <dolfinx::scalar T, std::floating_point U = dolfinx::scalar_value_t<T>>
class CUDAForm {

public:
/// Create GPU copies of data needed for assembly
///
/// @param[in] cuda_context A context for a CUDA device
/// @param[in] form Pointer to the variational form
CUDAForm(
const CUDA::Context& cuda_context,
Form<T,U>* form,
ufcx_form* ufcx_form,
std::vector<std::string>& tabulate_tensor_names,
std::vector<std::string>& tabulate_tensor_sources,
std::vector<int>& integral_tensor_indices
)
: _coefficients(cuda_context, form, _dofmap_store)
, _constants(cuda_context, form)
, _form(form)
, _ufcx_form(ufcx_form)
, _compiled(false)
{
_coefficients = CUDAFormCoefficients<T,U>(cuda_context, form, _dofmap_store);
const int* integral_offsets = ufcx_form->form_integral_offsets;
CUDAForm(const CUDA::Context &cuda_context, Form<T, U> *form,
ufcx_form *ufcx_form,
std::vector<std::string> &tabulate_tensor_names,
std::vector<std::string> &tabulate_tensor_sources,
std::vector<int> &integral_tensor_indices)
: _coefficients(cuda_context, form, _dofmap_store),
_constants(cuda_context, form), _form(form), _ufcx_form(ufcx_form),
_compiled(false) {
_coefficients =
CUDAFormCoefficients<T, U>(cuda_context, form, _dofmap_store);
const int *integral_offsets = ufcx_form->form_integral_offsets;
if (integral_offsets[3] != integral_tensor_indices.size()) {
throw std::runtime_error("UFCx form has " + std::to_string(integral_offsets[3])
+ " integrals, but only " + std::to_string(tabulate_tensor_names.size())
+ " tabulate tensor sources provided to CUDAForm!"
);
throw std::runtime_error(
"UFCx form has " + std::to_string(integral_offsets[3]) +
" integrals, but only " +
std::to_string(tabulate_tensor_names.size()) +
" tabulate tensor sources provided to CUDAForm!");
}
for (int i = 0; i < 3; i++) {
for (int offset = integral_offsets[i]; offset < integral_offsets[i+1]; offset++) {
for (int offset = integral_offsets[i]; offset < integral_offsets[i + 1];
offset++) {
int idx = offset - integral_offsets[i];
int tensor_offset = integral_tensor_indices[offset];
_cuda_kernel_sources[i].insert({idx, {tabulate_tensor_names[tensor_offset], tabulate_tensor_sources[tensor_offset]}});
_cuda_kernel_sources[i].insert(
{idx,
{tabulate_tensor_names[tensor_offset],
tabulate_tensor_sources[tensor_offset]}});
}
}

// initialize the CUDAFormIntegral objects
// this mostly just intitializes the mesh entities for each integral
// however there is an MPI component to determine the active ghost entities
// and so the constructor needs to be called on all processes
for (IntegralType integral_type : {IntegralType::cell, IntegralType::exterior_facet, IntegralType::interior_facet,
IntegralType::vertex, IntegralType::ridge})
for (IntegralType integral_type :
{IntegralType::cell, IntegralType::exterior_facet,
IntegralType::interior_facet, IntegralType::vertex,
IntegralType::ridge})

{
// TODO: add full mixed-topology support when it becomes more mature in dolfinx
// TODO: add full mixed-topology support when it becomes more mature in
// dolfinx
int num_integrals = form->num_integrals(integral_type, 0);
if (num_integrals > 0) {
std::vector<CUDAFormIntegral<T,U>>& cuda_integrals =
_integrals[integral_type];
std::vector<CUDAFormIntegral<T, U>> &cuda_integrals =
_integrals[integral_type];
for (int i = 0; i < num_integrals; i++)
cuda_integrals.emplace_back(*form, integral_type, i);
cuda_integrals.emplace_back(*form, integral_type, i);
}
}
}
}

/// Compile form on GPU
/// Under the hood, this creates the integrals
void compile(
const CUDA::Context& cuda_context,
int32_t max_threads_per_block,
int32_t min_blocks_per_multiprocessor,
std::string cachedir,
enum assembly_kernel_type assembly_kernel_type)
{
void compile(const CUDA::Context &cuda_context, int32_t max_threads_per_block,
int32_t min_blocks_per_multiprocessor, std::string cachedir,
enum assembly_kernel_type assembly_kernel_type) {
// Get the number of vertices and coordinates
const mesh::Mesh<U>& mesh = *_form->mesh();
std::int32_t num_vertices_per_cell = mesh::num_cell_vertices(mesh.geometry().cmap().cell_shape());
//std::int32_t num_coordinates_per_vertex = mesh.geometry().dim();
const mesh::Mesh<U> &mesh = *_form->mesh();
#if DOLFINX_VERSION_MINOR >= 11
if (mesh.geometry().cmaps().size() > 1)
throw std::runtime_error(
"CUDAForm does not yet support meshes with multiple cell types.");
std::int32_t num_vertices_per_cell =
mesh::num_cell_vertices(mesh.geometry().cmaps()[0].cell_shape());
#else
std::int32_t num_vertices_per_cell =
mesh::num_cell_vertices(mesh.geometry().cmap().cell_shape());
#endif
// std::int32_t num_coordinates_per_vertex = mesh.geometry().dim();
std::int32_t num_coordinates_per_vertex = 3;

// Find the number of degrees of freedom per cell
int32_t num_dofs_per_cell0 = 1;
int32_t num_dofs_per_cell1 = 1;
if (_form->rank() > 0) {
const DofMap& dofmap0 = *_form->function_spaces()[0]->dofmap();
num_dofs_per_cell0 = dofmap0.element_dof_layout().num_dofs() * dofmap0.element_dof_layout().block_size();
const DofMap &dofmap0 = *_form->function_spaces()[0]->dofmap();
num_dofs_per_cell0 = dofmap0.element_dof_layout().num_dofs() *
dofmap0.element_dof_layout().block_size();
}
if (_form->rank() > 1) {
const DofMap& dofmap1 = *_form->function_spaces()[1]->dofmap();
num_dofs_per_cell1 = dofmap1.element_dof_layout().num_dofs() * dofmap1.element_dof_layout().block_size();
const DofMap &dofmap1 = *_form->function_spaces()[1]->dofmap();
num_dofs_per_cell1 = dofmap1.element_dof_layout().num_dofs() *
dofmap1.element_dof_layout().block_size();
}

for (auto& [integral_type, integrals_for_type] : _integrals) {
auto sources = _cuda_kernel_sources[static_cast<std::size_t>(integral_type)];

for (auto &[integral_type, integrals_for_type] : _integrals) {
auto sources =
_cuda_kernel_sources[static_cast<std::size_t>(integral_type)];
for (int i = 0; i < integrals_for_type.size(); i++) {
std::string name = "";
auto it = sources.find(i);
if (it == sources.end())
throw std::runtime_error("No kernel for requested domain index.");
// name is modified by this function
auto assembly_module = compile_form_integral_kernel(
cuda_context,
cachedir,
_form->rank(),
integral_type,
it->second,
max_threads_per_block,
min_blocks_per_multiprocessor,
num_vertices_per_cell,
num_coordinates_per_vertex,
num_dofs_per_cell0,
num_dofs_per_cell1,
_form->coefficient_offsets().back(),
assembly_kernel_type,
false,
false,
name);
cuda_context, cachedir, _form->rank(), integral_type, it->second,
max_threads_per_block, min_blocks_per_multiprocessor,
num_vertices_per_cell, num_coordinates_per_vertex,
num_dofs_per_cell0, num_dofs_per_cell1,
_form->coefficient_offsets().back(), assembly_kernel_type, false,
false, name);
integrals_for_type[i].set_kernels(std::move(assembly_module), name);
}
}
_compiled = true;
}

/// Copy constructor
CUDAForm(const CUDAForm& form) = delete;
CUDAForm(const CUDAForm &form) = delete;

/// Move constructor
CUDAForm(CUDAForm&& form) = default;
CUDAForm(CUDAForm &&form) = default;

/// Destructor
virtual ~CUDAForm() = default;

bool compiled() { return _compiled; }

bool restricted() { return _restricted_dofmaps.size() > 0; }
std::map<IntegralType, std::vector<CUDAFormIntegral<T,U>>>& integrals() {

std::map<IntegralType, std::vector<CUDAFormIntegral<T, U>>> &integrals() {
if (!_compiled) {
throw std::runtime_error("Cannot access integrals for uncompiled cuda form!");
throw std::runtime_error(
"Cannot access integrals for uncompiled cuda form!");
}
return _integrals;
}

CUDAFormCoefficients<T,U>& coefficients() { return _coefficients; }
CUDAFormCoefficients<T, U> &coefficients() { return _coefficients; }

const CUDAFormConstants<T>& constants() { return _constants; }
const CUDAFormConstants<T> &constants() { return _constants; }

std::shared_ptr<const CUDADofMap> unrestricted_dofmap(size_t i) {
if (i >= _form->function_spaces().size()) throw std::runtime_error("Dofmap index out of bounds!");
return _dofmap_store.get_device_object(_form->function_spaces()[i]->dofmap().get());
if (i >= _form->function_spaces().size())
throw std::runtime_error("Dofmap index out of bounds!");
return _dofmap_store.get_device_object(
_form->function_spaces()[i]->dofmap().get());
}

std::shared_ptr<const CUDADofMap> dofmap(size_t i) {
if (!restricted()) return unrestricted_dofmap(i);
if (i >= _restricted_dofmaps.size()) throw std::runtime_error("Dofmap index out of bounds!");
if (!restricted())
return unrestricted_dofmap(i);
if (i >= _restricted_dofmaps.size())
throw std::runtime_error("Dofmap index out of bounds!");
return _restricted_dofmaps[i];
}

Form<T,U>* form() { return _form; }
Form<T, U> *form() { return _form; }

CUDADirichletBC<T,U> bc(
const CUDA::Context& cuda_context, size_t i,
std::vector<std::shared_ptr<const DirichletBC<T,U>>> bcs)
{
return CUDADirichletBC<T,U>(cuda_context, *_form->function_spaces()[i], bcs);
CUDADirichletBC<T, U>
bc(const CUDA::Context &cuda_context, size_t i,
std::vector<std::shared_ptr<const DirichletBC<T, U>>> bcs) {
return CUDADirichletBC<T, U>(cuda_context, *_form->function_spaces()[i],
bcs);
}

/// Copy the coefficient and constant data to the device
/// This can be necessary if either changes on the host
void to_device(const CUDA::Context& cuda_context)
{
void to_device(const CUDA::Context &cuda_context) {
_coefficients.copy_coefficients_to_device(cuda_context);
_constants.update_constant_values();
_constants.update_constant_values();
}

void set_restriction(
std::vector<std::int32_t> offsets,
std::vector<std::int32_t> ghost_offsets,
std::vector<std::shared_ptr<std::map<std::int32_t, std::int32_t>>> restriction)
{
std::vector<std::int32_t> offsets,
std::vector<std::int32_t> ghost_offsets,
std::vector<std::shared_ptr<std::map<std::int32_t, std::int32_t>>>
restriction) {
if (restriction.size() != _form->function_spaces().size()) {
throw std::runtime_error("Number of restrictions must equal arity of form (1 for vector, 2 for matrix)!");
throw std::runtime_error("Number of restrictions must equal arity of "
"form (1 for vector, 2 for matrix)!");
}
_restriction = restriction;
if (_restricted_dofmaps.size()) {
// need to update the restriction
for (int i = 0; i < _restricted_dofmaps.size(); i++) {
_restricted_dofmaps[i]->update(
offsets[i],
ghost_offsets[i],
restriction[i].get()
);
}
}
else {
_restricted_dofmaps[i]->update(offsets[i], ghost_offsets[i],
restriction[i].get());
}
} else {
for (int i = 0; i < restriction.size(); i++) {
_restricted_dofmaps.push_back(
std::make_shared<CUDADofMap>(
_form->function_spaces()[i]->dofmap().get(),
offsets[i],
ghost_offsets[i],
restriction[i].get()
)
);
_restricted_dofmaps.push_back(std::make_shared<CUDADofMap>(
_form->function_spaces()[i]->dofmap().get(), offsets[i],
ghost_offsets[i], restriction[i].get()));
}
}
}

const std::vector<std::shared_ptr<std::map<std::int32_t, std::int32_t>>> get_restriction()
{
const std::vector<std::shared_ptr<std::map<std::int32_t, std::int32_t>>>
get_restriction() {
return _restriction;
}

std::shared_ptr<dolfinx::common::IndexMap> restriction_index_map(size_t i) {
std::vector<std::int32_t> restricted_inds;
for (auto const& pair: *_restriction[i]) restricted_inds.push_back(pair.first);
for (auto const &pair : *_restriction[i])
restricted_inds.push_back(pair.first);
auto [sub_imap, inds] = dolfinx::common::create_sub_index_map(
*_form->function_spaces()[0]->dofmap()->index_map,
restricted_inds,
dolfinx::common::IndexMapOrder::preserve, false
);
*_form->function_spaces()[0]->dofmap()->index_map, restricted_inds,
dolfinx::common::IndexMapOrder::preserve, false);
return std::make_shared<dolfinx::common::IndexMap>(std::move(sub_imap));
}

Expand All @@ -251,18 +246,20 @@ class CUDAForm
// Restricted dofmaps
std::vector<std::shared_ptr<CUDADofMap>> _restricted_dofmaps;
// Restriction
std::vector<std::shared_ptr<std::map<std::int32_t, std::int32_t>>> _restriction;
std::vector<std::shared_ptr<std::map<std::int32_t, std::int32_t>>>
_restriction;
// Form coefficients
CUDAFormCoefficients<T, U> _coefficients;
// Form Constants
CUDAFormConstants<T> _constants;
// Compiled CUDA kernels
std::map<IntegralType, std::vector<CUDAFormIntegral<T,U>>> _integrals;
// CUDA tabulate tensors
std::array<std::map<int, std::pair<std::string, std::string>>, 4> _cuda_kernel_sources;
std::map<IntegralType, std::vector<CUDAFormIntegral<T, U>>> _integrals;
// CUDA tabulate tensors
std::array<std::map<int, std::pair<std::string, std::string>>, 4>
_cuda_kernel_sources;
bool _compiled;
Form<T,U>* _form;
ufcx_form* _ufcx_form;
Form<T, U> *_form;
ufcx_form *_ufcx_form;
};

} // end namespace fem
Expand Down
Loading