Fix VAE regression

Signed-off-by: MaheshRavishankar <[email protected]>

compiler/src/iree/compiler/DispatchCreation/FormDispatchRegions.cpp

@@ -410,10 +410,10 @@ static bool areOpsFusableAfterInterchangeOfConsumer(
     return false;
   }
 
-  // The input map must be a permutation (i.e. this is not a broadcasting
-  // access), but not identity.
+  // If input map is identity already, nothing to do here (and we dont care
+  // about maps that are not projected permutations)
   AffineMap inputMap = genericOp.getMatchingIndexingMap(&fusableOperand);
-  if (!inputMap.isPermutation() || inputMap.isIdentity()) {
+  if (!inputMap.isProjectedPermutation() || inputMap.isIdentity()) {
     return false;
   }
 
@@ -429,6 +429,12 @@ static bool areOpsFusableAfterInterchangeOfConsumer(
       llvm::map_to_vector(inputMap.getResults(), [](AffineExpr e) -> unsigned {
         return cast<AffineDimExpr>(e).getPosition();
       });
+  llvm::SmallDenseSet<unsigned int> permSet(perm.begin(), perm.end());
+  for (auto origDims : llvm::seq(0u, inputMap.getNumDims())) {
+    if (!permSet.contains(origDims)) {
+      perm.push_back(origDims);
+    }
+  }
   IRRewriter rewriter(consumer->getContext());
   FailureOr<linalg::GenericOp> interchangedOp =
       linalg::interchangeGenericOp(rewriter, genericOp, perm);