This is an archive of the discontinued LLVM Phabricator instance.

[MLIR][TOSA] First lowerings from Tosa to Linalg
ClosedPublic

Authored by rsuderman on Jan 7 2021, 10:58 AM.

Details

Summary

Initial commit to add support for lowering from TOSA to Linalg. The focus is on
the essential infrastructure for these lowerings and integration with existing
passes.

Includes lowerings for a subset of operations including:

abs, add, sub, pow, and, or, xor, left shift, right shift, tanh

Lit tests are used to validate correctness.

Diff Detail

Event Timeline

rsuderman created this revision.Jan 7 2021, 10:58 AM
rsuderman requested review of this revision.Jan 7 2021, 10:58 AM
rsuderman updated this revision to Diff 315187.Jan 7 2021, 11:03 AM

Some minor fixes before code review.

Should this be in the lib/Conversion directory with all the other dialect conversions?

mlir/include/mlir/Dialect/Tosa/Transforms/Passes.h
30 ↗(On Diff #315185)

I don't know what a "dispatch function" is, if this is in the context of IREE please remove as it won't help anyone working upstream without this context.

30 ↗(On Diff #315185)

Is it true? You're using a partial conversion, not a full one which would enforce this.

32 ↗(On Diff #315185)

Remove reference to IREE from upstream please.

mlir/lib/Dialect/Tosa/Transforms/TosaToLinalgOnTensors.cpp
35 ↗(On Diff #315185)

Nit: I think function are starting with a lower case.

41 ↗(On Diff #315185)

The TODO here looks like a copy paste, that seems like something to refactor instead of copying, feel free to split this refactoring in a NFC commit that you can rebase on.

56 ↗(On Diff #315185)

I'm a bit concerned about the code bloat for the template here.

Using some sort of table / interface may just be more appropriate.

Alternatively, refactoring the matchAndRewrite to make it "trivial" and mostly calling non-templated helper can probably work as well.

65 ↗(On Diff #315185)

Nit: in general, please use auto when it make the type is obvious and/or the code more clear (everywhere else as well).

78 ↗(On Diff #315185)

Seem like you could do:

Type uniqueTy = operation->getOperand(0).getType();

bool all_input_types_equal = llvm::all_of(operation->getOperandTypes(), [] (Type operandTy) { return operandTy == uniqueTy; });
if (!all_input_types_equal) 
  return notifyMatchFailure("All operands must have the same type");
bool all_result_types_equal = llvm::all_of(operation->getResultTypes(), [] (Type resultTy) { return resultTy == uniqueTy; });
if (!all_result_types_equal) 
  return notifyMatchFailure("All results must have the same type as the input");

Note that return notifyMatchFailure is useful when debugging failed conversion (it shows up in tracing mode), it can be a good thing to use in place of returning failure().
(and anectodically the code here isn't O(n2) like the nested loops).

81 ↗(On Diff #315185)

Please name the lambda according to what it is checking, like checkValidShapedType or something like that.

93 ↗(On Diff #315185)

Please use notifyMatchFailure, conversion in general are supposed to fail gracefully because other patterns could kick in.

102 ↗(On Diff #315185)

Nit: remove trivial braces.

105 ↗(On Diff #315185)
124 ↗(On Diff #315185)

The rewriter.getMultiDimIdentityMap(nloops) does not work with the 0 case?

140 ↗(On Diff #315185)

Remove this, failed is never modified

198 ↗(On Diff #315185)

patterns->insert is a variadic template you can do:

patterns->insert<PointwiseConverter<tosa::AbsOp, FloatType, mlir::AbsFOp>,

PointwiseConverter<tosa::AddOp, FloatType, mlir::AddFOp>,
PointwiseConverter<tosa::PowOp, FloatType, mlir::PowFOp>,
...>(context);
mlir/test/Dialect/Tosa/to_linalg_on_tensors.mlir
223 ↗(On Diff #315185)

Can you add variety of rank and test for the non-conversion of the non-supported cases?

mehdi_amini added a comment.EditedJan 7 2021, 11:27 AM

Looks like I commented on the original diff and you updated in-between, my comments aren't necessarily on the right line, ping me on chat if some of them don't make sense!

(They seem mostly offset by 5 lines down)

silvas added inline comments.Jan 7 2021, 1:50 PM
mlir/lib/Dialect/Tosa/Transforms/TosaToLinalgOnTensors.cpp
47 ↗(On Diff #315187)

I would recommend to constrain this pass to only work for static shapes for now.

The code that this pass currently produces will miscompile in the case of dynamic shapes. E.g.

"tosa.add"(%arg0, %arg1) : (tensor<?xf32>, tensor<?xf32>) -> tensor<?xf32>

If the dynamic sizes are tensor<1xf32>, tensor<2xf32> -> tensor<2xf32> then the code emitted by the current pass here will miscompile.

78 ↗(On Diff #315187)

I belive that tosa ops are always ranked. Also I think the verifier for those ops ensure that the rank of all inputs and outputs match.

mlir/test/Dialect/Tosa/to_linalg_on_tensors.mlir
2 ↗(On Diff #315187)

Put this pass / test in ConvertTosaToLinalg and name the pass convert-tosa-to-linalg

48 ↗(On Diff #315187)

Add a test exhibiting that we don't handle broadcasting.

rsuderman updated this revision to Diff 315502.Jan 8 2021, 1:26 PM

Made corresponding changes for the code review comments.

ExtractDynamicDimensions is still TBD in another CL.

rriddle requested changes to this revision.Jan 8 2021, 1:29 PM
rriddle added inline comments.
mlir/include/mlir/Conversion/TosaToLinalg/TosaToLinalg.h
18

This include doesn't look necessary.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
24

Please make sure that only classes are in namespaces, functions should be at the top level.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp
27

Same here.

This revision now requires changes to proceed.Jan 8 2021, 1:29 PM
rriddle added inline comments.Jan 8 2021, 1:35 PM
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
40

Remove user names from TODO comments.

47–51
70

nit: Only prefer auto if it makes the code more readable. Many of this uses could just use the proper type instead.

101

nit: Can you move this functor out-of-line into a variable before the if? Makes the control flow here pretty hard to read.

150

Can you just replace the use of operation here with op directly?

176

Is the to_vector actually necessary? Does the builder of ScalarOp not take a ValueRange?

silvas added inline comments.Jan 8 2021, 1:54 PM
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
30

use getParallelIteratorTypeName() instead of kParallelIterType

61

Just merge SimplePointwiseConverter into this pattern.

70

please restrict this to the static shape case for now.

95

hasRank and getRank != nloops is redundant

97

do we actually need ComplexType? I don't see any test covering it.

103

move the dyn_cast into checkValidShapedType

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
3

This testing seems really redundant. I think that things like:

// CHECK: absf
%0 = "tosa.abs"(%arg0) : (tensor<f32>) -> tensor<f32>

is sufficient. (with one test that does the more in-depth check of the full linalg op.

7

also check the indexing maps.

And check higher-rank cases as well.

319

test this error in only one place.

rsuderman updated this revision to Diff 315550.Jan 8 2021, 4:32 PM

Handled remaining comments from Mehdi/Sean

silvas added inline comments.Jan 8 2021, 4:53 PM
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
28

these should be static functions not inside an anonymous namespace

45

Does this have to be a conversion pattern? It isn't doing a type conversion, so a regular RewritePattern would work and reduce complexity I think.

50

I don't think there's any value in having this be virtual.

57

can this be a static helper function outside the class?

Also, you might want to comment that it is just this way due to code the amount of templated code.

84

you can use cast here instead of dyn_cast.

93

can either of these checks ever fail? I think all these conditions are guaranteed by the tosa op verifiers. If not, can you add a test case that hits this?

108

hasStaticShape is already guaranteed by the above code.

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
49

check at least the full op name (addf in this case)

also can just put all the tosa.* ops in a single function to reduce boilerplate

rsuderman updated this revision to Diff 315560.Jan 8 2021, 5:42 PM
rsuderman marked 6 inline comments as done.

Change from OpConversionPattern to RewritePattern

mehdi_amini added inline comments.Jan 8 2021, 6:00 PM
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
81

Nit: if t can be null, you can't call cast or dyn_cast on it right above I believe, so I guess this test can be removed?

sjarus added inline comments.Jan 11 2021, 10:43 AM
mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
29 ↗(On Diff #315560)

Shouldn't this go in include/mlir/Conversion/Passes.td ?

rsuderman updated this revision to Diff 315964.Jan 11 2021, 4:59 PM
rsuderman marked 2 inline comments as done.

Additional round of comments.

rsuderman updated this revision to Diff 315969.Jan 11 2021, 5:27 PM
rsuderman marked 23 inline comments as done.

Updates to tosa to linalg testing

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
45

Made a RewritePattern.

47–51

Removed entirely.

50

The main advantage is the rewrite helper is identical between instantiations and won't be duplicated decreasing binary size.

61

The intention is for multi op versions to instead implement the virtual function directly. I have some example lowerings in the pipeline to do that.

93

I am attempting to defensively code these practices given how early TOSA support is. Better to verify before lowering than assume valid lowerings.

rriddle added inline comments.Jan 11 2021, 5:32 PM
mlir/include/mlir/Dialect/Tosa/Transforms/Passes.td
9 ↗(On Diff #315969)

Can changes to this file be reverted now?

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
34

Remove the outer namespaces here. Anonymous namespaces should be top-level.

62

nit: Can you use Type here instead.

167

Prefer using full name resolution for methods. Only classes should really be in namespaces in .cpp files.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalgPass.cpp
44

nit: auto -> FuncOp

rsuderman updated this revision to Diff 315970.Jan 11 2021, 5:37 PM

River Riddle comments

rsuderman marked 4 inline comments as done.Jan 11 2021, 5:38 PM
silvas added inline comments.Jan 11 2021, 5:41 PM
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
61
rsuderman updated this revision to Diff 316279.Jan 12 2021, 5:07 PM

Changed to a single function for generating generic op contents.

rsuderman marked an inline comment as done.Jan 12 2021, 5:08 PM
silvas added inline comments.Jan 12 2021, 5:17 PM
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
100

I think we can hit this with the way we have the code structured now. E.g. an integer add op would currently slip through. Maybe a rewriter.notifyMatchFailure instead?

113

Can you put this in a static helper? I think it will still be duplicated code-size-wise due to the enclosing template.

120

You can remove this check and use cast<ShapedType>. The verifier for tosa ops guarantees it.

144

!st is redundant
!st.hasRank() is also redundant (TOSA ops require it)
!st.getElementType().isSignlessIntOrFloat() is redundant I think (TOSA ops require it)

191

I think you need to handle the case where createLinalgBodyCalculationForElementwiseOp doesn't succeed. (I think just set a bool in the enclosing function and return rewriter.notifyMatchFailure)

mlir/test/Conversion/TosaToLinalg/tosa-to-linalg.mlir
125

Can you test some cases don't get converted?

  • mismatching ranks (rank broadcasting)
  • dynamic shapes
rsuderman updated this revision to Diff 316501.Jan 13 2021, 1:21 PM

Removed redundant checks, added more failure test cases.

rsuderman updated this revision to Diff 316502.Jan 13 2021, 1:22 PM
rsuderman marked 2 inline comments as done.

Added missing static qualifier.

rsuderman marked 4 inline comments as done.Jan 13 2021, 1:23 PM
rsuderman added inline comments.
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
113

I moved the whole thing to a static helper. Should guarantee no replication.

silvas accepted this revision.Jan 13 2021, 4:03 PM
silvas added inline comments.
mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
137

nit: remove mention of "buffers" here. initTensors seems like a better name?

149

nit: auto bodyResultTypes = and remove the declaration above.

167

This name "error" seems very confusing. I read it as "did encounter an error", so setting it to true here is confusing. Suggest calling it didEncounterError and setting it to false here.

mehdi_amini accepted this revision.Jan 13 2021, 4:14 PM

LGTM in general.

mlir/lib/Conversion/TosaToLinalg/TosaToLinalg.cpp
23

(FYI: We don't need to specify the SmallVector size anymore.)

32

This kind of helper make me think that they look like a missing build method to be added on the InitTensorOp instead.

45

Remove trivial braces, here and elsewhere.

rsuderman marked 6 inline comments as done.Jan 14 2021, 11:16 AM
rriddle accepted this revision.Jan 14 2021, 11:16 AM

Awesome!

This revision is now accepted and ready to land.Jan 14 2021, 11:16 AM
This revision was landed with ongoing or failed builds.Jan 14 2021, 11:25 AM
This revision was automatically updated to reflect the committed changes.