This is an archive of the discontinued LLVM Phabricator instance.

[mlir] Extend BufferAssignmentTypeConverter with result conversion callbacks
ClosedPublic

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

Details

Summary

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

Input:

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.

mlir/include/mlir/Transforms/BufferPlacement.h
93

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?

219–243

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

257

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

mlir/lib/Transforms/BufferPlacement.cpp
872

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
mlir/include/mlir/Transforms/BufferPlacement.h
93

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.
mlir/include/mlir/Transforms/BufferPlacement.h
224

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

226

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

254–255

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
mlir/lib/Transforms/BufferPlacement.cpp
734–735

Nit: Can fold these two into one line.

837

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

837

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

841–843

Comments for the member variables please.

878

Avoid trailing comment - please move to next line.

895

Typo: it -> its

900

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

903–907

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

mlir/test/lib/Dialect/Test/TestOps.td
1625

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

mlir/test/lib/Transforms/TestBufferPlacement.cpp
184–186

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

185

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!

mlir/test/Transforms/buffer-placement-preparation-allowed-memref-results.mlir
122–133

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

163

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
mlir/include/mlir/Transforms/BufferPlacement.h
107–108

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
mlir/include/mlir/Transforms/BufferPlacement.h
254

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

See https://buildkite.com/mlir/mlir-core/builds/7580#ad1609b5-eedd-4d04-9fd0-361402d40f36.

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: https://github.com/tensorflow/mlir-hlo/blob/master/lib/Dialect/mhlo/transforms/hlo_legalize_to_lhlo.cc#L441

BufferAssignmentTypeConverter converter;
...
if (results_escape_function) {
  populateWithBufferAssignmentOpConversionPatterns<
      mlir::ReturnOp, mlir::ReturnOp, lmhlo::CopyOp,
      /*allowMemrefFunctionResults=*/true>(&context, &bufferAssignment,
                                           &converter, &patterns);
} else {
  populateWithBufferAssignmentOpConversionPatterns<
      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>(
    kind);
...
populateWithBufferAssignmentOpConversionPatterns<
    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.