wip

Author: AllanZyne

CMakeLists.txt

@@ -152,7 +152,7 @@ include_directories(
   ${LLVM_INCLUDE_DIRS}
   ${MLIR_INCLUDE_DIRS}
   ${IMEX_INCLUDE_DIRS}
-  "/export/users/yzhao/work/sharpy_ws/builds/debug/tools/Imex/include")
+  "/localdisk2/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")

setup.py

@@ -43,8 +43,10 @@ def build_cmake(self, ext):
         build_args = ["--config", config]
 
         os.chdir(str(build_temp))
+        print('!!!!!!!!!!', ["cmake", str(cwd)] + cmake_args)
         self.spawn(["cmake", str(cwd)] + cmake_args)
         if not self.dry_run:
+            print('!!!!!!!!!!', ["cmake", "--build", ".", f"-j{multiprocessing.cpu_count()}"] + build_args)
             self.spawn(
                 ["cmake", "--build", ".", f"-j{multiprocessing.cpu_count()}"]
                 + build_args

src/ManipOp.cpp

@@ -219,7 +219,7 @@ struct DeferredPermuteDims : public Deferred {
                      jit::DepManager &dm) override {
     auto arrayValue = dm.getDependent(builder, Registry::get(_array));
 
-    auto axesAttr = builder.getI64ArrayAttr(_axes);
+    auto axesAttr = builder.getDenseI64ArrayAttr(_axes);
 
     auto aTyp =
         ::mlir::cast<::imex::ndarray::NDArrayType>(arrayValue.getType());

src/idtr.cpp

@@ -718,20 +718,26 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
                                    SHARPY::Transceiver *tc, int64_t iNDims,
                                    int64_t *iGShapePtr, int64_t *iOffsPtr,
                                    void *iDataPtr, int64_t *iDataShapePtr,
-                                   int64_t *iDataStridesPtr, int64_t oNDims,
-                                   int64_t *oGShapePtr, int64_t *oOffsPtr,
+                                   int64_t *iDataStridesPtr, int64_t *oOffsPtr,
                                    void *oDataPtr, int64_t *oDataShapePtr,
                                    int64_t *oDataStridesPtr, int64_t *axesPtr) {
 #ifdef NO_TRANSCEIVER
   initMPIRuntime();
   tc = SHARPY::getTransceiver();
 #endif
   if (!iGShapePtr || !iOffsPtr || !iDataPtr || !iDataShapePtr ||
-      !iDataStridesPtr || !oGShapePtr || !oOffsPtr || !oDataPtr ||
-      !oDataShapePtr || !oDataStridesPtr || !tc) {
+      !iDataStridesPtr || !oOffsPtr || !oDataPtr || !oDataShapePtr ||
+      !oDataStridesPtr || !tc) {
     throw std::invalid_argument("Fatal: received nullptr in reshape");
   }
 
+  std::vector<int64_t> oGShape(iNDims);
+  for (int64_t i = 0; i < iNDims; ++i) {
+    oGShape[i] = iGShapePtr[axesPtr[i]];
+  }
+  auto *oGShapePtr = oGShape.data();
+  const auto oNDims = iNDims;
+
   assert(std::accumulate(&iGShapePtr[0], &iGShapePtr[iNDims], 1,
                          std::multiplies<int64_t>()) ==
          std::accumulate(&oGShapePtr[0], &oGShapePtr[oNDims], 1,
@@ -817,21 +823,21 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
     });
 
     int lastOffset = 0;
-    for (size_t i = 0; i < nRanks; i++) {
-      sendSizes[i] = sendRankBuffer[i].size();
-      sendOffsets[i] = lastOffset;
-      sendBuffer.insert(sendBuffer.end(), sendRankBuffer[i].begin(),
-                        sendRankBuffer[i].end());
-      lastOffset += sendSizes[i];
+    for (size_t rank = 0; rank < nRanks; rank++) {
+      sendSizes[rank] = sendRankBuffer[rank].size();
+      sendOffsets[rank] = lastOffset;
+      sendBuffer.insert(sendBuffer.end(), sendRankBuffer[rank].begin(),
+                        sendRankBuffer[rank].end());
+      lastOffset += sendSizes[rank];
     }
 
     output.localIndices([&](const id &outputIndex) {
       id inputIndex = outputIndex.permute(axes);
       auto rank = getInputRank(parts, inputIndex[0]);
       ++receiveSizes[rank];
     });
-    for (size_t i = 1; i < nRanks; i++) {
-      receiveOffsets[i] = receiveOffsets[i - 1] + receiveSizes[i - 1];
+    for (size_t rank = 1; rank < nRanks; rank++) {
+      receiveOffsets[rank] = receiveOffsets[rank - 1] + receiveSizes[rank - 1];
     }
   }
 
@@ -842,7 +848,7 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
 
   {
     std::vector<std::vector<T>> receiveRankBuffer(nRanks);
-    for (int64_t rank = 0; rank < nRanks; ++rank) {
+    for (size_t rank = 0; rank < nRanks; ++rank) {
       auto &rankBuffer = receiveRankBuffer[rank];
       rankBuffer.insert(
           rankBuffer.end(), receiveBuffer.begin() + receiveOffsets[rank],
@@ -866,12 +872,12 @@ template <typename T>
 WaitHandleBase *
 _idtr_copy_permute(SHARPY::Transceiver *tc, int64_t iNSzs, void *iGShapeDescr,
                    int64_t iNOffs, void *iLOffsDescr, int64_t iNDims,
-                   void *iDataDescr, int64_t oNSzs, void *oGShapeDescr,
-                   int64_t oNOffs, void *oLOffsDescr, int64_t oNDims,
-                   void *oDataDescr, int64_t axesSzs, void *axesDescr) {
+                   void *iDataDescr, int64_t oNOffs, void *oLOffsDescr,
+                   int64_t oNDims, void *oDataDescr, int64_t axesSzs,
+                   void *axesDescr) {
 
-  if (!iGShapeDescr || !iLOffsDescr || !iDataDescr || !oGShapeDescr ||
-      !oLOffsDescr || !oDataDescr || !axesDescr) {
+  if (!iGShapeDescr || !iLOffsDescr || !iDataDescr || !oLOffsDescr ||
+      !oDataDescr || !axesDescr) {
     throw std::invalid_argument(
         "Fatal error: received nullptr in update_halo.");
   }
@@ -882,15 +888,14 @@ _idtr_copy_permute(SHARPY::Transceiver *tc, int64_t iNSzs, void *iGShapeDescr,
   MRIdx1d iGShape(iNSzs, iGShapeDescr);
   MRIdx1d iOffs(iNOffs, iLOffsDescr);
   SHARPY::UnrankedMemRefType<T> iData(iNDims, iDataDescr);
-  MRIdx1d oGShape(oNSzs, oGShapeDescr);
   MRIdx1d oOffs(oNOffs, oLOffsDescr);
   SHARPY::UnrankedMemRefType<T> oData(oNDims, oDataDescr);
   MRIdx1d axes(axesSzs, axesDescr);
 
-  return _idtr_copy_permute<T>(
-      sharpyType, tc, iNDims, iGShape.data(), iOffs.data(), iData.data(),
-      iData.sizes(), iData.strides(), oNDims, oGShape.data(), oOffs.data(),
-      oData.data(), oData.sizes(), oData.strides(), axes.data());
+  return _idtr_copy_permute<T>(sharpyType, tc, iNDims, iGShape.data(),
+                               iOffs.data(), iData.data(), iData.sizes(),
+                               iData.strides(), oOffs.data(), oData.data(),
+                               oData.sizes(), oData.strides(), axes.data());
 }
 
 extern "C" {
@@ -919,12 +924,11 @@ TYPED_COPY_RESHAPE(i1, bool);
   void *_idtr_copy_permute_##_sfx(                                             \
       SHARPY::Transceiver *tc, int64_t iNSzs, void *iGShapeDescr,              \
       int64_t iNOffs, void *iLOffsDescr, int64_t iNDims, void *iLDescr,        \
-      int64_t oNSzs, void *oGShapeDescr, int64_t oNOffs, void *oLOffsDescr,    \
-      int64_t oNDims, void *oLDescr, int64_t axesSzs, void *axesDescr) {       \
-    return _idtr_copy_permute<_typ>(tc, iNSzs, iGShapeDescr, iNOffs,           \
-                                    iLOffsDescr, iNDims, iLDescr, oNSzs,       \
-                                    oGShapeDescr, oNOffs, oLOffsDescr, oNDims, \
-                                    oLDescr, axesSzs, axesDescr);              \
+      int64_t oNOffs, void *oLOffsDescr, int64_t oNDims, void *oLDescr,        \
+      int64_t axesSzs, void *axesDescr) {                                      \
+    return _idtr_copy_permute<_typ>(                                           \
+        tc, iNSzs, iGShapeDescr, iNOffs, iLOffsDescr, iNDims, iLDescr, oNOffs, \
+        oLOffsDescr, oNDims, oLDescr, axesSzs, axesDescr);                     \
   }                                                                            \
   _Pragma(STRINGIFY(weak _mlir_ciface__idtr_copy_permute_##_sfx =              \
                         _idtr_copy_permute_##_sfx))