commit 3e2c4f31477b8ac6b7181007251f34021f234449
author Guray Ozen <guray.ozen@gmail.com> Fri Nov 18 16:36:51 2022 +0100
committer Guray Ozen <guray.ozen@gmail.com> Fri Nov 18 16:36:51 2022 +0100
tree 512ab4c91f48cc900e78ba4bb9b6b92aff5d78b3
parent 20048f775f0043f6b8f38441771f6742dec0c9ad

Set vectorSize back to 4 for non-divisible by 4 dimensions

For issue #10003, the `WorkgroupSpecializationPass` nicely implements combination codegen of fastpath and the slowpath. But the tile size is still set incorrectly, this causes load stores on fastpath are still not vectorized.

The example given in the PR has a tensor of size 102401. `TileAndDistributeToWorkgroups` sets the tile size to "<tile_sizes = [[64]]>". However, it should be 256, because the inner loop is unrolled with 4, and the 4 loads/stores are vectorized. See the generated ir below after this pass.

```
scf.for %arg0 = blockIdx.x to 102401 step blockDim.x  {
...
	linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%6, %7 : tensor<?xf32>, tensor<?xf32>) outs(%8 : tensor<?xf32>) attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[64]]>} ...
}
```
This leds load/store are fastpath do not get vectorized, see `vector<1xf32>`.

```
if(fastpath) {
    %7 = affine.apply #map3(%threadIdx.x)[%blockIdx]
    %8 = vector.transfer_read %0[%7], %cst {in_bounds = [true]} : memref<102401xf32>, vector<1xf32>
    ...
} else {
	//slowpath
}
```

This assumes that fastpath+slowpath is implemented. Even if loop trip count is not divisable by 4, it still sets vectorSize to 4. This helps generating correct tile size, see `tile_sizes = [[256]]` in code below.

```
scf.for %arg0 = blockIdx.x to 102401 step blockDim.x  {
        %9 = linalg.generic {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>], iterator_types = ["parallel"]} ins(%6, %7 : tensor<?xf32>, tensor<?xf32>) outs(%8 : tensor<?xf32>) attrs =  {lowering_config = #iree_codegen.lowering_config<tile_sizes = [[256]]>} ...
}
```
Setting tile size correctly helps to vectorize load/store, see `vector<4xf32>`.

```
 if(fastpath) {
    %7 = affine.apply #map3(%threadIdx.x)[%blockIdx]
    %8 = vector.transfer_read %0[%7], %cst {in_bounds = [true]} : memref<102401xf32>, vector<4xf32>
    ...
} else {
	//slowpath
}
```