[mlir][Affine] Generalize the linearize(delinearize()) simplifications (llvm#117637)

The existing canonicalization patterns would only cancel out cases where
the entire result list of an affine.delineraize_index was passed to an
affine.lineraize_index and the basis elements matched exactly (except
possibly for the outer bounds).

This was correct, but limited, and left open many cases where a
delinearize_index would take a series of divisions and modulos only for
a subsequent linearize_index to use additions and multiplications to
undo all that work.

This sort of simplification is reasably easy to observe at the level of
splititng and merging indexes, but difficult to perform once the
underlying arithmetic operations have been created.

Therefore, this commit generalizes the existing simplification logic.

Now, any run of two or more delinearize_index results that appears
within the argument list of a linearize_index operation with the same
basis (or where they're both at the outermost position and so can be
unbonded, or when `linearize_index disjoint` implies a bound not present
on the `delinearize_index`) will be reduced to one signle
delinearize_index output, whose basis element (that is, size or length)
is equal to the product of the sizes that were simplified away.

That is, we can now simplify

    %0:2 = affine.delinearize_index %n into (8, 8) : inde, index
%1 = affine.linearize_index [%x, %0#0, %0#1, %y] by (3, 8, 8, 5) : index

to the simpler

    %1 = affine.linearize_index [%x, %n, %y] by (3, 64, 5) : index

This new pattern also works with dynamically-sized basis values.

While I'm here, I fixed a bunch of typos in existing tests, and added a
new getPaddedBasis() method to make processing
potentially-underspecified basis elements simpler in some cases.

Author: krzysz00

mlir/include/mlir/Dialect/Affine/IR/AffineOps.td

@@ -1083,6 +1083,9 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> {
     %indices_2 = affine.apply #map2()[%linear_index]
     ```
 
+    In other words, `%0:3 = affine.delinearize_index %x into (B, C)` produces
+    `%0 = {%x / (B * C), (%x mod (B * C)) / C, %x mod C}`.
+
     The basis may either contain `N` or `N-1` elements, where `N` is the number of results.
     If there are N basis elements, the first one will not be used during computations,
     but may be used during analysis and canonicalization to eliminate terms from
@@ -1098,7 +1101,12 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> {
     %0:3 = affine.delinearize_index %linear_index into (244, 244) : index, index
     ```
 
-    Note that, due to the constraints of affine maps, all the basis elements must
+    Note that, for symmetry with `getPaddedBasis()`, if `hasOuterBound` is `true`
+    when one of the `OpFoldResult` builders is called but the first element of the
+    basis is `nullptr`, that first element is ignored and the builder proceeds as if
+    there was no outer bound.
+
+    Due to the constraints of affine maps, all the basis elements must
     be strictly positive. A dynamic basis element being 0 or negative causes
     undefined behavior.
   }];
@@ -1136,6 +1144,11 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index", [Pure]> {
     /// Return a vector that contains the basis of the operation, removing
     /// the outer bound if one is present.
     SmallVector<OpFoldResult> getEffectiveBasis();
+
+    /// Return the vector with one basis element per result of the operation. If
+    /// there is no outer bound specified, the leading entry of this result will be
+    /// nullptr.
+    SmallVector<OpFoldResult> getPaddedBasis();
   }];
 
   let hasVerifier = 1;
@@ -1160,6 +1173,9 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
     sum(i = 0 to N-1) %idx_i * product(j = i + 1 to N-1) B_j
     ```
 
+    In other words, `%0 = affine.linearize_index [%z, %y, %x] by (Z, Y, X)`
+    gives `%0 = %x + %y * X + %z * X * Y`, or `%0 = %x + X * (%y + Y * (%z))`.
+
     The basis may either have `N` or `N-1` elements, where `N` is the number of
     inputs to linearize_index. If `N` inputs are provided, the first one is not used
     in computation, but may be used during analysis or canonicalization as a bound
@@ -1168,6 +1184,10 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
     If all `N` basis elements are provided, the linearize_index operation is said to
     "have an outer bound".
 
+    As a convenience, and for symmetry with `getPaddedBasis()`, ifg the first
+    element of a set of `OpFoldResult`s passed to the builders of this operation is
+    `nullptr`, that element is ignored.
+
     If the `disjoint` property is present, this is an optimization hint that,
     for all `i`, `0 <= %idx_i < B_i` - that is, no index affects any other index,
     except that `%idx_0` may be negative to make the index as a whole negative.
@@ -1224,6 +1244,11 @@ def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
     /// Return a vector that contains the basis of the operation, removing
     /// the outer bound if one is present.
     SmallVector<OpFoldResult> getEffectiveBasis();
+
+    /// Return the vector with one basis element per index operand of the operation.
+    /// If there is no outer bound specified, the leading entry of this basis will be
+    /// nullptr.
+    SmallVector<OpFoldResult> getPaddedBasis();
   }];
 
   let hasVerifier = 1;