Tile and fuse support

[ofri_frishman]

When working on adding support on a downstream project for creating groups of ops that are tiled together via using TilingInterface and scf::tileConsumerAndFuseProducersUsingSCF I’ve encountered some questions regarding how best to add such support:

1. Only part of the upstream ops implement TilingInterface, but there are other ops that I would like to add to a group of ops that should be tiled together which currently don’t implement the interface. For example tensor.expand_shape and tosa.transpose. Is this the case simply because there was no need for supporting tiling on these ops up till now, or is there a reason why not to implement the interface for those ops?
2. It seems that the TilingInterface provides everything needed for implementing the change in the IR from untiled ops to tiled ops, but I find it hard to use the interface for planning which ops to tile together.
   For example - I’m trying to avoid cases where data is recomputed. This could happen in cases where some ops are dependent on all the induction variables of the iteration domain, but other ops aren’t dependent on all of them. The ops that aren’t dependent on all induction variables will recompute data if placed in the inner most loop. A way to avoid recomputing data could be to not place all ops in the same nesting level or creating a different grouping of ops.
   There are methods in the LinalgStructuredInterface that seem to provide the missing information. The mapIterationSpaceDimToAllOperandDims provides a mapping between iteration domain dims and args/return values dims. Using this mapping can be used to understand which ops will depend on which induction variables and plan groups of ops to be tiled together accordingly.
   Currently this mapping exists only in the linalg dialect and is limited to cases where the mapping is a projected permutation, but could it make sense to add such a method as an addition to the TilingInterface to assist in planning of groups of ops to tile together?

I haven’t found much information in the upstream project regarding how such planning is done, if there are other cases where downstream projects implemented such planning, it would be interesting to understand how that was approached.

[MaheshRavishankar]

@ofri_frishman

Only part of the upstream ops implement TilingInterface, but there are other ops that I would like to add to a group of ops that should be tiled together which currently don’t implement the interface. For example tensor.expand_shape and tosa.transpose. Is this the case simply because there was no need for supporting tiling on these ops up till now, or is there a reason why not to implement the interface for those ops?

In general it makes sense to only implement TilingInterface for “real compute ops”. expand_shape has copy semantics, but in general these should become a no-op (in a full compiler stack). I think there are very few ops in the tensor dialect that would match the semantics of having a TilingInterface. These are driven by experience using these transformations, and a some exploration. I dont think I can give you a hard-and-fast rule to go by. I can say that expand_shape definitely does not fit what would be a good candidate to implement TilingInterface in my book.

For ops from dialects like Tosa and Torch, the idea is to lower them to Linalg and rely on Linalgs implementation of TilingInterface since otherwise its a bit of redundancy in implementation.

@ofri_frishman

It seems that the TilingInterface provides everything needed for implementing the change in the IR from untiled ops to tiled ops, but I find it hard to use the interface for planning which ops to tile together.
For example - I’m trying to avoid cases where data is recomputed. This could happen in cases where some ops are dependent on all the induction variables of the iteration domain, but other ops aren’t dependent on all of them. The ops that aren’t dependent on all induction variables will recompute data if placed in the inner most loop. A way to avoid recomputing data could be to not place all ops in the same nesting level or creating a different grouping of ops.

These are meant to be a transformation, without taking any stance on heuristics or profitability. Those vary based on use cases. For example the scf::tileConsumersAndFuseProducers takes an options struct that contains a call back that allows caller to control what is fused and what is not. You can do some pre-analysis to make that decision and just use that analysis result in the callback. There is a space from some pre-existing callbacks that implement certain strategies (like no redundant computations), but the core transformation “does what it is told to do” without really thinking about profitability.

@ofri_frishman

For example - I’m trying to avoid cases where data is recomputed. This could happen in cases where some ops are dependent on all the induction variables of the iteration domain, but other ops aren’t dependent on all of them. The ops that aren’t dependent on all induction variables will recompute data if placed in the inner most loop. A way to avoid recomputing data could be to not place all ops in the same nesting level or creating a different grouping of ops.
There are methods in the LinalgStructuredInterface that seem to provide the missing information. The mapIterationSpaceDimToAllOperandDims provides a mapping between iteration domain dims and args/return values dims. Using this mapping can be used to understand which ops will depend on which induction variables and plan groups of ops to be tiled together accordingly.

I am not able to map this exactly. Maybe if you provide a more representative example it will be clearer what you are trying to do.

[ofri_frishman]

@MaheshRavishankar

For example the scf::tileConsumersAndFuseProducers takes an options struct that contains a call back that allows caller to control what is fused and what is not. You can do some pre-analysis to make that decision and just use that analysis result in the callback

That is exactly what I am trying to do, split the fuse and tile process into 2 parts:

1. First part is just doing an analysis (I called this planning in my original post) and deciding which ops to tile and what tile size to use without any changes to the IR
2. Second part is to make the changes in the IR based on the decisions from the analysis.

