Bufferization with ControlFlow Asserts

MLIR
1

Hello team,

I am still a beginning so please correct me on any misunderstandings and help me with gaps. That being said I am trying to test out some lowerings in MLIR particularly with this sample test.mlir file:

module {
    func.func @main() {
        %cst = arith.constant dense<14> : tensor<1xi64>
        %cst_0 = arith.constant dense<15> : tensor<1xi64>
        %cst_1 = arith.constant dense<[[[-1], [-2], [10]]]> : tensor<1x3x1xi64>

        %800 = arith.constant 0 : index
        %801 = arith.cmpi slt, %800, %800 : index
        cf.assert %801, "index must be smaller than dim size"

        %0:2 = linalg.generic {
            indexing_maps = [
                affine_map<(d0, d1) -> (d0)>,
                affine_map<(d0, d1) -> (d0)>,
                affine_map<(d0, d1) -> (d0, d1, d0)>], 
            iterator_types = ["reduction", "reduction"]} 
        ins(%cst : tensor<1xi64>) outs(%cst_0, %cst_1 : tensor<1xi64>, tensor<1x3x1xi64>) {
        ^bb0(%in: i64, %out: i64, %out_3: i64):
            %1 = arith.minsi %out_3, %out_3 : i64
            linalg.yield %1, %out : i64, i64
        } -> (tensor<1xi64>, tensor<1x3x1xi64>)
        return
    }
}

I am running:

mlir-opt test.mlir -one-shot-bufferize -buffer-deallocation-pipeline

And I am getting this error on the cf.assert op:

error: ops with unknown memory side effects are not supported

Based on the documentation in OwnershipBasedBufferDeallocation.cpp, cf.assert sounds like it should pass verifyOperationPreconditions and be ignored as it does not operate on buffers and has no side effects. Is this a bug, or is there a pass I am missing that is needed before bufferization or some other requirements?

I want to lower MLIR files that need to be bufferized and may have cf.assert. I can create a pass that will convert cf.assert into the supported cf.cond_br however I would prefer not to make this pass for this specific use case.

Thank you for you time and support.

2

We should probably change the order of the checks in LogicalResult BufferDeallocation::verifyOperationPreconditions(Operation *op). If an op does not have buffer semantics, we can allow the op. (Does not matter if it implements the MemoryEffectOpInterface or not.) Do you want to prepare a PR on Github? Otherwise, I can take a look over the weekend.

3

Playing around with this proposal, I notice that reordering the checks does not solve my issue as cf.assert does indeed use buffer semantics. While I am not saying your change is not useful, would the solution to my problem be implementing the MemoryEffectOpInterface for cf.assert? I see it is the only op in that dialect that does not implement this interface. Is there a reason?

4

The op does not have memref operands, result or block arguments. So it should not have buffer semantics.

That would also work. I’m just wondering if the fact that the program may stop executing if the assertion fails is a side effect that should be modeled.

5

Thanks so much for all of the help. Just to follow up then, will there be or is there already a fix for this issue?

6

Feel free to send a PR on Github. I can review it.

7

Here is the pr: Bufferization with ControlFlow Asserts by mccowanzhang · Pull Request #95868 · llvm/llvm-project · GitHub
I implemented the changes we spoke about and additionally found a bug in a helper function. I have verified these changes have solved my issue. Thanks again for all your support.