Extend AffineMin canonicalization to support scf.parallel (#5289)

iree-org · Apr 2, 2021 · c61fefe · c61fefe
1 parent e0ee3f3
commit c61fefe
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 8 deletions.
diff --git a/iree/compiler/Conversion/Common/Transforms.cpp b/iree/compiler/Conversion/Common/Transforms.cpp
@@ -457,23 +457,38 @@ static bool affineMinOpDivisible(AffineMinOp minOp, int64_t dividend) {
   if (!minOp.getSymbolOperands().empty() ||
       minOp.getAffineMap().getNumResults() != 2)
     return {};
-
-  scf::ForOp forOp;
   Value iv;
-  // Check if any of the dimensions are a ForOp induction variable.
+  Value ub;
+  Value lb;
+  Value step;
+  // Check if any of the dimensions is a ForOp or ParallelOp induction variable.
   for (auto dim : minOp.getDimOperands()) {
     auto ivArg = dim.dyn_cast<BlockArgument>();
     if (!ivArg) continue;
     Operation *containingOp = ivArg.getOwner()->getParentOp();
-    forOp = dyn_cast_or_null<scf::ForOp>(containingOp);
+    auto forOp = dyn_cast_or_null<scf::ForOp>(containingOp);
     if (forOp && forOp.getInductionVar() == dim) {
       iv = dim;
+      ub = forOp.upperBound();
+      lb = forOp.lowerBound();
+      step = forOp.step();
       break;
     }
+    auto parallelOp = dyn_cast_or_null<scf::ParallelOp>(containingOp);
+    if (!parallelOp) continue;
+    for (auto inductionVar : llvm::enumerate(parallelOp.getInductionVars())) {
+      if (inductionVar.value() == dim) {
+        iv = dim;
+        ub = parallelOp.upperBound()[inductionVar.index()];
+        lb = parallelOp.lowerBound()[inductionVar.index()];
+        step = parallelOp.step()[inductionVar.index()];
+        break;
+      }
+    }
+    if (iv) break;
   }
   if (!iv) return false;
   // Calculate the affine map representing `%ub - %iv`.
-  Value ub = forOp.upperBound();
   AffineExpr ivDim;
   AffineExpr ubDim;
   for (auto dim : llvm::enumerate(minOp.getDimOperands())) {
@@ -504,9 +519,8 @@ static bool affineMinOpDivisible(AffineMinOp minOp, int64_t dividend) {
   // Now check that for every value of the induction variable `%ub - %iv` is
   // divisible by `dividend`. It is true if the lower bounder, the upper bound
   // and the step are all divisible by `dividend`.
-  std::array<Value, 3> forOperands = {forOp.lowerBound(), forOp.step(),
-                                      forOp.upperBound()};
-  return llvm::all_of(forOperands,
+  std::array<Value, 3> loopOperands = {lb, step, ub};
+  return llvm::all_of(loopOperands,
                       [dividend](Value v) { return isDivisible(v, dividend); });
 }
 

diff --git a/iree/compiler/Conversion/Common/test/affinemin_canonicalization.mlir b/iree/compiler/Conversion/Common/test/affinemin_canonicalization.mlir
@@ -45,5 +45,23 @@ func @scf_for_distributed(%A : memref<i64>, %id1 : index, %count1 : index,
     }
   }
 
+  // Same case but using scf.parallel.
+  //      CHECK: scf.for
+  //      CHECK:   %[[MIN:.*]] = affine.min
+  //      CHECK:   scf.parallel {{.*}} to (%[[MIN]])
+  // CHECK-NEXT:     %[[C4:.*]] = constant 4 : index
+  // CHECK-NEXT:     %[[C4I64:.*]] = index_cast %[[C4:.*]]
+  // CHECK-NEXT:     memref.store %[[C4I64]], %{{.*}}[] : memref<i64>
+  scf.for %arg0 = %0 to %c1020 step %1 {
+    %2 = affine.min affine_map<(d0) -> (32, -d0 + 1020)>(%arg0)
+    %3 = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%id2]
+    %4 = affine.apply affine_map<()[s0] -> (s0 * 4)>()[%count2]
+    scf.parallel (%arg1) = (%3) to (%2) step (%4) {
+      %5 = affine.min affine_map<(d0, d1) -> (4, d0 - d1)>(%2, %arg1)
+      %6 = index_cast %5: index to i64
+      memref.store %6, %A[]: memref<i64>
+    }
+  }
+
   return
 }