Skip to content

Commit

Permalink
task_system: Rehaul of the interface to go back to std::function... T…
Browse files Browse the repository at this point in the history 
…his will eventually be replaced by something better which provides better control to how the capture data is allocated.

The idéa is to introduce a granularity size to task_system to ensure that task_sets are processed in chunks where the user can control the size.
  • Loading branch information
@scanberg
scanberg committed Nov 19, 2024
1 parent bdfa8dd commit 4281bbf
Show file tree
Hide file tree
Showing 6 changed files with 141 additions and 182 deletions.
18 changes: 10 additions & 8 deletions src/components/ramachandran/ramachandran.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1271,11 +1271,14 @@ struct Ramachandran : viamd::EventHandler {
uint32_t frame_end;
uint32_t frame_stride;
float sigma;
md_allocator_i* alloc;
};

md_allocator_i* alloc = md_get_heap_allocator();

uint64_t tex_size = sizeof(vec4_t) * density_tex_dim * density_tex_dim;
uint64_t alloc_size = sizeof(UserData) + tex_size + alignof(vec4_t);
UserData* user_data = (UserData*)md_alloc(md_get_heap_allocator(), sizeof(UserData) + tex_size);
UserData* user_data = (UserData*)md_alloc(alloc, alloc_size);
vec4_t* density_tex = (vec4_t*)NEXT_ALIGNED_ADDRESS(user_data + 1, alignof(vec4_t));
memset(density_tex, 0, tex_size);

Expand All @@ -1291,9 +1294,9 @@ struct Ramachandran : viamd::EventHandler {
user_data->frame_end = frame_end;
user_data->frame_stride = frame_stride;
user_data->sigma = blur_sigma;
user_data->alloc = alloc;

task_system::ID async_task = task_system::create_pool_task(STR_LIT("Rama density"), [](void* user_data) {
UserData* data = (UserData*)user_data;
task_system::ID async_task = task_system::create_pool_task(STR_LIT("Rama density"), [data = user_data]() {
const float angle_to_coord_scale = 1.0f / (2.0f * PI);
const float angle_to_coord_offset = 0.5f;

Expand Down Expand Up @@ -1333,13 +1336,12 @@ struct Ramachandran : viamd::EventHandler {
data->rep->den_sum[1] = (float)sum[1];
data->rep->den_sum[2] = (float)sum[2];
data->rep->den_sum[3] = (float)sum[3];
}, user_data);
});

task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update rama texture"), [](void* user_data) {
UserData* data = (UserData*)user_data;
task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update rama texture"), [data = user_data]() {
gl::set_texture_2D_data(data->rep->den_tex, data->density_tex, GL_RGBA32F);
md_free(md_get_heap_allocator(), data, data->alloc_size);
}, user_data);
md_free(data->alloc, data, data->alloc_size);
});

task_system::set_task_dependency(main_task, async_task);
task_system::enqueue_task(async_task);
Expand Down
5 changes: 2 additions & 3 deletions src/components/shapespace/shapespace.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -379,9 +379,8 @@ struct Shapespace : viamd::EventHandler {
md_array_resize(coords, num_frames * num_structures, arena);
MEMSET(weights, 0, md_array_bytes(weights));
MEMSET(coords, 0, md_array_bytes(coords));
evaluate_task = task_system::create_pool_task(STR_LIT("Eval Shape Space"), 0, (uint32_t)num_frames, [](uint32_t range_beg, uint32_t range_end, void* user_data, uint32_t thread_num) {
evaluate_task = task_system::create_pool_task(STR_LIT("Eval Shape Space"), (uint32_t)num_frames, [shapespace = this](uint32_t range_beg, uint32_t range_end, uint32_t thread_num) {
(void)thread_num;
Shapespace* shapespace = (Shapespace*)user_data;
ApplicationState* app_state = shapespace->app_state;
const size_t stride = ALIGN_TO(app_state->mold.mol.atom.count, 8);
const size_t bytes = stride * 3 * sizeof(float);
Expand Down Expand Up @@ -423,7 +422,7 @@ struct Shapespace : viamd::EventHandler {
}
}
md_array_free(xyzw, md_get_heap_allocator());
}, this);
});

task_system::enqueue_task(evaluate_task);
}
Expand Down
35 changes: 13 additions & 22 deletions src/components/veloxchem/veloxchem.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -805,8 +805,7 @@ struct VeloxChem : viamd::EventHandler {
task_system::ID async_task = evaluate_gto_on_grid_async(&payload->args);

// Launch task for main (render) thread to update the volume texture
task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update Volume"), [](void* user_data) {
Payload* data = (Payload*)user_data;
task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update Volume"), [data = payload]() {

// Ensure that the dimensions of the texture have not changed during evaluation
int dim[3];
Expand All @@ -816,7 +815,7 @@ struct VeloxChem : viamd::EventHandler {

md_free(data->alloc, data->args.gtos, data->args.num_gtos * sizeof(md_gto_t));
md_free(data->alloc, data->mem, data->mem_size);
}, payload);
});

task_system::set_task_dependency(main_task, async_task);
task_system::enqueue_task(async_task);
Expand Down Expand Up @@ -890,9 +889,7 @@ struct VeloxChem : viamd::EventHandler {
task_system::ID async_task = evaluate_gto_on_grid_async(&payload->args);

// Launch task for main (render) thread to update the volume texture
task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update Volume"), [](void* user_data) {
Payload* data = (Payload*)user_data;

task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update Volume"), [data = payload]() {
// Ensure that the dimensions of the texture have not changed during evaluation
int dim[3];
if (gl::get_texture_dim(dim, data->tex) && MEMCMP(dim, data->args.grid.dim, sizeof(dim)) == 0) {
Expand All @@ -901,7 +898,7 @@ struct VeloxChem : viamd::EventHandler {

md_free(data->alloc, data->args.gtos, data->args.num_gtos * sizeof(md_gto_t));
md_free(data->alloc, data->mem, data->mem_size);
}, payload);
});

