Bufferizing linalg.tensor_reshape?

Following the recent changes in HLO->Linalg lowering, I’m able to almost fully bufferize a rather complex specification such as resnet without using code of my own.

However, one operation remains unconverted: linalg.tensor_reshape (which is always applied in a trivial way, e.g. tensor<1x112x112x64xf32> into tensor<112x112x64xf32>).

Here is an example of generated code where tensors subsist:

    %536 = memref.tensor_load %535 : memref<1x112x112x64xf32>
    %537 = linalg.tensor_reshape %536 [#map1, #map2, #map3] : tensor<1x112x112x64xf32> into tensor<112x112x64xf32>
    %539 = memref.buffer_cast %537 : memref<112x112x64xf32>

Ideally, this should be replaced with:

    %539 = memref.reshape %535 : memref<1x112x112x64xf32> into tensor<112x112x64xf32>

Is there some automated way of converting it, too?

I’m currently using a combination of tf-opt --linalg-bufferize, which does almost all the work and iree-opt --iree-codegen-hlo-to-linalg-on-tensors --iree-linalg-on-tensors-path to remove linalg.pad_tensor. I’ve also tried tf-opt --linalg-detensorize (it didn’t seem to do much).

D.

This is really fragile. That option was added strictly for lit testing and is not going to survive when IREE moves to use Linalg on tensors by default (soon). You are probably looking for just the pattern that lowers pad_tensor to a fill + subtensor_insert. You can just copy that over to your local pass.

You can try to enhance the bufferize pass to handle this. One caveat to consider, linalg.tensor_reshape has copy semantics, while linalg.reshape has aliasing semantics, i.e. it returns a different view of the same buffer. If you know that %536 is dead after this use, you can just do a direct replacement of the linalg.tensor_reshapelinalg.reshape.

W.R.T to automated way of doing this @nicolasvasilache has been looking at bufferization and is looking at upstreaming some of it, but its really WIP. Sorry, I dont have a better answer here.

1 Like

Is there some documentation on this? And what does “copy/aliasing” means when applied to tensors, which have value semantics?
D.

It is documented in op-semantics of each of those operations. More concretely,

%1 = linalg.tensor_reshape %0 ....

Here %1 is a copy of %0. tensor are SSA values, so there is no other semantics possible here.

On the other hand linalg.reshape operates on memrefs.

%1 = linalg.reshape %0 ....

Here %1 aliases %0, i.e. if you store to a location in %1 and then read the same location using %0 you will get the value that was stored.

I found out that IREE has some pass that is executed for replacing linalg.tensor_reshape by linalg.reshape.

In the dialect those two ops have different semantic and operands type. linalg.tensor_reshape accepts tensors while linalg.reshape accepts memrefs only.

I am wondering is there something similar in MLIR origin repo?
I tried to bufferize the following very simple code:

#map0 = affine_map<(d0, d1) -> (d0, d1)>
#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
module  {
  func @generated_model_inference(%arg0: tensor<4x3xf32>) -> tensor<3x4x1xf32> {
    %0 = linalg.tensor_reshape %arg0 [#map0] : tensor<4x3xf32> into tensor<12xf32>
    %1 = linalg.tensor_reshape %0 [#map1] : tensor<12xf32> into tensor<3x4x1xf32>
    return %1 : tensor<3x4x1xf32>
  }
}

after mlir-opt --split-input-file --hgfc-to-linalg-on-tensors -func-bufferize --linalg-bufferize --finalizing-bufferize I’ve got the folowing error:

within split at test/lit-test/linalg_on_tensors/test.ll:1 offset :5:10: error: failed to legalize operation 'linalg.tensor_reshape'
    %0 = linalg.tensor_reshape %arg0 [#map0] : tensor<1x2x3x2x1x1xf32> into tensor<12xf32>
         ^
within split at test/lit-test/linalg_on_tensors/test.ll:1 offset :5:10: note: see current operation: %1 = "linalg.tensor_reshape"(%0) {reassociation = [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>]} : (tensor<1x2x3x2x1x1xf32>) -> tensor<12xf32>

So, when I remove the finalization pass I have the following output:

#map0 = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>
#map1 = affine_map<(d0, d1) -> (d0, d1)>
module  {
  func @generated_model_inference(%arg0: memref<1x2x3x2x1x1xf32>) -> memref<4x3xf32> {
    %0 = memref.tensor_load %arg0 : memref<1x2x3x2x1x1xf32>
    %1 = linalg.tensor_reshape %0 [#map0] : tensor<1x2x3x2x1x1xf32> into tensor<12xf32>
    %2 = linalg.tensor_reshape %1 [#map1] : tensor<12xf32> into tensor<4x3xf32>
    %3 = memref.buffer_cast %2 : memref<4x3xf32>
    return %3 : memref<4x3xf32>
  }
}

It seems to me that necessary implementation of bufferization pass for linalg.tensor_reshape is not developed yet as it was for linalg.Fill.

Could you please guide me where I am wrong or may be I missed something?

Sorry, but in the end our solution was to write our pass which replaces the remaining linalg.tensor_reshape operations with a sequence of operations including memref.reshape. We didn’t find a way do do this using existing tools.

@dpotop I ran into the same issue. Is there any chance you can point me to the code for the pass you wrote?

Do you make a copy of tensor before memref.reshape or you always use replacement as you said?

Do you run your pass before finalizing-bufferize?

For some reason, I don’t need finalizing-bufferize. Bufferization is complete without this pass.