CW initial commit of controller-worker (CW) mode

Author: fschlimb

ddptensor/__init__.py

@@ -25,6 +25,8 @@
     UINT32 as uint32,
     UINT16 as uint16,
     UINT8 as uint8,
+    BOOL as bool,
+    init,
     fini,
     sync
 )
@@ -55,7 +57,7 @@
         )
     elif func == "full":
         exec(
-            f"{func} = lambda shape, val, dtype: dtensor(_cdt.Creator.full(_cdt.{FUNC}, shape, val, dtype))"
+            f"{func} = lambda shape, val, dtype: dtensor(_cdt.Creator.full(shape, val, dtype))"
         )
     elif func == "arange":
         exec(

ddptensor/ddptensor.py

@@ -54,4 +54,4 @@ def __getitem__(self, *args):
         return dtensor(self._t.__getitem__(*args))
 
     def __setitem__(self, key, value):
-        x = self._t.__setitem__(key, value._t) # if isinstance(value, dtensor) else value)
+         self._t = self._t.__setitem__(key, value._t) # if isinstance(value, dtensor) else value)

scripts/code_gen.py

@@ -17,7 +17,7 @@
 namespace py = pybind11;
 """)
 
-prev = "0"
+prev = 0
 for cat, lst in api.api_categories.items():
     print(f"enum {cat}Id : int {{")
     for x in lst:
@@ -28,7 +28,6 @@
     print("};\n")
 
 print("static void def_enums(py::module_ & m)\n{")
-
 for cat, lst in api.api_categories.items():
     print(f'    py::enum_<{cat}Id>(m, "{cat}Id")')
     for x in lst:

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 'Release'  # 'RelWithDebInfo' #'Release'
+        config = 'Debug' # if self.debug else 'Release'  # 'RelWithDebInfo' #'Release'
         cmake_args = [
             '-DCMAKE_LIBRARY_OUTPUT_DIRECTORY=' + str(extdir.parent.absolute()),
             '-DCMAKE_BUILD_TYPE=' + config

src/Creator.cpp

@@ -2,6 +2,7 @@
 #include "ddptensor/TypeDispatch.hpp"
 #include "ddptensor/x.hpp"
 #include "ddptensor/Deferred.hpp"
+#include "ddptensor/Factory.hpp"
 
 namespace x {
 
@@ -15,7 +16,7 @@ namespace x {
         static ptr_type op(CreatorId c, const shape_type & shp)
         {
             PVSlice pvslice(shp);
-            shape_type shape(std::move(pvslice.shape_of_rank()));
+            shape_type shape(std::move(pvslice.tile_shape()));
             switch(c) {
             case EMPTY:
                 return operatorx<T>::mk_tx(std::move(pvslice), std::move(xt::empty<T>(std::move(shape))));
@@ -28,14 +29,19 @@ namespace x {
             };
         };
 
-        template<typename V>
-        static ptr_type op(CreatorId c, const shape_type & shp, const V & v)
+        static ptr_type op(CreatorId c, const shape_type & shp, PyScalar v)
         {
+            T val;
+            if constexpr (std::is_integral<T>::value) val = static_cast<T>(v._int);
+            else if constexpr (std::is_floating_point<T>::value) val = static_cast<T>(v._float);
             if(c == FULL) {
+                if(VPROD(shp) <= 1) {
+                    return operatorx<T>::mk_tx(val, REPLICATED);
+                }
                 PVSlice pvslice(shp);
-                shape_type shape(std::move(pvslice.shape_of_rank()));
+                shape_type shape(std::move(pvslice.tile_shape()));
                 auto a = xt::empty<T>(std::move(shape));
-                a.fill(to_native<T>(v));
+                a.fill(val);
                 return operatorx<T>::mk_tx(std::move(pvslice), std::move(a));
             }
             throw std::runtime_error("Unknown creator");
@@ -47,24 +53,39 @@ namespace x {
             auto lslc = pvslice.slice_of_rank();
             const auto & l1dslc = lslc.dim(0);
             auto a = xt::arange<T>(start + l1dslc._start*step, start + l1dslc._end * step, l1dslc._step);
-            return operatorx<T>::mk_tx(std::move(pvslice), std::move(a));
+            auto r = operatorx<T>::mk_tx(std::move(pvslice), std::move(a));
+            return r;
         }
     }; // class creatorx
 } // namespace x
 
 struct DeferredFromShape : public Deferred
 {
-    CreatorId _op;
     shape_type _shape;
     DTypeId _dtype;
+    CreatorId _op;
 
+    DeferredFromShape() = default;
     DeferredFromShape(CreatorId op, const shape_type & shape, DTypeId dtype)
-        : _op(op), _shape(shape), _dtype(dtype)
+        : _shape(shape), _dtype(dtype), _op(op)
     {}
 
     void run()
     {
-        set_value(TypeDispatch<x::Creator>(_dtype, _op, _shape));
+        set_value(std::move(TypeDispatch<x::Creator>(_dtype, _op, _shape)));
+    }
+
+    FactoryId factory() const
+    {
+        return F_FROMSHAPE;
+    }
+
+    template<typename S>
+    void serialize(S & ser)
+    {
+        ser.template container<sizeof(shape_type::value_type)>(_shape, 8);
+        ser.template value<sizeof(_dtype)>(_dtype);
+        ser.template value<sizeof(_op)>(_op);
     }
 };
 
@@ -76,30 +97,46 @@ tensor_i::future_type Creator::create_from_shape(CreatorId op, const shape_type
 struct DeferredFull : public Deferred
 {
     shape_type _shape;
-    const py::object & _val;
+    PyScalar _val;
     DTypeId _dtype;
 
-    DeferredFull(const shape_type & shape, const py::object & val, DTypeId dtype)
+    DeferredFull() = default;
+    DeferredFull(const shape_type & shape, PyScalar val, DTypeId dtype)
         : _shape(shape), _val(val), _dtype(dtype)
     {}
 
     void run()
     {
         auto op = FULL;
-        set_value(TypeDispatch<x::Creator>(_dtype, op, _shape, _val));
+        set_value(std::move(TypeDispatch<x::Creator>(_dtype, op, _shape, _val)));
+    }
+
+    FactoryId factory() const
+    {
+        return F_FULL;
+    }
+
+    template<typename S>
+    void serialize(S & ser)
+    {
+        ser.template container<sizeof(shape_type::value_type)>(_shape, 8);
+        ser.template value<sizeof(_val)>(_val._int);
+        ser.template value<sizeof(_dtype)>(_dtype);
     }
 };
 
 tensor_i::future_type Creator::full(const shape_type & shape, const py::object & val, DTypeId dtype)
 {
-    return defer<DeferredFull>(shape, val, dtype);
+    auto v = mk_scalar(val, dtype);
+    return defer<DeferredFull>(shape, v, dtype);
 }
 
 struct DeferredArange : public Deferred
 {
     uint64_t _start, _end, _step;
     DTypeId _dtype;
 
+    DeferredArange() = default;
     DeferredArange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype)
         : _start(start), _end(end), _step(step), _dtype(dtype)
     {}
@@ -108,9 +145,39 @@ struct DeferredArange : public Deferred
     {
         set_value(std::move(TypeDispatch<x::Creator>(_dtype, _start, _end, _step)));
     };
+
+    FactoryId factory() const
+    {
+        return F_ARANGE;
+    }
+
+    template<typename S>
+    void serialize(S & ser)
+    {
+        ser.template value<sizeof(_start)>(_start);
+        ser.template value<sizeof(_end)>(_end);
+        ser.template value<sizeof(_step)>(_step);        
+        ser.template value<sizeof(_dtype)>(_dtype);
+    }
 };
 
 tensor_i::future_type Creator::arange(uint64_t start, uint64_t end, uint64_t step, DTypeId dtype)
 {
     return defer<DeferredArange>(start, end, step, dtype);
 }
+
+tensor_i::future_type Creator::mk_future(const py::object & b)
+{
+    if(py::isinstance<tensor_i::future_type>(b)) {
+        return b.cast<tensor_i::future_type>();
+    } else if(py::isinstance<py::float_>(b)) {
+        return Creator::full({1}, b, FLOAT64);
+    } else if(py::isinstance<py::int_>(b)) {
+        return Creator::full({1}, b, INT64);
+    }
+    throw std::runtime_error("Invalid right operand to elementwise binary operation");
+};
+
+FACTORY_INIT(DeferredFromShape, F_FROMSHAPE);
+FACTORY_INIT(DeferredFull, F_FULL);
+FACTORY_INIT(DeferredArange, F_ARANGE);

src/Deferred.cpp

@@ -1,16 +1,33 @@
-#include "include/ddptensor/Deferred.hpp"
 #include <oneapi/tbb/concurrent_queue.h>
+#include "include/ddptensor/Deferred.hpp"
+#include "include/ddptensor/Transceiver.hpp"
+#include "include/ddptensor/Mediator.hpp"
+#include "include/ddptensor/Registry.hpp"
 
 static tbb::concurrent_bounded_queue<Deferred::ptr_type> _deferred;
 
-Deferred::future_type Deferred::defer(Deferred::ptr_type && d)
+Deferred::future_type Deferred::get_future()
+{
+    return {std::move(tensor_i::promise_type::get_future()), _guid};
+}
+
+void Deferred::set_value(tensor_i::ptr_type && v)
 {
+    if(_guid != Registry::NOGUID) {
+        Registry::put(_guid, v);
+    }
+    tensor_i::promise_type::set_value(std::forward<tensor_i::ptr_type>(v));
+}
+
+Deferred::future_type Deferred::defer(Deferred::ptr_type && d, bool is_global)
+{
+    if(is_global) {
+        if(is_cw() && theTransceiver->rank() == 0) theMediator->to_workers(d);
+        if(d) d->_guid = Registry::get_guid();
+    }
     auto f = d ? d->get_future() : Deferred::future_type();
     _deferred.push(std::move(d));
     return f;
-    /* auto aa = Deferred::undefer_next();
-    aa->run();
-    return aa->get_future(); */
 }
 
 Deferred::ptr_type Deferred::undefer_next()
@@ -25,7 +42,7 @@ void process_promises()
     while(true) {
         Deferred::ptr_type d;
         _deferred.pop(d);
-        // auto d = std::move(Deferred::undefer_next());
+        std::cerr << "Executing something" << std::endl;
         if(d) d->run();
         else break;
         d.reset();
@@ -34,6 +51,7 @@ void process_promises()
 
 void sync()
 {
+    // FIXME this does not wait for the last deferred to complete
     while(!_deferred.empty()) {
         std::this_thread::sleep_for(std::chrono::milliseconds(1));
     }

src/EWBinOp.cpp

@@ -2,6 +2,9 @@
 #include "ddptensor/LinAlgOp.hpp"
 #include "ddptensor/TypeDispatch.hpp"
 #include "ddptensor/x.hpp"
+#include "ddptensor/Factory.hpp"
+#include "ddptensor/Registry.hpp"
+#include "ddptensor/Creator.hpp"
 
 namespace x {
 
@@ -135,27 +138,43 @@ namespace x {
 
 struct DeferredEWBinOp : public Deferred
 {
-    tensor_i::future_type _a;
-    tensor_i::future_type _b;
+    id_type _a;
+    id_type _b;
     EWBinOpId _op;
 
+    DeferredEWBinOp() = default;
     DeferredEWBinOp(EWBinOpId op, const tensor_i::future_type & a, const tensor_i::future_type & b)
-        : _a(a), _b(b), _op(op)
+        : _a(a.id()), _b(b.id()), _op(op)
     {}
 
     void run()
     {
-        const auto a = std::move(_a.get());
-        const auto b = std::move(_b.get());
+        const auto a = std::move(Registry::get(_a));
+        const auto b = std::move(Registry::get(_b));
         set_value(std::move(TypeDispatch<x::EWBinOp>(a, b, _op)));
     }
+
+    FactoryId factory() const
+    {
+        return F_EWBINOP;
+    }
+
+    template<typename S>
+    void serialize(S & ser)
+    {
+        ser.template value<sizeof(_a)>(_a);
+        ser.template value<sizeof(_b)>(_b);
+        ser.template value<sizeof(_op)>(_op);
+    }
 };
 
 tensor_i::future_type EWBinOp::op(EWBinOpId op, const tensor_i::future_type & a, const py::object & b)
 {
-    auto bb = x::mk_ftx(b);
+    auto bb = Creator::mk_future(b);
     if(op == __MATMUL__) {
         return LinAlgOp::vecdot(a, bb, 0);
     }
     return defer<DeferredEWBinOp>(op, a, bb);
 }
+
+FACTORY_INIT(DeferredEWBinOp, F_EWBINOP);

src/EWUnyOp.cpp

@@ -1,6 +1,7 @@
 #include "ddptensor/EWUnyOp.hpp"
 #include "ddptensor/TypeDispatch.hpp"
 #include "ddptensor/x.hpp"
+#include "ddptensor/Factory.hpp"
 
 namespace x {
 
@@ -109,22 +110,37 @@ namespace x {
 } //namespace x
 
 struct DeferredEWUnyOp : public Deferred
- {
-    tensor_i::future_type _a;
+{
+    id_type _a;
     EWUnyOpId _op;
 
+    DeferredEWUnyOp() = default;
     DeferredEWUnyOp(EWUnyOpId op, const tensor_i::future_type & a)
-        : _a(a), _op(op)
+        : _a(a.id()), _op(op)
     {}
 
     void run()
     {
-        const auto a = std::move(_a.get());
+        const auto a = std::move(Registry::get(_a));
         set_value(std::move(TypeDispatch<x::EWUnyOp>(a, _op)));
     }
+
+    FactoryId factory() const
+    {
+        return F_EWUNYOP;
+    }
+    
+    template<typename S>
+    void serialize(S & ser)
+    {
+        ser.template value<sizeof(_a)>(_a);
+        ser.template value<sizeof(_op)>(_op);
+    }
 };
 
 tensor_i::future_type EWUnyOp::op(EWUnyOpId op, const tensor_i::future_type & a)
 {
     return defer<DeferredEWUnyOp>(op, a);
 }
+
+FACTORY_INIT(DeferredEWUnyOp, F_EWUNYOP);