task_system::set_task_dependency(main_task, async_task);
task_system::enqueue_task(async_task);
Expand Down Expand Up @@ -1006,9 +1003,7 @@ struct VeloxChem : viamd::EventHandler {
task_system::ID async_task = evaluate_gto_on_grid_async(&payload->args);

// Launch task for main (render) thread to update the volume texture
task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update Volume"), [](void* user_data) {
Payload* data = (Payload*)user_data;

task_system::ID main_task = task_system::create_main_task(STR_LIT("##Update Volume"), [data = payload]() {
// Ensure that the dimensions of the texture have not changed during evaluation
int dim[3];
if (gl::get_texture_dim(dim, data->tex) && MEMCMP(dim, data->args.grid.dim, sizeof(dim)) == 0) {
Expand All @@ -1017,7 +1012,7 @@ struct VeloxChem : viamd::EventHandler {

md_free(data->alloc, data->args.gtos, data->args.num_gtos * sizeof(md_gto_t));
md_free(data->alloc, data, data->alloc_size);
}, payload);
});

task_system::set_task_dependency(main_task, async_task);
task_system::enqueue_task(async_task);
Expand Down Expand Up @@ -1554,9 +1549,7 @@ struct VeloxChem : viamd::EventHandler {
task_system::ID eval_task = evaluate_gto_on_grid_async(&payload->args);

// @TODO: This should be performed as a range task in parallel
task_system::ID segment_task = task_system::create_pool_task(STR_LIT("##Segment Volume"), [](void* user_data) {
Payload* data = (Payload*)user_data;

task_system::ID segment_task = task_system::create_pool_task(STR_LIT("##Segment Volume"), [data = payload]() {
#if DEBUG
double sum = 0.0;
size_t len = data->args.grid.dim[0] * data->args.grid.dim[1] * data->args.grid.dim[2];
Expand All @@ -1572,7 +1565,7 @@ struct VeloxChem : viamd::EventHandler {

md_free(data->alloc, data->args.gtos, data->args.num_gtos * sizeof(md_gto_t));
md_free(data->alloc, data, data->alloc_size);
}, payload);
});

task_system::set_task_dependency(segment_task, eval_task);

Expand All @@ -1592,9 +1585,8 @@ struct VeloxChem : viamd::EventHandler {

// We evaluate the in parallel over smaller NxNxN blocks
const uint32_t num_blocks = (args->grid.dim[0] / BLK_DIM) * (args->grid.dim[1] / BLK_DIM) * (args->grid.dim[2] / BLK_DIM);
task_system::ID async_task = task_system::create_pool_task(STR_LIT("Evaluate Orbital"), 0, num_blocks, [](uint32_t range_beg, uint32_t range_end, void* user_data, uint32_t thread_num) {
task_system::ID async_task = task_system::create_pool_task(STR_LIT("Evaluate Orbital"), num_blocks, [data = args](uint32_t range_beg, uint32_t range_end, uint32_t thread_num) {
(void)thread_num;
AsyncGridEvalArgs* data = (AsyncGridEvalArgs*)user_data;
MD_LOG_DEBUG("Starting async eval of orbital grid [%i][%i][%i]", data->grid.dim[0], data->grid.dim[1], data->grid.dim[2]);

// Number of NxNxN blocks in each dimension
Expand Down Expand Up @@ -1657,7 +1649,7 @@ struct VeloxChem : viamd::EventHandler {
}

md_temp_set_pos_back(temp_pos);
}, args);
});

return async_task;
}
Expand Down Expand Up @@ -2500,7 +2492,7 @@ struct VeloxChem : viamd::EventHandler {
rsp.first_plot_rot_ecd = false;
ImGui::TreePop();
}
#if 0
#if 1
if (rsp.first_plot_vib) { ImGui::SetNextItemOpen(true); }
if (ImGui::TreeNode("Vibrational Analysis")) {
// draw the vibrational analysis
Expand Down Expand Up @@ -4484,10 +4476,9 @@ struct VeloxChem : viamd::EventHandler {
if (compute_transition_group_values_async(&eval_attach, &seg_attach, nto.transition_density_part, nto.group.count, nto.atom_group_idx, nto.atom_xyzr, nto.num_atoms, nto_idx, MD_VLX_NTO_TYPE_PARTICLE, MD_GTO_EVAL_MODE_PSI_SQUARED, samples_per_angstrom) &&
compute_transition_group_values_async(&eval_detach, &seg_detach, nto.transition_density_hole, nto.group.count, nto.atom_group_idx, nto.atom_xyzr, nto.num_atoms, nto_idx, MD_VLX_NTO_TYPE_HOLE, MD_GTO_EVAL_MODE_PSI_SQUARED, samples_per_angstrom))
{
task_system::ID compute_matrix_task = task_system::create_main_task(STR_LIT("##Compute Transition Matrix"), [](void* user_data) {
VeloxChem::Nto* nto = (VeloxChem::Nto*)user_data;
task_system::ID compute_matrix_task = task_system::create_main_task(STR_LIT("##Compute Transition Matrix"), [nto = &nto]() {
compute_transition_matrix(nto->transition_matrix, nto->group.count, nto->transition_density_hole, nto->transition_density_part);
}, &nto);
});

task_system::set_task_dependency(compute_matrix_task, seg_attach);
task_system::set_task_dependency(compute_matrix_task, seg_detach);
Expand Down
Loading

0 comments on commit 4281bbf

Please sign in to comment.