[mlir] fix off-by-one error in collapseParallelLoops

[gysit]

Summary

The patch fixes an off by one error in the method collapseParallelLoops. It ensures the same normalized bound is used for the computation of the division and the remainder.

[herhut]

Thanks for finding and fixing this!

mlir/lib/Transforms/Utils/LoopUtils.cpp

auto newPloop = outsideBuilder.create<scf::ParallelOp>(
             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 + 1]);
+                loc, previous, normalizedUpperBounds[idx]);
           }
 
           // The final induction value is just the remaining value.
           unsigned idx = combinedDimensions[i][numberCombinedDimensions - 1];
           replaceAllUsesInRegionWith(loops.getBody()->getArgument(idx),
                                      previous, loops.region());
         }
       });

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

func @parallel_many_dims() {
   }
   return
 }
 
 // CHECK-LABEL: func @parallel_many_dims() {
 // CHECK:         [[C6:%.*]] = constant 6 : index
 // CHECK:         [[C9:%.*]] = constant 9 : index
 // CHECK:         [[C10:%.*]] = constant 10 : index
-// CHECK:         [[C12:%.*]] = constant 12 : 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:           [[V18:%.*]] = muli [[NEW_I0]], [[C10]] : index
-// CHECK:           [[I3:%.*]] = addi [[V18]], [[C9]] : 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

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

func @collapse_to_single() {
   return
 }
 
 // CHECK-LABEL: func @collapse_to_single() {
 // CHECK:         [[C7:%.*]] = constant 7 : 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]], [[C6]] : index
-// CHECK:           [[VAL_10:%.*]] = muli [[I1_COUNT]], [[C4]] : index
-// CHECK:           [[I1:%.*]] = addi [[VAL_10]], [[C7]] : index
-// CHECK:           [[VAL_12:%.*]] = muli [[I0_COUNT]], [[C3]] : index
-// CHECK:           [[I0:%.*]] = addi [[VAL_12]], [[C3]] : index
+// CHECK:           [[I1_COUNT:%.*]] = divi_signed [[NEW_I]], [[C3]] : index
+// CHECK:           [[V0:%.*]] = muli [[I1_COUNT]], [[C4]] : index
+// CHECK:           [[I1:%.*]] = addi [[V0]], [[C7]] : index
+// CHECK:           [[V1:%.*]] = muli [[I0_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