This is an archive of the discontinued LLVM Phabricator instance.

[MLIR][LinAlg] Detensorize interal function control flow.
ClosedPublic

Authored by ergawy on Feb 21 2021, 6:56 AM.

Details

Summary

This patch continues detensorizing implementation by detensoring
internal control flow in functions.

In order to detensorize functions, all the non-entry block's arguments
are detensored and branches between such blocks are properly updated to
reflect the detensored types as well. Function entry block (signature)
is left intact.

This continues work towards handling github/google/iree#1159.

Diff Detail

Event Timeline

ergawy created this revision.Feb 21 2021, 6:56 AM
ergawy requested review of this revision.Feb 21 2021, 6:56 AM
ergawy updated this revision to Diff 325302.Feb 21 2021, 7:02 AM

Assert on 0 inputs.

ergawy added inline comments.Feb 21 2021, 7:22 AM
mlir/include/mlir/Dialect/Linalg/Passes.h
65 ↗(On Diff #325302)

Eventually, I will not implement this as a separate pass. Both passes will be merged into one. I am doing this temporarily to make understanding how dialect conversion works a bit better (i.e. to have a clear separation point between internal conversion within a BB and conversion across the BB boundary in a function).

ergawy updated this revision to Diff 326010.Feb 24 2021, 1:09 AM

Properly handle non-entry function blocks.

ergawy retitled this revision from {WIP: PLZ DON'T REVIEW YET}[MLIR][LinAlg] Detensorize interal CF. to [MLIR][LinAlg] Detensorize interal function control flow..Feb 24 2021, 1:10 AM
ergawy added reviewers: silvas, rriddle.
ergawy updated this revision to Diff 326011.Feb 24 2021, 1:13 AM

Update docs.

ergawy edited the summary of this revision. (Show Details)Feb 24 2021, 1:45 AM
silvas added inline comments.Feb 24 2021, 11:00 AM
mlir/include/mlir/Dialect/Linalg/Passes.td
155

nit: indent is weird

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h
37

"to return" reads weird to me

mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
185

more specifically: "detensoring can't happen along external calling convention boundaries, which we conservatively approximate as all function signatures"

196

seems like there is some duplication with func-bufferize. can you put a helper next to the populate*Pattern function (populate*Legality)?

mlir/test/Dialect/Linalg/detensorized_while.mlir
3

remove iree.module.export

7

can you save these attributes in a #attr_alias (see other uses in the codebase). that would make the code somewhat simpler. E.g. https://github.com/llvm/llvm-project/blob/e339bba637b941c8e78057319b7654c4babf18cb/mlir/test/Dialect/Linalg/tile.mlir#L343

27

check: tensor.extract

ergawy updated this revision to Diff 326184.Feb 24 2021, 12:56 PM
ergawy marked 7 inline comments as done.

Handle review comments.

silvas added inline comments.Feb 24 2021, 1:48 PM
mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
124

I don't think that hasTrait<ReturnLike> is relevant to the branch pattern.

rriddle added inline comments.Feb 24 2021, 4:11 PM
mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h
42

Just a note but, conversion target only supports one "unknown legality" callback function.

mlir/lib/Dialect/Linalg/Transforms/Detensorize.cpp
18

Is this really necessary?

91

Functions should not be placed inside of anonymous namespaces. They should be marked as static and at the top level.

104

The .getResult() shouldn't be necessary.

ergawy updated this revision to Diff 326308.Feb 25 2021, 12:13 AM
ergawy marked 3 inline comments as done.

Handle some review comments.

mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h
42

I noticed that and and I think that might be an argument against adding this utility function. I mean it is not general enough, it's logic just happens to be needed in 2 different places. If the needs of one pass diverges from the other we might end up inlining the code again.

I guess for now we can keep the util function and see if the need arises to inline it again at call sites in the future. Let me know if you disagree.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
124

If we remove this condition, then ops that are ReturnLike but not BranchOpInterface will be marked as illegal.

I don't disagree that this should ideally be removed. However, I guess we will have to change the expected behavior a little bit.

I think we can rewrite the entire lambda as follows:

target.markUnknownOpDynamicallyLegal([&](Operation *op) {
  // All successor operands of branch like operations must be rewritten.
  if (auto branchOp = dyn_cast<BranchOpInterface>(op)) {
    for (int p = 0, e = op->getBlock()->getNumSuccessors(); p < e; ++p) {
      auto successorOperands = branchOp.getSuccessorOperands(p);
      if (successorOperands.hasValue() &&
          !converter.isLegal(successorOperands.getValue().getTypes()))
        return false;
    }
  }

  return true;
});

With the above, all bufferization and detensorization tests pass fine except for one: https://github.com/llvm/llvm-project/blob/main/mlir/test/Dialect/Standard/func-bufferize.mlir#L65. In particular, test.terminator won't be marked as illegal as before.

I am not talking about this particular test of course. In general, any op that is not explicitly BranchOpInterface but might have IsTerminator will be legal. Would that be fine?

ergawy added inline comments.Feb 25 2021, 12:48 AM
mlir/include/mlir/Dialect/StandardOps/Transforms/FuncConversions.h
42

Oh by the way, is it by design that conversion target supports exactly one unknown legality callback? Or is it something that might be extended to support multiple ones?

I think multiple ones can open the door for somewhat general utils like the above and allow customization/additional behavior on top of that. But it would complicate things by quite a bit as well I guess and there might be no need for it.

If you think extension to multiple callbacks is worth doing, I would gladly give this a try.

silvas added inline comments.Feb 25 2021, 12:06 PM
mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
124

I think the core issue here is that we are conflating legality for populateReturnOpTypeConversionPattern with legality for populateBranchOpInterfaceTypeConversionPattern.

I think we should split this into two functions bool isLegalForBranchOpInterfaceTypeConversionPattern(Operation *op, TypeConverter &converter) and bool isLegalForReturnOpTypeConversionPattern(Operation *op) (it is important that these be split so that they correspond very tightly with their corresponding populate*Pattern function). It i then up to the user to call manually if they want both:

target.markUnknownOpDynamicallyLegal([&](Operation *op) {
  return isLegalForBranchOpInterfaceTypeConversionPattern(op, converter) || isLegalForReturnOpTypeConversionPattern(op);
})

to work around the fact that only one legality callback is allowed in markUnknownOpDynamicallyLegal.

ergawy updated this revision to Diff 326614.Feb 26 2021, 12:00 AM

Handle review comments.

ergawy marked 2 inline comments as done.Feb 26 2021, 12:02 AM
ergawy added inline comments.
mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
124

I adopted your suggestion with a slight change. I broke up the previous lambda into 3 functions to make the purpose of each one more fine-grained. Let me know if this is not a good approach.

ergawy updated this revision to Diff 326615.Feb 26 2021, 12:04 AM
ergawy marked an inline comment as done.

Clean-up.

ergawy updated this revision to Diff 326617.Feb 26 2021, 12:11 AM

Rename function.

ergawy edited the summary of this revision. (Show Details)Feb 26 2021, 12:17 AM
silvas accepted this revision.Feb 26 2021, 1:58 PM

One conceptual concern, but otherwise this LGTM.

mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
111

The isNotBranchOpInterfaceOrReturnLikeOp is kind of annoying, and I'm sorry for getting the logic here backward in my suggestion. I realize now that for this to compose properly it basically needs to be:

target.markUnknownOpDynamicallyLegal([&](Operation *op) {
  if (needsToBeLegalizedByBranchOpInterfaceTypeConversion(op, converter))
    return false;
  if (needsToBeLegalizedByReturnOpTypeConversion(op, converter))
    return false;
  return true;
}

With a slight modification to the current code, we can write it as !isLegalForBranchOpInterfaceTypeConversionPattern(...)) return false;. The current code returns "false" for "needs to convert BranchOpInterface" and also for "i have no idea about this op". It should only return "false" for "needs to convert BranchOpInterface".

This dialect conversion / legality stuff is super tricky, sorry.

133

I don't understand why this function has to worry about ReturnLike. populateReturnOpTypeConversionPattern only handles std.return, and this function should only be concerned with that one.

This revision is now accepted and ready to land.Feb 26 2021, 1:58 PM
ergawy added inline comments.Feb 28 2021, 3:54 AM
mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
111

No problem, happy to discuss until we get it right. And apologies if I am not getting what you suggest sometimes.

With your suggestion in the above comment, we won't be able to report an error whenever a terminator neither implements BranchOpInterface nor is a ReturnOp. Such terminator would be marked legal. In order to handle this, I think we would need to:

  • either centralize the legality check in one place for return-like ops and branch op interface in one place like you had before.
  • or provide a catch-all utility like isNotBranchOpInterfaceOrReturnLikeOp that would mark unknown terminator ops as illegal and mark all other ops as legal.
ergawy updated this revision to Diff 327394.Mar 2 2021, 2:13 AM
  • Rebase.
  • Add a TODO.
mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
111

I added a TODO to revisit this and try getting rid of isNotBranchOpInterfaceOrReturnLikeOp. I will commit the patch and if you really don't like, I will handle this in a follow up patch right away.

This revision was landed with ongoing or failed builds.Mar 2 2021, 2:49 AM
This revision was automatically updated to reflect the committed changes.
silvas added inline comments.Mar 2 2021, 12:21 PM
mlir/lib/Dialect/StandardOps/Transforms/FuncConversions.cpp
111

ah, I see. yeah that's tricky and would have the issue you describe.

Nit: fix the typo in the commit title, "interal"