fix test

Author: AllanZyne

examples/transposed3d.py

@@ -0,0 +1,55 @@
+import sharpy as sp
+import numpy as np
+
+
+def sp_tranposed3d_1():
+    a = sp.arange(0,2*3*4,1)
+    a = sp.reshape(a,[2,3,4])
+
+    # b = a.swapaxes(1,0).swapaxes(1,2)
+    b = sp.permute_dims(a, (1,0,2)) # 2x4x4 -> 4x2x4 || 4x4x4
+    b = sp.permute_dims(b, (0,2,1)) # 4x2x4 -> 4x4x2 || 4x4x4
+
+    # c = b.swapaxes(1,2).swapaxes(1,0)
+    c = sp.permute_dims(b, (0,2,1))
+    c = sp.permute_dims(c, (1,0,2))
+
+    assert(np.allclose(sp.to_numpy(a), sp.to_numpy(c)))
+    return b
+
+def sp_tranposed3d_2():
+    a = sp.arange(0,2*3*4,1)
+    a = sp.reshape(a,[2,3,4])
+
+    # b = a.swapaxes(2,1).swapaxes(2,0)
+    b = sp.permute_dims(a, (0,2,1))
+    b = sp.permute_dims(b, (2,1,0))
+
+    # c = b.swapaxes(2,1).swapaxes(0,1)
+    c = sp.permute_dims(b, (0,2,1))
+    c = sp.permute_dims(c, (1,0,2))
+    
+    return c
+
+def np_tranposed3d_1():
+    a = np.arange(0,2*3*4,1)
+    a = np.reshape(a,[2,3,4])
+    b = a.swapaxes(1,0).swapaxes(1,2)
+    return b
+
+def np_tranposed3d_2():
+    a = np.arange(0,2*3*4,1)
+    a = np.reshape(a,[2,3,4])
+    b = a.swapaxes(2,1).swapaxes(2,0)
+    c = b.swapaxes(2,1).swapaxes(0,1)
+    return c
+
+sp.init(False)
+
+b1 = sp_tranposed3d_1()
+assert(np.allclose(sp.to_numpy(b1), np_tranposed3d_1()))
+
+b2 = sp_tranposed3d_2()
+assert(np.allclose(sp.to_numpy(b2), np_tranposed3d_2()))
+
+sp.fini()

setup.py

@@ -43,10 +43,8 @@ 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/idtr.cpp

@@ -646,6 +646,15 @@ template <typename T> class ndarray {
     }
   }
 
+  void globalIndices(const std::function<void(const id &)> &callback) const {
+    size_t size = gSize();
+    id idx(_nDims);
+    while (size--) {
+      callback(idx);
+      idx.next(_gShape);
+    }
+  }
+
   int64_t getLocalDataOffset(const id &idx) const {
     auto localIdx = idx - _gOffsets;
     int64_t offset = 0;
@@ -711,14 +720,16 @@ size_t getOutputRank(const std::vector<Parts> &parts, int64_t dim0) {
 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, std::vector<int64_t> oGShape,
+              SHARPY::rank_type cRank, SHARPY::rank_type nRanks,
+              std::vector<Parts> &&parts, std::vector<int64_t> &&axes,
+              std::vector<int64_t> oGShape, ndarray<T> &&input,
               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)),
+      : tc(tc), hdl(hdl), cRank(cRank), nRanks(nRanks), parts(std::move(parts)),
         axes(std::move(axes)), oGShape(std::move(oGShape)),
-        output(std::move(output)), receiveBuffer(std::move(receiveBuffer)),
+        input(std::move(input)), output(std::move(output)),
+        receiveBuffer(std::move(receiveBuffer)),
         receiveOffsets(std::move(receiveOffsets)),
         receiveSizes(std::move(receiveSizes)) {}
 
@@ -733,9 +744,12 @@ template <typename T> class WaitPermute {
     }
 
     std::vector<size_t> receiveRankBufferCount(nRanks, 0);
