[Mlir-commits] [mlir] [mlir][linalg] Relax tensor.extract vectorization (PR #99299)

[Han-Chung Wang]

hanhanW wrote:

Hi, this breaks downstream IREE project, and I'm going to revert it. Below is the upstream repro (and I will attach the case to the revert commit).

Run: `mlir-opt -transform-interpreter ~/repro.mlir`

```mlir
#map = affine_map<(d0, d1) -> (d0)>
#map1 = affine_map<(d0, d1) -> (d1)>
#map2 = affine_map<(d0, d1) -> (d0, d1)>
#map3 = affine_map<(d0, d1) -> (d0 + d1)>
module {
  func.func @foo(%arg0: index, %arg1: tensor<2xf32>, %arg2: tensor<4xf32>, %arg3: tensor<1xf32>) -> tensor<4x1xf32> {
    %c0 = arith.constant 0 : index
    %cst = arith.constant 1.000000e+00 : f32
    %cst_0 = arith.constant 0.000000e+00 : f32
    %0 = tensor.empty() : tensor<4x1xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel"]} ins(%arg2, %arg3 : tensor<4xf32>, tensor<1xf32>) outs(%0 : tensor<4x1xf32>) {
    ^bb0(%in: f32, %in_1: f32, %out: f32):
      %2 = linalg.index 0 : index
      %3 = linalg.index 1 : index
      %4 = affine.apply #map3(%3, %arg0)
      %extracted = tensor.extract %arg1[%c0] : tensor<2xf32>
      %5 = arith.cmpi eq, %2, %c0 : index
      %6 = arith.cmpi ult, %2, %c0 : index
      %7 = arith.select %5, %cst, %in : f32
      %8 = arith.select %6, %cst_0, %7 : f32
      %9 = arith.cmpi eq, %4, %c0 : index
      %10 = arith.cmpi ult, %4, %c0 : index
      %11 = arith.select %9, %cst, %in_1 : f32
      %12 = arith.select %10, %cst_0, %11 : f32
      %13 = arith.mulf %8, %12 : f32
      %14 = arith.mulf %13, %extracted : f32
      %15 = arith.cmpi eq, %2, %4 : index
      %16 = arith.select %15, %cst, %cst_0 : f32
      %17 = arith.subf %16, %14 : f32
      linalg.yield %17 : f32
    } -> tensor<4x1xf32>
    return %1 : tensor<4x1xf32>
  }
}

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 : !transform.any_op
    transform.yield
  }
}
```

Error:

```
mlir-opt: llvm-project/mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp:820: bool isLoopInvariantIdx(mlir::linalg::LinalgOp&, mlir::Value&): Assertion `targetShape.back() != 1 && "1-D vectors with the trailing dim eqaual 1 are not yet supp
orted"' failed.                                                                                                                                                                                                                                              PLEASE submit a bug report to https://github.com/llvm/llvm-project/issues/ and include the crash backtrace.
Stack dump:
...
```

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