@MaheshRavishankar

I am not able to map this exactly. Maybe if you provide a more representative example it will be clearer what you are trying to do.

Sure.
If you consider the following IR

%A = linalg.elemwise_unary ins(%arg0 : tensor<8x16x32xf32>)
                         outs(%arg0: tensor<8x16x32xf32>) -> tensor<8x16x32xf32>
%B = linalg.elemwise_unary ins(%arg1 : tensor<8x32x14xf32>)
                         outs(%arg1: tensor<8x32x14xf32>) -> tensor<8x32x14xf32>
%1 = linalg.batch_matmul ins(%A , %B:  tensor<8x16x32xf32>, tensor<8x32x14xf32>) 
                         outs(%C : tensor<8x16x14xf32>)  -> tensor<8x16x14xf32>

There is several different possibilities regarding how to construct the loop nesting if we would like to tile and fuse these ops.
One option could be to tile all 3 dimensions of the linalg.batch_matmul output, such a decisions could lead to the following IR:

%0 = scf.for %arg3 = %c0 to %c8 step %c1 iter_args(%arg4 = %arg2) -> (tensor<8x16x14xf32>) {
  %1 = scf.for %arg5 = %c0 to %c16 step %c4 iter_args(%arg6 = %arg4) -> (tensor<8x16x14xf32>) {
    %2 = scf.for %arg7 = %c0 to %c14 step %c7 iter_args(%arg8 = %arg6) -> (tensor<8x16x14xf32>) {
      %extracted_slice = tensor.extract_slice %arg0[%arg3, %arg5, 0] [1, 4, 32] [1, 1, 1] : tensor<8x16x32xf32> to tensor<1x4x32xf32>
      %extracted_slice_0 = tensor.extract_slice %arg0[%arg3, %arg5, 0] [1, 4, 32] [1, 1, 1] : tensor<8x16x32xf32> to tensor<1x4x32xf32>
      %3 = linalg.elemwise_unary ins(%extracted_slice : tensor<1x4x32xf32>) outs(%extracted_slice_0 : tensor<1x4x32xf32>) -> tensor<1x4x32xf32>
      %extracted_slice_1 = tensor.extract_slice %arg1[%arg3, 0, %arg7] [1, 32, 7] [1, 1, 1] : tensor<8x32x14xf32> to tensor<1x32x7xf32>
      %extracted_slice_2 = tensor.extract_slice %arg1[%arg3, 0, %arg7] [1, 32, 7] [1, 1, 1] : tensor<8x32x14xf32> to tensor<1x32x7xf32>
      %4 = linalg.elemwise_unary ins(%extracted_slice_1 : tensor<1x32x7xf32>) outs(%extracted_slice_2 : tensor<1x32x7xf32>) -> tensor<1x32x7xf32>
      %extracted_slice_3 = tensor.extract_slice %arg8[%arg3, %arg5, %arg7] [1, 4, 7] [1, 1, 1] : tensor<8x16x14xf32> to tensor<1x4x7xf32>
      %5 = linalg.batch_matmul ins(%3, %4 : tensor<1x4x32xf32>, tensor<1x32x7xf32>) outs(%extracted_slice_3 : tensor<1x4x7xf32>) -> tensor<1x4x7xf32>
      %inserted_slice = tensor.insert_slice %5 into %arg8[%arg3, %arg5, %arg7] [1, 4, 7] [1, 1, 1] : tensor<1x4x7xf32> into tensor<8x16x14xf32>
      scf.yield %inserted_slice : tensor<8x16x14xf32>
    }
    scf.yield %2 : tensor<8x16x14xf32>
  }
  scf.yield %1 : tensor<8x16x14xf32>
}

When looking at the produced IR you can see that the linalg.batch_matmul uses all of the induction variables, but that each of the linalg.elemwise_unary uses only 2 of the induction variables. This could lead to each of linalg.elemwise_unary performing recomputes. The first one would be computed twice since it doesn’t use the inner loop induction variable, the second one would be computed four times since it doesn’t use the middle loop induction variable.
If we want to avoid such recomputations, one way to do this could be to tile only the first dimension of the linalg.batch_matmul output.
That could lead us to the following IR:

