[Mlir-commits] [mlir] [mlir][affine] Define `affine.linearize_index` (PR #114480)

[Jakub Kuderski]

================
@@ -1099,4 +1099,73 @@ def AffineDelinearizeIndexOp : Affine_Op<"delinearize_index",
   let hasCanonicalizer = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// AffineLinearizeIndexOp
+//===----------------------------------------------------------------------===//
+def AffineLinearizeIndexOp : Affine_Op<"linearize_index",
+    [Pure, AttrSizedOperandSegments]> {
+  let summary = "linearize an index";
+  let description = [{
+    The `affine.linearize_index` operation takes a sequence of index values and a
+    basis of the same length and linearizes the indices using that basis.
+
+    That is, for indices %idx_1 through %i_N and basis elements b_1 through b_N,
+    it computes
+
+    ```
+    sum(i = 1 to N) %idx_i * product(j = i + 1 to N) B_j
+    ```
+
+    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.
+
+    Note that the outputs of `affine.delinearize_index` are, by definition, `disjoint`.
+
+    Example:
+
+    ```
+    %linear_index = affine.delinearize_index [%index_0, %index_1, %index_2] (16, 224, 224) : index
+    ```
+
+    In the above example, `%linear_index` conceptually holds the following:
+
+    ```
+    #map = affine_map<()[s0, s1, s2] -> (s0 * 50176 + s1 * 224 + s2)>
+    %linear_index = affine.apply #map()[%index_0, %index_1, %index_2]
+    ```
----------------
kuhar wrote:

```suggestion
    ```mlir
    #map = affine_map<()[s0, s1, s2] -> (s0 * 50176 + s1 * 224 + s2)>
    %linear_index = affine.apply #map()[%index_0, %index_1, %index_2]
    ```
```

https://github.com/llvm/llvm-project/pull/114480