Should I merge this now?

Hi @matthias-springer . Yes, it is ready to merge. Thank you.

This change breaks bufferization of ops that implement the FunctionOpInterface but have a terminator that doesn't have the ReturnLike trait, such as transform.named_sequence and transform.yield. For example, the following snippet now breaks:

// mlir-opt test.mlir -one-shot-bufferize="bufferize-function-boundaries=1"
module attributes {transform.with_named_sequence} {
  transform.named_sequence @__transform_main() {
    transform.yield 
  }
}

This produes the following error:

test.mlir:1:54: error: cannot bufferize a FuncOp with tensors and without a unique ReturnOp
module attributes {transform.with_named_sequence} {  "transform.named_sequence"() <{function_type = () -> (), sym_name = "__transform_main"}> ({
                                                     ^
test.mlir:1:54: note: see current operation: 
"transform.named_sequence"() <{function_type = () -> (), sym_name = "__transform_main"}> ({
  "transform.yield"() : () -> ()
}) : () -> ()

Also note that the error message still hardcodes FuncOp.

Finally, there might be other cases where people didn't expect their FunctionOpInterfaces to be bufferized and now be surprised but I guess that that's acceptable breakage, for which we get better genericity in return.

I think the fix is as easy to add the ReturnLike trait to transform.yield. I'll create a PR and run in through the test suite to check...

I have just created #111408, but I am getting second thoughts: AFAIU, FunctionOpInterface does not prescribe the ops implementing it to have ReturnLike terminators, so the pass should not assume this kind of terminators.

@tzunghanjuang Can you check how we are using the ReturnLike op in this file? I suspect it's to gather aliasing + equivalence information. If FunctionOpInterface is really not tied to the ReturnLike trait (i.e., the trait is not mentioned in the FunctionOpInterface documentation as a required terminator), then this analysis can not really be done just with the interface.

In that case, we could also make the FunctionOpInterface stricter: by requiring that function ops that have a result value must have a ReturnLike op. And verify result/operand types. I think we already do that for CallOpInterface + FunctionOpInterface today.

@matthias-springer Yes, ReturnLike is used during aliasing/equivalence analysis and getting the order of FunctionOpInterfaces. I have also checked that FunctionOpInterface is indeed not tied to ReturnLike. I will create another PR to update it.

Hi, @matthias-springer. I just want to make sure if I update the ReturnLike restriction correctly for FunctionOpInterface. I would like to make the verifyBody in FunctionInterfaces.td check the ReturnLike trait and the types. Or is there any other place more appropriate for this change? Thank you.

This change breaks bufferization of ops that implement the FunctionOpInterface but have a terminator that doesn't have the ReturnLike trait

Linaro's bots have been very behind but we finally noticed that this has caused SME test failures on AArch64.

Last good build: https://lab.llvm.org/staging/#/builders/125/builds/388
First bad: https://lab.llvm.org/staging/#/builders/125/builds/389

$ git bisect good
2026501cf107fcb3cbd51026ba25fda3af823941 is the first bad commit
commit 2026501cf107fcb3cbd51026ba25fda3af823941
Author: Tzung-Han Juang <tzunghan.juang@gmail.com>
Date:   Tue Oct 1 09:58:52 2024 -0400

    [MLIR] Make `OneShotModuleBufferize` use `OpInterface` (#110322)

    **Description:**
    This PR replaces a part of `FuncOp` and `CallOp` with
    `FunctionOpInterface` and `CallOpInterface` in `OneShotModuleBufferize`.
    Also fix the error from an integration test in the a previous PR
    attempt. (https://github.com/llvm/llvm-project/pull/107295)

This bot is running SME tests using qemu for emulation. Check the cmake stage to see how to enable that. You will need a recent version of qemu-aarch64 installed.

@banach-space I'll leave it to you to to handle, the linked PR may already solve the issue.

Thanks for flagging this @DavidSpickett.

This can be reproduced without requiring an emulator. Plain mlir-opt is fine. From you build dir:

bin/mlir-opt ../../mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir -transform-interpreter -test-transform-dialect-erase-schedule
../../mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir:56:3: error: cannot bufferize a FuncOp with tensors and without a unique ReturnOp
  transform.named_sequence @__transform_main(%module : !transform.any_op {transform.consumed}) {
  ^
../../mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir:56:3: note: see current operation:
"transform.named_sequence"() <{arg_attrs = [{transform.consumed}], function_type = (!transform.any_op) -> (), sym_name = "__transform_main"}> ({
^bb0(%arg0: !transform.any_op):
  %0 = "transform.structured.match"(%arg0) <{ops = ["linalg.matmul"]}> : (!transform.any_op) -> !transform.any_op
  %1:4 = "transform.structured.tile_using_for"(%0) <{scalable_sizes = array<i1: true, true, false>, static_sizes = array<i64: 4, 4, 1>}> : (!transform.any_op) -> (!transform.any_op, !transform.any_op, !transform.any_op, !transform.any_op)
  "transform.structured.vectorize"(%1#0) <{scalable_sizes = array<i1: true, true, false>, static_vector_sizes = array<i64: 4, 4, 1>}> : (!transform.any_op) -> ()
  %2 = "transform.bufferization.one_shot_bufferize"(%arg0) <{allow_return_allocs_from_loops = false, allow_unknown_ops = false, bufferize_function_boundaries = true, check_parallel_regions = true, dump_alias_sets = false, memcpy_op = "memref.copy", print_conflicts = false, test_analysis_only = false}> : (!transform.any_op) -> !transform.any_op
  %3 = "transform.structured.match"(%2) <{ops = ["func.func"]}> : (!transform.any_op) -> !transform.any_op
  "transform.apply_patterns"(%3) <{max_iterations = -1 : i64, max_num_rewrites = -1 : i64}> ({
    "transform.apply_patterns.vector.lower_masked_transfers"() : () -> ()
    "transform.apply_patterns.vector.transfer_permutation_patterns"() : () -> ()
    "transform.apply_patterns.vector.reduction_to_contract"() : () -> ()
  }) : (!transform.any_op) -> ()
  "transform.apply_patterns"(%3) <{max_iterations = -1 : i64, max_num_rewrites = -1 : i64}> ({
    "transform.apply_patterns.vector.lower_contraction"() <{lowering_strategy = 2 : i32}> : () -> ()
    "transform.apply_patterns.vector.lower_masks"() : () -> ()
    "transform.apply_patterns.vector.rank_reducing_subview_patterns"() : () -> ()
    "transform.apply_patterns.canonicalization"() : () -> ()
  }) : (!transform.any_op) -> ()
  %4 = "transform.structured.hoist_redundant_vector_transfers"(%3) : (!transform.any_op) -> !transform.any_op
  %5 = "transform.structured.match"(%2) <{interface = 2 : i32}> : (!transform.any_op) -> !transform.any_op
  "transform.apply_licm"(%5) : (!transform.any_op) -> ()
  "transform.loop.hoist_loop_invariant_subsets"(%5) : (!transform.any_op) -> ()
  "transform.yield"() : () -> ()
}) : () -> ()
../../mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir:71:18: error: bufferization failed
    %bufferize = transform.bufferization.one_shot_bufferize %module
                 ^

@tzunghanjuang or @matthias-springer , could you help us with this? This is not really my area of expertise 😅 Link to the test:

• https://github.com/llvm/llvm-project/blob/main/mlir/test/Integration/Dialect/Linalg/CPU/ArmSME/matmul.mlir

@banach-space , if you want to have your test pass while the PR above is discussed: you can wrap matmul and main in a payload module (i.e., module @payload attributes { transform.target_tag = "payload" } and use the debug-payload-root-tag in your RUN command. RUN: mlir-opt --transform-interpreter="debug-payload-root-tag=payload". You may also need to separate your RUN command into two mlir-opts similar to this other test after erasing the transform module: https://github.com/tzunghanjuang/llvm-project/blob/4a60687b1f744f3444855b39ae457c8efe6ae1bb/mlir/test/Examples/transform/ChH/full.mlir

@matthias-springer would it be desirable to have bufferization skip nested modules that do not contain a tensor values in their bodies? Does the builtin module implement hasTensorSemantics? Or at what point is hasTensorSemantics run / used? Because this error comes from the fact that bufferization is being run on the transform schedule but the functions there do not contain an op with a ReturnLike trait.

you can wrap matmul and main in a payload module (i.e., module @payload attributes { transform.target_tag = "payload" } and use the debug-payload-root-tag in your RUN command. RUN: mlir-opt --transform-interpreter="debug-payload-root-tag=payload". You may also need to separate your RUN command into two mlir-opts similar to this other test after erasing the transform module:

Thanks for these hints. I noticed that you have also rewritten the pass pipeline to use -pass-pipeline and used this nesting builtin.module(builtin.module(. That's quite a workaround 😅

Unfortunately, the steps outlined above weren't sufficient for the SME e2e tests to pass (#113117). I am running out of time and both Matthias and I are travelling this week, hence I am suggesting a revert:

• Revert "[MLIR] Make OneShotModuleBufferize use OpInterface (#110322)" #113124