[mlir] Improve async parallel for tests + fix typos

[ezhulenev]

Summary

Do load and store to verify that we process each element of the iteration space once.

[ezhulenev]

Revert accidental changes

[cota]

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-1d.mlir

80

Please add the same comment as in the commit log here:

// Do load and store to verify that we process each element of the iteration space once.

I wouldn't add it everywhere we do this, but at least once per file should help a reader quickly understand why we're doing this.

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-2d.mlir

74–76

store %5

[ezhulenev]

Address comments

mlir/lib/Dialect/Async/Transforms/AsyncParallelFor.cpp

createParallelComputeFunction(scf::ParallelOp op, PatternRewriter &rewriter) {
   // Remaining arguments are implicit captures of the parallel operation.
   ArrayRef<Value> captures = getArguments(block->getNumArguments() - offset);
 
   // Find one-dimensional iteration bounds: [blockFirstIndex, blockLastIndex]:
   //   blockFirstIndex = blockIndex * blockSize
   Value blockFirstIndex = b.create<arith::MulIOp>(blockIndex, blockSize);
 
   // The last one-dimensional index in the block defined by the `blockIndex`:
+  //   blockLastIndex = min(blockFirstIndex + blockSize, tripCount) - 1
   Value blockEnd0 = b.create<arith::AddIOp>(blockFirstIndex, blockSize);
   Value blockEnd1 = b.create<arith::MinSIOp>(blockEnd0, tripCount);
   Value blockLastIndex = b.create<arith::SubIOp>(blockEnd1, c1);
 
   // Convert one-dimensional indices to multi-dimensional coordinates.
   auto blockFirstCoord = delinearize(b, blockFirstIndex, tripCounts);
   auto blockLastCoord = delinearize(b, blockLastIndex, tripCounts);
 
createParallelComputeFunction(scf::ParallelOp op, PatternRewriter &rewriter) {
   using LoopBodyBuilder =
       std::function<void(OpBuilder &, Location, Value, ValueRange)>;
   using LoopNestBuilder = std::function<LoopBodyBuilder(size_t loopIdx)>;
 
   // Parallel region induction variables computed from the multi-dimensional
   // iteration coordinate using parallel operation bounds and step:
   //
   //   computeBlockInductionVars[loopIdx] =
-  //       lowerBound[loopIdx] + blockCoord[loopIdx] * step[loopDdx]
+  //       lowerBound[loopIdx] + blockCoord[loopIdx] * step[loopIdx]
   SmallVector<Value> computeBlockInductionVars(op.getNumLoops());
 
   // We need to know if we are in the first or last iteration of the
   // multi-dimensional loop for each loop in the nest, so we can decide what
   // loop bounds should we use for the nested loops: bounds defined by compute
   // block interval, or bounds defined by the parallel operation.
   //
   // Example: 2d parallel operation
return [&, loopIdx](OpBuilder &nestedBuilder, Location loc, Value iv,
         isBlockFirstCoord[loopIdx] = nb.create<arith::AndIOp>(
             isBlockFirstCoord[loopIdx], isBlockFirstCoord[loopIdx - 1]);
         isBlockLastCoord[loopIdx] = nb.create<arith::AndIOp>(
             isBlockLastCoord[loopIdx], isBlockLastCoord[loopIdx - 1]);
       }
 
       // Keep building loop nest.
       if (loopIdx < op.getNumLoops() - 1) {
-        // Select nested loop lower/upper bounds depending on out position in
+        // Select nested loop lower/upper bounds depending on our position in
         // the multi-dimensional iteration space.
         auto lb = nb.create<SelectOp>(isBlockFirstCoord[loopIdx],
                                       blockFirstCoord[loopIdx + 1], c0);
 
         auto ub = nb.create<SelectOp>(isBlockLastCoord[loopIdx],
                                       blockEndCoord[loopIdx + 1],
                                       tripCounts[loopIdx + 1]);
 

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-1d.mlir

func @entry() {
   %c3 = arith.constant 3 : index
 
   %lb = arith.constant 0 : index
   %ub = arith.constant 9 : index
 
   %A = memref.alloc() : memref<9xf32>
   %U = memref.cast %A :  memref<9xf32> to memref<*xf32>
 
+  // Initialize memref with zeros because we do load and store to in every test
+  // to verify that we process each element of the iteration space once.
+  scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
+    memref.store %c0, %A[%i] : memref<9xf32>
+  }
+
   // 1. %i = (0) to (9) step (1)
   scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
     %0 = arith.index_cast %i : index to i32
     %1 = arith.sitofp %0 : i32 to f32
-    memref.store %1, %A[%i] : memref<9xf32>
+    %2 = memref.load %A[%i] : memref<9xf32>
+    %3 = arith.addf %1, %2 : f32
+    memref.store %3, %A[%i] : memref<9xf32>
   }
   // CHECK: [0, 1, 2, 3, 4, 5, 6, 7, 8]
   call @print_memref_f32(%U): (memref<*xf32>) -> ()
 
   scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
     memref.store %c0, %A[%i] : memref<9xf32>
   }
 
   // 2. %i = (0) to (9) step (2)
   scf.parallel (%i) = (%lb) to (%ub) step (%c2) {
     %0 = arith.index_cast %i : index to i32
     %1 = arith.sitofp %0 : i32 to f32
-    memref.store %1, %A[%i] : memref<9xf32>
+    %2 = memref.load %A[%i] : memref<9xf32>
+    %3 = arith.addf %1, %2 : f32
+    memref.store %3, %A[%i] : memref<9xf32>
   }
   // CHECK:  [0, 0, 2, 0, 4, 0, 6, 0, 8]
   call @print_memref_f32(%U): (memref<*xf32>) -> ()
 
   scf.parallel (%i) = (%lb) to (%ub) step (%c1) {
     memref.store %c0, %A[%i] : memref<9xf32>
   }
 
   // 3. %i = (-20) to (-11) step (3)
   %lb0 = arith.constant -20 : index
   %ub0 = arith.constant -11 : index
   scf.parallel (%i) = (%lb0) to (%ub0) step (%c3) {
     %0 = arith.index_cast %i : index to i32
     %1 = arith.sitofp %0 : i32 to f32
     %2 = arith.constant 20 : index
     %3 = arith.addi %i, %2 : index
-    memref.store %1, %A[%3] : memref<9xf32>
+    %4 = memref.load %A[%3] : memref<9xf32>
+    %5 = arith.addf %1, %4 : f32
+    memref.store %5, %A[%3] : memref<9xf32>
   }
   // CHECK: [-20, 0, 0, -17, 0, 0, -14, 0, 0]
   call @print_memref_f32(%U): (memref<*xf32>) -> ()
 
   // 4. Check that loop with zero iterations doesn't crash at runtime.
   %lb1 = call @zero(): () -> (index)
   %ub1 = call @zero(): () -> (index)
 

mlir/test/Integration/Dialect/Async/CPU/test-async-parallel-for-2d.mlir

func @entry() {
   %c8 = arith.constant 8 : index
 
   %lb = arith.constant 0 : index
   %ub = arith.constant 8 : index
 
   %A = memref.alloc() : memref<8x8xf32>
   %U = memref.cast %A :  memref<8x8xf32> to memref<*xf32>
 
+  // Initialize memref with zeros because we do load and store to in every test
+  // to verify that we process each element of the iteration space once.
+  scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%c1, %c1) {
+    memref.store %c0, %A[%i, %j] : memref<8x8xf32>
+  }
+
   // 1. (%i, %i) = (0, 8) to (8, 8) step (1, 1)
   scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%c1, %c1) {
     %0 = arith.muli %i, %c8 : index
     %1 = arith.addi %j, %0  : index
     %2 = arith.index_cast %1 : index to i32
     %3 = arith.sitofp %2 : i32 to f32
-    memref.store %3, %A[%i, %j] : memref<8x8xf32>
+    %4 = memref.load %A[%i, %j] : memref<8x8xf32>
+    %5 = arith.addf %3, %4 : f32
+    memref.store %5, %A[%i, %j] : memref<8x8xf32>
   }
 
   // CHECK:      [0, 1, 2, 3, 4, 5, 6, 7]
   // CHECK-NEXT: [8, 9, 10, 11, 12, 13, 14, 15]
   // CHECK-NEXT: [16, 17, 18, 19, 20, 21, 22, 23]
   // CHECK-NEXT: [24, 25, 26, 27, 28, 29, 30, 31]
   // CHECK-NEXT: [32, 33, 34, 35, 36, 37, 38, 39]
   // CHECK-NEXT: [40, 41, 42, 43, 44, 45, 46, 47]
   // CHECK-NEXT: [48, 49, 50, 51, 52, 53, 54, 55]
   // CHECK-NEXT: [56, 57, 58, 59, 60, 61, 62, 63]
   call @print_memref_f32(%U): (memref<*xf32>) -> ()
 
   scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%c1, %c1) {
     memref.store %c0, %A[%i, %j] : memref<8x8xf32>
   }
 
   // 2. (%i, %i) = (0, 8) to (8, 8) step (2, 1)
   scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%c2, %c1) {
     %0 = arith.muli %i, %c8 : index
     %1 = arith.addi %j, %0  : index
     %2 = arith.index_cast %1 : index to i32
     %3 = arith.sitofp %2 : i32 to f32
-    memref.store %3, %A[%i, %j] : memref<8x8xf32>
+    %4 = memref.load %A[%i, %j] : memref<8x8xf32>
+    %5 = arith.addf %3, %4 : f32
+    memref.store %5, %A[%i, %j] : memref<8x8xf32>
   }
 
   // CHECK:      [0, 1, 2, 3, 4, 5, 6, 7]
   // CHECK-NEXT: [0, 0, 0, 0, 0, 0, 0, 0]
   // CHECK-NEXT: [16, 17, 18, 19, 20, 21, 22, 23]
   // CHECK-NEXT: [0, 0, 0, 0, 0, 0, 0, 0]
   // CHECK-NEXT: [32, 33, 34, 35, 36, 37, 38, 39]
   // CHECK-NEXT: [0, 0, 0, 0, 0, 0, 0, 0]
   // CHECK-NEXT: [48, 49, 50, 51, 52, 53, 54, 55]
   // CHECK-NEXT: [0, 0, 0, 0, 0, 0, 0, 0]
   call @print_memref_f32(%U): (memref<*xf32>) -> ()
 
   scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%c1, %c1) {
     memref.store %c0, %A[%i, %j] : memref<8x8xf32>
   }
 
   // 3. (%i, %i) = (0, 8) to (8, 8) step (1, 2)
   scf.parallel (%i, %j) = (%lb, %lb) to (%ub, %ub) step (%c1, %c2) {
     %0 = arith.muli %i, %c8 : index
     %1 = arith.addi %j, %0  : index
     %2 = arith.index_cast %1 : index to i32
     %3 = arith.sitofp %2 : i32 to f32
-    memref.store %3, %A[%i, %j] : memref<8x8xf32>
+    %4 = memref.load %A[%i, %j] : memref<8x8xf32>
+    %5 = arith.addf %3, %4 : f32
+    memref.store %5, %A[%i, %j] : memref<8x8xf32>
   }
 
   // CHECK:      [0, 0, 2, 0, 4, 0, 6, 0]
   // CHECK-NEXT: [8, 0, 10, 0, 12, 0, 14, 0]
   // CHECK-NEXT: [16, 0, 18, 0, 20, 0, 22, 0]
   // CHECK-NEXT: [24, 0, 26, 0, 28, 0, 30, 0]
   // CHECK-NEXT: [32, 0, 34, 0, 36, 0, 38, 0]
   // CHECK-NEXT: [40, 0, 42, 0, 44, 0, 46, 0]