Support memref::SubviewOp in FlattenMemRefSubspanPass.cpp

[matthias-springer]

Request description

There is a TODO in DispatchLinalgOnTensors.cpp that can be resolved now that #8637 is fixed:

  // TODO(#8637): `tensor.collapse_shape` and `tensor.expand_shape` are
  // trivially clonable too, but they cause problems
  // with bufferization. Make them clonable when fixed.
   if (isa<arith::IndexCastOp, linalg::InitTensorOp, tensor::CastOp,
          tensor::ExtractOp, tensor::ExtractSliceOp, tensor::PadOp>(op)) {

Making these two ops clonable works mostly, but triggered a bug in FlattenMemRefSubspanPass.cpp. This bug is fixed with #9828. Another test case (iree/tests/e2e/models/unidirectional_lstm.mlir.test) is failing because the following test case fails FlattenMemRefSubspanPass. It looks like memref::SubviewOps are not supported by the pass.

func.func @expand_shape3(%offset : index, %i0: index, %i1: index) -> f32 {
  %subspan = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) offset(%offset) : memref<1x40xf32>
  %subview = memref.subview %subspan[0, 30] [1, 10] [1, 1] : memref<1x40xf32> to memref<10xf32, affine_map<(d0) -> (d0 + 30)>>
  %expand = memref.expand_shape %subview [[0, 1]] : memref<10xf32, affine_map<(d0) -> (d0 + 30)>> into memref<1x10xf32, affine_map<(d0, d1) -> (d0 * 10 + d1 + 30)>>
  %value = memref.load %expand[%i0, %i1] : memref<1x10xf32, affine_map<(d0, d1) -> (d0 * 10 + d1 + 30)>>
  return %value : f32
}

fails with:

within split at /usr/local/google/home/springerm/iree/compiler/src/iree/compiler/Codegen/Common/test/flatten_memref_subspan.mlir:360 offset :4:14: error: failed to materialize conversion for result #0 of operation 'hal.interface.binding.subspan' that remained live after conversion
  %subspan = hal.interface.binding.subspan set(0) binding(0) type(storage_buffer) offset(%offset) : memref<1x40xf32>
             ^
within split at /usr/local/google/home/springerm/iree/compiler/src/iree/compiler/Codegen/Common/test/flatten_memref_subspan.mlir:360 offset :4:14: note: see current operation: %2 = "hal.interface.binding.subspan"(%arg0) {binding = 0 : index, descriptor_type = 7 : i32, operand_segment_sizes = dense<[1, 0]> : vector<2xi32>, set = 0 : index} : (index) -> memref<1x40xf32>
within split at /usr/local/google/home/springerm/iree/compiler/src/iree/compiler/Codegen/Common/test/flatten_memref_subspan.mlir:360 offset :5:14: note: see existing live user here: %3 = "memref.subview"(%2) {operand_segment_sizes = dense<[1, 0, 0, 0]> : vector<4xi32>, static_offsets = [0, 30], static_sizes = [1, 10], static_strides = [1, 1]} : (memref<1x40xf32>) -> memref<10xf32, affine_map<(d0) -> (d0 + 30)>>
  %subview = memref.subview %subspan[0, 30] [1, 10] [1, 1] : memref<1x40xf32> to memref<10xf32, affine_map<(d0) -> (d0 + 30)>>

What component(s) does this issue relate to?

Compiler

Additional context

No response

[allieculp]

@MaheshRavishankar new codegen labeled issue. Can you review and add priority?

[MaheshRavishankar]

Thanks @matthias-springer . The pass expects subviews to be folded away. Here the folding is thwarted by the memref.expand_shape between the subview and load. Its a known issue with the pass. I am going to keep this bug with low priority. There is no easy way to fix it, but is not a blocker at this point.