working dist-coalesce

Author: fschlimb

CMakeLists.txt

@@ -169,6 +169,7 @@ target_link_libraries(_ddpt_rt PRIVATE
     IMEXPTensorTransforms
     IMEXPTensorToLinalg
     IMEXDistDialect
+    IMEXDistTransforms
     IMEXDistToStandard
 	  MLIROptLib
     MLIRExecutionEngine

src/Creator.cpp

@@ -81,7 +81,7 @@ struct DeferredFromShape : public Deferred
 
     DeferredFromShape() = default;
     DeferredFromShape(CreatorId op, const shape_type & shape, DTypeId dtype)
-        : Deferred(dtype, shape.size()),
+        : Deferred(dtype, shape.size(), true),
           _shape(shape), _dtype(dtype), _op(op)
     {}
 
@@ -119,7 +119,7 @@ struct DeferredFull : public Deferred
 
     DeferredFull() = default;
     DeferredFull(const shape_type & shape, PyScalar val, DTypeId dtype)
-        : Deferred(dtype, shape.size()),
+        : Deferred(dtype, shape.size(), true),
           _shape(shape), _val(val), _dtype(dtype)
     {}
 
@@ -158,11 +158,11 @@ struct DeferredFull : public Deferred
 
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::CreateOp>(loc, shp, dtyp, val, nullptr, team),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
+                  [this](Transceiver * transceiver, uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned, uint64_t balanced) {
             assert(rank == this->_shape.size());
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
-                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
+            this->set_value(std::move(mk_tnsr(transceiver, _dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned, balanced)));
         });
         return false;
     }
@@ -193,7 +193,7 @@ struct DeferredArange : public Deferred
 
     DeferredArange() = default;
     DeferredArange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype, uint64_t team = 0)
-        : Deferred(dtype, 1),
+        : Deferred(dtype, 1, true),
           _start(start), _end(end), _step(step), _team(team)
     {}
 
@@ -211,12 +211,12 @@ struct DeferredArange : public Deferred
         auto team = ::imex::createIndex(loc, builder, reinterpret_cast<uint64_t>(getTransceiver()));
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::ARangeOp>(loc, start, stop, step, nullptr, team),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
+                  [this](Transceiver * transceiver, uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned, uint64_t balanced) {
             assert(rank == 1);
             assert(strides[0] == 1);
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
-                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
+            this->set_value(std::move(mk_tnsr(transceiver, _dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned, balanced)));
         });
         return false;
     }

src/DDPTensorImpl.cpp

@@ -10,11 +10,12 @@
 #include <algorithm>
 
 
-DDPTensorImpl::DDPTensorImpl(DTypeId dtype, uint64_t ndims,
-                void * allocated, void * aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned,
-                rank_type owner)
-    : _owner(owner),
+DDPTensorImpl::DDPTensorImpl(Transceiver * transceiver, DTypeId dtype, uint64_t ndims, 
+                  void * allocated, void * aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                  uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned,
+                  uint64_t balanced, rank_type owner)
+    : _transceiver(transceiver),
+      _owner(owner),
       _allocated(allocated),
       _aligned(aligned),
       _gs_allocated(gs_allocated),
