adding getitem and simple single-node sort

fschlimb · fschlimb · commit 5dbf1d203424 · 2022-11-23T02:54:44.000-06:00
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -111,7 +111,6 @@ set(IDTRSrcs
     ${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
@@ -173,6 +172,7 @@ target_link_libraries(idtr PRIVATE
     MLIRLinalgDialect
     MLIRLinalgToLLVM
     MLIRLinalgTransforms
+    MLIRLLVMDialect
     MLIRMathDialect
     MLIRMathToLLVM
     MLIRMathTransforms
diff --git a/ddptensor/__init__.py b/ddptensor/__init__.py
@@ -72,7 +72,7 @@ def to_numpy(a):
         )
     elif func == "arange":
         exec(
-            f"{func} = lambda start, end, step, dtype: dtensor(_cdt.Creator.arange(start, end, step, dtype))"
+            f"{func} = lambda start, end, step, dtype, team=0: dtensor(_cdt.Creator.arange(start, end, step, dtype, team))"
         )
 
 for func in api.api_categories["ReduceOp"]:
diff --git a/ddptensor/ddptensor.py b/ddptensor/ddptensor.py
@@ -50,8 +50,8 @@ def _inplace(self, t):
             f"{att} = property(lambda self: self._t.{att})"
         )
 
-    def __getitem__(self, *args):
-        return dtensor(self._t.__getitem__(*args))
+    def __getitem__(self, key):
+        return dtensor(self._t.__getitem__(key if isinstance(key, list) else [key,]))
 
     def __setitem__(self, key, value):
          self._t = self._t.__setitem__(key, value._t) # if isinstance(value, dtensor) else value)
diff --git a/src/CollComm.cpp b/src/CollComm.cpp
@@ -15,6 +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)
 {
+#if 0
     auto nr = getTransceiver()->nranks();
     std::vector<int> counts_send(nr, 0);
     std::vector<int> disp_send(nr, 0);
@@ -63,4 +64,6 @@ std::vector<std::vector<int>> CollComm::map(const PVSlice & n_slc, const PVSlice
         disp_send[r] = soverlap._start - o_llslc._start;
     }
     return {counts_send, disp_send, counts_recv, disp_recv};
+#endif // if 0
+    return {};
 }
diff --git a/src/Creator.cpp b/src/Creator.cpp
@@ -152,12 +152,12 @@ ddptensor * Creator::full(const shape_type & shape, const py::object & val, DTyp
 
 struct DeferredArange : public Deferred
 {
-    uint64_t _start, _end, _step;
+    uint64_t _start, _end, _step, _team;
 
     DeferredArange() = default;
-    DeferredArange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype)
+    DeferredArange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype, uint64_t team = 0)
         : Deferred(dtype, 1),
-          _start(start), _end(end), _step(step)
+          _start(start), _end(end), _step(step), _team(team)
     {}
 
     void run() override
@@ -171,9 +171,10 @@ struct DeferredArange : public Deferred
         auto stop = ::imex::createInt(loc, builder, _end);
         auto step = ::imex::createInt(loc, builder, _step);
         auto dtype = builder.getI64Type(); // FIXME
-        auto artype = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get({-1}, dtype), false);
+        auto artype = ::imex::ptensor::PTensorType::get(builder.getContext(), 1, dtype, false);
         auto dmy = ::imex::createInt<1>(loc, builder, 0);
-        auto team = ::imex::createInt(loc, builder, reinterpret_cast<uint64_t>(getTransceiver()));
+        // ::mlir::Value 
+        auto team = ::imex::createIndex(loc, builder, reinterpret_cast<uint64_t>(getTransceiver()));
         dm.addVal(this->guid(),
                   builder.create<::imex::ptensor::ARangeOp>(loc, artype, start, stop, step, dmy, team),
                   [this](uint64_t rank, void *allocated, void *aligned, intptr_t offset, const intptr_t * sizes, const intptr_t * strides,
@@ -201,9 +202,9 @@ struct DeferredArange : public Deferred
     }
 };
 
-ddptensor * Creator::arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype)
+ddptensor * Creator::arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype, uint64_t team)
 {
-    return new ddptensor(defer<DeferredArange>(start, end, step, dtype));
+    return new ddptensor(defer<DeferredArange>(start, end, step, dtype, team));
 }
 
 ddptensor * Creator::mk_future(const py::object & b)
