adding spmd.get_slice, more attributes, comments, cleanup

Author: fschlimb

ddptensor/__init__.py

@@ -2,6 +2,7 @@
 from .ddptensor import float64, int64, fini, dtensor
 from os import getenv
 from . import array_api as api
+from . import spmd
 
 #__impl_str = getenv("DDPNP_ARRAY", 'numpy')
 #exec(f"import {__impl_str} as __impl")

ddptensor/array_api.py

@@ -152,5 +152,20 @@
     "__dlpack_device__",  # (self, /)
     "__float__",  # (self, /)
     "__int__",  # (self, /)
+    "__index__",
     "__len__",  # (self, /)
 ]
+
+misc_methods = [
+    "__getitem__",
+    "__setitem__",
+]
+
+attributes = [
+    "dtype",
+    "shape",
+    "device",
+    "ndim",
+    "size",
+    "T"
+]

ddptensor/ddptensor.py

@@ -2,8 +2,6 @@
 from ._ddptensor import float64, int64, fini
 from . import array_api as api
 
-t_attributes = ["dtype", "shape", ]  #"device", "ndim", "size", "T"]
-
 #def try_except(func, *args, **kwargs):
 #    try:
 #        return func(*args, **kwargs)
@@ -37,19 +35,16 @@ def __repr__(self):
 
     for method in api.unary_methods:
         exec(
-            f"{method} = lambda self: _cdt.UnyOp.{method}(self._t)"
+            f"{method} = lambda self: self._t.{method}()"
         )
 
-    # for att in t_attributes:
-    #     exec(
-    #         f"{att} = property(lambda self: self._t.{att})"
-    #     )
+    for att in api.attributes:
+        exec(
+            f"{att} = property(lambda self: self._t.{att})"
+        )
 
     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)
-
-    # def get_slice(self, *args):
-    #     return self._t.get_slice(*args)

ddptensor/spmd.py

@@ -0,0 +1,4 @@
+from . import _ddptensor as _cdt
+
+def get_slice(self, *args):
+    return _cdt._get_slice(self._t, *args)

scripts/code_gen.py

@@ -1,6 +1,12 @@
 import array_api as api
 
