Skip to content

Commit

Permalink
Generate PackInfo
Browse files Browse the repository at this point in the history
  • Loading branch information
@jngrad
jngrad committed Sep 19, 2024
1 parent 5e39e84 commit 8ea7111
Show file tree
Hide file tree
Showing 13 changed files with 3,622 additions and 12 deletions.
24 changes: 24 additions & 0 deletions maintainer/walberla_kernels/generate_lb_kernels.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,6 +193,30 @@ def paramlist(parameters, keys):
ctx, config, method, templates
)

# generate PackInfo
assignments = pystencils_espresso.generate_pack_info_pdfs_field_assignments(
fields, streaming_pattern="pull")
spec = pystencils_espresso.generate_pack_info_vector_field_specifications(
config, stencil, force_field.layout)
for params, target_suffix in paramlist(parameters, ["CPU"]):
pystencils_walberla.generate_pack_info_from_kernel(
ctx, f"PackInfoPdf{precision_prefix}{target_suffix}", assignments,
kind="pull", **params)
pystencils_walberla.generate_pack_info(
ctx, f"PackInfoVec{precision_prefix}{target_suffix}", spec, **params)
if target_suffix == "CUDA":
continue
token = "\n //TODO: optimize by generating kernel for this case\n"
for field_suffix in ["Pdf", "Vec"]:
class_name = f"PackInfo{field_suffix}{precision_prefix}{target_suffix}" # nopep8
with open(f"{class_name}.h", "r+") as f:
content = f.read()
assert token in content
content = content.replace(token, "\n")
f.seek(0)
f.truncate()
f.write(content)

# boundary conditions
ubb_dynamic = lbmpy_espresso.UBB(
lambda *args: None, dim=3, data_type=config.data_type.default_factory())
Expand Down
54 changes: 54 additions & 0 deletions maintainer/walberla_kernels/pystencils_espresso.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -207,6 +207,60 @@ def generate_fields(config, stencil, field_layout='fzyx'):
return fields


def generate_pack_info_pdfs_field_assignments(fields, streaming_pattern):
"""
Visualize the stencil directions with::
import lbmpy
import matplotlib.pyplot as plt
stencil = lbmpy.LBStencil(lbmpy.Stencil.D3Q19)
stencil.plot(data=[i for i in range(19)])
plt.show()
"""
stencil = lbmpy.enums.Stencil.D3Q19
lbm_config = lbmpy.LBMConfig(stencil=stencil,
method=lbmpy.Method.CUMULANT,
compressible=True,
zero_centered=False,
weighted=True,
streaming_pattern=streaming_pattern,
relaxation_rate=sp.Symbol("omega_shear"),
)
lbm_opt = lbmpy.LBMOptimisation(
symbolic_field=fields["pdfs" if streaming_pattern ==
"pull" else "pdfs_tmp"],
symbolic_temporary_field=fields["pdfs" if streaming_pattern ==
"push" else "pdfs_tmp"],
field_layout=fields['pdfs'].layout)
lbm_update_rule = lbmpy.create_lb_update_rule(
lbm_config=lbm_config,
lbm_optimisation=lbm_opt)
return lbm_update_rule.all_assignments


def generate_pack_info_vector_field_specifications(config, stencil, layout):
import collections
import itertools
field = ps.Field.create_generic(
"field",
3,
data_type_np[config.data_type.default_factory().c_name],
index_dimensions=1,
layout=layout,
index_shape=(3,)
)
q = len(stencil)
coord = itertools.product(*[(-1, 0, 1)] * 3)
if q == 19:
dirs = tuple((i, j, k) for i, j, k in coord if i**2 + j**2 + k**2 != 3)
else:
dirs = tuple((i, j, k) for i, j, k in coord)
spec = collections.defaultdict(set)
spec[dirs] = {field[0, 0, 0](i) for i in range(3)}
return spec


def generate_config(ctx, params):
return pystencils_walberla.utility.config_from_context(ctx, **params)

