HoistIntoGlobals fails to hoist constantOp

[ziereis]

What happened?

running the pass --iree-util-hoist-into-globals fails to hoist constants into globals.

Changing one NON-const op, namely tensor. expand_shape to tensor.reshape make the constants get hoisted.

Steps to reproduce your issue

input ir:

util.func @hoist(%arg0: tensor<2xindex>) -> tensor<1x2xf32> {
  %const_t = arith.constant dense<[0.0, 1.0]> : tensor<2xf32>
  %one = arith.constant 1.0 : f32
  %empty = tensor.empty() : tensor<2xf32>
  %add_one = linalg.generic {
    indexing_maps = [
      affine_map<(d0) -> (d0)>,
      affine_map<(d0) -> (d0)>
    ],
    iterator_types = ["parallel"]
  } ins(%const_t : tensor<2xf32>) outs(%empty: tensor<2xf32>) {
    ^bb0(%in : f32, %out : f32):
      %added = arith.addf %in, %one : f32
      linalg.yield %added : f32
  } -> tensor<2xf32>
  %empty2 = tensor.empty() : tensor<2xf32>
  %loaded = linalg.generic {
    indexing_maps = [
      affine_map<(d0) -> (d0)>,
      affine_map<(d0) -> (d0)>
    ],
    iterator_types = ["parallel"]
  } ins(%arg0 : tensor<2xindex>) outs(%empty2 : tensor<2xf32>) {
    ^bb0(%in: index, %out: f32):
      %extracted = tensor.extract %add_one[%in] : tensor<2xf32>
      linalg.yield %extracted : f32
  } -> tensor<2xf32>
  %reshaped = tensor.expand_shape %loaded [[0, 1]] output_shape[1, 2]: tensor<2xf32> into tensor<1x2xf32>
  // %shape = arith.constant dense<[1, 2]> : tensor<2xi32>
  // %reshaped = tensor.reshape %loaded(%shape) : (tensor<2xf32>, tensor<2xi32>) -> tensor<1x2xf32>

  util.return %reshaped : tensor<1x2xf32>
}

running the following command fails to hoist the constant linalg.generic op %add_one

./build/tools/iree-opt --iree-util-hoist-into-globals --mlir-disable-threading

Changing the following lines from using a tensor.expand_shape to a tensor.reshape, will make the linalg.generic op get hoisted.

%reshaped = tensor.expand_shape %loaded [[0, 1]] output_shape[1, 2]: tensor<2xf32> into tensor<1x2xf32>

to

%shape = arith.constant dense<[1, 2]> : tensor<2xi32>
%reshaped = tensor.reshape %loaded(%shape) : (tensor<2xf32>, tensor<2xi32>) -> tensor<1x2xf32>

What component(s) does this issue relate to?

Compiler

Version information

573eed59c8eae6529757de3918f73ac8e8395e53

Additional context

In the ConstExpr analysis with the IR using expand_shape, the variable %loaded will not be assigned as a user of %add_one since it is only used inside its nested region, which leads to the constant %8 not having a consumer. This will then lead to the problem that in the ConstExprHoistingPolicy it will get marked as non-hoistable as it has no consumer. This won't happen when using reshape since reshape has %loaded as its producer, so %loaded will eventually be added as a consumer of %add_one. I have some small fix that seems to work for me where I just add the parent op as a user; however, I'm not sure if this will break other things.

Is it possible to have the ConstExpr pass be changed so it will actually capture all consumers even if they will be flagged as non-const? For analysis which exprs are const it won't matter, but for the ConstExprHoistingPolicy it's important to know about the consumers of the constantOps.