[mlir] properly support min/max in affine parallelization

[ftynse]

Summary

The existing implementation of the affine parallelization silently copies over
the lower and upper bound maps from affine.for to affine.parallel. However, the
semantics of these maps differ between these two ops: in affine.for, a max(min)
of results is taken for the lower(upper) bound; in affine.parallel, multiple
induction variables can be defined an each result corresponds to one induction
variable. Thus the existing implementation could generate invalid IR or IR that
passes the verifier but has different semantics than the original code. Fix the
parallelization utility to emit dedicated min/max operations before the
affine.parallel in such cases. Disallow parallelization if min/max would have
been in an operation without the AffineScope trait, e.g., in another loop,
since the result of these operations is not considered a valid affine dimension
identifier and may not be properly handled by the affine analyses.

[wsmoses]

mlir/lib/Dialect/Affine/Utils/Utils.cpp

145

While seeming correct, I am a bit worried that checking for num results > 1 here won't be equivalent to checking if there's a max (or min) in the future if AffineForOp is extended, resulting in this code getting overlooked.

[ftynse]

mlir/lib/Dialect/Affine/Utils/Utils.cpp

145

This has been a part of the op's semantics since it was created - https://mlir.llvm.org/docs/Dialects/Affine/#affinefor-affineforop - there's no other way of checking for this (well, unless you want me to introduce the condition in the op itself and update the existing users).

[wsmoses]

mlir/lib/Dialect/Affine/Utils/Utils.cpp

145

I think that may be useful for clarity to add a helper wrapper to AffineForOp, but not necessary for this patch.

mlir/lib/Dialect/Affine/Utils/Utils.cpp

static AffineIfOp hoistAffineIfOp(AffineIfOp ifOp, Operation *hoistOverOp) {
   return hoistedIfOp;
 }
 
 /// Replace affine.for with a 1-d affine.parallel and clone the former's body
 /// into the latter while remapping values.
 void mlir::affineParallelize(AffineForOp forOp) {
   Location loc = forOp.getLoc();
   OpBuilder outsideBuilder(forOp);
+
+  // If a loop has a 'max' in the lower bound, emit it outside the parallel loop
+  // as it does not have implicit 'max' behavior.
+  AffineMap lowerBoundMap = forOp.getLowerBoundMap();
+  ValueRange lowerBoundOperands = forOp.getLowerBoundOperands();
+  AffineMap upperBoundMap = forOp.getUpperBoundMap();
+  ValueRange upperBoundOperands = forOp.getUpperBoundOperands();
+
+  bool needsMax = lowerBoundMap.getNumResults() > 1;
+  bool needsMin = upperBoundMap.getNumResults() > 1;
+  AffineMap identityMap;
+  if (needsMax || needsMin) {
+    if (forOp->getParentOp() &&
+        !forOp->getParentOp()->hasTrait<OpTrait::AffineScope>())
+      return;
+
+    identityMap = AffineMap::getMultiDimIdentityMap(1, loc->getContext());
+  }
+  if (needsMax) {
+    auto maxOp = outsideBuilder.create<AffineMaxOp>(loc, lowerBoundMap,
+                                                    lowerBoundOperands);
+    lowerBoundMap = identityMap;
+    lowerBoundOperands = maxOp->getResults();
+  }
+
+  // Same for the upper bound.
+  if (needsMin) {
+    auto minOp = outsideBuilder.create<AffineMinOp>(loc, upperBoundMap,
+                                                    upperBoundOperands);
+    upperBoundMap = identityMap;
+    upperBoundOperands = minOp->getResults();
+  }
+
   // Creating empty 1-D affine.parallel op.
   AffineParallelOp newPloop = outsideBuilder.create<AffineParallelOp>(
-      loc, llvm::None, llvm::None, forOp.getLowerBoundMap(),
-      forOp.getLowerBoundOperands(), forOp.getUpperBoundMap(),
-      forOp.getUpperBoundOperands());
+      loc, llvm::None, llvm::None, lowerBoundMap, lowerBoundOperands,
+      upperBoundMap, upperBoundOperands);
   // Steal the body of the old affine for op and erase it.
   newPloop.region().takeBody(forOp.region());
   forOp.erase();
 }
 
 // Returns success if any hoisting happened.
 LogicalResult mlir::hoistAffineIfOp(AffineIfOp ifOp, bool *folded) {
   // Bail out early if the ifOp returns a result.  TODO: Consider how to

mlir/test/Dialect/Affine/parallelize.mlir

 func @non_affine_load() {
   %0 = alloc() : memref<100 x f32>
   affine.for %i = 0 to 100 {
 // CHECK:  affine.for %{{.*}} = 0 to 100 {
     load %0[%i] : memref<100 x f32>
   }
   return
 }
+
+// CHECK-LABEL: for_with_minmax
+func @for_with_minmax(%m: memref<?xf32>, %lb0: index, %lb1: index,
+                      %ub0: index, %ub1: index) {
+  // CHECK: %[[lb:.*]] = affine.max
+  // CHECK: %[[ub:.*]] = affine.min
+  // CHECK: affine.parallel (%{{.*}}) = (%[[lb]]) to (%[[ub]])
+  affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %lb1)
+          to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
+    affine.load %m[%i] : memref<?xf32>
+  }
+  return
+}
+
+// CHECK-LABEL: nested_for_with_minmax
+func @nested_for_with_minmax(%m: memref<?xf32>, %lb0: index,
+                             %ub0: index, %ub1: index) {
+  // CHECK: affine.parallel
+  affine.for %j = 0 to 10 {
+    // Cannot parallelize the inner loop because we would need to compute
+    // affine.max for its lower bound inside the loop, and that is not (yet)
+    // considered as a valid affine dimension.
+    // CHECK: affine.for
+    affine.for %i = max affine_map<(d0, d1) -> (d0, d1)>(%lb0, %j)
+            to min affine_map<(d0, d1) -> (d0, d1)>(%ub0, %ub1) {
+      affine.load %m[%i] : memref<?xf32>
+    }
+  }
+  return
+}