[mlir] fix wrong symbol order in AffineApplyNormalizer

[ftynse]

Summary

AffineApplyNormalizer provides common logic for folding affine maps that appear
in affine.apply into other affine operations that use the result of said
affine.apply. In the process, affine maps of both operations are composed.
During the composition A.compose(B) the symbols from the map A are placed
before those of the map B in a single concatenated symbol list. However,
AffineApplyNormalizer was ordering the operands of the operation being
normalized by iteratively appending the symbols into a single list accoridng to
the operand order, regardless of whether these operands are symbols of the
current operation or of the map that is being folded into it. This could lead
to wrong order of symbols and, when the symbols were bound to constant values,
to visibly incorrect folding of constants into affine maps as reported in
PR45031. Make sure symbols operands to the current operation are always placed
before symbols coming from the folded maps.

Update the test that was exercising the incorrect folder behavior. For some
reason, the order of symbol operands was swapped in the test input compared to
the previous operations, making it easy to assume the correct maps were
produced whereas they were swapping the symbols back due to the problem
described above.

Closes https://bugs.llvm.org/show_bug.cgi?id=45031

[nicolasvasilache]

Thanks for tracking this down Alex!

Note that a lot of the complexity of AffineApplyNormalizer comes from its supporting chain of affine applies and I strongly advocated for being unnecessary cognitive overhead.
We had a window of opportunity to clean all this up about 1 year ago but it got a lot of pushback for very bad reasons (e.g. a few tests with single affine_apply and composed maps would have been more complex).
Are we now in a position were we could say that chains of ops with affine map composition are outlawed so we could simplify this logic?

[nicolasvasilache]

Well actually the example you added is affine.min(affine.apply()) and goes against what I wrote above and would ideally have been composed in the first place.

I'd recommend we start putting verifiers that disallow such chains everywhere and simplify the logic, but I have not followed deeply what new assumptions have crept inside affine.

[ftynse]

I only replicated the IR that exhibited the buggy behavior. Same problem would have happened with affine.load(affine.apply) or any other composition. Same problem would have also happened if parts of this code were used for folding-on-construction.

[herhut]

Thanks for fixing this!

Nit: accoridng in the CL description.

Regarding Nicholas' comment: By such chains, you mean chains of affine.apply? Or also the cases of an affine.min on an affine.apply?

[bondhugula]

In D75247#1895307, @ftynse wrote:

I only replicated the IR that exhibited the buggy behavior. Same problem would have happened with affine.load(affine.apply) or any other composition. Same problem would have also happened if parts of this code were used for folding-on-construction.

+1 to @ftynse. Absolutely! One needs in-memory composition and compose/fold in by construction - both for creating IR and for analysis. Making and pushing IR design decisions because there is buggy code based on incorrect assumptions (along with matching incorrect test cases) sounds hilarious to me! :-) And in any case, the fundamental need for compositions of maps with operands isn't avoided by just restricting what the printed/parsed IR represents.

[nicolasvasilache]

Making and pushing IR design decisions because there is buggy code based on incorrect assumptions (along with matching incorrect test cases) sounds hilarious to me! :-)

I think you misspelled "unnecessary cognitive overhead" :)
In any case I don't see what convoluted logic results in the implication bug => IR design decision.

The non-fallacious logic is IR design => unnecessary cognitive overhead => consider changing IR design :)
There are many places in affine that have been guilty of that historically and some are remaining today.
This particular one has been around for a lot of time.

The bug fix is one thing that goes in regardless.

Adding others to visibility: @albertcohen @andydavis1

mlir/include/mlir/Dialect/AffineOps/AffineOps.h

private:
   /// Maps of Value to position in `affineMap`.
   DenseMap<Value, unsigned> dimValueToPosition;
 
   /// Ordered dims and symbols matching positional dims and symbols in
   /// `affineMap`.
   SmallVector<Value, 8> reorderedDims;
   SmallVector<Value, 8> concatenatedSymbols;
 
+  /// The number of symbols in concatenated symbols that belong to the original
+  /// map as opposed to those concatendated during map composition.
+  unsigned numProperSymbols;
+
   AffineMap affineMap;
 
   /// Used with RAII to control the depth at which AffineApply are composed
   /// recursively. Only accepts depth 1 for now to allow a behavior where a
   /// newly composed AffineApplyOp does not increase the length of the chain of
   /// AffineApplyOps. Full composition is implemented iteratively on top of
   /// this behavior.
   static unsigned &affineApplyDepth() {
     static thread_local unsigned depth = 0;
     return depth;
   }
   static constexpr unsigned kMaxAffineApplyDepth = 1;
 
-  AffineApplyNormalizer() { affineApplyDepth()++; }
+  AffineApplyNormalizer() : numProperSymbols(0) { affineApplyDepth()++; }
 
 public:
   ~AffineApplyNormalizer() { affineApplyDepth()--; }
 };
 
 } // end namespace mlir
 
 #endif

mlir/lib/Dialect/AffineOps/AffineOps.cpp