diff --git a/src/EWBinOp.cpp b/src/EWBinOp.cpp
@@ -8,7 +8,6 @@
 #include "ddptensor/Creator.hpp"
 #include "ddptensor/TypePromotion.hpp"
 #include "ddptensor/CollComm.hpp"
-#include "ddptensor/Chunker.hpp"
 #include "ddptensor/DDPTensorImpl.hpp"
 
 #include <imex/Dialect/PTensor/IR/PTensorOps.h>
diff --git a/src/MPIMediator.cpp b/src/MPIMediator.cpp
@@ -9,7 +9,6 @@
 #include "ddptensor/CppTypes.hpp"
 #include "ddptensor/MPIMediator.hpp"
 #include "ddptensor/MPITransceiver.hpp"
-#include "ddptensor/NDSlice.hpp"
 #include "ddptensor/Factory.hpp"
 
 constexpr static int REQ_TAG = 14711;
@@ -50,6 +49,7 @@ MPIMediator::~MPIMediator()
     }
 }
 
+#if 0
 void MPIMediator::pull(rank_type from, id_type guid, const NDSlice & slice, void * rbuff)
 {
     MPI_Request request[2];
@@ -81,6 +81,7 @@ void MPIMediator::pull(rank_type from, id_type guid, const NDSlice & slice, void
     MPI_Get_count(&status[1], MPI_CHAR, &cnt);
     if(cnt != sz) throw(std::runtime_error("Received unexpected message size."));
 }
+#endif
 
 void send_to_workers(const Runable * dfrd, bool self, MPI_Comm comm)
 {
@@ -162,6 +163,7 @@ void MPIMediator::listen()
             uptr.get()->defer(std::move(uptr)); // grmpf
             break;
         }
+#if 0
         case PULL_TAG: {
             uint64_t id;
             ser.value8b(id);
@@ -182,6 +184,7 @@ void MPIMediator::listen()
             MPI_Isend(rbuff.data(), rbuff.size(), MPI_CHAR, requester, PUSH_TAG, _comm, &request_out);
             break;
         }
+#endif
         case EXIT_TAG:
             defer(nullptr);
             return;
diff --git a/src/ReduceOp.cpp b/src/ReduceOp.cpp
@@ -124,9 +124,9 @@ struct DeferredReduceOp : public Deferred
         auto av = dm.getDependent(builder, _a);
         auto aPtTyp = ::imex::dist::getPTensorType(av);
         assert(aPtTyp);
-        ::mlir::Type dtype = aPtTyp.getRtensor().getElementType();
+        ::mlir::Type dtype = aPtTyp.getElementType();
         // return type 0d with same dtype as input
-        auto retPtTyp = ::imex::ptensor::PTensorType::get(builder.getContext(), ::mlir::RankedTensorType::get({}, dtype), false);
+        auto retPtTyp = ::imex::ptensor::PTensorType::get(builder.getContext(), 0, dtype, false);
         // reduction op
         auto mop = ddpt2mlir(_op);
         auto op = builder.getIntegerAttr(builder.getIntegerType(sizeof(mop)*8), mop);
diff --git a/src/SetGetItem.cpp b/src/SetGetItem.cpp
@@ -5,6 +5,11 @@
 #include "ddptensor/Mediator.hpp"
 #include "ddptensor/Factory.hpp"
 
+#include <imex/Dialect/PTensor/IR/PTensorOps.h>
+#include <imex/Dialect/Dist/IR/DistOps.h>
+#include <imex/internal/PassUtils.h>
+#include <mlir/IR/Builders.h>
+
 #if 0
 namespace x {
 
@@ -223,14 +228,48 @@ struct DeferredGetItem : public Deferred
 
     DeferredGetItem() = default;
     DeferredGetItem(const tensor_i::future_type & a, const std::vector<py::slice> & v)
-        : _a(a.id()), _slc(v)
+        : Deferred(a.dtype(), a.rank()), _a(a.id()), _slc(v)
     {}
 
     void run()
     {
         //const auto a = std::move(Registry::get(_a).get());
         //set_value(std::move(TypeDispatch<x::GetItem>(a, _slc)));
     }
+    
+    bool generate_mlir(::mlir::OpBuilder & builder, ::mlir::Location loc, jit::DepManager & dm) override
+    {
+        // get params and extract offsets/sizes/strides
+        const auto dtype = this->dtype();
+        auto av = dm.getDependent(builder, _a);
+        auto & offs = _slc.offsets();
+        auto & sizes = _slc.sizes();
+        auto & strides = _slc.strides();
+        auto nd = offs.size();
+        // convert C++ slices into vectors of MLIR Values
+        std::vector<::mlir::Value> offsV(nd);
+        std::vector<::mlir::Value> sizesV(nd);
+        std::vector<::mlir::Value> stridesV(nd);
+        for(auto i = 0; i<nd; ++i) {
+            offsV[i] = ::imex::createIndex(loc, builder, offs[i]);
+            sizesV[i] = ::imex::createIndex(loc, builder, sizes[i]);
+            stridesV[i] = ::imex::createIndex(loc, builder, strides[i]);
+        }
+        // now we can create the PTensor op using the above Values
+        dm.addVal(this->guid(),
+                  builder.create<::imex::ptensor::ExtractSliceOp>(loc,
+                      ::imex::dist::getPTensorType(av),
+                      av,
+                      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)));
+        });
+        return false;
+    }
 
     FactoryId factory() const
     {
diff --git a/src/include/ddptensor/Creator.hpp b/src/include/ddptensor/Creator.hpp
@@ -10,6 +10,6 @@ struct Creator
 {
     static ddptensor * create_from_shape(CreatorId op, const shape_type & shape, DTypeId dtype=FLOAT64);
     static ddptensor * full(const shape_type & shape, const py::object & val, DTypeId dtype=FLOAT64);
-    static ddptensor * arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype=INT64);
+    static ddptensor * arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype=INT64, uint64_t team=0);
     static ddptensor * mk_future(const py::object & b);
 };
