commit 59532d30c0b8b3a0cdf261eaa4e69784c9fddeee
author Andrzej Warzyński <andrzej.warzynski@arm.com> Wed Apr 24 09:10:07 2024 +0100
committer GitHub <noreply@github.com> Wed Apr 24 09:10:07 2024 +0100
tree e035dd33654de32bbf7fe3772ca7f1e96c27c6bf
parent 1b4c76f931e40ba665930e09d7a6e3f0dc2e0853

Extend DecomposeConvolutionToLowerDimOpsPass (#17069)

Updates `DecomposeConvolutionToLowerDimOpsPass` so that when a
convolution Op is decomposed, the lowering config attribute is updated
("decomposed") accordingly. This change leverages the fact that we know
which dimension is being "removed" when the decomposition happens. The
corresponding dim/size is simply deleted when updating the lowering
config.

With the logic added in this PR, we are making sure that
`DecomposeConvolutionToLowerDimOpsPass` will no longer "drop" the
lowering config. This is beneficial - it removes the need to restore the
tiling config (for e.g. the vectoriser) from the loop nest.

BEFORE `DecomposeConvolutionToLowerDimOpsPass`:
```mlir
  #config = #iree_codegen.lowering_config<tile_sizes = [[0, 0, 0, 0, 0, 0], [1, 1, 1, 4, 0, 0], [0, 0, 0, 0, 1, 4], [0, 0, 0, 0, 0, 0]]>
  linalg.depthwise_conv_2d_nhwc_hwc lowering_config = #config
```

AFTER `DecomposeConvolutionToLowerDimOpsPass` _without_ this change:
```mlir
   linalg.depthwise_conv_1d_nwc_wc
```

AFTER `DecomposeConvolutionToLowerDimOpsPass` _with_ this change:
```mlir
   #config = #iree_codegen.lowering_config<tile_sizes = [[0, 0, 0, 0], [1, 1, 4, 0], [0, 0, 0, 1], [0, 0, 0, 0]]>
   linalg.depthwise_conv_1d_nwc_wc lowering_config = #config
```

ATM, the logic implemented here is limited to deptwhise HWC
convolutions. It can easily be extended to other Convs when that's
required.

For simplicity, it is also assumed that there's only one Conv Op per
function. This seems like the most common case. Adding support for
multiple conv ops per function is left as a TODO. It shouldn't be too
difficult to implement - it will require building a map of all
convs/lowering configs.