commit 8128e961889a804598787b17b25d411b182c6db5
author Han-Chung Wang <hanchung@google.com> Sat Sep 25 06:22:47 2021 +0800
committer GitHub <noreply@github.com> Sat Sep 25 06:22:47 2021 +0800
tree db23c7d2d3f7926c83516766120e69216bab9d81
parent 95ad84c242250e7356af843bac6ce0989c2cfe07

Implement tiling method for linalg_ext.reverse ops. (#7159)

To tile a linalg_ext.reverse op, we have to store the tiled result to
mirror offsets. E.g.,

Before:

```
[T_0], [T_1], [T_2], [T_3]
```

After:

```
[Rev_t3], [Rev_t2], [Rev_t1], [Rev_t0]
```

The snippet after DispatchLinalgOnTensors, which shows that the op is
tiled and distributed:

```mlir
func private @_foo(%arg0: tensor<?xi32>, %arg1: !shapex.ranked_shape<[?]>) -> (tensor<?xi32>, !shapex.ranked_shape<[?]>) {
  %c0 = constant 0 : index
  %c1 = constant 1 : index
  %0 = shapex.ranked_dim %arg1[0] : !shapex.ranked_shape<[?]> -> index
  %1 = shapex.tie_shape %arg0, %arg1 : tensor<?xi32>, !shapex.ranked_shape<[?]>
  %2 = linalg.init_tensor [%0] : tensor<?xi32>
  %3 = tensor.dim %1, %c0 : tensor<?xi32>
  %4 = tensor.dim %2, %c0 : tensor<?xi32>
  %5 = tensor.dim %1, %c0 : tensor<?xi32>
  %6 = flow.dispatch.workgroups[%3, %c1, %c1](%1, %0) : (tensor<?xi32>{%5}, index) -> tensor<?xi32>{%4} =
      (%arg2: !flow.dispatch.tensor<readonly:?xi32>, %arg3: index, %arg4: !flow.dispatch.tensor<writeonly:?xi32>) {
    %c0_0 = constant 0 : index
    %8 = flow.dispatch.tensor.load %arg2, offsets = [], sizes = [], strides = [] : !flow.dispatch.tensor<readonly:?xi32> -> tensor<?xi32>
    %9 = linalg.init_tensor [%arg3] : tensor<?xi32>
    %workgroup_size_0 = flow.dispatch.workgroup.size[0] : index
    %10 = tensor.dim %8, %c0_0 : tensor<?xi32>
    %workgroup_id_0 = flow.dispatch.workgroup.id[0] : index
    %workgroup_count_0 = flow.dispatch.workgroup.count[0] : index
    %11 = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%workgroup_id_0)[%workgroup_size_0]
    %12 = affine.apply affine_map<(d0)[s0] -> (d0 * s0)>(%workgroup_count_0)[%workgroup_size_0]
    scf.for %arg5 = %11 to %10 step %12 {
      %13 = affine.min affine_map<(d0)[s0, s1] -> (s0, -d0 + s1)>(%arg5)[%workgroup_size_0, %10]
      %14 = flow.dispatch.tensor.load %arg2, offsets = [%arg5], sizes = [%13], strides = [1] : !flow.dispatch.tensor<readonly:?xi32> -> tensor<?xi32>
      %15 = tensor.dim %8, %c0_0 : tensor<?xi32>
      %16 = subi %15, %arg5 : index
      %17 = subi %16, %13 : index
      %18 = tensor.extract_slice %9[%arg5] [%13] [1] : tensor<?xi32> to tensor<?xi32>
      %19 = linalg_ext.reverse dimensions(dense<0> : tensor<1xi64>) {__internal_linalg_transform__ = "workgroup"} ins(%14 : tensor<?xi32>) outs(%18 : tensor<?xi32>) : tensor<?xi32>
      flow.dispatch.tensor.store %19, %arg4, offsets = [%17], sizes = [%13], strides = [1] : tensor<?xi32> -> !flow.dispatch.tensor<writeonly:?xi32>
    }
    flow.return
  }
  %7 = shapex.get_ranked_shape %6 : tensor<?xi32> -> !shapex.ranked_shape<[?]>
  return %6, %7 : tensor<?xi32>, !shapex.ranked_shape<[?]>
}
```

The `print-ir-after-all` log can be found at:
https://gist.githubusercontent.com/hanhanW/bba1d36149f37d948e8cea5fe95a6287/raw

This is a step toward https://github.com/google/iree/issues/5045