%0 = scf.for %arg3 = %c0 to %c8 step %c1 iter_args(%arg4 = %arg2) -> (tensor<8x16x14xf32>) {
  %extracted_slice = tensor.extract_slice %arg0[%arg3, 0, 0] [1, 16, 32] [1, 1, 1] : tensor<8x16x32xf32> to tensor<1x16x32xf32>
  %extracted_slice_0 = tensor.extract_slice %arg0[%arg3, 0, 0] [1, 16, 32] [1, 1, 1] : tensor<8x16x32xf32> to tensor<1x16x32xf32>
  %1 = linalg.elemwise_unary ins(%extracted_slice : tensor<1x16x32xf32>) outs(%extracted_slice_0 : tensor<1x16x32xf32>) -> tensor<1x16x32xf32>
  %extracted_slice_1 = tensor.extract_slice %arg1[%arg3, 0, 0] [1, 32, 14] [1, 1, 1] : tensor<8x32x14xf32> to tensor<1x32x14xf32>
  %extracted_slice_2 = tensor.extract_slice %arg1[%arg3, 0, 0] [1, 32, 14] [1, 1, 1] : tensor<8x32x14xf32> to tensor<1x32x14xf32>
  %2 = linalg.elemwise_unary ins(%extracted_slice_1 : tensor<1x32x14xf32>) outs(%extracted_slice_2 : tensor<1x32x14xf32>) -> tensor<1x32x14xf32>
  %extracted_slice_3 = tensor.extract_slice %arg4[%arg3, 0, 0] [1, 16, 14] [1, 1, 1] : tensor<8x16x14xf32> to tensor<1x16x14xf32>
  %3 = linalg.batch_matmul ins(%1, %2 : tensor<1x16x32xf32>, tensor<1x32x14xf32>) outs(%extracted_slice_3 : tensor<1x16x14xf32>) -> tensor<1x16x14xf32>
  %inserted_slice = tensor.insert_slice %3 into %arg4[%arg3, 0, 0] [1, 16, 14] [1, 1, 1] : tensor<1x16x14xf32> into tensor<8x16x14xf32>
  scf.yield %inserted_slice : tensor<8x16x14xf32>
}

I prefer that the outcome of tile and fuse for these ops will be the second option and am trying to find the right tools the achieve that outcome.
One way to achieve it could be to try and tile all the dimensions of the output of the linalg.batch_matmul and then look at the generated IR (the first example) and see which induction variables are used by which ops. Such an analysis would lead to understanding that only the induction variable of the outermost loop is used by all operations. Given that information, I could erase/rollback the generated IR and then tile the original ops again only on the first dimension of the output.
But, I’m trying to achieve this outcome without making any changes to the IR and then analyzing to outcome. Instead I would like to understand how the outcome would look just from analyzing the connections between operation values and operation iteration domain. linalg dialect seems to provide a way to do that via the mapIterationSpaceDimToAllOperandDims interface method and another one or two methods. If I use the information from this mapping for all the ops in the example I could understand that only the outer iteration dimension of the matmul is also an iteration dimension for the other 2 operations, and that each of the two inner iteration dimensions of the matmul is used only by one of the other operations.
Currently using this approach is limited to pretiling analysis for linalg dialect operations only, but if such mapping method/s would be added as part of the tiling interface then the approach could be used for all operations that implement the interface.

[MaheshRavishankar]

@ofri_frishman

I prefer that the outcome of tile and fuse for these ops will be the second option and am trying to find the right tools the achieve that outcome.

You should get that by just setting tile size to [8] instead of [8, 16, 14]. [8] is essentially [8, 0, 0] . So only tiles the first loop of batch-matmul and then greedy fusion pulls in everything else.

The tile and fuse implementation is really tile and then fuse. The loop structure you want is decided by the tiling of the last operation in the chain. Once that is done, all fusion does is pulls it into the loop (greedily, but controllable using the callback).

[ofri_frishman]

@MaheshRavishankar

You should get that by just setting tile size to [8] instead of [8, 16, 14]. [8] is essentially [8, 0, 0] . So only tiles the first loop of batch-matmul and then greedy fusion pulls in everything else.

I should clarify what I meant - I understand how to use the tile and fuse utility to tile only the first dimension.
What I’m working on is trying to be able to analyze the IR to understand what valid options exist regarding tiling and fusing ops together and the tradeoffs between the different options.
One such tradeoff is if data would be recomputed or not. That is a tradeoff I referenced in the example I gave, and for that example and that decision about tradeoffs, the right tile to use would be [8, 0, 0]. But there could be other tradeoffs such as estimated memory footprint required or various other information that could be relevant based on the target HW.
But I would like to do all this analysis before performing the actual tiling, and don’t currently see how I could use the tiling interface to do this analysis prior to tiling. I was thinking that perhaps some extension to the tiling interface could be useful for this analysis phase.

[MaheshRavishankar]

@ofri_frishman

But I would like to do all this analysis before performing the actual tiling, and don’t currently see how I could use the tiling interface to do this analysis prior to tiling. I was thinking that perhaps some extension to the tiling interface could be useful for this analysis phase.

While I dont have any strong opinions on what analysis could be performed (I assume each downstream user of these transformations have their own metrics to evaluate in terms how, how much recompute they can tolerate depending on their individual compute vs bandwidth balance point, etc.), the only thing I really have a strong opinion is on keeping heuristics and transformations separate. MLIR could have some “in-built heuristic” options, and contributions there would be welcome.

Maybe a useful example is the LinalgExtFusionInterface that we are using in IREE. We’ve been developing/using it in IREE, but in my view this should be folded into TilingInterface.