[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

@@ -448,13 +448,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) {