Add to_device method (#74)

* add array.to_device method
* test fail because remote temporaries pass removes memref.copy thus
effectively producing a view

Author: tkarna

sharpy/ndarray.py

@@ -54,6 +54,9 @@ def __repr__(self):
     def astype(self, dtype, copy=False):
         return ndarray(self._t.astype(dtype, copy))
 
+    def to_device(self, device=""):
+        return ndarray(self._t.to_device(device))
+
     def __getitem__(self, key):
         key = key if isinstance(key, tuple) else (key,)
         key = [slicefy(x) for x in key]

src/ManipOp.cpp

@@ -142,6 +142,65 @@ struct DeferredAsType : public Deferred {
   }
 };
 
+// ***************************************************************************
+
+struct DeferredToDevice : public Deferred {
+  id_type _a;
+
+  DeferredToDevice() = default;
+  DeferredToDevice(const array_i::future_type &a, const std::string &device)
+      : Deferred(a.dtype(), a.shape(), device, a.team()), _a(a.guid()) {}
+
+  bool generate_mlir(::mlir::OpBuilder &builder, const ::mlir::Location &loc,
+                     jit::DepManager &dm) override {
+    auto av = dm.getDependent(builder, _a);
+
+    auto srcType = av.getType().dyn_cast<::imex::ndarray::NDArrayType>();
+    assert(srcType);
+    // copy envs, drop gpu env (if any)
+    auto srcEnvs = srcType.getEnvironments();
+    ::mlir::SmallVector<::mlir::Attribute> envs;
+    for (auto e : srcEnvs) {
+      if (!::mlir::isa<::imex::region::GPUEnvAttr>(e)) {
+        envs.emplace_back(e);
+      }
+    }
+    // append device attr
+    if (!_device.empty()) {
+      envs.emplace_back(
+          ::imex::region::GPUEnvAttr::get(builder.getStringAttr(_device)));
+    }
+    auto outType = ::imex::ndarray::NDArrayType::get(srcType.getShape(),
+                                                     srcType.getElementType(),
+                                                     envs, srcType.getLayout());
+    auto res = builder.create<::imex::ndarray::CopyOp>(loc, outType, av);
+    dm.addVal(this->guid(), res,
+              [this](uint64_t rank, void *l_allocated, void *l_aligned,
+                     intptr_t l_offset, const intptr_t *l_sizes,
+                     const intptr_t *l_strides, void *o_allocated,
+                     void *o_aligned, intptr_t o_offset,
+                     const intptr_t *o_sizes, const intptr_t *o_strides,
+                     void *r_allocated, void *r_aligned, intptr_t r_offset,
+                     const intptr_t *r_sizes, const intptr_t *r_strides,
+                     uint64_t *lo_allocated, uint64_t *lo_aligned) {
+                auto t = mk_tnsr(reinterpret_cast<Transceiver *>(this->team()),
+                                 _dtype, this->shape(), l_allocated, l_aligned,
+                                 l_offset, l_sizes, l_strides, o_allocated,
+                                 o_aligned, o_offset, o_sizes, o_strides,
+                                 r_allocated, r_aligned, r_offset, r_sizes,
+                                 r_strides, lo_allocated, lo_aligned);
+                this->set_value(std::move(t));
+              });
+    return false;
+  }
+
+  FactoryId factory() const { return F_TODEVICE; }
+
+  template <typename S> void serialize(S &ser) {
+    ser.template value<sizeof(_a)>(_a);
+  }
+};
+
 FutureArray *ManipOp::reshape(const FutureArray &a, const shape_type &shape,
                               const py::object &copy) {
   auto doCopy = copy.is_none()
@@ -155,12 +214,18 @@ FutureArray *ManipOp::reshape(const FutureArray &a, const shape_type &shape,
   return new FutureArray(defer<DeferredReshape>(a.get(), shape, doCopy));
 }
 
-FutureArray *AsType::astype(const FutureArray &a, DTypeId dtype,
-                            const py::object &copy) {
+FutureArray *ManipOp::astype(const FutureArray &a, DTypeId dtype,
+                             const py::object &copy) {
   auto doCopy = copy.is_none() ? false : copy.cast<bool>();
   return new FutureArray(defer<DeferredAsType>(a.get(), dtype, doCopy));
 }
 
+FutureArray *ManipOp::to_device(const FutureArray &a,
+                                const std::string &device) {
+  return new FutureArray(defer<DeferredToDevice>(a.get(), device));
+}
+
 FACTORY_INIT(DeferredReshape, F_RESHAPE);
 FACTORY_INIT(DeferredAsType, F_ASTYPE);
+FACTORY_INIT(DeferredToDevice, F_TODEVICE);
 } // namespace SHARPY

src/_sharpy.cpp

@@ -203,7 +203,8 @@ PYBIND11_MODULE(_sharpy, m) {
       .def("__index__",
            [](const FutureArray &f) { REPL_SYNC_RETURN(f, __int__); })
       // attributes returning a new FutureArray
-      .def("astype", &AsType::astype)
+      .def("astype", &ManipOp::astype)
+      .def("to_device", &ManipOp::to_device)
       .def("__getitem__", &GetItem::__getitem__)
       .def("__setitem__", &SetItem::__setitem__)
       .def("map", &SetItem::map);

src/include/sharpy/CppTypes.hpp

@@ -343,6 +343,7 @@ enum FactoryId : int {
   F_SERVICE,
   F_SETITEM,
   F_ASTYPE,
+  F_TODEVICE,
   FACTORY_LAST
 };
 

src/include/sharpy/ManipOp.hpp

@@ -14,10 +14,11 @@ namespace SHARPY {
 struct ManipOp {
   static FutureArray *reshape(const FutureArray &a, const shape_type &shape,
                               const py::object &copy);
-};
 
-struct AsType {
   static FutureArray *astype(const FutureArray &a, DTypeId dtype,
                              const py::object &copy);
+
+  static FutureArray *to_device(const FutureArray &a,
+                                const std::string &device);
 };
 } // namespace SHARPY

src/jit/mlir.cpp

@@ -354,15 +354,14 @@ void DepManager::deliver(std::vector<intptr_t> &outputV, uint64_t sz) {
         // no sizes/stride needed, just skip
         pos += memref_sz(1);
         // call finalization callback
-        v->second(
-            rank, t_allocated[0], t_aligned[0], t_offset[0], t_sizes[0],
-            t_strides[0], // lhsHalo
-            t_allocated[1], t_aligned[1], t_offset[1], t_sizes[1],
-            t_strides[1], // lData
-            t_allocated[2], t_aligned[2], t_offset[2], t_sizes[2],
-            t_strides[2], // rhsHalo
-            lo_allocated,
-            lo_aligned + lo_offset // local offset is 1d array of uint64_t
+        v->second(rank, t_allocated[0], t_aligned[0], t_offset[0], t_sizes[0],
+                  t_strides[0], // lhsHalo
+                  t_allocated[1], t_aligned[1], t_offset[1], t_sizes[1],
+                  t_strides[1], // lData
+                  t_allocated[2], t_aligned[2], t_offset[2], t_sizes[2],
+                  t_strides[2], // rhsHalo
+                  lo_allocated,
+                  lo_aligned + lo_offset // local offset is 1d array of uint64_t
         );
       } else { // 0d array or non-dist
         pos += getMR(rank, &output[pos], t_allocated[1], t_aligned[1],
@@ -631,6 +630,7 @@ JIT::JIT()
   _context.getOrLoadDialect<::mlir::arith::ArithDialect>();
   _context.getOrLoadDialect<::mlir::func::FuncDialect>();
   _context.getOrLoadDialect<::mlir::linalg::LinalgDialect>();
+  _context.getOrLoadDialect<::imex::region::RegionDialect>();
   // create the pass pipeline from string
   if (::mlir::failed(::mlir::parsePassPipeline(pass_pipeline, _pm)))
     throw std::runtime_error("failed to parse pass pipeline");

test/test_manip.py

@@ -68,3 +68,15 @@ def test_astype_suite(self):
             b = a.astype(to_type)
             assert b.dtype == to_type
             assert numpy.allclose(sp.to_numpy(b), [0, 1, 2, 3, 4, 5, 6, 7])
+
+    @pytest.mark.skip(reason="FIXME copy is eliminated")
+    def test_todevice_host2host(self):
+        a = sp.arange(0, 8, 1, sp.int32)
+        b = a.to_device()
+        assert numpy.allclose(sp.to_numpy(b), [0, 1, 2, 3, 4, 5, 6, 7])
+
+    @pytest.mark.skip(reason="FIXME copy is eliminated, device support")
+    def test_todevice_host2gpu(self):
+        a = sp.arange(0, 8, 1, sp.int32)
+        b = a.to_device(device="GPU")
+        assert numpy.allclose(sp.to_numpy(b), [0, 1, 2, 3, 4, 5, 6, 7])