Page MenuHomePhabricator

[MLIR] Update the FunctionAndBlockSignatureConverter and NonVoidToVoidReturnOpConverter of Buffer Assignment
ClosedPublic

Authored by dfki-ehna on Mon, May 4, 7:39 AM.

Details

Summary

Making these two converters more generic. FunctionAndBlockSignatureConverter now moves only memref results (after type conversion) to the function argument and keeps other legal function results unchanged. NonVoidToVoidReturnOpConverter is renamed to NoBufferOperandsReturnOpConverter`. It removes only the buffer operands from the operands of the converted ReturnOp and inserts CopyOps to copy each buffer to the target function argument.

Input:

func @example(%arg0: tensor<4x8xf32>, %arg1: i1) -> (tensor<4x8xf32>, i1) {
    return %arg0, %arg1 : i1, tensor<4x8xf32>
}

Output:

func @example(%arg0: memref<4x8xf32>, %arg1: i1, %arg2: memref<4x8xf32>) -> (i1) {
    copy(%arg0, %arg2) : memref<4x8xf32>, memref<4x8xf32>
    return %arg1 : i1
}

Diff Detail

Event Timeline

dfki-ehna created this revision.Mon, May 4, 7:39 AM
Herald added a project: Restricted Project. · View Herald TranscriptMon, May 4, 7:40 AM
dfki-ehna edited the summary of this revision. (Show Details)Mon, May 4, 7:43 AM
dfki-ehna added reviewers: mehdi_amini, pifon2a, herhut.
dfki-ehna added a project: Restricted Project.
pifon2a edited the summary of this revision. (Show Details)Mon, May 4, 11:35 PM

nice! thanks!

mlir/include/mlir/Transforms/BufferPlacement.h
84
/// Converts the signature of the function using the type converter.
/// It adds an extra argument for each illegally-typed function
/// result to the function arguments. `BufferAssignmentTypeConverter`
/// is a helper `TypeConverter for this purpose. All the non-shaped types
/// of the input function will be converted to memref.
98
/// Converts the source `ReturnOp` to target `ReturnOp`, removes all
/// the buffer operands from the operands list, and inserts `CopyOp`s
/// for all buffer operands instead.
121

these 5 lines are just llvm::erase_if, right?

125
// For each memref type operand of the source `ReturnOp`, a new `CopyOp` is
// inserted that copies the buffer content from the operand to the target.
131

can this be moved outside of the loop?

mlir/test/Transforms/buffer-placement-preparation.mlir
15

please, reformat it so that it fits

herhut requested changes to this revision.Tue, May 5, 1:10 AM

Please add a test for the case that is currently asserted.

mlir/include/mlir/Transforms/BufferPlacement.h
114–115

In the case where the return has non-memref operands, this can fail even for correct code as the non-memref operands do not require a function argument. Maybe add a test for such a case?

This revision now requires changes to proceed.Tue, May 5, 1:10 AM
dfki-ehna updated this revision to Diff 262089.Tue, May 5, 6:06 AM

Addressing the comments.

dfki-ehna marked 8 inline comments as done.Tue, May 5, 6:53 AM
dfki-ehna added inline comments.
mlir/include/mlir/Transforms/BufferPlacement.h
114–115

Thanks, you are right. I changed the assert and updated the target.addDynamicallyLegalOp<FuncOp>. A few new test cases were also added.

pifon2a accepted this revision.Tue, May 5, 6:57 AM
herhut requested changes to this revision.Thu, May 7, 4:50 AM
herhut added inline comments.
mlir/lib/Transforms/BufferPlacement.cpp
454

it is remained -> it remains unchanged as a function result?

mlir/test/Transforms/buffer-placement-preparation.mlir
34

If this returned a memref before, why would it not return a memref now? We should only rewrite the return values that got rewritten from tensor to memref. Not the ones that already were a memref. I'd assume those would be left alone.

This revision now requires changes to proceed.Thu, May 7, 4:50 AM
dfki-ehna marked 2 inline comments as done.Thu, May 7, 5:30 AM
dfki-ehna added inline comments.
mlir/test/Transforms/buffer-placement-preparation.mlir
34

The idea of removing all Memref types from function results is proposed to avoid placing a DeallocOp after the return operation using buffer-placement pass. This case would happen if there is a defining operation for a return operand.

func @memref_in_function_results(%arg0: i1, %arg1: f16, %arg2:memref<2xf32>) -> (i1, memref<2xf32>, f16){
  %0 = alloc() : memref<2xf32>,
  lhlo.exp(%arg2, %0)
  return %arg0, %0, %arg1  : i1, memref<2xf32>, f16
}

Making it fail would be good enough for landing.

mlir/test/Transforms/buffer-placement-preparation.mlir
34

I think it is better to just not allow this case. If one has a memref that escapes in the original program, it should still escape in the rewritten program. We cannot support this use case at the moment it seems, so it would be better to fail.

I wonder what it would take, though, to handle this case. Would it suffice to know that the last use is a return and then treat the return like a dealloc from the local scope? This could be an alternative version of this rewriting that turns tensor typed returns into memref typed returns (and not arguments) and lets allocations escape. Would the main logic of placement still work in this setting? What would it take to support both modes?

I see a use for the latter mode in scenarios where the buffers for results are not provided from the context, which is particularly useful in the presence of dynamic shapes where the context might not know sizes.

Such a rewrite would not be globally correct (the caller would need to free the buffer after all and that cannot be decided statically) but at least locally per function it would be. For global correctness we could think in the direction of reference counting but that would be a next step.

dfki-ehna updated this revision to Diff 264225.EditedFri, May 15, 6:30 AM

BufferAssignmentPlacer now fails in cases that function returns buffers already. A note has also been added to the description of BufferPlacement.cpp about not supporting these cases. The related TODO in TensorsToBuffers.cpp has also been removed.

dfki-ehna marked 2 inline comments as done.Fri, May 15, 6:31 AM
mehdi_amini added inline comments.Fri, May 15, 4:23 PM
mlir/test/Transforms/buffer-placement-preparation.mlir
34

+1: I would really prefer that we keep passes semantically preserving and be conservative with respect to not changing the original behavior.

herhut requested changes to this revision.Mon, May 18, 2:35 AM
herhut added inline comments.
mlir/lib/Transforms/BufferPlacement.cpp
435

Instead of an assert, can you fail the rewrite and report an error? That way, the program does not just crash and this also fails outside of debug builds.

Also, please add a test that the error is generated,

This revision now requires changes to proceed.Mon, May 18, 2:35 AM
dfki-ehna updated this revision to Diff 264824.Tue, May 19, 2:02 AM

Change Assert to emitError and add a test case for the generated error.

dfki-ehna marked an inline comment as done.Tue, May 19, 2:03 AM
pifon2a accepted this revision.Tue, May 19, 2:34 AM
pifon2a added inline comments.
mlir/test/Transforms/buffer-placement-preparation.mlir
1

typo in the filename /buffer-placement-prepration.mlir

herhut accepted this revision.Tue, May 19, 3:19 AM

Thanks. Please fix the filename and land.

This revision is now accepted and ready to land.Tue, May 19, 3:19 AM
This revision was automatically updated to reflect the committed changes.