[Mlir-commits] [mlir] [mlir][linalg] Add folder for `linalg.index` (PR #136640)

[Andrzej Warzyński]

================
@@ -305,6 +305,86 @@ func.func @self_copy(%arg0 : memref<2x3x?x4xf32>) {
 }
 
 // -----
+
+// CHECK: func @fold_linalg_index_tensor_static
+func.func @fold_linalg_index_tensor_static(%0: tensor<4x16xi32>, %1: tensor<1x16xi32>,
+                                           %2: tensor<4x1xi32>) -> tensor<4x1xi32> {
+// CHECK-NEXT: linalg.generic
+// CHECK:        %[[IDX_0:.+]] = linalg.index 0 : index
+// CHECK-NOT:    linalg.index 1
----------------
banach-space wrote:

AFAIK, this will only ensure that linalg.index 1 doesn't appear **between** the following lines:

mlir
Copy
Edit
```mlir
// CHECK:        %[[IDX_0:.+]] = linalg.index 0 : index
// CHECK:        %[[IDX_2:.+]] = linalg.index 2 : index
```

Instead, you probably want to check that `linalg.index 1` **does not appear at all inside the linalg.generic:
```mlir
// CHECK-DAG: linalg.generic
// CHECK-NOT:    linalg.index 1
// CHECK:        %[[IDX_0:.+]] = linalg.index 0 : index
// CHECK:        %[[IDX_2:.+]] = linalg.index 2 : index
// CHECK:        %[[ADD:.+]] = arith.addi %[[IDX_0]], %[[IDX_2]]
// CHECK:        %[[CAST:.+]] = arith.index_cast %[[ADD]]
// CHECK-DAG:        linalg.yield %[[CAST]]
```

(Reference: https://llvm.org/docs/CommandGuide/FileCheck.html#the-check-dag-directive)

I'm raising this because, in my experience, `CHECK-NOT` can be a weak check when used alongside other `CHECK` lines 0 it's easy for it to pass unintentionally. That said, it's also pretty clear here that `%idx1` isn't used in the yielded value, so feel free to skip this check entirely if you're confident it's not relevant.

Btw, why not test something trivial (you could even remove the other inputs):
```mlir
  %res = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d2)>,
                                          affine_map<(d0, d1, d2) -> (d1, d2)>,
                                          affine_map<(d0, d1, d2) -> (d0, d1)>],
                         iterator_types = ["parallel", "parallel", "reduction"]}
    ins(%0, %1 : tensor<4x16xi32>, tensor<1x16xi32>)
    outs(%2 : tensor<4x1xi32>) {
      ^bb0(%lhs: i32, %rhs: i32, %out: i32):
        %idx1 = linalg.index 1 : index
        %int = arith.index_cast %idx1 : index to i32
        linalg.yield %int : i32
    } -> tensor<4x1xi32>
  return %res : tensor<4x1xi32>
}
```


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