Linalg.tiled_loop does not bufferize

Hello all,

I’m trying to bufferize linalg.tiled_loop which takes two tensors, but it fails with failed to legalize operation 'linalg.tiled_loop' error message. :frowning:

The linalg.tiled_loop looks like this:

    %7 = linalg.tiled_loop (%i) = (%c0) to (%c18) step (%c1) ins (%arg3 = %arg0: tensor<18x131072xi64>, %arg4 = %5: tensor<18x1xi64>) outs (%arg5 = %6: tensor<18x131072xi64>) {
      %slice1 = tensor.extract_slice %arg3[%i, 0] [1, 131072] [1, 1] : tensor<18x131072xi64> to tensor<1x131072xi64>
      %slice2 = tensor.extract_slice %arg4[%i, 0] [1, 1] [1, 1] : tensor<18x1xi64> to tensor<1x1xi64>
      ... (has a call)
      %res = tensor.insert_slice %28 into %arg5[%i, 0] [1, 131072] [1, 1] : tensor<1x131072xi64> into tensor<18x131072xi64>
      linalg.yield %res : tensor<18x131072xi64>
    }

The loop satisfies the conditions described in the document (link) – args must be extract_sliced, the yielded result must be insert_slice.
I have no clue why it is failing. :confused: Does anyone have ideas?

This is the full code: bufferize-fail.mlir · GitHub

Hi @aqjune,

the IR that you provided seems fine. At the moment, bufferization for tiled_loop in MLIR Core is supported only by comprehensive bufferization pass. If you want to stick to the usual “partial” bufferization, then you can find the pass for it in TensorFlow link. Notice, that the bufferization in that case is done in two steps. You can see the invocation in the test.

2 Likes

Thank you very much! :slight_smile:

1 Like