diff --git a/src/include/ddptensor/DDPTensorImpl.hpp b/src/include/ddptensor/DDPTensorImpl.hpp
@@ -48,12 +48,14 @@ class DDPTensorImpl : public tensor_i
         : _owner(owner),
           _ndims(N)
     {
+        assert(_ndims <= 1);
     }
 
     // incomplete, useful for computing meta information
     DDPTensorImpl()
         : _owner(REPLICATED)
     {
+        assert(_ndims <= 1);
     }
 
     DDPTensorImpl::ptr_type clone(bool copy = true);
diff --git a/src/include/ddptensor/MPIMediator.hpp b/src/include/ddptensor/MPIMediator.hpp
@@ -14,7 +14,7 @@ class MPIMediator : public Mediator
 public:
     MPIMediator();
     virtual ~MPIMediator();
-    virtual void pull(rank_type from, id_type guid, const NDSlice & slice, void * buffer);
+    // virtual void pull(rank_type from, id_type guid, const NDSlice & slice, void * buffer);
     virtual void to_workers(const Runable * dfrd);
 
 protected:
diff --git a/src/include/ddptensor/Mediator.hpp b/src/include/ddptensor/Mediator.hpp
@@ -13,7 +13,7 @@ class Mediator
 public:
     enum : uint64_t {LOCAL_ONLY = 0};
     virtual ~Mediator() {}
-    virtual void pull(rank_type from, id_type guid, const NDSlice & slice, void * buffer) = 0;
+    // virtual void pull(rank_type from, id_type guid, const NDSlice & slice, void * buffer) = 0;
     virtual void to_workers(const Runable * dfrd) = 0;
 };
 
diff --git a/src/include/ddptensor/NDSlice.hpp b/src/include/ddptensor/NDSlice.hpp
@@ -7,29 +7,53 @@
 #include <bitsery/bitsery.h>
 #include <bitsery/traits/vector.h>
 
