diff --git a/src/language/MeshModule.cpp b/src/language/MeshModule.cpp index 696799455ebef773a52770848b296d7769a966dc..fe0c75d9c317dd53ae911570bc0a19b0ea1f4e2d 100644 --- a/src/language/MeshModule.cpp +++ b/src/language/MeshModule.cpp @@ -10,12 +10,44 @@ #include <mesh/Mesh.hpp> #include <utils/Exceptions.hpp> +#include <Kokkos_Core.hpp> +#include <cstdio> + template <> inline ASTNodeDataType ast_node_data_type_from<std::shared_ptr<IMesh>> = {ASTNodeDataType::type_id_t, "mesh"}; template <> inline ASTNodeDataType ast_node_data_type_from<FunctionSymbolId> = {ASTNodeDataType::function_t}; +struct TagA +{ +}; +struct TagC +{ +}; + +struct Foo +{ + static std::vector<int> m_ids; + KOKKOS_INLINE_FUNCTION + void + operator()(const TagA, const Kokkos::TeamPolicy<>::member_type& team) const + { + m_ids[team.league_rank()] = team.league_rank() * 3; + } + + KOKKOS_INLINE_FUNCTION + void + operator()(const TagC, const Kokkos::TeamPolicy<>::member_type& team) const + { + m_ids[team.league_rank()] = team.league_rank() * 2; + } + + Foo() {} +}; + +inline std::vector<int> Foo::m_ids; + MeshModule::MeshModule() { this->_addTypeDescriptor( @@ -31,37 +63,60 @@ MeshModule::MeshModule() )); - this - ->_addBuiltinFunction("transform", - std::make_shared< - BuiltinFunctionEmbedder<std::shared_ptr<IMesh>, std::shared_ptr<IMesh>, FunctionSymbolId>>( - std::function<std::shared_ptr<IMesh>(std::shared_ptr<IMesh>, FunctionSymbolId)>{ - - [](std::shared_ptr<IMesh> p_mesh, - FunctionSymbolId function_id) -> std::shared_ptr<IMesh> { - auto& symbol_table = function_id.symbolTable(); - auto& function_expression = - *symbol_table.functionTable()[function_id.id()].definitionNode().children[1]; - auto& function_context = function_expression.m_symbol_table->context(); - - switch (p_mesh->dimension()) { - case 1: { - throw NotImplementedError("not implemented in 1d"); - break; - } - case 2: { - throw NotImplementedError("not implemented in 2d"); - break; - } - case 3: { - using MeshType = Mesh<Connectivity3D>; - const MeshType& given_mesh = dynamic_cast<const MeshType&>(*p_mesh); - NodeValue<const TinyVector<3>> given_xr = given_mesh.xr(); - - NodeValue<TinyVector<3>> xr(given_mesh.connectivity()); - - parallel_for(given_mesh.numberOfNodes(), [=, &function_expression, - &function_context](NodeId r) { + this->_addBuiltinFunction("transform", + std::make_shared<BuiltinFunctionEmbedder<std::shared_ptr<IMesh>, std::shared_ptr<IMesh>, + FunctionSymbolId>>( + std::function<std::shared_ptr<IMesh>(std::shared_ptr<IMesh>, FunctionSymbolId)>{ + + [](std::shared_ptr<IMesh> p_mesh, + FunctionSymbolId function_id) -> std::shared_ptr<IMesh> { + auto& symbol_table = function_id.symbolTable(); + auto& function_expression = + *symbol_table.functionTable()[function_id.id()].definitionNode().children[1]; + auto& function_context = function_expression.m_symbol_table->context(); + + // Foo foo; + // Foo::m_ids.resize(Kokkos::DefaultExecutionSpace::impl_thread_pool_size()); + // Kokkos::parallel_for(Kokkos::TeamPolicy< + // TagA>(Kokkos::DefaultExecutionSpace::impl_thread_pool_size(), + // 1), + // foo); + // std::cout << "--------\n" << std::endl; + + // for (int i = 0; i < Kokkos::DefaultExecutionSpace::impl_thread_pool_size(); ++i) { + // std::cout << "m_ids[" << i << "] = " << Foo::m_ids[i] << " | " << i * 3 << + // std::endl; + // } + + // std::cout << "********\n" << std::endl; + + // Kokkos::parallel_for(Kokkos::TeamPolicy< + // TagC>(Kokkos::DefaultExecutionSpace::impl_thread_pool_size(), + // 1), + // foo); + // std::cout << "--------\n" << std::endl; + + // for (int i = 0; i < Kokkos::DefaultExecutionSpace::impl_thread_pool_size(); ++i) { + // std::cout << "m_ids[" << i << "] = " << Foo::m_ids[i] << " | " << i * 2 << + // std::endl; + // } + + switch (p_mesh->dimension()) { + case 1: { + throw NotImplementedError("not implemented in 1d"); + break; + } + case 2: { + throw NotImplementedError("not implemented in 2d"); + break; + } + case 3: { + using MeshType = Mesh<Connectivity3D>; + const MeshType& given_mesh = dynamic_cast<const MeshType&>(*p_mesh); + NodeValue<const TinyVector<3>> given_xr = given_mesh.xr(); + + NodeValue<TinyVector<3>> xr(given_mesh.connectivity()); + ExecutionPolicy::Context context{function_context.id(), std::make_shared<ExecutionPolicy::Context::Values>( function_context.size())}; @@ -69,21 +124,23 @@ MeshModule::MeshModule() ExecutionPolicy execution_policy; ExecutionPolicy context_execution_policy{execution_policy, context}; - context_execution_policy.currentContext()[0] = given_xr[r]; + parallel_for(given_mesh.numberOfNodes(), [=, &function_expression, + &context_execution_policy](NodeId r) { + context_execution_policy.currentContext()[0] = given_xr[r]; - auto&& value = function_expression.execute(context_execution_policy); + auto&& value = function_expression.execute(context_execution_policy); - AggregateDataVariant& v = std::get<AggregateDataVariant>(value); - xr[r] = {std::get<double>(v[0]), std::get<double>(v[1]), std::get<double>(v[2])}; - }); + AggregateDataVariant& v = std::get<AggregateDataVariant>(value); + xr[r] = {std::get<double>(v[0]), std::get<double>(v[1]), std::get<double>(v[2])}; + }); - return std::make_shared<MeshType>(given_mesh.shared_connectivity(), xr); - } - default: { - return nullptr; - } - } - }} + return std::make_shared<MeshType>(given_mesh.shared_connectivity(), xr); + } + default: { + return nullptr; + } + } + }} - )); + )); }