Cyclic dependencies between dialects when adding a new canonicalization pattern

MLIR
1

Hi!
I am trying to add a canonicalization pattern that will handle tensor.dim of tosa.transpose,
for example:

Input:

func.func @inserts_slice_dynamic(%arg0: tensor<1x1x?x1xsi16>, %arg1: tensor<1x8x1x?xsi16>) -> (tensor<1x8x1x?xsi16>) {
  %0 = "tosa.const"() <{value = dense<[0, 3, 1, 2]> : tensor<4xi32>}> : () -> tensor<4xi32>
  %1 = tosa.transpose %arg0, %0 : (tensor<1x1x?x1xsi16>, tensor<4xi32>) -> tensor<1x1x1x?xsi16>
  %c3 = arith.constant 3 : index
  %dim = tensor.dim %1, %c3 : tensor<1x1x1x?xsi16>
  %inserted_slice = tensor.insert_slice %1 into %arg1[0, 0, 0, 0] [1, 1, 1, %dim] [1, 1, 1, 1] : tensor<1x1x1x?xsi16> into tensor<1x8x1x?xsi16>
  return %inserted_slice : tensor<1x8x1x?xsi16>
}

Output:

func.func @inserts_slice_dynamic(%arg0: tensor<1x1x?x1xsi16>, %arg1: tensor<1x8x1x?xsi16>) -> (tensor<1x8x1x?xsi16>) {
  %c2 = arith.constant 2 : index
  %dim = tensor.dim %arg0, %c2 : <1x1x?x1xsi16>
  %inserted_slice = tensor.insert_slice %1 into %arg1[0, 0, 0, 0] [1, 1, 1, %dim] [1, 1, 1, 1] : tensor<1x1x1x?xsi16> into tensor<1x8x1x?xsi16>
  return %inserted_slice : tensor<1x8x1x?xsi16>
}

I think the pattern should be under the canonicalization patterns of tensor::DimOp.
This means Tensor dialect should depend on Tosa dialect.
Since Tosa dialect already depends on Tensor dialect, I get cyclic dependency.

In the meantime I added my pattern locally to Tosa canonicalization patterns, just to see it works.
But I really don’t think it should be there…

Does anybody have an idea how I can solve it?

Thanks in advance,
Maya

2

It is possible to add the canonicalization pattern to tosa, regardless of which ops it is matching. Iiuc, what you are trying to do is very much a “tosa optimization” in that it wouldn’t make sense at all without the tosa.transpose.

The usual advice applies of “are you sure this is a canonicalization?” vs some other kind of optimization, but that is a question for tosa.

3

Its good to consider indeed if this is the canonical form (vs an optimization part of another set of patterns). But as Stella said, you can add this pattern to TOSA even though it acts on tensor dialect one - canonicalization patterns are really just patterns that are consumed by the canonicalize pass. The convenience of expressing is also a downfall as too easy to add where it shouldn’t be :slight_smile: (which may also indicate we should make the others easier).

This feels like a limited form of dim propagation, feels like a general construct in the larger scheme of which this is a local pattern. Might be useful, but also seems like something that could be subsumed later.

4

You can have the tosa op implement the ReifyRankedShapedTypeOpInterface . Then the ResolveRankedShapeTypeResultDims pass will resolve the tensor.dim operation in terms of its operands. This will also avoid the circular dependency. (Its unfortunate that this pass is in memref dialect, but reason is that it can resolve dims of memref ops and tensor ops at the same time, and was just put there.

5

Thanks everybody! I uploaded a patch: [mlir] Add ReifyRankedShapedTypeOpInterface to tosa::TransposeOp by amrami · Pull Request #88890 · llvm/llvm-project · GitHub