-#include "Slice.hpp"
-#include "NDIndex.hpp"
-#include "CppTypes.hpp"
-
 ///
 /// A slice of n-dimensional range with utility features to extract nd-indices.
 /// Represented as a vector of triplets [start, end, step[.
 ///
 class NDSlice {
-    // type of elements, n-dimensional
-    typedef std::vector<Slice::value_type> value_type;
-    // the nd-slice, e.g. once slice per dimension
-    typedef std::vector<Slice> ndslice_t;
-
-
-    // vector with one Slice per dimension, e.g. our nd-slice
-    ndslice_t _slice_vec;
-    // helper, holding sizes of each dimensions's slice
-    mutable std::vector<uint64_t> _sizes;
+public:
+    using vec_t = std::vector<uint64_t>;
+private:
+    // vector with offsets per dimension
+    vec_t _offsets;
+    // vector with sizes per dimension
+    vec_t _sizes;
+    // vector with strides per dimension
+    vec_t _strides;
 
 public:
     NDSlice() = default;
     NDSlice(NDSlice &&) = default;
+    NDSlice(const NDSlice &) = default;
+    // assumes a function exists to extract values per slice (compute_slice(T, ...))
+    template<typename T>
+    NDSlice(const std::vector<T> & v)
+        : _offsets(v.size()), _sizes(v.size()), _strides(v.size())
+    {
+        auto nd = v.size();
+        auto i = 0;
+        for(auto s : v) {
+            compute_slice(s, _offsets[i], _sizes[i], _strides[i]);
+            ++i;
+        };
+    }
+
+    const vec_t & offsets() const { return _offsets; };
+    const vec_t & sizes()   const { return _sizes; };
+    const vec_t & strides() const { return _strides; };
+
+    template<typename S>
+    void serialize(S & ser)
+    {
+        ser.container8b(_offsets, 8);
+        ser.container8b(_sizes, 8);
+        ser.container8b(_strides, 8);
+    }
+};
+
+
+#if 0
     NDSlice(const Slice & slc) : _slice_vec(1, slc), _sizes() {}
     NDSlice(const NDSlice & o) : _slice_vec(o._slice_vec), _sizes() {}
     NDSlice(const std::initializer_list<Slice> & slc) : _slice_vec(slc), _sizes() {}
@@ -402,7 +426,7 @@ class NDSlice {
         return iterator();
     }
 };
-
+#endif // if 0
 
 #if 0
     // ###########################################################################
