Update normalizeAffineFor to canonicalize maps/operands before using them

[bondhugula]

Summary

Update normalizeAffineFor utility to canonicalize maps and operands
before using them.

[ftynse]

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

31

We don't care that the for is in the next line after the func declaration, so just CHECK should suffice here.

[bondhugula]

CHECK instead of CHECK-NEXT.

[bondhugula]

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

31

Sure.

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

newUbExprs.push_back(
         (origUbExprs[i] - origLbExprs[0]).ceilDiv(origLoopStep));
   }
 
   // Construct newUbMap.
   AffineMap newUbMap =
       AffineMap::get(origLbMap.getNumDims() + origUbMap.getNumDims(),
                      origLbMap.getNumSymbols() + origUbMap.getNumSymbols(),
                      newUbExprs, opBuilder.getContext());
+  canonicalizeMapAndOperands(&newUbMap, &ubOperands);
 
   // Normalize the loop.
   op.setUpperBound(ubOperands, newUbMap);
   op.setLowerBound({}, opBuilder.getConstantAffineMap(0));
   op.setStep(1);
 
   // Calculate the Value of new loopIV. Create affine.apply for the value of
   // the loopIV in normalized loop.
   opBuilder.setInsertionPointToStart(op.getBody());
   SmallVector<Value, 4> lbOperands(lb.getOperands().begin(),
                                    lb.getOperands().begin() +
                                        lb.getMap().getNumDims());
   // Add an extra dim operand for loopIV.
   lbOperands.push_back(op.getInductionVar());
   // Add symbol operands from lower bound.
   for (unsigned j = 0, e = origLbMap.getNumSymbols(); j < e; ++j)
     lbOperands.push_back(lb.getOperand(origLbMap.getNumDims() + j));
 
   AffineExpr origIVExpr = opBuilder.getAffineDimExpr(lb.getMap().getNumDims());
   AffineExpr newIVExpr = origIVExpr * origLoopStep + origLbMap.getResult(0);
   AffineMap ivMap = AffineMap::get(origLbMap.getNumDims() + 1,
                                    origLbMap.getNumSymbols(), newIVExpr);
+  canonicalizeMapAndOperands(&ivMap, &lbOperands);
   Operation *newIV = opBuilder.create<AffineApplyOp>(loc, ivMap, lbOperands);
   op.getInductionVar().replaceAllUsesExcept(newIV->getResult(0), newIV);
   return success();
 }
 
 /// Ensure that all operations that could be executed after `start`
 /// (noninclusive) and prior to `memOp` (e.g. on a control flow/op path
 /// between the operations) do not have the potential memory effect

mlir/test/Dialect/Affine/affine-loop-normalize.mlir

affine.parallel (%k) = (%i) to (%j) {
       affine.store %cst, %0[%i, %k] : memref<2x4xf32>
     }
   }
   return
 }
 
 // -----
 
