[mlir] Op needs producer extract_slice op to bufferize in place

[Max191]

The following two func ops have different results when running OneShotBufferize:

// RUN:  mlir-opt test.mlir --pass-pipeline="builtin.module(one-shot-bufferize{test-analysis-only=true print-conflicts=true}, canonicalize, cse, canonicalize)"
#map = affine_map<(d0) -> (d0 mod 256)>
module {
  func.func @slice_bufferize_inplace(%2: tensor<2xf32>) -> tensor<2xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %3 = scf.forall (%arg0) in (1) shared_outs(%arg2 = %2) -> (tensor<2xf32>) {
      %extracted_slice = tensor.extract_slice %arg2[%arg0] [2] [1] : tensor<2xf32> to tensor<2xf32>
      %fill = linalg.fill ins(%cst : f32) outs(%extracted_slice : tensor<2xf32>) -> tensor<2xf32>
      scf.forall.in_parallel {
        tensor.parallel_insert_slice %fill into %arg2[%arg0] [2] [1] : tensor<2xf32> into tensor<2xf32>
      }
    } {mapping = [#gpu.thread<linear_dim_0>]}
    return %3 : tensor<2xf32>
  }

  func.func @no_slice_bufferize_outplace(%2: tensor<2xf32>) -> tensor<2xf32> {
    %cst = arith.constant 0.000000e+00 : f32
    %3 = scf.forall (%arg0) in (1) shared_outs(%arg2 = %2) -> (tensor<2xf32>) {
      %fill = linalg.fill ins(%cst : f32) outs(%arg2 : tensor<2xf32>) -> tensor<2xf32>
      scf.forall.in_parallel {
        tensor.parallel_insert_slice %fill into %arg2[%arg0] [2] [1] : tensor<2xf32> into tensor<2xf32>
      }
    } {mapping = [#gpu.thread<linear_dim_0>]}
    return %3 : tensor<2xf32>
  }
}

The result of bufferization analysis for these func ops is the following:

run mlir-opt test.mlir --pass-pipeline="builtin.module(one-shot-bufferize{test-analysis-only=true print-conflicts=true}, canonicalize, cse, canonicalize)"

module {
  func.func @slice_bufferize_inplace(%arg0: tensor<2xf32>) -> tensor<2xf32> attributes {"W_1[NOT-WRITABLE: bbArg 0]"} {
    %cst = arith.constant 0.000000e+00 : f32
    %0 = scf.forall (%arg1) in (1) shared_outs(%arg2 = %arg0) -> (tensor<2xf32>) {
      %extracted_slice = tensor.extract_slice %arg2[%arg1] [2] [1] {__inplace_operands_attr__ = ["true", "none"]} : tensor<2xf32> to tensor<2xf32>
      %1 = linalg.fill {__inplace_operands_attr__ = ["none", "true"]} ins(%cst : f32) outs(%extracted_slice : tensor<2xf32>) -> tensor<2xf32>
      scf.forall.in_parallel {
        tensor.parallel_insert_slice %1 into %arg2[%arg1] [2] [1] {__inplace_operands_attr__ = ["true", "true", "none"]} : tensor<2xf32> into tensor<2xf32>
      }
    } {__inplace_operands_attr__ = ["false"], mapping = [#gpu.thread<linear_dim_0>]}
    return {__inplace_operands_attr__ = ["true"]} %0 : tensor<2xf32>
  }
  func.func @no_slice_bufferize_outplace(%arg0: tensor<2xf32>) -> tensor<2xf32> attributes {"W_0[NOT-WRITABLE: bbArg 0]"} {
    %cst = arith.constant 0.000000e+00 : f32
    %0 = scf.forall (%arg1) in (1) shared_outs(%arg2 = %arg0) -> (tensor<2xf32>) {
      %1 = linalg.fill {"C_0[CONFL-WRITE: 1]", __inplace_operands_attr__ = ["none", "false"]} ins(%cst : f32) outs(%arg2 : tensor<2xf32>) -> tensor<2xf32>
      scf.forall.in_parallel {
        tensor.parallel_insert_slice %1 into %arg2[%arg1] [2] [1] {"C_0[READ: 1]", __inplace_operands_attr__ = ["true", "true", "none"]} : tensor<2xf32> into tensor<2xf32>
      }
    } {"C_0[DEF: bbArg 1]", __inplace_operands_attr__ = ["false"], mapping = [#gpu.thread<linear_dim_0>]}
    return {__inplace_operands_attr__ = ["true"]} %0 : tensor<2xf32>
  }
}

The key difference is that the linalg.fill op bufferizes in place when there is an extract_slice producer, but it bufferizes out of place when the block argument is the init arg of the fill. The extract_slice is a full slice, so both IRs are equivalent, and they should bufferize the same way. This seems like a bug in the bufferization analysis somewhere.

[Max191]

@matthias-springer any ideas why this would be happening? I am not familiar with the inner workings of the bufferization analysis, so it is not easy for me to pinpoint the root issue.

[matthias-springer]

The analysis is looking for RaW conflicts in the form of "definition, write, read" triples. In your example:

• definition: %arg2
• write: linalg.fill operand 1
• read: parallel_insert_slice operand 1

There is a special case here where the entire tensor is overwritten by the source of the parallel_insert_slice, so the dest is never read. We could account for that here: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp#L932

Check if the op operand is the dest and if the size of the source equals the size of the dest. In that case, return false. That should fix it.

Generally speaking, we have some limitations in the analysis around extract_slice/insert_slice. We usually expect them to appear in pairs. Otherwise, we sometimes insert extra copies.

[matthias-springer]

Actually, we can just the same as for insert_slice: https://github.com/llvm/llvm-project/blob/main/mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp#L647. It's implemented there.

[Max191]

Check if the op operand is the dest and if the size of the source equals the size of the dest. In that case, return false. That should fix it.

I see, that makes sense. In this case, the offsets are not constant, and they come from a loop iteration variable that has a mapping, so it can't be folded. Is it still okay to consider such a case as a full slice? It would imply that parallel_insert_slice is guaranteed to be in-bounds. Do the semantics of the op support that assumption?