diff --git a/iree/compiler/Conversion/Common/Transforms.cpp b/iree/compiler/Conversion/Common/Transforms.cpp index 55e578a31589..3637b4091e7e 100644 --- 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(); if (!ivArg) continue; Operation *containingOp = ivArg.getOwner()->getParentOp(); - forOp = dyn_cast_or_null(containingOp); + auto forOp = dyn_cast_or_null(containingOp); if (forOp && forOp.getInductionVar() == dim) { iv = dim; + ub = forOp.upperBound(); + lb = forOp.lowerBound(); + step = forOp.step(); break; } + auto parallelOp = dyn_cast_or_null(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 forOperands = {forOp.lowerBound(), forOp.step(), - forOp.upperBound()}; - return llvm::all_of(forOperands, + std::array 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 index 975eee557e0c..a829ea310bf7 100644 --- 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, %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 + 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 + } + } + return }