diff --git a/src/include/ddptensor/PyTypes.hpp b/src/include/ddptensor/PyTypes.hpp
@@ -109,6 +109,12 @@ 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);
+}
+
 #if 0
 inline py::tuple _make_tuple(const NDSlice & v)
 {
diff --git a/src/include/ddptensor/Sorting.hpp b/src/include/ddptensor/Sorting.hpp
@@ -0,0 +1,12 @@
+// SPDX-License-Identifier: BSD-3-Clause
+
+#pragma once
+
+#include "UtilsAndTypes.hpp"
+#include "ddptensor.hpp"
+#include "p2c_ids.hpp"
+
+struct SortOp
+{
+    static ddptensor * sort(const ddptensor & a, bool descending = false);
+};
diff --git a/src/include/ddptensor/jit/mlir.hpp b/src/include/ddptensor/jit/mlir.hpp
@@ -32,18 +32,6 @@ using SetResFunc = std::function<void(
 // initialize jit
 void init();
 
-template<int W=64, typename T=int64_t>
-::mlir::Value createInt(const ::mlir::Location & loc, ::mlir::OpBuilder & builder, T val)
-{
-    auto attr = builder.getIntegerAttr(builder.getIntegerType(W), val);
-    return builder.create<::mlir::arith::ConstantOp>(loc, attr);
-}
-inline ::mlir::Value createIndex(const ::mlir::Location & loc, ::mlir::OpBuilder & builder, uint64_t val)
-{
-    auto attr = builder.getIndexAttr(val);
-    return builder.create<::mlir::arith::ConstantOp>(loc, attr);
-}
-
 /// Manages iput/output (tensor) dependences
 class DepManager
 {
diff --git a/src/jit/mlir.cpp b/src/jit/mlir.cpp

Original file line number	Diff line number	Diff line change
`@@ -72,7 +72,7 @@ def to_numpy(a):`
`72`	`72`	`)`
`73`	`73`	`elif func == "arange":`
`74`	`74`	`exec(`
`75`		`- f"{func} = lambda start, end, step, dtype: dtensor(_cdt.Creator.arange(start, end, step, dtype))"`
	`75`	`+ f"{func} = lambda start, end, step, dtype, team=0: dtensor(_cdt.Creator.arange(start, end, step, dtype, team))"`
`76`	`76`	`)`
`77`	`77`
`78`	`78`	`for func in api.api_categories["ReduceOp"]:`
Original file line number	Diff line number	Diff line change
`@@ -50,8 +50,8 @@ def _inplace(self, t):`
`50`	`50`	`f"{att} = property(lambda self: self._t.{att})"`
`51`	`51`	`)`
`52`	`52`
`53`		`- def __getitem__(self, *args):`
`54`		`- return dtensor(self._t.__getitem__(*args))`
	`53`	`+ def __getitem__(self, key):`
	`54`	`+ return dtensor(self._t.__getitem__(key if isinstance(key, list) else [key,]))`
`55`	`55`
`56`	`56`	`def __setitem__(self, key, value):`
`57`	`57`	`self._t = self._t.__setitem__(key, value._t) # if isinstance(value, dtensor) else value)`
Original file line number	Diff line number	Diff line change
`@@ -15,6 +15,7 @@`
`15`	`15`	`// This is not implemented: we need an extra mechanism to work with reshape-views or alike.`
`16`	`16`	`std::vector<std::vector<int>> CollComm::map(const PVSlice & n_slc, const PVSlice & o_slc)`
`17`	`17`	`{`
	`18`	`+#if 0`
`18`	`19`	`auto nr = getTransceiver()->nranks();`
`19`	`20`	`std::vector<int> counts_send(nr, 0);`
`20`	`21`	`std::vector<int> disp_send(nr, 0);`
`@@ -63,4 +64,6 @@ std::vector<std::vector<int>> CollComm::map(const PVSlice & n_slc, const PVSlice`
`63`	`64`	`disp_send[r] = soverlap._start - o_llslc._start;`
`64`	`65`	`}`
`65`	`66`	`return {counts_send, disp_send, counts_recv, disp_recv};`
	`67`	`+#endif // if 0`
	`68`	`+ return {};`
`66`	`69`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,7 +9,6 @@`
`9`	`9`	`#include "ddptensor/CppTypes.hpp"`
`10`	`10`	`#include "ddptensor/MPIMediator.hpp"`
`11`	`11`	`#include "ddptensor/MPITransceiver.hpp"`
`12`		`-#include "ddptensor/NDSlice.hpp"`
`13`	`12`	`#include "ddptensor/Factory.hpp"`
`14`	`13`
`15`	`14`	`constexpr static int REQ_TAG = 14711;`
`@@ -50,6 +49,7 @@ MPIMediator::~MPIMediator()`
`50`	`49`	`}`
`51`	`50`	`}`
`52`	`51`
	`52`	`+#if 0`
`53`	`53`	`void MPIMediator::pull(rank_type from, id_type guid, const NDSlice & slice, void * rbuff)`
`54`	`54`	`{`
`55`	`55`	`MPI_Request request[2];`
`@@ -81,6 +81,7 @@ void MPIMediator::pull(rank_type from, id_type guid, const NDSlice & slice, void`
`81`	`81`	`MPI_Get_count(&status[1], MPI_CHAR, &cnt);`
`82`	`82`	`if(cnt != sz) throw(std::runtime_error("Received unexpected message size."));`
`83`	`83`	`}`
	`84`	`+#endif`
`84`	`85`
`85`	`86`	`void send_to_workers(const Runable * dfrd, bool self, MPI_Comm comm)`
`86`	`87`	`{`
`@@ -162,6 +163,7 @@ void MPIMediator::listen()`
`162`	`163`	`uptr.get()->defer(std::move(uptr)); // grmpf`
`163`	`164`	`break;`
`164`	`165`	`}`
	`166`	`+#if 0`
`165`	`167`	`case PULL_TAG: {`
`166`	`168`	`uint64_t id;`
`167`	`169`	`ser.value8b(id);`
`@@ -182,6 +184,7 @@ void MPIMediator::listen()`
`182`	`184`	`MPI_Isend(rbuff.data(), rbuff.size(), MPI_CHAR, requester, PUSH_TAG, _comm, &request_out);`
`183`	`185`	`break;`
`184`	`186`	`}`
	`187`	`+#endif`
`185`	`188`	`case EXIT_TAG:`
`186`	`189`	`defer(nullptr);`
`187`	`190`	`return;`
Original file line number	Diff line number	Diff line change
`@@ -10,6 +10,6 @@ struct Creator`
`10`	`10`	`{`
`11`	`11`	`static ddptensor * create_from_shape(CreatorId op, const shape_type & shape, DTypeId dtype=FLOAT64);`
`12`	`12`	`static ddptensor * full(const shape_type & shape, const py::object & val, DTypeId dtype=FLOAT64);`
`13`		`- static ddptensor * arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype=INT64);`
	`13`	`+ static ddptensor * arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype=INT64, uint64_t team=0);`
`14`	`14`	`static ddptensor * mk_future(const py::object & b);`
`15`	`15`	`};`