-    output.localIndices([&](const id &outputIndex) {
-      id inputIndex = outputIndex.permute(axes);
-      auto rank = getInputRank(parts, inputIndex[0]);
+    input.globalIndices([&](const id &inputIndex) {
+      id outputIndex = inputIndex.permute(axes);
+      auto rank = getOutputRank(parts, outputIndex[0]);
+      if (rank != cRank)
+        return;
+      rank = getInputRank(parts, inputIndex[0]);
       auto &count = receiveRankBufferCount[rank];
       output[outputIndex] = receiveRankBuffer[rank][count++];
     });
@@ -744,10 +758,12 @@ template <typename T> class WaitPermute {
 private:
   SHARPY::Transceiver *tc;
   SHARPY::Transceiver::WaitHandle hdl;
+  SHARPY::rank_type cRank;
   SHARPY::rank_type nRanks;
   std::vector<Parts> parts;
   std::vector<int64_t> axes;
   std::vector<int64_t> oGShape;
+  ndarray<T> input;
   ndarray<T> output;
   std::vector<T> receiveBuffer;
   std::vector<int> receiveOffsets;
@@ -791,9 +807,9 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
   assert(std::accumulate(&oOffsPtr[1], &oOffsPtr[oNDims], 0,
                          std::plus<int64_t>()) == 0);
 
-  auto nRanks = tc->nranks();
-  auto rank = tc->rank();
-  if (nRanks <= rank) {
+  const auto nRanks = tc->nranks();
+  const auto cRank = tc->rank();
+  if (nRanks <= cRank) {
     throw std::out_of_range("Fatal: rank must be < number of ranks");
   }
 
@@ -833,10 +849,10 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
 
   // First we allgather the current and target partitioning
   std::vector<Parts> parts(nRanks);
-  parts[rank].iStart = iOffsPtr[0];
-  parts[rank].iEnd = iOffsPtr[0] + iDataShapePtr[0];
-  parts[rank].oStart = oOffsPtr[0];
-  parts[rank].oEnd = oOffsPtr[0] + oDataShapePtr[0];
+  parts[cRank].iStart = iOffsPtr[0];
+  parts[cRank].iEnd = iOffsPtr[0] + iDataShapePtr[0];
+  parts[cRank].oStart = oOffsPtr[0];
+  parts[cRank].oEnd = oOffsPtr[0] + oDataShapePtr[0];
   std::vector<int> counts(nRanks, 4);
   std::vector<int> dspl(nRanks);
   for (auto i = 0ul; i < nRanks; ++i) {
@@ -891,10 +907,10 @@ WaitHandleBase *_idtr_copy_permute(SHARPY::DTypeId sharpytype,
                           sendOffsets.data(), sharpytype, receiveBuffer.data(),
                           receiveSizes.data(), receiveOffsets.data());
 
-  auto wait = WaitPermute(tc, hdl, nRanks, std::move(parts), std::move(axes),
-                          std::move(oGShape), std::move(output),
-                          std::move(receiveBuffer), std::move(receiveOffsets),
-                          std::move(receiveSizes));
+  auto wait = WaitPermute(tc, hdl, cRank, nRanks, std::move(parts),
+                          std::move(axes), std::move(oGShape), std::move(input),
+                          std::move(output), std::move(receiveBuffer),
+                          std::move(receiveOffsets), std::move(receiveSizes));
 
   assert(parts.empty() && axes.empty() && receiveBuffer.empty() &&
          receiveOffsets.empty() && receiveSizes.empty());

test/test_manip.py

@@ -95,8 +95,8 @@ def test_todevice_host2gpu(self):
         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)
+        a = sp.arange(0, 10, 1, sp.int64)
+        b = sp.reshape(a, (2, 5))
+        c1 = sp.to_numpy(sp.permute_dims(b, [1, 0]))
+        c2 = sp.to_numpy(b).transpose(1, 0)
+        assert numpy.allclose(c1, c2)