for (unsigned i = 0, e = operands.size(); i < e; ++i) {
 
         auxiliaryExprs.push_back(renumberedMap.getResult(0));
       } else {
         if (i < numDimsBeforeRewrite) {
           // b. The mathematical composition of AffineMap composes dims.
           auxiliaryExprs.push_back(renumberOneDim(t));
         } else {
           // c. The mathematical composition of AffineMap concatenates symbols.
-          //    We do the same for symbol operands.
-          concatenatedSymbols.push_back(t);
+          //    Note that the map composition will put symbols already present
+          //    in the map before any symbols coming from the auxiliary map, so
+          //    we insert them before any symbols that are due to renumbering,
+          //    and after the proper symbols we have seen already.
+          concatenatedSymbols.insert(
+              std::next(concatenatedSymbols.begin(), numProperSymbols++), t);
         }
       }
     }
   }
 
   // Early exit if `map` is already composed.
   if (auxiliaryExprs.empty()) {
     affineMap = map;

mlir/test/Dialect/AffineOps/canonicalize.mlir

 
 // Affine map for test case: compose_affine_maps_d2_tile
 // CHECK-DAG: [[MAP8:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 4) * 4 - (d1 floordiv 4) * 4)>
 // CHECK-DAG: [[MAP8a:#map[0-9]+]] = affine_map<(d0, d1) -> (d1 + (d0 ceildiv 8) * 8 - (d1 floordiv 8) * 8)>
 
 // Affine maps for test case: compose_affine_maps_dependent_loads
 // CHECK-DAG: [[MAP9:#map[0-9]+]] = affine_map<(d0) -> (d0 + 3)>
 // CHECK-DAG: [[MAP10:#map[0-9]+]] = affine_map<(d0) -> (d0 * 3)>
-// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 7) ceildiv 3)>
+// CHECK-DAG: [[MAP11:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 3) ceildiv 3)>
 // CHECK-DAG: [[MAP12:#map[0-9]+]] = affine_map<(d0) -> (d0 * 7 - 49)>
 
 // Affine maps for test case: compose_affine_maps_diamond_dependency
 // CHECK-DAG: [[MAP13A:#map[0-9]+]] = affine_map<(d0) -> ((d0 + 6) ceildiv 8)>
 // CHECK-DAG: [[MAP13B:#map[0-9]+]] = affine_map<(d0) -> ((d0 * 4 - 4) floordiv 3)>
 
 // Affine maps for test case: partial_fold_map
 // CHECK-DAG: [[MAP15:#map[0-9]+]] = affine_map<()[s0] -> (s0 - 42)>
affine.for %i1 = 0 to 3 {
         %v1 = load %0[%x00, %x02] : memref<16x32xf32>
 
         // Swizzle %i0, %i1
         // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I1]], [[I0]]{{\]}}
         %v2 = load %0[%x01, %x00] : memref<16x32xf32>
 
         // Swizzle %x00, %x01 and %c3, %c7
         %x10 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d0 * s1)>
-           (%x01, %x00)[%c7, %c3]
+           (%x01, %x00)[%c3, %c7]
         %x11 = affine.apply affine_map<(d0, d1)[s0, s1] -> (d1 ceildiv s0)>
-           (%x01, %x00)[%c7, %c3]
+           (%x01, %x00)[%c3, %c7]
 
         // CHECK-NEXT: [[I2A:%[0-9]+]] = affine.apply [[MAP12]](%{{.*}})
         // CHECK-NEXT: [[I2B:%[0-9]+]] = affine.apply [[MAP11]](%{{.*}})
         // CHECK-NEXT: load %{{[0-9]+}}{{\[}}[[I2A]], [[I2B]]{{\]}}
         %v3 = load %0[%x10, %x11] : memref<16x32xf32>
       }
     }
   }
affine.for %j = 0 to %arg0 {
       %c2 = constant 2 : index
       // CHECK: affine.min #[[map]]
       %0 = affine.min affine_map<(d0,d1,d2)->(d0, d1 - d2)>(%i, %j, %c2)
       "consumer"(%0) : (index) -> ()
     }
   }
   return
 }
+
+// -----
+
+// Reproducer for PR45031. This used to fold into an incorrect map because
+// symbols were concatenated in the wrong order during map folding. Map
+// composition places the symbols of the original map before those of the map
+// it is composed with, e.g. A.compose(B) will first have all symbols of A,
+// then all symbols of B.
+
+#map1 = affine_map<(d0)[s0, s1] -> (d0 * s0 + s1)>
+#map2 = affine_map<(d0)[s0] -> (1024, -d0 + s0)>
+
+// CHECK: #[[MAP:.*]] = affine_map<()[s0, s1] -> (1024, s1 * -1024 + s0)>
+
+// CHECK: func @rep(%[[ARG0:.*]]: index, %[[ARG1:.*]]: index)
+func @rep(%arg0 : index, %arg1 : index) -> index {
+  // CHECK-NOT: constant
+  %c0 = constant 0 : index
+  %c1024 = constant 1024 : index
+  // CHECK-NOT: affine.apply
+  %0 = affine.apply #map1(%arg0)[%c1024, %c0]
+
+  // CHECK: affine.min #[[MAP]]()[%[[ARG1]], %[[ARG0]]]
+  %1 = affine.min #map2(%0)[%arg1]
+  return %1 : index
+}