wip

Author: AllanZyne

Diff for: CMakeLists.txt

@@ -149,8 +149,10 @@ include_directories(
   ${PROJECT_SOURCE_DIR}/third_party/bitsery/include
   ${MPI_INCLUDE_PATH}
   ${pybind11_INCLUDE_DIRS}
+  ${LLVM_INCLUDE_DIRS}
   ${MLIR_INCLUDE_DIRS}
-  ${IMEX_INCLUDE_DIRS})
+  ${IMEX_INCLUDE_DIRS}
+  "/export/users/yzhao/work/sharpy_ws/builds/debug/tools/Imex/include")
 
 if (CMAKE_SYSTEM_NAME STREQUAL Linux)
   target_link_options(_sharpy PRIVATE "LINKER:--version-script=${CMAKE_CURRENT_SOURCE_DIR}/export.txt")

Diff for: setup.py

@@ -1,3 +1,4 @@
+import multiprocessing
 import os
 import pathlib
 
@@ -44,7 +45,10 @@ def build_cmake(self, ext):
         os.chdir(str(build_temp))
         self.spawn(["cmake", str(cwd)] + cmake_args)
         if not self.dry_run:
-            self.spawn(["cmake", "--build", ".", "-j5"] + build_args)
+            self.spawn(
+                ["cmake", "--build", ".", f"-j{multiprocessing.cpu_count()}"]
+                + build_args
+            )
         # Troubleshooting: if fail on line above then delete all possible
         # temporary CMake files including "CMakeCache.txt" in top level dir.
         os.chdir(str(cwd))

Diff for: sharpy/__init__.py

@@ -130,6 +130,10 @@ def _validate_device(device):
         exec(
             f"{func} = lambda this, shape, cp=None: ndarray(_csp.ManipOp.reshape(this._t, shape, cp))"
         )
+    elif func == "permute_dims":
+        exec(
+            f"{func} = lambda this, axes: ndarray(_csp.ManipOp.permute_dims(this._t, axes))"
+        )
 
 for func in api.api_categories["ReduceOp"]:
     FUNC = func.upper()

Diff for: sharpy/array_api.py

@@ -179,6 +179,7 @@
             "roll",  # (x, /, shift, *, axis=None)
             "squeeze",  # (x, /, axis)
             "stack",  # (arrays, /, *, axis=0)
+            "permute_dims",  # (x: array, /, axes: Tuple[int, ...]) → array
         ],
         "LinAlgOp": [
             "matmul",  # (x1, x2, /)

Diff for: src/ManipOp.cpp

@@ -205,6 +205,57 @@ struct DeferredToDevice : public Deferred {
   }
 };
 
