enabling tests with 2 mpir ranks (#23)

enabling tests with 2 mpi ranks

Author: fschlimb

.github/workflows/ci.yml

@@ -132,6 +132,7 @@ jobs:
           cd test
           pytest .
           DDPT_FORCE_DIST=1 pytest .
+          mpirun -n 2 pytest .
           cd -
       - name: Run examples
         run: |

ddptensor/__init__.py

@@ -39,6 +39,8 @@
 
 _ddpt_cw = _bool(int(getenv("DDPT_CW", False)))
 
+pi = 3.1415926535897932384626433
+
 
 def init(cw=None):
     cw = _ddpt_cw if cw is None else cw

imex_version.txt

@@ -1 +1 @@
-865cea0c1af5fe7ac1458907637c28f69afd7146
+e9b8ef1bb2126a4aa544ae1b25eae33f7f9a2d80

src/DDPTensorImpl.cpp

@@ -115,6 +115,23 @@ void *DDPTensorImpl::data() {
   return ret;
 }
 
+bool DDPTensorImpl::is_sliced() const {
+  if (ndims() == 0)
+    return false;
+  auto d = ndims() - 1;
+  intptr_t tsz = _strides[d];
+  if (tsz == 1) {
+    for (; d > 0; --d) {
+      tsz *= _sizes[d];
+      if (tsz <= 0)
+        break;
+      if (_strides[d - 1] > tsz)
+        return true;
+    }
+  }
+  return false;
+}
+
 std::string DDPTensorImpl::__repr__() const {
   const auto nd = ndims();
   std::ostringstream oss;

src/IO.cpp

@@ -12,5 +12,5 @@
 
 GetItem::py_future_type IO::to_numpy(const ddptensor &a) {
   assert(!getTransceiver()->is_cw() || getTransceiver()->rank() == 0);
-  return GetItem::gather(a, 0);
+  return GetItem::gather(a, getTransceiver()->is_cw() ? 0 : REPLICATED);
 }

src/idtr.cpp

@@ -349,7 +349,7 @@ void _idtr_reshape(DTypeId ddpttype, int64_t lRank, int64_t *gShapePtr,
   for (auto i = 0; i < N; ++i) {
     dspl[i] = 4 * i;
   }
-  tc->gather(buff.data(), counts.data(), dspl.data(), ddpttype, REPLICATED);
+  tc->gather(buff.data(), counts.data(), dspl.data(), INT64, REPLICATED);
 
   // compute overlaps of current parts with requested parts
   // and store meta for alltoall
@@ -409,7 +409,7 @@ void _idtr_reshape(int64_t gShapeRank, void *gShapeDescr, int64_t lOffsRank,
 
   _idtr_reshape(ddpttype, lRank, MRIdx1d(gShapeRank, gShapeDescr).data(),
                 lData.data(), lData.sizes(), lData.strides(),
-                MRIdx1d(oOffsRank, oOffsDescr).data(), oRank,
+                MRIdx1d(lOffsRank, lOffsDescr).data(), oRank,
                 MRIdx1d(oGShapeRank, oGShapeDescr).data(), oData.data(),
                 oData.sizes(), MRIdx1d(oOffsRank, oOffsDescr).data(), tc);
 }
@@ -465,7 +465,7 @@ void _idtr_repartition(DTypeId ddpttype, int64_t rank, void *lDataPtr,
   for (auto i = 0; i < N; ++i) {
     dspl[i] = 2 * rank * i;
   }
-  tc->gather(buff.data(), counts.data(), dspl.data(), ddpttype, REPLICATED);
+  tc->gather(buff.data(), counts.data(), dspl.data(), INT64, REPLICATED);
 
   // compute overlap of my local data with each requested part
 

src/include/ddptensor/CppTypes.hpp

@@ -152,6 +152,9 @@ inline RedOpType red_op(const char *op) {
   throw std::logic_error("unsupported reduction operation");
 }
 
+/// denotes the full dimension, like when using '[:]' in array subscription
+constexpr auto ALL_SIZE = std::numeric_limits<int64_t>::max();
+
 inline shape_type reduce_shape(const shape_type &shape,
                                const dim_vec_type &dims) {
   auto ssz = shape.size();

src/include/ddptensor/DDPTensorImpl.hpp

@@ -76,10 +76,7 @@ class DDPTensorImpl : public tensor_i {
   void *data();
 
   /// @return true if tensor is a sliced
-  bool is_sliced() const {
-    assert(!"Not implemented");
-    return false;
-  }
+  bool is_sliced() const;
 
   /// python object's __repr__
   virtual std::string __repr__() const override;

src/include/ddptensor/PyTypes.hpp

@@ -126,8 +126,7 @@ template <typename T> T to_native(const py::object &o) { return o.cast<T>(); }
 inline void compute_slice(const py::slice &slc, uint64_t &offset,
                           uint64_t &size, uint64_t &stride) {
   uint64_t dmy = 0;
-  slc.compute(std::numeric_limits<int64_t>::max(), &offset, &dmy, &stride,
-              &size);
+  slc.compute(ALL_SIZE, &offset, &dmy, &stride, &size);
 }
 
 #if 0

test/test_ewb.py

@@ -18,7 +18,7 @@ def test_add1(self):
             a = dt.ones((6, 6), dtype=dtyp)
             b = dt.ones((6, 6), dtype=dtyp)
             c = a + b
-            r1 = dt.sum(c, [0, 1])
+            r1 = dt.sum(c)
             v = 6 * 6 * 2
             assert float(r1) == v