[Mlir-commits] [mlir] [mlir][linalg] Fix exponential walk of tensor.extract indices (PR #222319)

[Jelle Schühmacher]

jschuhmacher wrote:

Verified below test:

| `arith.muli` ops | before | after |
|---|---|---|
| 20 | 0.03 s | 0.02 s |
| 24 | 0.26 s | 0.02 s |
| 27 | 2.01 s | 0.02 s |
| 30 | 16.10 s | 0.01 s |
| 40 | >300 s, killed | 0.02 s |

which is probably not a good test to add to CI.

```
// RUN: mlir-opt %s -transform-interpreter | FileCheck %s

func.func @vectorize_nd_tensor_extract_deep_index_chain(
    %src: tensor<4xf32>,
    %seed: tensor<8xi32>,
    %output: tensor<8xf32>) -> tensor<8xf32> {
  %res = linalg.generic {
    indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],
    iterator_types = ["parallel"]
  } ins(%seed : tensor<8xi32>) outs(%output : tensor<8xf32>) {
  ^bb0(%in: i32, %out: f32):
    %0 = arith.muli %in, %in : i32
    %1 = arith.muli %0, %0 : i32
    %2 = arith.muli %1, %1 : i32
    %3 = arith.muli %2, %2 : i32
    %4 = arith.muli %3, %3 : i32
    %5 = arith.muli %4, %4 : i32
    %6 = arith.muli %5, %5 : i32
    %7 = arith.muli %6, %6 : i32
    %8 = arith.muli %7, %7 : i32
    %9 = arith.muli %8, %8 : i32
    %10 = arith.muli %9, %9 : i32
    %11 = arith.muli %10, %10 : i32
    %12 = arith.muli %11, %11 : i32
    %13 = arith.muli %12, %12 : i32
    %14 = arith.muli %13, %13 : i32
    %15 = arith.muli %14, %14 : i32
    %16 = arith.muli %15, %15 : i32
    %17 = arith.muli %16, %16 : i32
    %18 = arith.muli %17, %17 : i32
    %19 = arith.muli %18, %18 : i32
    %20 = arith.muli %19, %19 : i32
    %21 = arith.muli %20, %20 : i32
    %22 = arith.muli %21, %21 : i32
    %23 = arith.muli %22, %22 : i32
    %24 = arith.muli %23, %23 : i32
    %25 = arith.muli %24, %24 : i32
    %26 = arith.muli %25, %25 : i32
    %27 = arith.muli %26, %26 : i32
    %28 = arith.muli %27, %27 : i32
    %29 = arith.muli %28, %28 : i32
    %30 = arith.muli %29, %29 : i32
    %31 = arith.muli %30, %30 : i32
    %32 = arith.muli %31, %31 : i32
    %33 = arith.muli %32, %32 : i32
    %34 = arith.muli %33, %33 : i32
    %35 = arith.muli %34, %34 : i32
    %36 = arith.muli %35, %35 : i32
    %37 = arith.muli %36, %36 : i32
    %38 = arith.muli %37, %37 : i32
    %39 = arith.muli %38, %38 : i32
    %idx = arith.index_cast %39 : i32 to index
    %extracted = tensor.extract %src[%idx] : tensor<4xf32>
    linalg.yield %extracted : f32
  } -> tensor<8xf32>
  return %res : tensor<8xf32>
}

// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_deep_index_chain
// CHECK-COUNT-40: arith.muli {{.*}} : vector<8xi32>
// CHECK:          vector.gather

module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {
     %0 = transform.structured.match ops{["linalg.generic"]} in %arg1 : (!transform.any_op) -> !transform.any_op
     transform.structured.vectorize %0 vector_sizes [8] vectorize_nd_extract : !transform.any_op
     transform.yield
   }
}
```

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