@@ -23,6 +24,7 @@ DDPTensorImpl::DDPTensorImpl(DTypeId dtype, uint64_t ndims,
       _lo_aligned(lo_aligned),
       _offset(offset),
       _ndims(ndims),
+      _balanced(balanced),
       _dtype(dtype)
 {
     if(ndims > 0) {
@@ -74,8 +76,8 @@ DDPTensorImpl::ptr_type DDPTensorImpl::clone(bool copy)
     memcpy(lo_aligned, _lo_aligned, nd*sizeof(*lo_aligned));
 
     // strides and sizes are allocated/copied in constructor
-    return std::make_shared<DDPTensorImpl>(dtype(), nd, allocated, aligned, _offset, _sizes, _strides,
-                                           gs_allocated, gs_aligned, lo_allocated, lo_aligned, owner());
+    return std::make_shared<DDPTensorImpl>(transceiver(), dtype(), nd, allocated, aligned, _offset, _sizes, _strides,
+                                           gs_allocated, gs_aligned, lo_allocated, lo_aligned, balanced(), owner());
 }
 
 void DDPTensorImpl::alloc(bool all)
@@ -164,10 +166,12 @@ void DDPTensorImpl::add_to_args(std::vector<void*> & args, int ndims)
     memcpy(buff+3, _sizes, ndims*sizeof(intptr_t));
     memcpy(buff+3+ndims, _strides, ndims*sizeof(intptr_t));
     args.push_back(buff);
-    // second the team
-    args.push_back(reinterpret_cast<void*>(1));
+    // second the transceiver
+    args.push_back(&_transceiver);
+    // balanced third
+    args.push_back(&_balanced);
     if(ndims > 0) {
-        // global shape third
+        // global shape next
         buff = new intptr_t[dtensor_sz(1)];
         buff[0] = reinterpret_cast<intptr_t>(_gs_allocated);
         buff[1] = reinterpret_cast<intptr_t>(_gs_aligned);

src/Deferred.cpp

@@ -29,7 +29,7 @@ void _dist(const Runable * p)
 
 Deferred::future_type Deferred::get_future()
 {
-    return {std::move(promise_type::get_future().share()), _guid, _dtype, _rank};
+    return {std::move(promise_type::get_future().share()), _guid, _dtype, _rank, _balanced};
 }
 
 Deferred::future_type defer_tensor(Runable::ptr_type && _d, bool is_global)

src/EWBinOp.cpp

@@ -444,7 +444,7 @@ struct DeferredEWBinOp : public Deferred
 
     DeferredEWBinOp() = default;
     DeferredEWBinOp(EWBinOpId op, const tensor_i::future_type & a, const tensor_i::future_type & b)
-        : Deferred(a.dtype(), a.rank()),
+        : Deferred(a.dtype(), a.rank(), true),
           _a(a.id()), _b(b.id()), _op(op)
     {}
 
@@ -468,10 +468,10 @@ struct DeferredEWBinOp : public Deferred
 
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::EWBinOp>(loc, aPtTyp, builder.getI32IntegerAttr(ddpt2mlir(_op)), av, bv),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
-                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
+                  [this](Transceiver * transceiver, uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned, uint64_t balanced) {
+            this->set_value(std::move(mk_tnsr(transceiver, _dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned, balanced)));
         });
         return false;
     }

src/MPITransceiver.cpp

@@ -204,18 +204,19 @@ void MPITransceiver::gather(void* buffer,
                             DTypeId datatype,
                             rank_type root)
 {
+    auto dtype = to_mpi(datatype);
     if(root == REPLICATED) {
-        MPI_Allgatherv(MPI_IN_PLACE, 0, to_mpi(datatype),
-                       buffer, counts, displacements, to_mpi(datatype),
+        MPI_Allgatherv(MPI_IN_PLACE, 0, dtype,
+                       buffer, counts, displacements, dtype,
                        _comm);
     } else {
         if(root == _rank) {
-            MPI_Gatherv(MPI_IN_PLACE, 0, to_mpi(datatype),
-                        buffer, counts, displacements, to_mpi(datatype),
+            MPI_Gatherv(MPI_IN_PLACE, 0, dtype,
+                        buffer, counts, displacements, dtype,
                         root, _comm);
         } else {
-            MPI_Gatherv(buffer, counts[_rank], to_mpi(datatype),
-                        nullptr, nullptr, nullptr, to_mpi(datatype),
+            MPI_Gatherv(buffer, counts[_rank], dtype,
+                        nullptr, nullptr, nullptr, dtype,
                         root, _comm);
         }
     }

src/ReduceOp.cpp

@@ -106,7 +106,7 @@ struct DeferredReduceOp : public Deferred
 
     DeferredReduceOp() = default;
     DeferredReduceOp(ReduceOpId op, const tensor_i::future_type & a, const dim_vec_type & dim)
-        : Deferred(a.dtype(), 0), // FIXME rank
+        : Deferred(a.dtype(), 0, true), // FIXME rank
           _a(a.id()), _dim(dim), _op(op)
     {}
 
@@ -132,10 +132,10 @@ struct DeferredReduceOp : public Deferred
         auto op = builder.getIntegerAttr(builder.getIntegerType(sizeof(mop)*8), mop);
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::ReductionOp>(loc, retPtTyp, op, av),
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
-            this->set_value(std::move(mk_tnsr(_dtype, rank, allocated, aligned, offset, sizes, strides,
-                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
+                  [this](Transceiver * transceiver, uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                         uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned, uint64_t balanced) {
+            this->set_value(std::move(mk_tnsr(transceiver, _dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned, balanced)));
         });
         return false;
     }

src/SetGetItem.cpp

@@ -226,8 +226,8 @@ struct DeferredSetItem : public Deferred
         builder.create<::imex::ptensor::InsertSliceOp>(loc, av, bv, offsV, sizesV, stridesV);
         // ... and use av as to later create the ptensor
         dm.addVal(this->guid(), av,
-                  [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                                uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
+                  [this](Transceiver * transceiver, uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                                uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned, uint64_t balanced) {
             this->set_value(Registry::get(this->_a).get());
             // this->set_value(std::move(mk_tnsr(dtype, rank, allocated, aligned, offset, sizes, strides,
             //                                   gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
@@ -261,7 +261,7 @@ struct DeferredGetItem : public Deferred
 
     DeferredGetItem() = default;
     DeferredGetItem(const tensor_i::future_type & a, const std::vector<py::slice> & v)
-        : Deferred(a.dtype(), a.rank()), _a(a.id()), _slc(v)
+        : Deferred(a.dtype(), a.rank(), false), _a(a.id()), _slc(v)
     {}
 
     void run()
@@ -296,10 +296,10 @@ struct DeferredGetItem : public Deferred
                       offsV,
                       sizesV,
                       stridesV),
-                  [this, dtype](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
-                                uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned) {
-            this->set_value(std::move(mk_tnsr(dtype, rank, allocated, aligned, offset, sizes, strides,
-                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned)));
+                  [this, dtype](Transceiver * transceiver, uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
+                                uint64_t * gs_allocated, uint64_t * gs_aligned, uint64_t * lo_allocated, uint64_t * lo_aligned, uint64_t balanced) {
+            this->set_value(std::move(mk_tnsr(transceiver, dtype, rank, allocated, aligned, offset, sizes, strides,
+                                              gs_allocated, gs_aligned, lo_allocated, lo_aligned, balanced)));
         });
         return false;
     }

src/idtr.cpp

@@ -269,4 +269,104 @@ void _idtr_rebalance(uint64_t rank, const int64_t * gShape, const int64_t * lOff
     // Finally communicate elements
     getTransceiver()->alltoall(ptr, sszs.data(), soffs.data(), ddpttype, out, rszs.data(), roffs.data());
 }
+
+
+/// @brief repartition tensor
+/// We assume tensor is partitioned along the first dimension (only) and partitions are ordered by ranks
+/// @param rank 
+/// @param gShapePtr 
+/// @param dtype 
+/// @param lDataPtr 
+/// @param lOffsPtr 
+/// @param lShapePtr 
+/// @param lStridesPtr 
+/// @param offsPtr 
+/// @param szsPtr 
+/// @param outPtr 
+/// @param tc  
+void _idtr_repartition(int64_t rank, int64_t * gShapePtr, int dtype,
+                       void * lDataPtr, int64_t * lOffsPtr, int64_t * lShapePtr, int64_t * lStridesPtr,
+                       int64_t * offsPtr, int64_t * szsPtr, void * outPtr, Transceiver * tc)
+{
+    assert(is_contiguous(lShapePtr, lStridesPtr, rank));
+
+    auto N = tc->nranks();
+    auto me = tc->rank();
+    auto ddpttype = mlir2ddpt(static_cast<::imex::ptensor::DType>(dtype));
+    auto nSz = std::accumulate(&lShapePtr[1], &lShapePtr[rank], 1, std::multiplies<int64_t>());
+
+    // First we allgather the requested target partitioning
+
+    auto myBOff = 2 * rank * me;
+    ::std::vector<int64_t> buff(2*rank*N);
+    for(int64_t i=0; i<rank; ++i) {
+        buff[myBOff+i] = offsPtr[i];
+        buff[myBOff+i+rank] = szsPtr[i];
+    }
+    ::std::vector<int> counts(N, rank*2);
+    ::std::vector<int> dspl(N);
+    for(auto i=0; i<N; ++i) {
+        dspl[i] = 2*rank*i;
+    }
+    tc->gather(buff.data(), counts.data(), dspl.data(), ddpttype, REPLICATED);
+
+    // compute overlap of my local data with each requested part
+
+    auto myOff = lOffsPtr[0];
+    auto mySz = lShapePtr[0];
+    auto myEnd = myOff + mySz;
+
+    std::vector<int> soffs(N);
+    std::vector<int> sszs(N, 0);
+
+    for(auto i=0; i<N; ++i) {
+        auto tOff = buff[2*rank*i];
+        auto tSz = buff[2*rank*i+rank];
+        auto tEnd = tOff + tSz;
+        if(tEnd > myOff && tOff < myEnd) {
+            // We have a target partition which is inside my local data
+            // we now compute what data goes to this target partition
+            auto start = std::max(myOff, tOff);
+            auto end = std::min(myEnd, tEnd);
+            soffs[i] = (int)(start - myOff) * nSz;
+            sszs[i] = (int)(end - start) * nSz;
+        } else {
+            soffs[i] = i ? soffs[i-1] + sszs[i-1] : 0;
+        }
+    }
+    
+    // send our send sizes to others and receive theirs
+    std::vector<int> rszs(N);
+    getTransceiver()->alltoall(sszs.data(), 1, INT32, rszs.data());
+
+    // compute receive-displacements
+    std::vector<int> roffs(N);
+    roffs[0] = 0;
+    for(auto i=1; i<N; ++i) {
+        // compute for all i > 0
+        roffs[i] = roffs[i-1] + rszs[i-1];
+    }
+
+    // Finally communicate elements
+    getTransceiver()->alltoall(lDataPtr, sszs.data(), soffs.data(), ddpttype, outPtr, rszs.data(), roffs.data());
+}
+
+void _idtr_extractslice(int64_t * slcOffs,
+                            int64_t * slcSizes,
+                            int64_t * slcStrides,
+                            int64_t * tOffs,
+                            int64_t * tSizes,
+                            int64_t * lSlcOffsets,
+                            int64_t * lSlcSizes,
+                            int64_t * gSlcOffsets)
+{
+    std::cerr << "slcOffs: " << slcOffs[0] << " " << slcOffs[1] << std::endl;
+    std::cerr << "slcSizes: " << slcSizes[0] << " " << slcSizes[1] << std::endl;
+    std::cerr << "slcStrides: " << slcStrides[0] << " " << slcStrides[1] << std::endl;
+    std::cerr << "tOffs: " << tOffs[0] << " " << tOffs[1] << std::endl;
+    std::cerr << "tSizes: " << tSizes[0] << " " << tSizes[1] << std::endl;
+    std::cerr << "lSlcOffsets: " << lSlcOffsets[0] << " " << lSlcOffsets[1] << std::endl;
+    std::cerr << "lSlcSizes: " << lSlcSizes[0] << " " << lSlcSizes[1] << std::endl;
+    std::cerr << "gSlcOffsets: " << gSlcOffsets[0] << " " << gSlcOffsets[1] << std::endl;
+}
 } // extern "C"

src/include/ddptensor/CppTypes.hpp

@@ -158,4 +158,4 @@ enum FactoryId : int {
 
 // size of memreftype in number of intptr_t's
 inline uint64_t memref_sz(int rank) { return 3 + 2 * rank; }
-inline uint64_t dtensor_sz(int rank) { return 2 * memref_sz(1) + memref_sz(rank) + 1; };
+inline uint64_t dtensor_sz(int rank) { return 2 * memref_sz(1) + memref_sz(rank) + 2; };