[Mlir-commits] [mlir] [mlir][tosa] Lower concat to tensor.concat (PR #231726)

[Jack Frankland]

FranklandJack wrote:

Hey, this is very reasonable. I guess my only concern here is that this might have performance implications for downstream projects. I ran the resulting IR in both cases for concatenating two constant tensors through bufferization and the new rewrite results in two allocations for the input one for the result and a two copies into subviews, whereas the the IR in the current flow results in a single allocation for the result and then a pair of fills into subviews, something like this:
```
module {                                                                                                                                                                                                                                        func.func @slice_form() -> memref<10xf32> {
    %cst = arith.constant 1.000000e+00 : f32
    %cst_0 = arith.constant 2.000000e+00 : f32
    %alloc = memref.alloc() alignment = 64 : memref<10xf32>
    %subview = memref.subview %alloc[0] [4] [1] : memref<10xf32> to memref<4xf32, strided<[1]>>
    linalg.fill ins(%cst : f32) outs(%subview : memref<4xf32, strided<[1]>>)
    %subview_1 = memref.subview %alloc[4] [6] [1] : memref<10xf32> to memref<6xf32, strided<[1], offset: 4>>
    linalg.fill ins(%cst_0 : f32) outs(%subview_1 : memref<6xf32, strided<[1], offset: 4>>)
    return %alloc : memref<10xf32>
  }
  func.func @concat_form() -> memref<10xf32> {
    %cst = arith.constant 1.000000e+00 : f32
    %cst_0 = arith.constant 2.000000e+00 : f32
    %alloc = memref.alloc() alignment = 64 : memref<4xf32>
    linalg.fill ins(%cst : f32) outs(%alloc : memref<4xf32>)
    %alloc_1 = memref.alloc() alignment = 64 : memref<6xf32>
    linalg.fill ins(%cst_0 : f32) outs(%alloc_1 : memref<6xf32>)
    %alloc_2 = memref.alloc() alignment = 64 : memref<10xf32>
    %subview = memref.subview %alloc_2[0] [4] [1] : memref<10xf32> to memref<4xf32, strided<[1]>>
    memref.copy %alloc, %subview : memref<4xf32> to memref<4xf32, strided<[1]>>
    %subview_3 = memref.subview %alloc_2[4] [6] [1] : memref<10xf32> to memref<6xf32, strided<[1], offset: 4>>
    memref.copy %alloc_1, %subview_3 : memref<6xf32> to memref<6xf32, strided<[1], offset: 4>>
    return %alloc_2 : memref<10xf32>
  }
}
```

There is a rewrite pattern in `ConcatOpPatterns.cpp` called `DecomposeTensorConcatOp` that restores the current lowering (two insertions) but this isn't run as part of canonicalization so we might want to to make it clear that downstream compilers may want to run this rewrite. Otherwise LGTM!

https://github.com/llvm/llvm-project/pull/231726