adding linspace (#6)

adding linspace

Author: fschlimb

ddptensor/__init__.py

@@ -83,6 +83,10 @@ def to_numpy(a):
         exec(
             f"{func} = lambda start, end, step, dtype, team=0: dtensor(_cdt.Creator.arange(start, end, step, dtype, team))"
         )
+    elif func == "linspace":
+        exec(
+            f"{func} = lambda start, end, step, endpoint, dtype, team=0: dtensor(_cdt.Creator.linspace(start, end, step, endpoint, dtype, team))"
+        )
 
 for func in api.api_categories["ReduceOp"]:
     FUNC = func.upper()

src/Creator.cpp

@@ -79,13 +79,11 @@ namespace x {
 
 struct DeferredFromShape : public Deferred {
   shape_type _shape;
-  DTypeId _dtype;
   CreatorId _op;
 
   DeferredFromShape() = default;
   DeferredFromShape(CreatorId op, const shape_type &shape, DTypeId dtype)
-      : Deferred(dtype, shape.size(), true), _shape(shape), _dtype(dtype),
-        _op(op) {}
+      : Deferred(dtype, shape.size(), true), _shape(shape), _op(op) {}
 
   void run() {
     // set_value(std::move(TypeDispatch<x::Creator>(_dtype, _op, _shape)));
@@ -97,7 +95,6 @@ struct DeferredFromShape : public Deferred {
 
   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);
   }
 };
@@ -110,12 +107,10 @@ ddptensor *Creator::create_from_shape(CreatorId op, const shape_type &shape,
 struct DeferredFull : public Deferred {
   shape_type _shape;
   PyScalar _val;
-  DTypeId _dtype;
 
   DeferredFull() = default;
   DeferredFull(const shape_type &shape, PyScalar val, DTypeId dtype)
-      : Deferred(dtype, shape.size(), true), _shape(shape), _val(val),
-        _dtype(dtype) {}
+      : Deferred(dtype, shape.size(), true), _shape(shape), _val(val) {}
 
   void run() {
     // auto op = FULL;
@@ -189,6 +184,8 @@ ddptensor *Creator::full(const shape_type &shape, const py::object &val,
   return new ddptensor(defer<DeferredFull>(shape, v, dtype));
 }
 
+// ***************************************************************************
+
 struct DeferredArange : public Deferred {
   uint64_t _start, _end, _step, _team;
 
@@ -205,19 +202,25 @@ struct DeferredArange : public Deferred {
 
   bool generate_mlir(::mlir::OpBuilder &builder, ::mlir::Location loc,
                      jit::DepManager &dm) override {
-    auto start = ::imex::createInt(loc, builder, _start);
-    auto stop = ::imex::createInt(loc, builder, _end);
-    auto step = ::imex::createInt(loc, builder, _step);
     // ::mlir::Value
     auto team = /*getTransceiver()->nranks() <= 1
                 ? ::mlir::Value()
                 :*/
         ::imex::createIndex(loc, builder,
                             reinterpret_cast<uint64_t>(getTransceiver()));
 
+    auto _num = (_end - _start + _step + (_step < 0 ? 1 : -1)) / _step;
+
+    auto start = ::imex::createFloat(loc, builder, _start);
+    auto stop = ::imex::createFloat(loc, builder, _start + _num * _step);
+    auto num = ::imex::createIndex(loc, builder, _num);
+    auto rTyp = ::imex::ptensor::PTensorType::get(
+        ::llvm::ArrayRef<int64_t>{::mlir::ShapedType::kDynamic},
+        imex::ptensor::toMLIR(builder, jit::getPTDType(_dtype)));
+
     dm.addVal(this->guid(),
-              builder.create<::imex::ptensor::ARangeOp>(loc, start, stop, step,
-                                                        nullptr, team),
+              builder.create<::imex::ptensor::LinSpaceOp>(
+                  loc, rTyp, start, stop, num, false, nullptr, team),
               [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,
@@ -239,7 +242,6 @@ struct DeferredArange : public Deferred {
     ser.template value<sizeof(_start)>(_start);
     ser.template value<sizeof(_end)>(_end);
     ser.template value<sizeof(_step)>(_step);
-    ser.template value<sizeof(_dtype)>(_dtype);
   }
 };
 
@@ -248,6 +250,76 @@ ddptensor *Creator::arange(uint64_t start, uint64_t end, uint64_t step,
   return new ddptensor(defer<DeferredArange>(start, end, step, dtype, team));
 }
 
+// ***************************************************************************
+
+struct DeferredLinspace : public Deferred {
+  double _start, _end;
+  uint64_t _num, _team;
+  bool _endpoint;
+
+  DeferredLinspace() = default;
+  DeferredLinspace(double start, double end, uint64_t num, bool endpoint,
+                   DTypeId dtype, uint64_t team = 0)
+      : Deferred(dtype, 1, true), _start(start), _end(end), _num(num),
+        _team(team), _endpoint(endpoint) {}
+
+  void run() override{
+      // set_value(std::move(TypeDispatch<x::Creator>(_dtype, _start, _end,
+      // _num)));
+  };
+
+  bool generate_mlir(::mlir::OpBuilder &builder, ::mlir::Location loc,
+                     jit::DepManager &dm) override {
+    // ::mlir::Value
+    auto team = /*getTransceiver()->nranks() <= 1
+                ? ::mlir::Value()
+                :*/
+        ::imex::createIndex(loc, builder,
+                            reinterpret_cast<uint64_t>(getTransceiver()));
+
+    auto start = ::imex::createFloat(loc, builder, _start);
+    auto stop = ::imex::createFloat(loc, builder, _end);
+    auto num = ::imex::createIndex(loc, builder, _num);
+    auto rTyp = ::imex::ptensor::PTensorType::get(
+        ::llvm::ArrayRef<int64_t>{::mlir::ShapedType::kDynamic},
+        imex::ptensor::toMLIR(builder, jit::getPTDType(_dtype)));
+
+    dm.addVal(this->guid(),
+              builder.create<::imex::ptensor::LinSpaceOp>(
+                  loc, rTyp, start, stop, num, _endpoint, nullptr, team),
+              [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(transceiver, _dtype, rank, allocated, aligned,
+                            offset, sizes, strides, gs_allocated, gs_aligned,
+                            lo_allocated, lo_aligned, balanced)));
+              });
+    return false;
+  }
+
+  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(_num)>(_num);
+    ser.template value<sizeof(_endpoint)>(_endpoint);
+  }
+};
+
+ddptensor *Creator::linspace(double start, double end, uint64_t num,
+                             bool endpoint, DTypeId dtype, uint64_t team) {
+  return new ddptensor(
+      defer<DeferredLinspace>(start, end, num, endpoint, dtype, team));
+}
+
+// ***************************************************************************
+
 std::pair<ddptensor *, bool> Creator::mk_future(const py::object &b) {
   if (py::isinstance<ddptensor>(b)) {
     return {b.cast<ddptensor *>(), false};
@@ -263,3 +335,4 @@ std::pair<ddptensor *, bool> Creator::mk_future(const py::object &b) {
 FACTORY_INIT(DeferredFromShape, F_FROMSHAPE);
 FACTORY_INIT(DeferredFull, F_FULL);
 FACTORY_INIT(DeferredArange, F_ARANGE);
+FACTORY_INIT(DeferredLinspace, F_LINSPACE);

src/ddptensor.cpp

@@ -128,6 +128,7 @@ PYBIND11_MODULE(_ddptensor, m) {
   Factory::init<F_GETITEM>();
   Factory::init<F_IEWBINOP>();
   Factory::init<F_LINALGOP>();
+  Factory::init<F_LINSPACE>();
   Factory::init<F_MANIPOP>();
   Factory::init<F_RANDOM>();
   Factory::init<F_REDUCEOP>();
@@ -162,7 +163,8 @@ PYBIND11_MODULE(_ddptensor, m) {
   py::class_<Creator>(m, "Creator")
       .def("create_from_shape", &Creator::create_from_shape)
       .def("full", &Creator::full)
-      .def("arange", &Creator::arange);
+      .def("arange", &Creator::arange)
+      .def("linspace", &Creator::linspace);
 
   py::class_<EWUnyOp>(m, "EWUnyOp").def("op", &EWUnyOp::op);
   py::class_<IEWBinOp>(m, "IEWBinOp").def("op", &IEWBinOp::op);

src/include/ddptensor/CppTypes.hpp

@@ -203,6 +203,7 @@ enum FactoryId : int {
   F_GETLOCAL,
   F_IEWBINOP,
   F_LINALGOP,
+  F_LINSPACE,
   F_MANIPOP,
   F_MAP,
   F_RANDOM,

src/include/ddptensor/Creator.hpp

@@ -17,5 +17,7 @@ struct Creator {
                          DTypeId dtype = FLOAT64);
   static ddptensor *arange(uint64_t start, uint64_t end, uint64_t step,
                            DTypeId dtype = INT64, uint64_t team = 0);
+  static ddptensor *linspace(double start, double end, uint64_t num,
+                             bool endpoint, DTypeId dtype, uint64_t team);
   static std::pair<ddptensor *, bool> mk_future(const py::object &b);
 };

src/include/ddptensor/jit/mlir.hpp

@@ -63,6 +63,35 @@ template <> struct PT_DTYPE<bool> {
   constexpr static ::imex::ptensor::DType value = ::imex::ptensor::I1;
 };
 
+inline ::imex::ptensor::DType getPTDType(DTypeId dt) {
+  switch (dt) {
+  case FLOAT64:
+    return PT_DTYPE<TYPE<FLOAT64>::dtype>::value;
+  case INT64:
+    return PT_DTYPE<TYPE<INT64>::dtype>::value;
+  case FLOAT32:
+    return PT_DTYPE<TYPE<FLOAT32>::dtype>::value;
+  case INT32:
+    return PT_DTYPE<TYPE<INT32>::dtype>::value;
+  case INT16:
+    return PT_DTYPE<TYPE<INT16>::dtype>::value;
+  case INT8:
+    return PT_DTYPE<TYPE<INT8>::dtype>::value;
+  case UINT64:
+    return PT_DTYPE<TYPE<UINT64>::dtype>::value;
+  case UINT32:
+    return PT_DTYPE<TYPE<UINT32>::dtype>::value;
+  case UINT16:
+    return PT_DTYPE<TYPE<UINT16>::dtype>::value;
+  case UINT8:
+    return PT_DTYPE<TYPE<UINT8>::dtype>::value;
+  case BOOL:
+    return PT_DTYPE<TYPE<BOOL>::dtype>::value;
+  default:
+    throw std::runtime_error("unknown dtype");
+  }
+}
+
 // function type used for reporting back tensor results generated
 // by Deferred::generate_mlir
 using SetResFunc = std::function<void(