introducing idtr, separating of python-dependent features

Author: fschlimb

CMakeLists.txt

@@ -82,14 +82,45 @@ add_custom_command(
 # ============
 # Target
 # ============
-FILE(GLOB MyCppSources ${PROJECT_SOURCE_DIR}/src/*.cpp ${PROJECT_SOURCE_DIR}/src/include/ddptensor/*.hpp)
-set(MyCppSources ${MyCppSources} ${PROJECT_SOURCE_DIR}/src/jit/mlir.cpp ${P2C_HPP})
 
-pybind11_add_module(_ddptensor MODULE ${MyCppSources})
+FILE(GLOB Hpps ${PROJECT_SOURCE_DIR}/src/include/ddptensor/*.hpp)
+set(Hpps ${Hpps} ${P2C_HPP})
+
+set(DDPTSrcs
+    ${PROJECT_SOURCE_DIR}/src/ddptensor.cpp
+    ${PROJECT_SOURCE_DIR}/src/Creator.cpp
+    ${PROJECT_SOURCE_DIR}/src/EWBinOp.cpp
+    ${PROJECT_SOURCE_DIR}/src/EWUnyOp.cpp
+    ${PROJECT_SOURCE_DIR}/src/IEWBinOp.cpp
+    ${PROJECT_SOURCE_DIR}/src/IO.cpp
+    ${PROJECT_SOURCE_DIR}/src/LinAlgOp.cpp
+    ${PROJECT_SOURCE_DIR}/src/ManipOp.cpp
+    ${PROJECT_SOURCE_DIR}/src/Random.cpp
+    ${PROJECT_SOURCE_DIR}/src/ReduceOp.cpp
+    ${PROJECT_SOURCE_DIR}/src/Service.cpp
+    ${PROJECT_SOURCE_DIR}/src/SetGetItem.cpp
+)
+set(IDTRSrcs
+    ${PROJECT_SOURCE_DIR}/src/idtr.cpp
+    ${PROJECT_SOURCE_DIR}/src/CollComm.cpp
+    ${PROJECT_SOURCE_DIR}/src/Deferred.cpp
+    ${PROJECT_SOURCE_DIR}/src/Factory.cpp
+    ${PROJECT_SOURCE_DIR}/src/Mediator.cpp
+    ${PROJECT_SOURCE_DIR}/src/MPIMediator.cpp
+    ${PROJECT_SOURCE_DIR}/src/MPITransceiver.cpp
+    ${PROJECT_SOURCE_DIR}/src/PVSlice.cpp
+    ${PROJECT_SOURCE_DIR}/src/Registry.cpp
+    ${PROJECT_SOURCE_DIR}/src/Transceiver.cpp
+    ${PROJECT_SOURCE_DIR}/src/jit/mlir.cpp
+)
+
+pybind11_add_module(_ddptensor MODULE ${DDPTSrcs} ${Hpps})
+add_library(idtr SHARED ${IDTRSrcs} ${Hpps})
+set(AllTargets _ddptensor idtr)
 
-target_compile_definitions(_ddptensor PRIVATE USE_MKL=1 DDPT_2TYPES=1)
-target_compile_options(_ddptensor PRIVATE "-ftemplate-backtrace-limit=0")
-target_include_directories(_ddptensor PRIVATE
+add_compile_definitions(USE_MKL=1)
+add_compile_options("-ftemplate-backtrace-limit=0")
+include_directories(
   ${PROJECT_SOURCE_DIR}/src/include
   ${PROJECT_SOURCE_DIR}/third_party/bitsery/include
   ${MPI_INCLUDE_PATH}
@@ -103,20 +134,22 @@ if (CMAKE_SYSTEM_NAME STREQUAL Linux)
   target_link_options(_ddptensor PRIVATE "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/export.txt")
 endif()
 
-#target_compile_options(_ddptensor PRIVATE -fopenmp)
+#compile_options(_ddptensor PRIVATE -fopenmp)
 get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
 get_property(mlir_all_libs GLOBAL PROPERTY MLIR_ALL_LIBS)
 get_property(imex_all_libs GLOBAL PROPERTY IMEX_ALL_LIBS)
 
 #llvm_update_compile_flags(_ddpttensor)
-target_link_directories(_ddptensor PRIVATE
-    ${CONDA_PREFIX}/lib
-    ${IMEX_INSTALL_PREFIX}/lib
-)
+target_link_directories(_ddptensor PRIVATE ${CONDA_PREFIX}/lib)
+target_link_directories(idtr PRIVATE ${CONDA_PREFIX}/lib ${IMEX_INSTALL_PREFIX}/lib)
 
-message(${imex_all_libs})
 target_link_libraries(_ddptensor PRIVATE
+    # ${MKL_LIBRARIES}
+    # tbb
+    idtr
+)
+target_link_libraries(idtr PRIVATE
     ${MPI_C_LIBRARIES}
     # ${MKL_LIBRARIES}
     tbb

export.txt

@@ -0,0 +1,4 @@
+{
+	global: PyInit__ddptensor;
+	local: *;
+};

scripts/code_gen.py

@@ -12,9 +12,11 @@
 // #######################################################
 // SPDX-License-Identifier: BSD-3-Clause
 #pragma once
+#ifdef DEF_PY11_ENUMS
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
 namespace py = pybind11;
+#endif
 """)
 
 prev = 0
@@ -27,14 +29,16 @@
     print(f"    {prev}")
     print("};\n")
 
-print("static void def_enums(py::module_ & m)\n{")
+print("""#ifdef DEF_PY11_ENUMS
+static void def_enums(py::module_ & m)
+{""")
 for cat, lst in api.api_categories.items():
     print(f'    py::enum_<{cat}Id>(m, "{cat}Id")')
     for x in lst:
         print(f'        .value("{x.upper()}", {x.upper()})')
     print("        .export_values();\n")
 
-print("}")
+print("}\n#endif\n")
 
 # Close the file
 sys.stdout.close()

src/CollComm.cpp

@@ -15,7 +15,7 @@
 // This is not implemented: we need an extra mechanism to work with reshape-views or alike.
 std::vector<std::vector<int>> CollComm::map(const PVSlice & n_slc, const PVSlice & o_slc)
 {
-    auto nr = theTransceiver->nranks();
+    auto nr = getTransceiver()->nranks();
     std::vector<int> counts_send(nr, 0);
     std::vector<int> disp_send(nr, 0);
     std::vector<int> counts_recv(nr, 0);
@@ -26,14 +26,14 @@ std::vector<std::vector<int>> CollComm::map(const PVSlice & n_slc, const PVSlice
     // tilesize of my local partition of orig array
     auto o_tsz = o_slc.tile_size();
     // linearized local slice of orig array
-    auto o_llslc = Slice(o_ntsz * theTransceiver->rank(), o_ntsz * theTransceiver->rank() + o_tsz);
+    auto o_llslc = Slice(o_ntsz * getTransceiver()->rank(), o_ntsz * getTransceiver()->rank() + o_tsz);
 
     // norm tile-size of new (reshaped) array
     auto n_ntsz =  n_slc.tile_size(0);
     // tilesize of my local partition of new (reshaped) array
     auto n_tsz = n_slc.tile_size();
     // linearized/flattened/1d local slice of new (reshaped) array
-    auto n_llslc = Slice(n_ntsz * theTransceiver->rank(), n_ntsz * theTransceiver->rank() + n_tsz);
+    auto n_llslc = Slice(n_ntsz * getTransceiver()->rank(), n_ntsz * getTransceiver()->rank() + n_tsz);
 
     for(auto r=0; r<nr; ++r) {
         // determine what I receive from rank r

src/Deferred.cpp

@@ -22,8 +22,8 @@ void push_runable(Runable::ptr_type && r)
 
 void _dist(const Runable * p)
 {
-    if(is_cw() && theTransceiver->rank() == 0)
-        theMediator->to_workers(p);
+    if(getTransceiver()->is_cw() && getTransceiver()->rank() == 0)
+        getMediator()->to_workers(p);
 }
 
 Deferred::future_type Deferred::get_future()

src/EWBinOp.cpp

@@ -166,7 +166,7 @@ namespace x {
 
             // Step 1: Get the mapping of a and b to our resulting slice
 
-            rank_type rank = theTransceiver->rank();
+            rank_type rank = getTransceiver()->rank();
             // Slice for/of the result
             PVSlice r_slc(a_sptr->slice().size() >= b_sptr->slice().size() ? a_sptr->slice().shape() : b_sptr->slice().shape());
             // Size of local result tile

src/IO.cpp

@@ -19,7 +19,7 @@ struct DeferredToNumpy : public DeferredT<promise_type, future_type>
     void run()
     {
         const auto a = std::move(Registry::get(_a).get());
-        set_value(GetItem::do_gather(a, is_cw() ? 0 : REPLICATED));
+        set_value(GetItem::do_gather(a, getTransceiver()->is_cw() ? 0 : REPLICATED));
     }
 
     FactoryId factory() const
@@ -36,7 +36,7 @@ struct DeferredToNumpy : public DeferredT<promise_type, future_type>
 
 py::object IO::to_numpy(const ddptensor & a)
 {
-    assert(!is_cw() || theTransceiver->rank() == 0);
+    assert(!getTransceiver()->is_cw() || getTransceiver()->rank() == 0);
     auto f = defer<DeferredToNumpy>(a.get());
     auto x = f.get();
     return x;

src/LinAlgOp.cpp

@@ -48,7 +48,7 @@ namespace x {
         static ptr_type vecdot_1d(const T1 & a, const T2 & b, int axis)
         {
             auto d = xt::linalg::dot(a, b)();
-            theTransceiver->reduce_all(&d, DTYPE<decltype(d)>::value, 1, SUM);
+            getTransceiver()->reduce_all(&d, DTYPE<decltype(d)>::value, 1, SUM);
             return operatorx<decltype(d)>::mk_tx(d, REPLICATED);
         }
 
@@ -58,8 +58,8 @@ namespace x {
             if(a_ptr->slice().split_dim() != 0)
                 throw(std::runtime_error("vecdoc_2d supported for split_dim=0 only"));
 
-            auto nr = theTransceiver->nranks();
-            auto me = theTransceiver->rank();
+            auto nr = getTransceiver()->nranks();
+            auto me = getTransceiver()->rank();
             rank_type right = (me + 1) % nr;
             rank_type left = (nr + me - 1) % nr;
             auto tsz = b_ptr->slice().tile_size(0);
@@ -97,7 +97,7 @@ namespace x {
                 if(i > 1) {
                     // data exchange
                     // FIXME: optimize data transfer: last partition might contain unused data
-                    theTransceiver->send_recv(buff.data(),
+                    getTransceiver()->send_recv(buff.data(),
                                               tsz,
                                               DTYPE<A>::value,
                                               left,

src/MPIMediator.cpp

@@ -6,7 +6,7 @@
 #include <unordered_map>
 #include <mutex>
 
-#include "ddptensor/UtilsAndTypes.hpp"
+#include "ddptensor/CppTypes.hpp"
 #include "ddptensor/MPIMediator.hpp"
 #include "ddptensor/MPITransceiver.hpp"
 #include "ddptensor/NDSlice.hpp"
@@ -24,7 +24,7 @@ void send_to_workers(const Runable * dfrd, bool self, MPI_Comm comm);
 MPIMediator::MPIMediator()
     : _listener(nullptr)
 {
-    auto c = dynamic_cast<MPITransceiver*>(theTransceiver);
+    auto c = dynamic_cast<MPITransceiver*>(getTransceiver());
     if(c == nullptr) throw std::runtime_error("Expected Transceiver to be MPITransceiver.");
     _comm = c->comm();
     int sz;
@@ -40,9 +40,9 @@ MPIMediator::~MPIMediator()
     MPI_Comm_rank(_comm, &rank);
     MPI_Comm_size(_comm, &sz);
 
-    if(is_cw() && rank == 0) to_workers(nullptr);
+    if(getTransceiver()->is_cw() && rank == 0) to_workers(nullptr);
     MPI_Barrier(_comm);
-    if(!is_cw() || rank == 0) send_to_workers(nullptr, true, _comm);
+    if(!getTransceiver()->is_cw() || rank == 0) send_to_workers(nullptr, true, _comm);
     if(_listener) {
         _listener->join();
         delete _listener;

src/MPITransceiver.cpp

@@ -3,10 +3,11 @@
 #include <mpi.h>
 #include <limits>
 #include <sstream>
+#include <iostream>
 #include "ddptensor/MPITransceiver.hpp"
 
-MPITransceiver::MPITransceiver()
-    : _nranks(1), _rank(0), _comm(MPI_COMM_WORLD)
+MPITransceiver::MPITransceiver(bool is_cw)
+    : _nranks(1), _rank(0), _comm(MPI_COMM_WORLD), _is_cw(is_cw)
 {
     int flag;
     MPI_Initialized(&flag);