[GraphOptimizer] Optimize consecutive non-inverse transposes

docs/Optimizations.md

@@ -32,7 +32,10 @@ Below you can see the list of currently supported graph optimizations:
     This optimization removes computations whose results or side-effects are
     not used.
 
-  * Elimination of transpose operations reversing each other
+  * Optimization of transpose nodes
+
+    This optimization combines multiple consecutive transpose nodes
+    into a single node, and also eliminates identity transpose nodes.
 
   * Sinking of transpose operations below other operations
 

lib/Optimizer/GraphOptimizer.cpp

@@ -116,26 +116,6 @@ static bool isIdentityShuffle(llvm::ArrayRef<unsigned> shuffle) {
   return true;
 }
 
-/// \returns true if the masks \p shuffle1 and shuffle2 are
-/// the inverse of on another. Applying both masks should result in the identity
-/// shuffle.
-static bool isIdentityShuffle(llvm::ArrayRef<unsigned_t> shuffle1,
-                              llvm::ArrayRef<unsigned_t> shuffle2) {
-
-  if (shuffle1.size() != shuffle2.size()) {
-    return false;
-  }
-
-  // Check if the combined masks are the identity mask.
-  for (unsigned i = 0, e = shuffle1.size(); i < e; i++) {
-    unsigned_t idx = shuffle2[shuffle1[i]];
-    if (idx != i) {
-      return false;
-    }
-  }
-  return true;
-}
-
 /// \returns True if the node \p N always evaluates to \p val.
 bool isSplatOfVal(Node *N, float val) {
   SplatNode *Z = dyn_cast<SplatNode>(N);
@@ -448,13 +428,17 @@ static bool sinkCode(Function *F) {
       auto mask2 = TR2->getShuffle();
       assert(mask1.size() == mask2.size() && "Invalid mask size");
 
-      // The two transposes are reversing one another. We can skip both of
-      // them alltogether.
-      if (isIdentityShuffle(mask1, mask2)) {
-        TR1->getResult().replaceAllUsesOfWith(TR2->getInput());
-        changed = true;
-        continue;
+      llvm::SmallVector<unsigned_t, max_tensor_dimensions> newMask;
+      newMask.resize(mask2.size());
+
+      for (size_t i = 0, end = mask2.size(); i < end; i++) {
+        newMask[i] = mask2[mask1[i]];
       }
+
+      auto *newTR = F->createTranspose("tranpose", TR2->getInput(), newMask);
+      TR1->getResult().replaceAllUsesOfWith(newTR->getResult());
+      changed = true;
+      continue;
     }
 
     if (auto *RN = dyn_cast<ReshapeNode>(node)) {

tests/unittests/GraphOptzTest.cpp

@@ -806,30 +806,38 @@ TEST_F(GraphOptz, removeIdentityTransposeWithPredicate) {
   EXPECT_EQ(A->dims(), llvm::makeArrayRef(origDims));
 }
 
-TEST_F(GraphOptz, dontCancelTwoTransposesIfNotMatching) {
+/// Check that consecutive non-inverse transposes are merged
+/// into an equivalent single transpose node.
+TEST_F(GraphOptz, mergeNonInverseTransposes) {
   const size_t origDims[] = {1, 5, 10, 15};
-  const size_t afterFirstTransposeDims[] = {1, 10, 15, 5};
-  const size_t afterSecondTransposeDims[] = {1, 15, 5, 10};
+  const size_t finalDims[] = {5, 1, 15, 10};
+
   Node *A = mod_.createPlaceholder(ElemKind::FloatTy, origDims, "input", false);
-  TransposeNode *T1 = F_->createTranspose("transpose", A, NCHW2NHWC);
-  TransposeNode *T2 = F_->createTranspose("transpose", T1, NCHW2NHWC);
-  SaveNode *save = F_->createSave("ret", T2);
+  TransposeNode *T1 = F_->createTranspose("transpose", A, {0, 3, 2, 1});
+  TransposeNode *T2 = F_->createTranspose("transpose", T1, {0, 2, 3, 1});
+  TransposeNode *T3 = F_->createTranspose("transpose", T2, {1, 0, 3, 2});
+  TransposeNode *T4 = F_->createTranspose("transpose", T3, {3, 1, 2, 0});
 
-  EXPECT_EQ(F_->getNodes().size(), 3);
+  // Intermediate dims after each tranpose
+  // Initial : {1, 5, 10, 15}
+  // After T1: {1, 15, 10, 5}
+  // After T2: {1, 10, 5, 15}
+  // After T3: {10, 1, 15, 5}
+  // After T4: {5, 1, 15, 10}
+
+  SaveNode *save = F_->createSave("ret", T4);
+
+  EXPECT_EQ(F_->getNodes().size(), 5);
 
   ::glow::optimize(F_, CompilationMode::Infer);
 
-  EXPECT_EQ(F_->getNodes().size(), 3);
-  // Make sure the structure of the graph did not change.
-  Node *secondTranspose = save->getInput();
-  EXPECT_EQ(secondTranspose->dims(0),
-            llvm::makeArrayRef(afterSecondTransposeDims));
-  EXPECT_EQ(secondTranspose, T2);
-  Node *firstTranspose = T2->getInput();
-  EXPECT_EQ(firstTranspose, T1);
-  EXPECT_EQ(T1->dims(0), llvm::makeArrayRef(afterFirstTransposeDims));
-  EXPECT_EQ(T1->getInput().getNode(), A);
-  EXPECT_EQ(A->dims(0), llvm::makeArrayRef(origDims));
+  auto *TR = llvm::dyn_cast<TransposeNode>(save->getInput());
+  ASSERT_NE(TR, nullptr);
+
+  EXPECT_EQ(F_->getNodes().size(), 2);
+  EXPECT_EQ(TR->getResult().dims(), llvm::makeArrayRef(finalDims));
+  EXPECT_EQ(A->getNthResult(0).dims(), llvm::makeArrayRef(origDims));
+  EXPECT_EQ(TR->getInput().getNode(), A);
 }
 
 TEST_F(GraphOptz, sinkTransposeBelowArithmeticNodes) {