+// CHECK-LABEL: func @relative_bounds
+func @relative_bounds(%arg: index) {
+  // CHECK: affine.for %{{.*}} = 0 to 4
+  affine.for %i = affine_map<(d0) -> (d0)>(%arg) to affine_map<(d0) -> (d0 + 4)>(%arg) {
+  }
+  return
+}
+
+// -----
+
 // Check that single iteration loop is removed and its body is promoted to the
 // parent block.
 
 // CHECK-LABEL: func @single_iteration_loop
 func @single_iteration_loop(%in: memref<1xf32>, %out: memref<1xf32>) {
   affine.for %i = 0 to 1 {
     %1 = affine.load %in[%i] : memref<1xf32>
     affine.store %1, %out[%i] : memref<1xf32>
func @loop_with_unknown_upper_bound(%arg0: memref<?x?xf32>, %arg1: index) {
   return
 }
 
 // -----
 
 // CHECK-DAG: [[$OUTERIV:#map[0-9]+]] = affine_map<(d0) -> (d0 * 32 + 2)>
 // CHECK-DAG: [[$INNERIV:#map[0-9]+]] = affine_map<(d0) -> (d0 + 2)>
 // CHECK-DAG: [[$OUTERUB:#map[0-9]+]] = affine_map<()[s0] -> ((s0 - 2) ceildiv 32)>
-// CHECK-DAG: [[$INNERUB:#map[0-9]+]] = affine_map<(d0) -> (d0 - 2, 510)>
+// CHECK-DAG: [[$INNERUB:#map[0-9]+]] = affine_map<()[s0] -> (s0 - 2, 510)>
 
 // CHECK-LABEL: func @loop_with_multiple_upper_bounds
 // CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index)
 // CHECK-NEXT:  %{{.*}} = arith.constant 0 : index
 // CHECK-NEXT:  %[[DIM:.*]] = memref.dim %arg0, %c0 : memref<?x?xf32>
 // CHECK-NEXT:   affine.for %[[I:.*]] = 0 to [[$OUTERUB]]()[%[[DIM]]] {
 // CHECK-NEXT:     %[[IIV:.*]] = affine.apply [[$OUTERIV]](%[[I]])
-// CHECK-NEXT:     affine.for %[[II:.*]] = 0 to min [[$INNERUB]](%[[ARG1]]) {
+// CHECK-NEXT:     affine.for %[[II:.*]] = 0 to min [[$INNERUB]]()[%[[ARG1]]] {
 // CHECK-NEXT:       %[[IIIV:.*]] = affine.apply [[$INNERIV]](%[[II]])
 // CHECK-NEXT:       "test.foo"(%[[IIV]], %[[IIIV]])
 // CHECK-NEXT:     }
 // CHECK-NEXT:   }
 // CHECK-NEXT:   return
 // CHECK-NEXT: }
 func @loop_with_multiple_upper_bounds(%arg0: memref<?x?xf32>, %arg1 : index) {
   %c0 = arith.constant 0 : index
   %0 = memref.dim %arg0, %c0 : memref<?x?xf32>
   affine.for %i0 = 2 to %0 step 32{
     affine.for %i1 = 2 to min affine_map<(d0)[] -> (d0, 512)>(%arg1) {
       "test.foo"(%i0, %i1) : (index, index) -> ()
     }
   }
   return
 }
 
 // -----
 
 // CHECK-DAG: [[$INTERUB:#map[0-9]+]] = affine_map<()[s0] -> (s0 ceildiv 32)>
 // CHECK-DAG: [[$INTERIV:#map[0-9]+]] = affine_map<(d0) -> (d0 * 32)>
-// CHECK-DAG: [[$INTRAUB:#map[0-9]+]] = affine_map<(d0, d1)[s0] -> (32, -d0 + s0)>
+// CHECK-DAG: [[$INTRAUB:#map[0-9]+]] = affine_map<(d0)[s0] -> (32, -d0 + s0)>
 // CHECK-DAG: [[$INTRAIV:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + d0)>
 
 // CHECK-LABEL: func @tiled_matmul
 // CHECK-SAME: (%[[ARG0:.*]]: memref<1024x1024xf32>, %[[ARG1:.*]]: memref<1024x1024xf32>, %[[ARG2:.*]]: memref<1024x1024xf32>)
 // CHECK-NEXT:    %{{.*}} = arith.constant 0 : index
 // CHECK-NEXT:    %{{.*}} = arith.constant 1 : index
 // CHECK-NEXT:    %[[DIM0:.*]] = memref.dim %[[ARG0]], %{{.*}}
 // CHECK-NEXT:    %[[DIM1:.*]] = memref.dim %[[ARG1]], %{{.*}}
 // CHECK-NEXT:    %[[DIM2:.*]] = memref.dim %[[ARG0]], %{{.*}}
 // CHECK-NEXT:    affine.for %[[I:.*]] = 0 to [[$INTERUB]]()[%[[DIM0]]] {
 // CHECK-NEXT:      %[[IIV:.*]] = affine.apply [[$INTERIV]](%[[I]])
 // CHECK-NEXT:      affine.for %[[J:.*]] = 0 to [[$INTERUB]]()[%[[DIM1]]] {
 // CHECK-NEXT:        %[[JIV:.*]] = affine.apply [[$INTERIV]](%[[J]])
 // CHECK-NEXT:        affine.for %[[K:.*]] = 0 to [[$INTERUB]]()[%[[DIM2]]] {
 // CHECK-NEXT:          %[[KIV:.*]] = affine.apply [[$INTERIV]](%[[K]])
-// CHECK-NEXT:          affine.for %[[II:.*]] = 0 to min [[$INTRAUB]](%[[IIV]], %[[IIV]])[%[[DIM0]]] {
+// CHECK-NEXT:          affine.for %[[II:.*]] = 0 to min [[$INTRAUB]](%[[IIV]])[%[[DIM0]]] {
 // CHECK-NEXT:            %[[IIIV:.*]] = affine.apply [[$INTRAIV]](%[[IIV]], %[[II]])
-// CHECK-NEXT:            affine.for %[[JJ:.*]] = 0 to min [[$INTRAUB]](%[[JIV]], %[[JIV]])[%[[DIM1]]] {
+// CHECK-NEXT:            affine.for %[[JJ:.*]] = 0 to min [[$INTRAUB]](%[[JIV]])[%[[DIM1]]] {
 // CHECK-NEXT:              %[[JJIV:.*]] = affine.apply [[$INTRAIV]](%[[JIV]], %[[JJ]])
-// CHECK-NEXT:              affine.for %[[KK:.*]] = 0 to min [[$INTRAUB]](%[[KIV]], %[[KIV]])[%[[DIM2]]] {
+// CHECK-NEXT:              affine.for %[[KK:.*]] = 0 to min [[$INTRAUB]](%[[KIV]])[%[[DIM2]]] {
 // CHECK-NEXT:                %[[KKIV:.*]] = affine.apply [[$INTRAIV]](%[[KIV]], %[[KK]])
 // CHECK-NEXT:                %{{.*}} = affine.load %[[ARG0]][%[[IIIV]], %[[KKIV]]] : memref<1024x1024xf32>
 // CHECK-NEXT:                %{{.*}} = affine.load %[[ARG1]][%[[KKIV]], %[[JJIV]]] : memref<1024x1024xf32>
 // CHECK-NEXT:                %{{.*}} = affine.load %[[ARG2]][%[[IIIV]], %[[JJIV]]] : memref<1024x1024xf32>
 // CHECK-NEXT:                %{{.*}} = arith.mulf %9, %10 : f32
 // CHECK-NEXT:                %{{.*}} = arith.addf %11, %12 : f32
 // CHECK-NEXT:                affine.store %{{.*}}, %[[ARG2]][%6, %7] : memref<1024x1024xf32>
 // CHECK-NEXT:              }