[MLIR] Make OneShotModuleBufferize use OpInterface

[tzunghanjuang]

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. (#107295)

The below fixes skip CallOpInterface so that the assertions are not triggered.

llvm-project/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

Lines 254 to 259 in 8d78000

	static void equivalenceAnalysis(FunctionOpInterface funcOp,
	OneShotAnalysisState &state,
	FuncAnalysisState &funcState) {
	funcOp->walk([&](CallOpInterface callOp) {
	if (isa<func::CallIndirectOp>(callOp))
	return WalkResult::skip();

llvm-project/mlir/lib/Dialect/Bufferization/Transforms/OneShotModuleBufferize.cpp

Lines 311 to 315 in 8d78000

	// Collect function calls and populate the caller map.
	numberCallOpsContainedInFuncOp[funcOp] = 0;
	return funcOp.walk([&](CallOpInterface callOp) -> WalkResult {
	if (isa<func::CallIndirectOp>(callOp))
	return WalkResult::skip();

Related Discord Discussion: Link

[erick-xanadu]

Hi @tzunghanjuang , I left a couple of suggestions. The reason for this is to avoid hardcoding func::CallIndirectOp. But wait for the other maintainers to chime in, as I am not sure if this algorithm requires all functions to be statically known.

Thanks!

[matthias-springer]

Should I merge this now?

[tzunghanjuang]

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

[ingomueller-net]

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.

[ingomueller-net]

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...

[ingomueller-net]

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.

[matthias-springer]

@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.

[matthias-springer]

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.

[tzunghanjuang]

@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.

[tzunghanjuang]

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.

llvm-project/mlir/include/mlir/Interfaces/FunctionInterfaces.td

Lines 80 to 87 in 67160c5

	InterfaceMethod<[{
	Verify the contents of the body of this function.
	Note: The default implementation merely checks that if the entry block
	exists, it has the same number and type of arguments as the function type.
	}],
	"::llvm::LogicalResult", "verifyBody", (ins),
	/*methodBody=*/[{}], /*defaultImplementation=*/[{

[DavidSpickett]

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.

[banach-space]

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

[erick-xanadu]

@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

[erick-xanadu]

@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.

[banach-space]

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