Expand Down
70 changes: 58 additions & 12 deletions src/walberla_bridge/src/lattice_boltzmann/LBWalberlaImpl.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -119,6 +119,11 @@ class LBWalberlaImpl : public LBWalberlaBase {
using VectorField = field::GhostLayerField<FT, uint_t{3u}>;
template <class Field>
using PackInfo = field::communication::PackInfo<Field>;
template <class Field>
using PackInfoStreaming =
std::conditional_t<std::is_same_v<Field, PdfField>,
typename detail::KernelTrait<FT, AT>::PackInfoPdf,
typename detail::KernelTrait<FT, AT>::PackInfoVec>;
template <class Stencil>
using RegularCommScheme =
blockforest::communication::UniformBufferedScheme<Stencil>;
Expand All @@ -133,6 +138,8 @@ class LBWalberlaImpl : public LBWalberlaBase {
using VectorField = gpu::GPUField<FT>;
template <class Field>
using PackInfo = gpu::communication::MemcpyPackInfo<Field>;
template <class Field>
using PackInfoStreaming = gpu::communication::MemcpyPackInfo<Field>;
template <class Stencil>
using RegularCommScheme = gpu::communication::UniformGPUScheme<Stencil>;
template <class Stencil>
Expand Down Expand Up @@ -284,6 +291,7 @@ class LBWalberlaImpl : public LBWalberlaBase {

/** Flag for boundary cells. */
FlagUID const Boundary_flag{"boundary"};
bool m_has_boundaries{false};

/**
* @brief Full communicator.
Expand All @@ -307,6 +315,10 @@ class LBWalberlaImpl : public LBWalberlaBase {
template <class Field>
using PackInfo =
typename FieldTrait<FloatType, Architecture>::template PackInfo<Field>;
template <class Field>
using PackInfoStreaming =
typename FieldTrait<FloatType,
Architecture>::template PackInfoStreaming<Field>;

// communicators
std::shared_ptr<BoundaryFullCommunicator> m_boundary_communicator;
Expand Down Expand Up @@ -414,6 +426,24 @@ class LBWalberlaImpl : public LBWalberlaBase {
#endif
}

void setup_streaming_communicator() {
auto const setup = [this]<typename PdfPackInfo>() {
auto const &blocks = m_lattice->get_blocks();
m_pdf_streaming_communicator =
std::make_shared<PDFStreamingCommunicator>(blocks);
m_pdf_streaming_communicator->addPackInfo(
std::make_shared<PdfPackInfo>(m_pdf_field_id));
m_pdf_streaming_communicator->addPackInfo(
std::make_shared<PackInfoStreaming<VectorField>>(
m_last_applied_force_field_id));
};
if (m_has_boundaries or (m_collision_model and has_lees_edwards_bc())) {
setup.template operator()<PackInfo<PdfField>>();
} else {
setup.template operator()<PackInfoStreaming<PdfField>>();
}
}

public:
LBWalberlaImpl(std::shared_ptr<LatticeWalberla> lattice, double viscosity,
double density)
Expand Down Expand Up @@ -448,12 +478,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
reset_boundary_handling();

// Set up the communication and register fields
m_pdf_streaming_communicator =
std::make_shared<PDFStreamingCommunicator>(blocks);
m_pdf_streaming_communicator->addPackInfo(
std::make_shared<PackInfo<PdfField>>(m_pdf_field_id));
m_pdf_streaming_communicator->addPackInfo(
std::make_shared<PackInfo<VectorField>>(m_last_applied_force_field_id));
setup_streaming_communicator();

m_full_communicator = std::make_shared<RegularFullCommunicator>(blocks);
m_full_communicator->addPackInfo(
Expand Down Expand Up @@ -555,7 +580,9 @@ class LBWalberlaImpl : public LBWalberlaBase {
integrate_collide(blocks);
m_pdf_streaming_communicator->communicate();
// Handle boundaries
integrate_boundaries(blocks);
if (m_has_boundaries) {
integrate_boundaries(blocks);
}
// LB stream
integrate_stream(blocks);
// Mark pending ghost layer updates
Expand All @@ -569,7 +596,9 @@ class LBWalberlaImpl : public LBWalberlaBase {
void integrate_pull_scheme() {
auto const &blocks = get_lattice().get_blocks();
// Handle boundaries
integrate_boundaries(blocks);
if (m_has_boundaries) {
integrate_boundaries(blocks);
}
// LB stream
integrate_stream(blocks);
// LB collide
Expand Down Expand Up @@ -690,6 +719,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
omega_odd, omega, seed, uint32_t{0u});
m_collision_model = std::make_shared<CollisionModel>(std::move(obj));
m_run_collide_sweep = CollideSweepVisitor(blocks);
setup_streaming_communicator();
}

void set_collision_model(
Expand Down Expand Up @@ -734,6 +764,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
blocks, m_last_applied_force_field_id, m_vec_tmp_field_id,
n_ghost_layers, shear_direction, shear_plane_normal,
m_lees_edwards_callbacks->get_pos_offset);
setup_streaming_communicator();
}

void check_lebc(unsigned int shear_direction,
Expand Down Expand Up @@ -765,10 +796,12 @@ class LBWalberlaImpl : public LBWalberlaBase {
bool consider_ghosts = false) const override {
assert(not(consider_ghosts and m_pending_ghost_comm.test(GhostComm::VEL)));
assert(not(consider_ghosts and m_pending_ghost_comm.test(GhostComm::UBB)));
auto const is_boundary = get_node_is_boundary(node, consider_ghosts);
if (is_boundary) // is info available locally
if (*is_boundary) // is the node a boundary
if (m_has_boundaries) {
auto const is_boundary = get_node_is_boundary(node, consider_ghosts);
if (is_boundary and *is_boundary) {
return get_node_velocity_at_boundary(node, consider_ghosts);
}
}
auto const bc = get_block_and_cell(get_lattice(), node, consider_ghosts);
if (!bc)
return std::nullopt;
Expand Down Expand Up @@ -1266,6 +1299,7 @@ class LBWalberlaImpl : public LBWalberlaBase {

bool set_node_velocity_at_boundary(Utils::Vector3i const &node,
Utils::Vector3d const &velocity) override {
on_boundary_add();
m_pending_ghost_comm.set(GhostComm::UBB);
auto bc = get_block_and_cell(get_lattice(), node, true);
if (bc) {
Expand Down Expand Up @@ -1307,6 +1341,7 @@ class LBWalberlaImpl : public LBWalberlaBase {
void set_slice_velocity_at_boundary(
Utils::Vector3i const &lower_corner, Utils::Vector3i const &upper_corner,
std::vector<std::optional<Utils::Vector3d>> const &velocity) override {
on_boundary_add();
m_pending_ghost_comm.set(GhostComm::UBB);
if (auto const ci = get_interval(lower_corner, upper_corner)) {
auto const &lattice = get_lattice();
Expand Down Expand Up @@ -1388,19 +1423,30 @@ class LBWalberlaImpl : public LBWalberlaBase {

void reallocate_ubb_field() override { m_boundary->boundary_update(); }

void on_boundary_add() {
if (not m_has_boundaries) {
m_has_boundaries = true;
setup_streaming_communicator();
}
m_has_boundaries = true;
}

void clear_boundaries() override {
reset_boundary_handling();
m_pending_ghost_comm.set(GhostComm::UBB);
ghost_communication();
m_has_boundaries = false;
setup_streaming_communicator();
}

void
update_boundary_from_shape(std::vector<int> const &raster_flat,
std::vector<double> const &data_flat) override {
on_boundary_add();
m_pending_ghost_comm.set(GhostComm::UBB);
auto const grid_size = get_lattice().get_grid_dimensions();
auto const data = fill_3D_vector_array(data_flat, grid_size);
set_boundary_from_grid(*m_boundary, get_lattice(), raster_flat, data);
m_pending_ghost_comm.set(GhostComm::UBB);
ghost_communication();
reallocate_ubb_field();
}
Expand Down
2 changes: 2 additions & 0 deletions src/walberla_bridge/src/lattice_boltzmann/generated_kernels/CMakeLists.txt
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,8 @@
target_sources(
espresso_walberla
PRIVATE StreamSweepSinglePrecision.cpp StreamSweepDoublePrecision.cpp
PackInfoPdfSinglePrecision.cpp PackInfoPdfDoublePrecision.cpp
PackInfoVecSinglePrecision.cpp PackInfoVecDoublePrecision.cpp
InitialPDFsSetterSinglePrecision.cpp
InitialPDFsSetterDoublePrecision.cpp Dynamic_UBB_single_precision.cpp
Dynamic_UBB_double_precision.cpp)
Expand Down
Loading

0 comments on commit 8ea7111

Please sign in to comment.