+struct DeferredPermuteDims : public Deferred {
+  id_type _array;
+  shape_type _axes;
+
+  DeferredPermuteDims() = default;
+  DeferredPermuteDims(const array_i::future_type &array,
+                      const shape_type &shape, const shape_type &axes)
+      : Deferred(array.dtype(), shape, array.device(), array.team()),
+        _array(array.guid()), _axes(axes) {}
+
+  bool generate_mlir(::mlir::OpBuilder &builder, const ::mlir::Location &loc,
+                     jit::DepManager &dm) override {
+    auto arrayValue = dm.getDependent(builder, Registry::get(_array));
+
+    auto axesAttr = builder.getI64ArrayAttr(_axes);
+
+    auto aTyp =
+        ::mlir::cast<::imex::ndarray::NDArrayType>(arrayValue.getType());
+    auto outTyp = imex::dist::cloneWithShape(aTyp, shape());
+
+    auto op = builder.create<::imex::ndarray::PermuteDimsOp>(
+        loc, outTyp, arrayValue, axesAttr);
+
+    dm.addVal(
+        this->guid(), op,
+        [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, std::vector<int64_t> &&loffs) {
+          auto t = mk_tnsr(this->guid(), _dtype, this->shape(), this->device(),
+                           this->team(), 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, std::move(loffs));
+          this->set_value(std::move(t));
+        });
+
+    return false;
+  }
+
+  FactoryId factory() const override { return F_PERMUTEDIMS; }
+
+  template <typename S> void serialize(S &ser) {
+    ser.template value<sizeof(_array)>(_array);
+    // ser.template value<sizeof(_axes)>(_axes);
+  }
+};
+
 FutureArray *ManipOp::reshape(const FutureArray &a, const shape_type &shape,
                               const py::object &copy) {
   auto doCopy = copy.is_none()
@@ -229,7 +280,32 @@ FutureArray *ManipOp::to_device(const FutureArray &a,
   return new FutureArray(defer<DeferredToDevice>(a.get(), device));
 }
 
+FutureArray *ManipOp::permute_dims(const FutureArray &array,
+                                   const shape_type &axes) {
+  auto shape = array.get().shape();
+
+  // verifyPermuteArray
+  if (shape.size() != axes.size()) {
+    throw std::invalid_argument("axes must have the same length as the shape");
+  }
+  for (auto i = 0ul; i < shape.size(); ++i) {
+    if (std::find(axes.begin(), axes.end(), i) == axes.end()) {
+      throw std::invalid_argument("axes must contain all dimensions");
+    }
+  }
+
+  auto permutedShape = shape_type(shape.size());
+  for (auto i = 0ul; i < shape.size(); ++i) {
+    permutedShape[i] = shape[axes[i]];
+  }
+
+  return new FutureArray(
+      defer<DeferredPermuteDims>(array.get(), permutedShape, axes));
+}
+
 FACTORY_INIT(DeferredReshape, F_RESHAPE);
 FACTORY_INIT(DeferredAsType, F_ASTYPE);
 FACTORY_INIT(DeferredToDevice, F_TODEVICE);
+FACTORY_INIT(DeferredPermuteDims, F_PERMUTEDIMS);
+
 } // namespace SHARPY

Diff for: src/_sharpy.cpp

@@ -196,7 +196,9 @@ PYBIND11_MODULE(_sharpy, m) {
   py::class_<IEWBinOp>(m, "IEWBinOp").def("op", &IEWBinOp::op);
   py::class_<EWBinOp>(m, "EWBinOp").def("op", &EWBinOp::op);
   py::class_<ReduceOp>(m, "ReduceOp").def("op", &ReduceOp::op);
-  py::class_<ManipOp>(m, "ManipOp").def("reshape", &ManipOp::reshape);
+  py::class_<ManipOp>(m, "ManipOp")
+      .def("reshape", &ManipOp::reshape)
+      .def("permute_dims", &ManipOp::permute_dims);
   py::class_<LinAlgOp>(m, "LinAlgOp").def("vecdot", &LinAlgOp::vecdot);
 
   py::class_<FutureArray>(m, "SHARPYFuture")

Diff for: src/include/sharpy/CppTypes.hpp

@@ -339,6 +339,7 @@ enum FactoryId : int {
   F_REDUCEOP,
   F_REPLICATE,
   F_RESHAPE,
+  F_PERMUTEDIMS,
   F_SERVICE,
   F_SETITEM,
   F_ASTYPE,

Diff for: src/include/sharpy/ManipOp.hpp

@@ -20,5 +20,8 @@ struct ManipOp {
 
   static FutureArray *to_device(const FutureArray &a,
                                 const std::string &device);
+
+  static FutureArray *permute_dims(const FutureArray &array,
+                                   const shape_type &axes);
 };
 } // namespace SHARPY

Diff for: test/test_manip.py

@@ -86,3 +86,10 @@ 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])
+
+    def test_permute_dims(self):
+        def doit(aapi, **kwargs):
+            a = aapi.arange(0, 12 * 11, 1, aapi.int32, **kwargs)
+            return aapi.permite_dims(a, [1, 0])
+
+        assert runAndCompare(doit)