[Mlir-commits] [mlir] [MLIR] Missed Optimization: Incomplete elimination of a tensor.empty (PR #228032)

[Ashutosh Sharma]

ashu3103 wrote:

Included the running example in the test.

Input:

```mlir
func.func @shared_empty(%arg0: i1, %arg1: f32, %arg2: f32, %arg3: tensor<16xf32>) -> tensor<16xf32> {
  %0 = tensor.empty() : tensor<4xf32>
  %1 = scf.if %arg0 -> (tensor<4xf32>) {
    %2 = linalg.fill ins(%arg1 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
    scf.yield %2 : tensor<4xf32>
  } else {
    %2 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
    scf.yield %2 : tensor<4xf32>
  }
  %inserted_slice = tensor.insert_slice %1 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
  return %inserted_slice : tensor<16xf32>
}
```

Command:

`mlir-opt --eliminate-empty-tensors --empty-tensor-to-alloc-tensor -cse -canonicalize /tmp/input.mlir`

Before PR:

```mlir
func.func @shared_empty(%arg0: i1, %arg1: f32, %arg2: f32, %arg3: tensor<16xf32>) -> tensor<16xf32> {
  %extracted_slice = tensor.extract_slice %arg3[4] [4] [1] : tensor<16xf32> to tensor<4xf32>
  %0 = bufferization.alloc_tensor() : tensor<4xf32>
  %1 = scf.if %arg0 -> (tensor<4xf32>) {
    %2 = linalg.fill ins(%arg1 : f32) outs(%extracted_slice : tensor<4xf32>) -> tensor<4xf32>
    scf.yield %2 : tensor<4xf32>
  } else {
    %2 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<4xf32>) -> tensor<4xf32>
    scf.yield %2 : tensor<4xf32>
  }
  %inserted_slice = tensor.insert_slice %1 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
  return %inserted_slice : tensor<16xf32>
}
```

After PR:

```mlir
func.func @shared_empty(%arg0: i1, %arg1: f32, %arg2: f32, %arg3: tensor<16xf32>) -> tensor<16xf32> {
  %extracted_slice = tensor.extract_slice %arg3[4] [4] [1] : tensor<16xf32> to tensor<4xf32>
  %0 = scf.if %arg0 -> (tensor<4xf32>) {
    %1 = linalg.fill ins(%arg1 : f32) outs(%extracted_slice : tensor<4xf32>) -> tensor<4xf32>
    scf.yield %1 : tensor<4xf32>
  } else {
    %1 = linalg.fill ins(%arg2 : f32) outs(%extracted_slice : tensor<4xf32>) -> tensor<4xf32>
    scf.yield %1 : tensor<4xf32>
  }
  %inserted_slice = tensor.insert_slice %0 into %arg3[4] [4] [1] : tensor<4xf32> into tensor<16xf32>
  return %inserted_slice : tensor<16xf32>
}
```

Previously, if use of an empty tensor appeared in both branches of a conditional, the replacement mechanism considered only one occurrence for subview substitution. Consequently, the empty tensor could not be completely eliminated, resulting in spurious tensor allocations.

@matthias-springer PTAL

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