[Mlir-commits] [mlir] [mlir][bufferization] implement BufferizableOpInterface for concat op (PR #140171)

[Jeremy Kun]

j2kun wrote:

I completely forgot about dynamic tensors, but in trying to support this I may need some advice.

The current snapshot has my attempt. It transforms

```mlir
func.func @tensor.concat_dynamic(%f: tensor<8x?xf32>, %g: tensor<8x?xf32>) -> tensor<8x?xf32> {
  %t = tensor.concat dim(1) %f, %g : (tensor<8x?xf32>, tensor<8x?xf32>) -> tensor<8x?xf32>
  return %t : tensor<8x?xf32>
}
```

Into

```mlir
#map = affine_map<()[s0, s1] -> (s0 + s1)>
module {
  func.func @tensor.concat_dynamic(%arg0: tensor<8x?xf32>, %arg1: tensor<8x?xf32>) -> tensor<8x?xf32> {
    %0 = bufferization.to_memref %arg1 : tensor<8x?xf32> to memref<8x?xf32>
    %1 = bufferization.to_memref %arg0 : tensor<8x?xf32> to memref<8x?xf32>
    %c1 = arith.constant 1 : index
    %dim = memref.dim %1, %c1 : memref<8x?xf32>
    %dim_0 = memref.dim %0, %c1 : memref<8x?xf32>
    %2 = affine.apply #map()[%dim, %dim_0]
    %alloc = memref.alloc(%2) {alignment = 64 : i64} : memref<8x?xf32>
    %c-9223372036854775808 = arith.constant -9223372036854775808 : index
    %c0 = arith.constant 0 : index
    %subview = memref.subview %alloc[0, %c0] [8, %dim] [1, 1] : memref<8x?xf32> to memref<8x?xf32, strided<[?, 1], offset: ?>>
    memref.copy %1, %subview : memref<8x?xf32> to memref<8x?xf32, strided<[?, 1], offset: ?>>
    %3 = arith.addi %c0, %dim : index
    %subview_1 = memref.subview %alloc[0, %3] [8, %dim_0] [1, 1] : memref<8x?xf32> to memref<8x?xf32, strided<[?, 1], offset: ?>>
    memref.copy %0, %subview_1 : memref<8x?xf32> to memref<8x?xf32, strided<[?, 1], offset: ?>>
    %4 = bufferization.to_tensor %alloc : memref<8x?xf32> to tensor<8x?xf32>
    return %4 : tensor<8x?xf32>
  }
}
```

I am mostly ignorant of how dynamic dimensions are supposed to work in MLIR in general, and bufferization acutely. The alloc seems fine to me? (using the affine map with memref.dim)

But the existence of this `arith.constant -9223372036854775808` that I did not create explicitly is a bit concerning. It obviously goes away with `canonicalize`, but is that supposed to be a sentinel for `kDynamic`? I don't understand why it's materialized in the IR, and it feels like a minor red flag to me.

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