[Mlir-commits] [mlir] [mlir][mem2reg] Promote whole-buffer memref to a vector SSA value (PR #211880)

[Jianhui Li]

Jianhui-Li wrote:

@banach-space  Thanks for your review! Let me answer the questions below. 
> Rather than recovering the SSA representation after bufferization, could this problem be addressed at the tensor level, before aliasing is introduced?
> 
> MLIR already provides several potentially relevant transformations:
> 
> * https://mlir.llvm.org/docs/Dialects/Transform/#transformloophoist_loop_invariant_subsets-transformhoistloopinvariantsubsetsop
> * https://mlir.llvm.org/docs/Dialects/Transform/#transformbufferizationbuffer_loop_hoisting-transformbufferloophoistingop
> * https://mlir.llvm.org/docs/Dialects/Transform/#transformstructuredhoist_redundant_vector_transfers-transformhoistredundantvectortransfersop

Thanks for the pointers. I actually was not aware of these passes when I came up with this PR.  I just tried them out and they do work well for some of use cases, but I also identify use cases that these passes don't help.

My motivation is to replace these temporary buffers (after bufferization) to vectors, and mem2reg appears to me the right solution to the problem, probably just because mem2reg is like a fundamental concept to compiler guys. Also, the vector/memref IR is the right place to address the issue since it is the lowerest target-independent representation.  Other than being conceptually well-established, the mem2reg implementation does bring extra value since it traces back to definition so it actually deals with control flow well (see my example below).  

> Could you provide a minimal end-to-end example showing:
> 
> 1. the tensor-level IR before bufferization;
> 2. the relevant transformation pipeline;
> 3. the resulting allocation and transfers; and
> 4. why the existing hoisting transformations cannot eliminate them?

This is the minimum motivation example at the linalg level. It is a simple GEMM but with an early exit: the function takes a parameter dyn_k, if the %k index for the k-loop is larger than dyn_k, it ends the computation and exit early.  
```
func.func @gemm_k_early_exit(%A: tensor<4x16xf32>, %B: tensor<16x4xf32>,
                             %C: tensor<4x4xf32>, %dyn_k: index) -> tensor<4x4xf32> {
  %c0 = arith.constant 0 : index
  %c4 = arith.constant 4 : index
  %c16 = arith.constant 16 : index
  %cst = arith.constant 0.0 : f32
  // Local accumulator tile D for this GEMM, zero-initialized.
  %e = tensor.empty() : tensor<4x4xf32>
  %dinit = linalg.fill ins(%cst : f32) outs(%e : tensor<4x4xf32>) -> tensor<4x4xf32>
  %d = scf.for %k = %c0 to %c16 step %c4 iter_args(%acc = %dinit) -> (tensor<4x4xf32>) {
    // Early exit: only accumulate K tiles whose offset is below dyn_k.
    %inrange = arith.cmpi slt, %k, %dyn_k : index
    %dnew = scf.if %inrange -> (tensor<4x4xf32>) {
      %atile = tensor.extract_slice %A[0, %k] [4, 4] [1, 1] : tensor<4x16xf32> to tensor<4x4xf32>
      %btile = tensor.extract_slice %B[%k, 0] [4, 4] [1, 1] : tensor<16x4xf32> to tensor<4x4xf32>
      %mm = linalg.matmul ins(%atile, %btile : tensor<4x4xf32>, tensor<4x4xf32>)
                          outs(%acc : tensor<4x4xf32>) -> tensor<4x4xf32>
      scf.yield %mm : tensor<4x4xf32>
    } else {
      scf.yield %acc : tensor<4x4xf32>
    }
    scf.yield %dnew : tensor<4x4xf32>
  }
  // Accumulate the local tile D into the C parameter (destination-passing): C = C + D.
  %cout = linalg.add ins(%C, %d : tensor<4x4xf32>, tensor<4x4xf32>)
                     outs(%C : tensor<4x4xf32>) -> tensor<4x4xf32>
  return %cout : tensor<4x4xf32>
}
```

The pipeline is following: 
//   mlir-opt cond-accumulator-mem2reg.mlir -transform-interpreter | \
//   mlir-opt --canonicalize --cse \
//     --one-shot-bufferize="bufferize-function-boundaries allow-return-allocs-from-loops" \
//     --canonicalize --cse --promote-buffers-to-stack --mem2reg
module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main(%arg0: !transform.any_op {transform.readonly}) {
    %f = transform.structured.match ops{["func.func"]} in %arg0 : (!transform.any_op) -> !transform.any_op
    %f2 = transform.structured.vectorize_children_and_apply_patterns %f {vectorize_padding}
        : (!transform.any_op) -> !transform.any_op
    transform.yield
  }
}

The result will show the temporary buffer introduced by bufferization is replaced by vector (appeared as loop iter arg). But if you take out the last pass mem2reg, the memref are not replaced, even if I extend the pipeline to include the passes you mentioned. 

My understanding is that the existing hoisting mechanism only hoists the buffer out of loop if the buffer is loop-invariant memory location. But the case above has a branch inside so it breaks the assumption.  The limitation sounds to me non-trivial and it is questionable whether we want to patch those, rather than extending a well-known mechanism.  

> More generally, this transformation appears fairly specialised, and some of its functionality may overlap with existing tensor-level or vector-transfer hoisting. Perhaps I am missing an important distinction, but based on the current motivation I am not yet convinced that the upstream community should maintain this additional mechanism.

I agree that the motivation part is not well described. I hope that my explanation and example above can help clarify. 

I am also interested to know specific limitation that you are aware for this approach by expanding mem2reg to cover vector. So we can discuss whether these limitations are fundamental or it is due to this initial PR. 

> As a side note, I find the summary quite long and difficult to follow. It contains implementation details that distract from the central design question. 

I will update the summary.  Thanks for the feedback.


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