wip

Author: AllanZyne

Diff for: src/idtr.cpp

@@ -641,6 +641,7 @@ template <typename T> class ndarray {
     size_t size = lSize();
     id idx = firstLocalIndex();
     while (size--) {
+      std::cout << "idx: " << idx[0] << ", " << idx[1] << std::endl;
       callback(idx);
       idx.next(_gShape);
     }
@@ -708,6 +709,52 @@ size_t getOutputRank(const std::vector<Parts> &parts, int64_t dim0) {
   return 0;
 }
 
+template <typename T> class WaitPermute {
+public:
+  WaitPermute(SHARPY::Transceiver *tc, SHARPY::Transceiver::WaitHandle hdl,
+              SHARPY::rank_type nRanks, std::vector<Parts> &&parts,
+              std::vector<int64_t> &&axes, ndarray<T> &&output,
+              std::vector<T> &&receiveBuffer, std::vector<int> &&receiveOffsets,
+              std::vector<int> &&receiveSizes)
+      : tc(tc), hdl(hdl), nRanks(nRanks), parts(std::move(parts)),
+        axes(std::move(axes)), output(std::move(output)),
+        receiveBuffer(std::move(receiveBuffer)),
+        receiveOffsets(std::move(receiveOffsets)),
+        receiveSizes(std::move(receiveSizes)) {}
+
+  void operator()() {
+    tc->wait(hdl);
+    std::vector<std::vector<T>> receiveRankBuffer(nRanks);
+    for (size_t rank = 0; rank < nRanks; ++rank) {
+      auto &rankBuffer = receiveRankBuffer[rank];
+      rankBuffer.insert(
+          rankBuffer.end(), receiveBuffer.begin() + receiveOffsets[rank],
+          receiveBuffer.begin() + receiveOffsets[rank] + receiveSizes[rank]);
+    }
+
+    std::vector<size_t> receiveRankBufferCount(nRanks, 0);
+    output.localIndices([&](const id &outputIndex) {
+      id inputIndex = outputIndex.permute(axes);
+      std::cout << "inputIndex: " << inputIndex[0] << ", " << inputIndex[1]
+                << std::endl;
+      auto rank = getInputRank(parts, inputIndex[0]);
+      auto &count = receiveRankBufferCount[rank];
+      output[outputIndex] = receiveRankBuffer[rank][count++];
+    });
+  }
+
+private:
+  SHARPY::Transceiver *tc;
+  SHARPY::Transceiver::WaitHandle hdl;
+  SHARPY::rank_type nRanks;
+  std::vector<Parts> parts;
+  std::vector<int64_t> axes;
+  ndarray<T> output;
+  std::vector<T> receiveBuffer;
+  std::vector<int> receiveOffsets;
+  std::vector<int> receiveSizes;
+};
+
 } // namespace
 
 /// @brief permute array
@@ -844,27 +891,20 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
   auto hdl = tc->alltoall(sendBuffer.data(), sendSizes.data(),
                           sendOffsets.data(), sharpytype, receiveBuffer.data(),
                           receiveSizes.data(), receiveOffsets.data());
-  tc->wait(hdl);
 
-  {
-    std::vector<std::vector<T>> receiveRankBuffer(nRanks);
-    for (size_t rank = 0; rank < nRanks; ++rank) {
-      auto &rankBuffer = receiveRankBuffer[rank];
-      rankBuffer.insert(
-          rankBuffer.end(), receiveBuffer.begin() + receiveOffsets[rank],
-          receiveBuffer.begin() + receiveOffsets[rank] + receiveSizes[rank]);
-    }
+  auto wait = WaitPermute(tc, hdl, nRanks, std::move(parts), std::move(axes),
+                          std::move(output), std::move(receiveBuffer),
+                          std::move(receiveOffsets), std::move(receiveSizes));
 
-    std::vector<size_t> receiveRankBufferCount(nRanks);
-    output.localIndices([&](const id &outputIndex) {
-      id inputIndex = outputIndex.permute(axes);
-      auto rank = getInputRank(parts, inputIndex[0]);
-      auto &count = receiveRankBufferCount[rank];
-      output[outputIndex] = receiveRankBuffer[rank][count++];
-    });
+  assert(parts.empty() && axes.empty() && receiveBuffer.empty() &&
+         receiveOffsets.empty() && receiveSizes.empty());
+
+  if (no_async) {
+    wait();
+    return nullptr;
   }
 
-  return nullptr;
+  return mkWaitHandle(std::move(wait));
 }
 
 /// @brief permute array

Diff for: src/jit/mlir.cpp

@@ -691,7 +691,7 @@ static const std::string cpu_pipeline =
     "one-shot-bufferize,"
     "canonicalize,"
     "imex-remove-temporaries,"
-    "func.func(buffer-deallocation),"
+    "buffer-deallocation-pipeline,"
     "func.func(convert-linalg-to-parallel-loops),"
     "func.func(scf-parallel-loop-fusion),"
     "drop-regions,"