supporting shifted ew binary ops (except scalar only)

Author: fschlimb

CMakeLists.txt

@@ -42,7 +42,7 @@ set(MyCppSources ${MyCppSources} ${P2C_HPP})
 
 pybind11_add_module(_ddptensor MODULE ${MyCppSources})
 
-target_compile_definitions(_ddptensor PRIVATE XTENSOR_USE_XSIMD=1 XTENSOR_USE_TBB=1 DDPT_2TYPES=1 USE_MKL=1)
+target_compile_definitions(_ddptensor PRIVATE XTENSOR_USE_XSIMD=1 XTENSOR_USE_TBB=1 USE_MKL=1 DDPT_2TYPES=1)
 target_include_directories(_ddptensor PRIVATE
   ${PROJECT_SOURCE_DIR}/src/include
   ${PROJECT_SOURCE_DIR}/third_party/xtl/include

setup.py

@@ -29,7 +29,7 @@ def build_cmake(self, ext):
         extdir.parent.mkdir(parents=True, exist_ok=True)
 
         # example of cmake args
-        config = 'Debug' if self.debug else 'RelWithDebInfo' #'Release'
+        config = 'Debug'# if self.debug else 'RelWithDebInfo' #'Release'
         cmake_args = [
             '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + str(extdir.parent.absolute()),
             '-DCMAKE_BUILD_TYPE=' + config

src/CollComm.cpp

@@ -0,0 +1,66 @@
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "ddptensor/CollComm.hpp"
+
+// Compute offset and displacements when mapping n_slc to o_slc. This is necessary when
+// slices are not equally partitioned.
+//
+// We assume we split in first dimension.
+// We also assume partitions are assigned to ranks in sequence from 0-N.
+// With this we know that our buffers (old and new) get data in the
+// same order. The only thing which might have changed is the tile-size.
+// Actually, the tile-size might change only if old or new shape does not evenly
+// distribute data (e.g. last partition is smaller).
+// In theory we could re-shape in-place when the norm-tile-size does not change.
+// 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();
+    std::vector<int> counts_send(nr, 0);
+    std::vector<int> disp_send(nr, 0);
+    std::vector<int> counts_recv(nr, 0);
+    std::vector<int> disp_recv(nr, 0);
+        
+    // norm tile-size of  orig array
+    auto o_ntsz =  o_slc.tile_size(0);
+    // 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);
+            
+    // 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);
+
+    for(auto r=0; r<nr; ++r) {
+        // determine what I receive from rank r
+        // e.g. which parts of my new slice overlap with rank r's old slice
+        // Get local slice of rank r of orig array
+        auto o_rslc = o_slc.tile_slice(r);
+        // Flatten to 1d
+        auto o_lrslc = Slice(o_ntsz * r, o_ntsz * r + o_rslc.size());
+        // Determine overlap with local partition of linearized new array
+        auto roverlap = n_llslc.overlap(o_lrslc);
+        // number of elements to be received from rank r
+        counts_recv[r] = roverlap.size();
+        // displacement in new array where elements from rank r get copied to
+        disp_recv[r] = roverlap._start - n_llslc._start;
+
+        // determine what I send to rank r
+        // e.g. which parts of my old slice overlap with rank r's new slice
+        // Get local slice of rank r of new array
+        auto n_rslc = n_slc.tile_slice(r);
+        // Flatten to 1d
+        auto n_lrslc = Slice(n_ntsz * r, n_ntsz * r + n_rslc.size());
+        // Determine overlap with local partition of linearized orig array
+        auto soverlap = o_llslc.overlap(n_lrslc);
+        // number of elements to be send to rank r
+        counts_send[r] = soverlap.size();
+        // displacement in orig array where elements from rank r get copied from
+        disp_send[r] = soverlap._start - o_llslc._start;
+    }
+    return {counts_send, disp_send, counts_recv, disp_recv};
+}

src/Creator.cpp

@@ -50,7 +50,7 @@ namespace x {
         static ptr_type op(uint64_t start, uint64_t end, uint64_t step)
         {
             PVSlice pvslice({static_cast<uint64_t>(Slice(start, end, step).size())});
-            auto lslc = pvslice.slice_of_rank();
+            auto lslc = pvslice.local_slice();
             const auto & l1dslc = lslc.dim(0);
             auto a = xt::arange<T>(start + l1dslc._start*step, start + l1dslc._end * step, l1dslc._step);
             auto r = operatorx<T>::mk_tx(std::move(pvslice), std::move(a));