-print("""// SPDX-License-Identifier: BSD-3-Clause
+print("""// Auto-generated file
+// #######################################################
+// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+// !! DO NOT EDIT, USE ../scripts/code_gen.py TO UPDATE !!
+// !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
+// #######################################################
+// SPDX-License-Identifier: BSD-3-Clause
 #pragma once
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>

src/ddptensor.cpp

@@ -2,11 +2,16 @@
 
 /*
   A Distributed Data-Parallel Tensor for Python, following the array API.
-  We have a 3-level hierachy
-    1. tensor_i: the abstract interface, not bound to types (like numpy)
-    2. dtensor_impl: a typed template layer with the actual functionality
-    3. dtensor: the PYthon API delegating to untyped tensor_i
-  We use pybind11.
+  
+  XTensor handles the actual functionality on each process.
+  pybind11 handles the bridge to Python.
+
+  We bridge dynamic dtypes of the Python array through dynamic type dispatch (TypeDispatch).
+  This means the compiler will instantiate the full functionality for all elements types.
+  Within kernels we dispatch the operation type by enum values (see x.hpp).
+  tensor_i is an abstract class to hide the element type which of the actual tensor.
+  The concrete tensor implementation (DPTensorX, x.hpp) requires the element type 
+  as a template parameter.
  */
 
 #include <pybind11/pybind11.h>
@@ -18,17 +23,17 @@ using namespace pybind11::literals; // to bring _a
 #include "ddptensor/MPIMediator.hpp"
 #include "ddptensor/x.hpp"
 
-// ###################################################################
-// ###################################################################
-// ###################################################################
-
+// Dependent on dt, dispatch arguments to a operation class.
+// The operation must
+//    * be a template class accepting the element type as argument
+//    * implement one or more "op" methods matching the given arguments (args)
+// All arguments other than dt are opaquely passed to the operation.
 template<template<typename OD> class OpDispatch, typename... Ts>
 auto TypeDispatch(DType dt, Ts&&... args)
 {
     switch(dt) {
     case DT_FLOAT64:
         return OpDispatch<double>::op(std::forward<Ts>(args)...);
-#if 0
     case DT_INT64:
         return OpDispatch<int64_t>::op(std::forward<Ts>(args)...);
     case DT_FLOAT32:
@@ -43,7 +48,6 @@ auto TypeDispatch(DType dt, Ts&&... args)
         return OpDispatch<uint32_t>::op(std::forward<Ts>(args)...);
     case DT_UINT16:
         return OpDispatch<uint16_t>::op(std::forward<Ts>(args)...);
-#endif
         /* FIXME
     case DT_BOOL:
         return OpDispatch<bool>::op(std::forward<Ts>(args)...);
@@ -53,6 +57,9 @@ auto TypeDispatch(DType dt, Ts&&... args)
     }
 }
 
+// #########################################################################
+// The following classes are wrappers bridging pybind11 defs to TypeDispatch
+
 struct Creator
 {
 
@@ -67,7 +74,7 @@ struct Creator
         return TypeDispatch<x::Creator>(dtype, op, std::forward<shape_type>(shape), std::forward<py::object>(val));
     }
 };
-#if 0
+
 struct IEWBinOp
 {
     static auto op(IEWBinOpId op, x::DPTensorBaseX::ptr_type a, x::DPTensorBaseX::ptr_type b)
@@ -92,43 +99,29 @@ struct EWUnyOp
     }
 };
 
-struct UnyOp
-{
-    static bool __bool__(x::DPTensorBaseX::ptr_type a)
-    {
-        return TypeDispatch<x::UnyOp>(a->dtype(), a, true);
-    }
-
-    static double __float__(x::DPTensorBaseX::ptr_type a)
-    {
-        return TypeDispatch<x::UnyOp>(a->dtype(), a, double(1));
-    }
-
-    static int64_t __int__(x::DPTensorBaseX::ptr_type a)
-    {
-        return TypeDispatch<x::UnyOp>(a->dtype(), a, int64_t(1));
-    }
-};
-
 struct ReduceOp
 {
     static auto op(ReduceOpId op, x::DPTensorBaseX::ptr_type a, const dim_vec_type & dim)
     {
         return TypeDispatch<x::ReduceOp>(a->dtype(), op, a, dim);
     }
 };
-#endif
 
 struct GetItem
 {
-    static auto op(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v)
+    static auto __getitem__(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v)
     {
         return TypeDispatch<x::GetItem>(a->dtype(), a, NDSlice(v));
     }
+    static auto get_slice(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v)
+    {
+        return TypeDispatch<x::SPMD>(a->dtype(), a, NDSlice(v));
+    }
 };
+
 struct SetItem
 {
-    static auto op(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v, x::DPTensorBaseX::ptr_type b)
+    static auto __setitem__(x::DPTensorBaseX::ptr_type a, const std::vector<py::slice> & v, x::DPTensorBaseX::ptr_type b)
     {
         return TypeDispatch<x::SetItem>(a->dtype(), a, NDSlice(v), b);
     }
@@ -142,6 +135,7 @@ rank_type myrank()
 Transceiver * theTransceiver = nullptr;
 Mediator * theMediator = nullptr;
 
+// users currently need to call fini to make MPI terminate gracefully
 void fini()
 {
     delete theMediator;
@@ -150,6 +144,8 @@ void fini()
     theTransceiver = nullptr;
 }
 
+// #########################################################################
+// Finally our Python module
 PYBIND11_MODULE(_ddptensor, m) {
     theTransceiver = new MPITransceiver();
     theMediator = new MPIMediator();
@@ -164,21 +160,17 @@ PYBIND11_MODULE(_ddptensor, m) {
         .value("bool", DT_BOOL)
         .export_values();
 
-    m.def("fini", &fini);
-    m.def("myrank", &myrank);
+    m.def("fini", &fini)
+        .def("myrank", &myrank)
+        .def("_get_slice", &GetItem::get_slice);
 
     py::class_<Creator>(m, "Creator")
         .def("create_from_shape", &Creator::create_from_shape)
         .def("full", &Creator::full);
-#if 0
+
     py::class_<EWUnyOp>(m, "EWUnyOp")
         .def("op", &EWUnyOp::op);
 
-    py::class_<UnyOp>(m, "UnyOp")
-        .def("__bool__", &UnyOp::__bool__)
-        .def("__float__", &UnyOp::__float__)
-        .def("__int__", &UnyOp::__int__);
-
     py::class_<IEWBinOp>(m, "IEWBinOp")
         .def("op", &IEWBinOp::op);
 
@@ -187,12 +179,20 @@ PYBIND11_MODULE(_ddptensor, m) {
 
     py::class_<ReduceOp>(m, "ReduceOp")
         .def("op", &ReduceOp::op);
-#endif
 
     py::class_<x::DPTensorBaseX, x::DPTensorBaseX::ptr_type>(m, "DPTensorX")
+        .def_property_readonly("dtype", &x::DPTensorBaseX::dtype)
+        .def_property_readonly("shape", &x::DPTensorBaseX::shape)
+        .def_property_readonly("size", &x::DPTensorBaseX::size)
+        .def_property_readonly("ndim", &x::DPTensorBaseX::ndim)
+        .def("__bool__", &x::DPTensorBaseX::__bool__)
+        .def("__float__", &x::DPTensorBaseX::__float__)
+        .def("__int__", &x::DPTensorBaseX::__int__)
+        .def("__index__", &x::DPTensorBaseX::__int__)
+        .def("__len__", &x::DPTensorBaseX::__len__)
         .def("__repr__", &x::DPTensorBaseX::__repr__)
-        .def("__getitem__", &GetItem::op)
-        .def("__setitem__", &SetItem::op);
+        .def("__getitem__", &GetItem::__getitem__)
+        .def("__setitem__", &SetItem::__setitem__);
 
 #if 0
     py::class_<dtensor>(m, "dtensor")