Page MenuHomePhabricator

[mlir] Extend BufferAssignmentTypeConverter with result conversion callbacks

Authored by dfki-ehna on Aug 3 2020, 7:04 AM.



In this PR, the users of BufferPlacement can configure
BufferAssginmentTypeConverter. These new configurations would give the user more
freedom in the process of converting function signature, and return and call
operation conversions.

These are the new features:

  • Accepting callback functions for decomposing types (i.e. 1 to N type conversion such as unpacking tuple types).
  • Defining ResultConversionKind for specifying whether a function result with a certain type should be appended to the function arguments list or should be kept as function result. (Usage: converter.setResultConversionKind<MemRefType>(AppendToArgumentList))
  • Accepting callback functions for composing or decomposing values (i.e. N to 1 and 1 to N value conversion).


func @test(%arg0: tuple<tensor<2xf32>,i1>) -> (tuple<tensor<2xf32>,i1>){
  return %arg0 : tuple<tensor<2xf32>,i1>

Output with [converter.setResultConversionKind<MemRefType>(AppendToArgumentList)]:

func @test(%arg0: memref<2xf32>, %arg1: i1, %arg2: memref<2xf32>) -> i1 {
  %0 = "test.make_tuple"(%arg0, %arg1) : (memref<2xf32>, i1) -> tuple<memref<2xf32>, i1>
  %1 = "test.get_tuple_element"(%0) {index = 0 : i32} : (tuple<memref<2xf32>, i1>) -> memref<2xf32>
  %2 = "test.get_tuple_element"(%0) {index = 1 : i32} : (tuple<memref<2xf32>, i1>) -> i1
  linalg.copy(%1, %arg2) : memref<2xf32>, memref<2xf32>
  return %2 : i1

Output with [converter.setResultConversionKind<MemRefType>(KeepAsFunctionResult)]:

func @test(%arg0: memref<2xf32>, %arg1: i1) -> (memref<2xf32>, i1) {
  %0 = "test.make_tuple"(%arg0, %arg1) : (memref<2xf32>, i1) -> tuple<memref<2xf32>, i1>
  %1 = "test.get_tuple_element"(%0) {index = 0 : i32} : (tuple<memref<2xf32>, i1>) -> memref<2xf32>
  %2 = "test.get_tuple_element"(%0) {index = 1 : i32} : (tuple<memref<2xf32>, i1>) -> i1
  return %1, %2 : memref<2xf32>, i1

Diff Detail

Event Timeline

dfki-ehna created this revision.Aug 3 2020, 7:04 AM
Herald added a project: Restricted Project. · View Herald TranscriptAug 3 2020, 7:04 AM
dfki-ehna requested review of this revision.Aug 3 2020, 7:04 AM
dfki-ehna edited the summary of this revision. (Show Details)Aug 3 2020, 7:07 AM

Just some quick comments.


Currently this could just be a map from Type to ResultConversionKind. Callbacks only help if the user can provide them and hence have some extra logic in them. Maybe by also passing the originating type?


This should be an assert, best in the super class' constructor.


Here reporting an error instead would be nicer. You could also append a note on the error showing the function signature,


If you had the original and new type available here, you could still differentiate between a rewritten memref and a newly created memref.

pifon2a added inline comments.Aug 7 2020, 3:19 AM

I agree with Stephan. It would be nice to make it a map.
If we need more flexibility in future, then we would need to provide a method that allows to set a map from functor(Type) -> bool to ResultConversionKind so that users would what and how to convert. For now, a map would be sufficient.

The functionality added looks great! Thanks. I'll shortly submit a review.

bondhugula requested changes to this revision.Aug 7 2020, 10:34 PM
bondhugula added inline comments.

You can move this declaration to right under the comment below.


Nit: " ... depending on whether they need a ..."


int -> unsigned ?

This revision now requires changes to proceed.Aug 7 2020, 10:34 PM
bondhugula added inline comments.Aug 7 2020, 10:55 PM

Nit: Can fold these two into one line.


auto -> const std::pair<unsigned, Value> &


I think you have more copy here due to lack of a const ref.


Comments for the member variables please.


Avoid trailing comment - please move to next line.


Typo: it -> its


Use i = 0, e = ...; i < e; ... form to avoid repeated evaluation.


Nit: Use braces all blocks since you need it for the else block below.


Nit: the i-th element -> a specified element


Aren't you already in the mlir namespace? Drop mlir:: ?


Avoid repeated evaluation of .size() - use i = 0, e = ... form.

dfki-ehna updated this revision to Diff 284620.Aug 11 2020, 2:35 AM

Resolve comments, change setResultConversionKind to a mapping (i.e. setResultConversionKind<RankedTensorType, MemRefType>(AppendToArguementList)), and introduce addDecomposeTypeConversion.

dfki-ehna marked 19 inline comments as done.Aug 11 2020, 2:42 AM
bondhugula accepted this revision.Aug 19 2020, 3:44 PM

This looks really good - thanks!


Could you add a line or two here and further below to summarize what this is testing? (although obvious to you at this point).


Please name these functions a bit more descriptively. Here and above.

This revision is now accepted and ready to land.Aug 19 2020, 3:44 PM
bondhugula added inline comments.Aug 30 2020, 12:06 PM

You could use llvm::is_one_of instead.

Resolving comments

dfki-ehna marked an inline comment as done.Sep 1 2020, 12:00 AM
mehdi_amini added inline comments.Sep 1 2020, 12:36 AM

Seems like the error message isn't terminated?

dfki-ehna updated this revision to Diff 289392.Sep 2 2020, 4:30 AM

Resolve the minor comment before landing

This revision was landed with ongoing or failed builds.Sep 2 2020, 4:32 AM
This revision was automatically updated to reflect the committed changes.

FYI: this change fails the following tests:

Failed Tests (3):
  MLIR :: Dialect/Linalg/tensors-to-buffers.mlir
  MLIR :: Transforms/buffer-placement-preparation-allowed-memref-results.mlir
  MLIR :: Transforms/buffer-placement-preparation.mlir


I'll revert this. Please resubmit after fixing the tests. :)

dfki-ehna reopened this revision.Sep 2 2020, 8:25 AM

Reopen to resolve the failure.

This revision is now accepted and ready to land.Sep 2 2020, 8:25 AM
dfki-ehna updated this revision to Diff 289452.Sep 2 2020, 8:29 AM

Resolve the failure

When trying to integrate this downstream, I'm having trouble figuring out what to do with this block:

BufferAssignmentTypeConverter converter;
if (results_escape_function) {
      mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp,
      /*allowMemrefFunctionResults=*/true>(&context, &bufferAssignment,
                                           &converter, &patterns);
} else {
      mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp,
      /*allowMemrefFunctionResults=*/false>(&context, &bufferAssignment,
                                            &converter, &patterns);

As the allowMemrefFunctionResults argument went away, I tried using setResultConversionKind as suggested to have the equivalent functionality, however I'm still seeing test failures:

BufferAssignmentTypeConverter converter;
auto kind = results_escape_function
                ? BufferAssignmentTypeConverter::KeepAsFunctionResult
                : BufferAssignmentTypeConverter::AppendToArgumentsList;
converter.setResultConversionKind<RankedTensorType, MemRefType>(kind);
converter.setResultConversionKind<UnrankedTensorType, UnrankedMemRefType>(
    mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp>(
    &context, &bufferAssignment, &converter, &patterns);

Is there anything wrong with the fix applied? (cc @antiagainst)

@rupprecht I am investigating the downstream issue.