[mlir] Preserve lexicographic order after loop collapsing.

[pifon2a]

Summary

Currently, for a scf.parallel (i,j,k) after the loop collapsing to 1D is done, the
IVs would be traversed as for an scf.parallel(k,j,i).

[bkramer]

looks good

[pifon2a]

Fix single-parallel-loop-collapsing.mlir.test

mlir/lib/Transforms/Utils/LoopUtils.cpp

void mlir::coalesceLoops(MutableArrayRef<scf::ForOp> loops) {
   second.erase();
 }
 
 void mlir::collapseParallelLoops(
     scf::ParallelOp loops, ArrayRef<std::vector<unsigned>> combinedDimensions) {
   OpBuilder outsideBuilder(loops);
   Location loc = loops.getLoc();
 
+  // Presort combined dimensions.
+  auto sortedDimensions = llvm::to_vector<3>(combinedDimensions);
+  for (auto &dims : sortedDimensions)
+    std::sort(dims.begin(), dims.end());
+
   // Normalize ParallelOp's iteration pattern.
-  SmallVector<Value, 3> normalizedLowerBounds;
-  SmallVector<Value, 3> normalizedSteps;
-  SmallVector<Value, 3> normalizedUpperBounds;
+  SmallVector<Value, 3> normalizedLowerBounds, normalizedSteps,
+      normalizedUpperBounds;
   for (unsigned i = 0, e = loops.getNumLoops(); i < e; ++i) {
     OpBuilder insideLoopBuilder = OpBuilder::atBlockBegin(loops.getBody());
     auto resultBounds =
         normalizeLoop(outsideBuilder, insideLoopBuilder, loc,
                       loops.lowerBound()[i], loops.upperBound()[i],
                       loops.step()[i], loops.getBody()->getArgument(i));
 
     normalizedLowerBounds.push_back(resultBounds.lowerBound);
     normalizedUpperBounds.push_back(resultBounds.upperBound);
     normalizedSteps.push_back(resultBounds.step);
   }
 
   // Combine iteration spaces.
-  SmallVector<Value, 3> lowerBounds;
-  SmallVector<Value, 3> steps;
-  SmallVector<Value, 3> upperBounds;
+  SmallVector<Value, 3> lowerBounds, upperBounds, steps;
   auto cst0 = outsideBuilder.create<ConstantIndexOp>(loc, 0);
   auto cst1 = outsideBuilder.create<ConstantIndexOp>(loc, 1);
-  for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
+  for (unsigned i = 0, e = sortedDimensions.size(); i < e; ++i) {
     Value newUpperBound = outsideBuilder.create<ConstantIndexOp>(loc, 1);
-    for (auto idx : combinedDimensions[i]) {
+    for (auto idx : sortedDimensions[i]) {
       newUpperBound = outsideBuilder.create<MulIOp>(loc, newUpperBound,
                                                     normalizedUpperBounds[idx]);
     }
     lowerBounds.push_back(cst0);
     steps.push_back(cst1);
     upperBounds.push_back(newUpperBound);
   }
 
   // Create new ParallelLoop with conversions to the original induction values.
   // The loop below uses divisions to get the relevant range of values in the
   // new induction value that represent each range of the original induction
   // value. The remainders then determine based on that range, which iteration
   // of the original induction value this represents. This is a normalized value
   // that is un-normalized already by the previous logic.
   auto newPloop = outsideBuilder.create<scf::ParallelOp>(
       loc, lowerBounds, upperBounds, steps,
       [&](OpBuilder &insideBuilder, Location, ValueRange ploopIVs) {
         for (unsigned i = 0, e = combinedDimensions.size(); i < e; ++i) {
           Value previous = ploopIVs[i];
           unsigned numberCombinedDimensions = combinedDimensions[i].size();
           // Iterate over all except the last induction value.
-          for (unsigned j = 0, e = numberCombinedDimensions - 1; j < e; ++j) {
+          for (unsigned j = numberCombinedDimensions - 1; j > 0; --j) {
             unsigned idx = combinedDimensions[i][j];
 
             // Determine the current induction value's current loop iteration
             Value iv = insideBuilder.create<SignedRemIOp>(
                 loc, previous, normalizedUpperBounds[idx]);
             replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx), iv,
                                        loops.region());
 
             // Remove the effect of the current induction value to prepare for
             // the next value.
             previous = insideBuilder.create<SignedDivIOp>(
                 loc, previous, normalizedUpperBounds[idx]);
           }
 
           // The final induction value is just the remaining value.
-          unsigned idx = combinedDimensions[i][numberCombinedDimensions - 1];
+          unsigned idx = combinedDimensions[i][0];
           replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx),
                                      previous, loops.region());
         }
       });
 
   // Replace the old loop with the new loop.
   loops.getBody()->back().erase();
   newPloop.getBody()->getOperations().splice(

mlir/test/Transforms/parallel-loop-collapsing.mlir

 // RUN: mlir-opt -allow-unregistered-dialect %s -pass-pipeline='func(parallel-loop-collapsing{collapsed-indices-0=0,3 collapsed-indices-1=1,4 collapsed-indices-2=2}, canonicalize)' | FileCheck %s
 
+// CHECK-LABEL: func @parallel_many_dims() {
 func @parallel_many_dims() {
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   %c4 = constant 4 : index
   %c5 = constant 5 : index
   %c6 = constant 6 : index
   %c7 = constant 7 : index
   %c8 = constant 8 : index
   %c9 = constant 9 : index
   %c10 = constant 10 : index
   %c11 = constant 11 : index
   %c12 = constant 12 : index
   %c13 = constant 13 : index
   %c14 = constant 14 : index
   %c15 = constant 15 : index
   %c26 = constant 26 : index
 
-  scf.parallel (%i0, %i1, %i2, %i3, %i4) = (%c0, %c3, %c6, %c9, %c12) to (%c2, %c5, %c8, %c11, %c14)
-                                          step (%c1, %c4, %c7, %c10, %c13) {
-    %result = "magic.op"(%i0, %i1, %i2, %i3, %i4): (index, index, index, index, index) -> index
+  scf.parallel (%i0, %i1, %i2, %i3, %i4) = (%c0, %c3, %c6, %c9, %c12)
+    to (%c2, %c5, %c8, %c26, %c14) step (%c1, %c4, %c7, %c10, %c13) {
+    %result = "magic.op"(%i0, %i1, %i2, %i3, %i4)
+        : (index, index, index, index, index) -> index
   }
   return
 }
 
-// CHECK-LABEL: func @parallel_many_dims() {
-// CHECK:         [[C6:%.*]] = constant 6 : index
-// CHECK:         [[C9:%.*]] = constant 9 : index
-// CHECK:         [[C10:%.*]] = constant 10 : index
-// CHECK:         [[C0:%.*]] = constant 0 : index
-// CHECK:         [[C1:%.*]] = constant 1 : index
-// CHECK:         [[C2:%.*]] = constant 2 : index
-// CHECK:         [[C3:%.*]] = constant 3 : index
-// CHECK:         [[C12:%.*]] = constant 12 : index
-// CHECK:         scf.parallel ([[NEW_I0:%.*]]) = ([[C0]]) to ([[C2]]) step ([[C1]]) {
-// CHECK:           [[I0:%.*]] = remi_signed [[NEW_I0]], [[C2]] : index
-// CHECK:           [[V0:%.*]] = divi_signed [[NEW_I0]], [[C2]] : index
-// CHECK:           [[V2:%.*]] = muli [[V0]], [[C10]] : index
-// CHECK:           [[I3:%.*]] = addi [[V2]], [[C9]] : index
-// CHECK:           "magic.op"([[I0]], [[C3]], [[C6]], [[I3]], [[C12]]) : (index, index, index, index, index) -> index
-// CHECK:           scf.yield
-// CHECK-NEXT:    }
-// CHECK-NEXT:    return
+// CHECK: [[C3:%.*]] = constant 3 : index
+// CHECK: [[C6:%.*]] = constant 6 : index
+// CHECK: [[C9:%.*]] = constant 9 : index
+// CHECK: [[C10:%.*]] = constant 10 : index
+// CHECK: [[C4:%.*]] = constant 4 : index
+// CHECK: [[C12:%.*]] = constant 12 : index
+// CHECK: [[C0:%.*]] = constant 0 : index
+// CHECK: [[C1:%.*]] = constant 1 : index
+// CHECK: [[C2:%.*]] = constant 2 : index
+// CHECK: scf.parallel ([[NEW_I0:%.*]]) = ([[C0]]) to ([[C4]]) step ([[C1]]) {
+// CHECK:   [[V0:%.*]] = remi_signed [[NEW_I0]], [[C2]] : index
+// CHECK:   [[I0:%.*]] = divi_signed [[NEW_I0]], [[C2]] : index
+// CHECK:   [[V2:%.*]] = muli [[V0]], [[C10]] : index
+// CHECK:   [[I3:%.*]] = addi [[V2]], [[C9]] : index
+// CHECK:   "magic.op"([[I0]], [[C3]], [[C6]], [[I3]], [[C12]]) : (index, index, index, index, index) -> index
+// CHECK:   scf.yield

mlir/test/Transforms/single-parallel-loop-collapsing.mlir

func @collapse_to_single() {
   scf.parallel (%i0, %i1) = (%c0, %c1) to (%c2, %c3) step (%c4, %c5) {
     %result = "magic.op"(%i0, %i1): (index, index) -> index
   }
   return
 }
 
 // CHECK-LABEL: func @collapse_to_single() {
 // CHECK:         [[C7:%.*]] = constant 7 : index
+// CHECK:         [[C3:%.*]] = constant 3 : index
 // CHECK:         [[C4:%.*]] = constant 4 : index
 // CHECK:         [[C18:%.*]] = constant 18 : index
-// CHECK:         [[C3:%.*]] = constant 3 : index
+// CHECK:         [[C6:%.*]] = constant 6 : index
 // CHECK:         [[C0:%.*]] = constant 0 : index
 // CHECK:         [[C1:%.*]] = constant 1 : index
 // CHECK:         scf.parallel ([[NEW_I:%.*]]) = ([[C0]]) to ([[C18]]) step ([[C1]]) {
-// CHECK:           [[I0_COUNT:%.*]] = remi_signed [[NEW_I]], [[C3]] : index
-// CHECK:           [[I1_COUNT:%.*]] = divi_signed [[NEW_I]], [[C3]] : index
-// CHECK:           [[V0:%.*]] = muli [[I1_COUNT]], [[C4]] : index
+// CHECK:           [[I0_COUNT:%.*]] = remi_signed [[NEW_I]], [[C6]] : index
+// CHECK:           [[I1_COUNT:%.*]] = divi_signed [[NEW_I]], [[C6]] : index
+// CHECK:           [[V0:%.*]] = muli [[I0_COUNT]], [[C4]] : index
 // CHECK:           [[I1:%.*]] = addi [[V0]], [[C7]] : index
-// CHECK:           [[V1:%.*]] = muli [[I0_COUNT]], [[C3]] : index
+// CHECK:           [[V1:%.*]] = muli [[I1_COUNT]], [[C3]] : index
 // CHECK:           [[I0:%.*]] = addi [[V1]], [[C3]] : index
 // CHECK:           "magic.op"([[I0]], [[I1]]) : (index, index) -> index
 // CHECK:           scf.yield
 // CHECK-NEXT:    }
 // CHECK-NEXT:    return