This is an archive of the discontinued LLVM Phabricator instance.

[mlir][linalg] Vectorize 1D pooling ops
ClosedPublic

Authored by vmurali on Dec 15 2022, 9:13 PM.

Diff Detail

Event Timeline

vmurali created this revision.Dec 15 2022, 9:13 PM
Herald added a project: Restricted Project. · View Herald Transcript
vmurali requested review of this revision.Dec 15 2022, 9:13 PM
rsuderman requested changes to this revision.Dec 16 2022, 11:13 AM
rsuderman added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1879–1897

This block is quite difficult to parse. I would recommend factoring it out into a helper or reworking to avoid multiple nested if/else statements. This could be as simple as getting your three block operands. Checking for casts. Seeing if there is a mul between input/kernel, then finding the combiner. with the accumulator. Then special checks (e.g. finding the combiner) can be done by helpers to avoid nesting the list inside of the entire block.

It would provide a more step-by-step process and avoid parsing deep parts of the implementation.

1940–1941

There appears to be a decent amount of shared implementation with the convolution implementation. I would recommend maintaining shared lines for shared implementation.

1970–1971

Ditto about the shared implementation.

2066–2077

Often if you have an if-else block in a simple loop (like here) it is simpler to do

for (...) {
  if (condition) {
    do work;
    continue;
  }
  do other work;
}

This avoids the extra else statement and demonstrates there is only one task.

2390–2393

Why are these being handled as globals vs being included a parameters to the builder?

This revision now requires changes to proceed.Dec 16 2022, 11:13 AM
vmurali added inline comments.Dec 16 2022, 11:45 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
2390–2393

The code works in two steps - the constructor checks for validity of the op and sets these values, and the transformer (conv) performs the transformation post construction. To keep the state between the construction and transformation, I need these values. (They are inside a class, so not really global.)

vmurali added inline comments.Dec 16 2022, 11:51 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
2390–2393

Another thing to note is that we never needed to specify the cast op or the reduction op for convolution because (a) convolution was always a sum of products, and (b) contraction automatically took care of casting. But for pooling, neither of these is true (as we have multiple kinds of pooling, and we don't use contraction for the equivalent of reduction). Finally, the isPool state is simply to distinguish between pooling and convolution. The only state that can potentially be eliminated is isPoolExt and rely on null strings for no reduction - I find that distasteful.

vmurali updated this revision to Diff 483628.Dec 16 2022, 12:06 PM
vmurali marked 3 inline comments as done.
vmurali marked an inline comment as done.
vmurali updated this revision to Diff 483652.Dec 16 2022, 1:27 PM

Addressed Rob's comments

vmurali marked an inline comment as done.Dec 16 2022, 1:28 PM

Addressed all of Rob's comments

vmurali updated this revision to Diff 483656.Dec 16 2022, 1:37 PM

Clang format

It looks good! Doing a first pass.

Design question: Wondering if it would make more sense to decouple conv and pooling implementation while sharing the common code. The fact that we have to preserve the isPool state and add conditional code all over the place is suggesting that something is off. Not sure how feasible this is but it would be good to give it a thought.

A couple of style comments:

  • Single statement ifs and first shouldn't have { }.
  • Period at the end of a comment.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1825

Nit: If they are not widely used, I would write these predicates in place to prevent an explosion of utilities with different combinations. If you need a local predicate, you can write a local lambda.

For example, since !isa<arith::MulIOp, arith::MulFOp>(op) is only used once, I would write it as:

if (!isa<arith::MulIOp, arith::MulFOp>(op) && 
    llvm::any_of(op->getOperands(), ...))
1878–1887

Is it possible to simplify this code? If I parse this correctly, the loop body beyond the isa<BlockArgument check will be executed only once as feedOp = getDefiningOp() will never be null. I would perhaps add a few simpler checks using llvm::count_if to check that there is only one operand that is not a block argument, then llvm::find_if to get it, and move the remaining code out of the loop. Just a quick suggestion, there must be other ways.

Perhaps is also a good idea to separate the two If described in the comments if that simplifies the code. I would prioritize simplicity over performance as the number of operands is very small.

1879–1897

what about something like

if (!maybeKind || !isSupportedConvKind(isPool))
  return;
1899

nit: !(A == B && C == D) -> A != B || C !=D ? It seems easier to parse.

Thanks for the review, I will address some of it and submit again for review.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1825

Rob was suggesting that I create specific functions for these. I can replace the for-loop with llvm::all_of. Maybe I can inline the multiply part alone after switching to llvm::all_of

1878–1887

I just modified what existed earlier - look at mulOp in the earlier code. But yes, I am just checking if there's at most one non-BlockArgument, and if I find it, I make sure it's either a MulOp of block arguments or casts of block arguments (convolution), or a cast of a block argument (pooling).

1879–1897

I think there are other CombiningKinds that are not used in Convolution or Pooling. We should have an explicit list for convolution and pooling and not rely on not-convolution being pooling.

1899

For me, !(A && B) is a lot more readable because I enumerate the valid conditions in my head and simply negate it for invalid.

vmurali updated this revision to Diff 483716.Dec 16 2022, 8:24 PM

Addressing Diego's comments

vmurali marked 4 inline comments as done.Dec 16 2022, 8:25 PM

It looks good! Doing a first pass.

Design question: Wondering if it would make more sense to decouple conv and pooling implementation while sharing the common code. The fact that we have to preserve the isPool state and add conditional code all over the place is suggesting that something is off. Not sure how feasible this is but it would be good to give it a thought.

A couple of style comments:

  • Single statement ifs and first shouldn't have { }.
  • Period at the end of a comment.

Big +1. We should decouple conv and pooling implementation. We should also factor common codes out and avoid having isPool everywhere?

Also, can we use an enum instead having isPool and isPoolExt? Using enum is much simpler because the boolean combination can be exponentially large and we'll have lacking documentations about the combinations.

(I did not look into too much implementation because the design would change.)

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1801–1825

llvm style nit: Don’t Use Braces on Simple Single-Statement Bodies of if/else/loop Statements

https://llvm.org/docs/CodingStandards.html#don-t-use-braces-on-simple-single-statement-bodies-of-if-else-loop-statements

1802

A CastOpInterface is allowed to have multiple inputs. However, we unconditionally check the first operand. That should be taken into account.

1816–1822

[optional] I think using case switch captures this better. That's also what we used in VectorOps.cpp. E.g.,

static bool isSupportedCombiningKind(CombiningKind combiningKind,
                                     Type elementType) {
  switch (combiningKind) {
  case CombiningKind::ADD:
  case CombiningKind::MUL:
    return elementType.isIntOrIndexOrFloat();
  case CombiningKind::MINUI:
  case CombiningKind::MINSI:
  case CombiningKind::MAXUI:
  case CombiningKind::MAXSI:
  case CombiningKind::AND:
  case CombiningKind::OR:
  case CombiningKind::XOR:
    return elementType.isIntOrIndex();
  case CombiningKind::MINF:
  case CombiningKind::MAXF:
    return elementType.isa<FloatType>();
  }
  return false;
}
1861–1864

should we just check if the op is a LinalgConvolutionOpInterface?

Maybe @antiagainst should weigh in because he's the author of the implementation.

1899

I think we should follow the pattern that most of people use in MLIR/LLVM. My experience tells me that I've seen A != B || C != D form a lot in the code browsing and code review. I seldom see !(A && B) form in the codebase.

vmurali added a comment.EditedDec 19 2022, 3:33 PM

It looks good! Doing a first pass.

Design question: Wondering if it would make more sense to decouple conv and pooling implementation while sharing the common code. The fact that we have to preserve the isPool state and add conditional code all over the place is suggesting that something is off. Not sure how feasible this is but it would be good to give it a thought.

A couple of style comments:

  • Single statement ifs and first shouldn't have { }.
  • Period at the end of a comment.

Big +1. We should decouple conv and pooling implementation. We should also factor common codes out and avoid having isPool everywhere?

Also, can we use an enum instead having isPool and isPoolExt? Using enum is much simpler because the boolean combination can be exponentially large and we'll have lacking documentations about the combinations.

(I did not look into too much implementation because the design would change.)

There's contradictory feedback here about separating out pooling (the earlier review, for instance). There are only a few places where isPool is checked (8 places after the constructor, 4 of which can be removed without affecting functionality, so effectively just 4 places) and the amount of shared code is big. When we had brainstormed the design, we explicitly decided to not create a new transformation for pooling. Pooling ops implement the convolution interface, precisely to combine it with convolution.

Just to reiterate, the only place where a distinction exists is in binding the kernel dimensions and in creating the final reduction op.

vmurali added a comment.EditedDec 19 2022, 3:42 PM

Also, can we use an enum instead having isPool and isPoolExt? Using enum is much simpler because the boolean combination can be exponentially large and we'll have lacking documentations about the combinations.

isPool and isPoolExt are logically almost independent. Yes you will remove one state because it cannot be isPoolExt without it being isPool. The hackier solution is to use the size of the poolExtOp to decide if extension has to be applied

vmurali marked an inline comment as done.Dec 19 2022, 3:50 PM
vmurali added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1802

Aah then we have to explicitly make sure there's only one operand. I will change that

1861–1864

I believe this check has already been performed

1899

Okay maybe I am weird in thinking about conditions this way 🤣. But, once we have more complicated conditions like what we have here, having an OR of negations is not enough. One has to negate complicated expressions leading to some complex expression containing ANDs of some negations and some literals as is, all ORed together. At that point it's just impossible to reason about. Instead, just listing all the valid options as OR, with the valid options containing a bunch of ANDed expressions for that option

I did not mean to create a new transformation for pooling. I'd suggest that we should consider going with enum and refactor things out. After review more details, I still feel that it's better for having enum. It helps people like me start thinking which part can be refactored out or be organized together. It also hints others which part should be modified if there are more new kinds of LinalgConvolutionOpInterface op.

Anyway, I'm not marking it "require to be fixed". It's more like wanna making sure it has been considered. (Because I do not the see the contradictory feedback and the discussions in the review thread.)

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1861–1864

I don't know if the check has been performed or not. I look into the struct members and I can only say that it is a LinalgOp. A LinalgOp does not have to be LinalgConvolutionOpInterface.

But now I figure out why we check this. It is because the generator works for Conv1D. And now we are updating the Conv1D definition to consider Pool1D cases. In this context, should we add back the check about resShapedType.getRank() != 3?

(IMO, we should have LinalgConvolutionOpInterface::is1D method and move the check to there. But I think it should not be a blocker of the pooling vectorization patch.)

1877

having the check is better than comment it.. And we can use switch-case for this kind of code.

1899

I can point out the if-cond that I think can be the way; mark them an optional comment..

1905–1907

we can just update the error message to conv/pool.. The notification already includes the op. People can look into the op name and figure out what's happening.

2017

Here is an example that having enum is better. We can have better documentation and write code like

// The reduce op of a convolution op is a binary op.
if (enum == ConvEnum)
  rhs = ...

It's more meaningful to me and consider the cases if there are more kinds of LinalgConvolutionOpInterface ops.

2042–2048

I believe the compiler is smart for handling it, but can we just swap for if to if for? :)

vmurali marked 5 inline comments as done.Dec 20 2022, 9:38 AM

(Because I do not the see the contradictory feedback and the discussions in the review thread.)

I meant the feedback about keeping more code common than separate.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1861–1864

resShapedType.getRank() != 3 is performed after I figure out it's convolution and not pooling (line 1918). I suppose I can refactor it in and make the check earlier, since both pooling and convolution have rank 3.

2017

Aah I didn't realize you meant Conv vs Pool enum. I assumed you meant an enum combining isPool and isPoolExt, which seems unnecessary.

hanchung added inline comments.Dec 20 2022, 10:51 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
2017

I meant an enum which may have conv, pool, poolext or just conv, pool.

vmurali updated this revision to Diff 484383.Dec 20 2022, 2:36 PM
vmurali marked 4 inline comments as done.

Addressing all of Hanhan's comments

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
2017

poolExt is a subset of pool, just to see if there's a cast during the pooling operation. And it's used in exactly one place, so it shouldn't be conflated with conv and pool

vmurali marked an inline comment as done.Dec 20 2022, 2:40 PM
vmurali updated this revision to Diff 484387.Dec 20 2022, 2:46 PM

Clang format

Thanks for addressing the feedback! I did another pass.

There's contradictory feedback here about separating out pooling (the earlier review, for instance). There are only a few places where isPool is checked (8 places after the constructor, 4 of which can be removed without affecting functionality, so effectively just 4 places) and the amount of shared code is big. When we had brainstormed the design, we explicitly decided to not create a new transformation for pooling. Pooling ops implement the convolution interface, precisely to combine it with convolution.

I think @hanchung's feedback is similar to mine. We both feel that something is off but we haven't written the original code so we can't make a specific suggestion without spending much more time on the code. The fact that we have code that conditionally depends on the Conv or Pool enum values indicates that we are conflating disjoint implementations into a single one and that the design would benefit from using hierarchy or overloading to separate the Conv or Pool specific code from the shared main algorithm. I'm ok if this is not addressed as part of this patch but it's very likely that more specific code is added in the near future as we add support for more Conv and Pool cases so it would be better to do some refactoring now that the implementation is simpler. If you think that having a separate transformation for pooling, with common utilities shared with conv, should be reconsider... that should be fine! I personally can't suggest anything more specific without diving more into the code and the structured generator infra.

Hopefully that helps!

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1825

Sorry, I missed Rob's comment. In general, utility functions make sense. However, for these specific and simple checks that are used all over the compiler there are some trade-offs that I would consider. If the check is used only once, it's probably simpler to parse the check directly (developers are very used to these kind of isa checks) than perhaps a complex name (e.g., isBlockArgumentOrCastOfBlockArgument). If you still feel that giving this a name is more convenient for a single use, a local named lambda might work better. The other problem I see with adding static utilities for these checks is that there could easily be a combinatorial explosion of them and, most importantly, we may create a style or API for these checks local to this specific file, which wouldn't be aligned with the rest of the compiler and it would probably be used inconsistently by those not aware of it.

Anyways, too much discussion for this subjective nit comment :). I hope with all this information in place you can make a good call.

1861–1864

Per other discussion: !(A ==B && C == D) -> A != B || C !=D

Can we also add a comment about what these checks mean?

1873–1874

Probably better to save reduceOp and avoid a pool specific name to keep it generic

1875–1911

Can we move this switch case above to a utility function? It looks like this is matching the convolution or pool kind so naming that accordingly would help readability.

1882

I found this check confusing: how could oper == Pool and isPoolingOp return false? Should we rename isPoolingOp to something like isSupportedPoolCombiner?

1906
2042–2048

Not necessarily in Debug mode, which is a compiler mode we should also optimize for efficient debugging :)

2043–2048

Per LLVM guidelines, please finish comments with period.
https://llvm.org/docs/CodingStandards.html#id8

mlir/test/Dialect/Linalg/vectorize-convolution.mlir
575

Missing -----?

585

We usually use CHECK-LABEL to match the function name. It can make the test more resilient to accidental checks.

599

Please, use CHECK-DAG only when absolutely necessary. It makes the mapping more expensive and can make debugging an actual nightmare. For example, It's common to use DAG for constants, maps, globals but I don't think we need it here for anything else.

vmurali marked 6 inline comments as done.Jan 4 2023, 1:59 PM

Thanks for addressing the feedback! I did another pass.

There's contradictory feedback here about separating out pooling (the earlier review, for instance). There are only a few places where isPool is checked (8 places after the constructor, 4 of which can be removed without affecting functionality, so effectively just 4 places) and the amount of shared code is big. When we had brainstormed the design, we explicitly decided to not create a new transformation for pooling. Pooling ops implement the convolution interface, precisely to combine it with convolution.

I think @hanchung's feedback is similar to mine. We both feel that something is off but we haven't written the original code so we can't make a specific suggestion without spending much more time on the code. The fact that we have code that conditionally depends on the Conv or Pool enum values indicates that we are conflating disjoint implementations into a single one and that the design would benefit from using hierarchy or overloading to separate the Conv or Pool specific code from the shared main algorithm. I'm ok if this is not addressed as part of this patch but it's very likely that more specific code is added in the near future as we add support for more Conv and Pool cases so it would be better to do some refactoring now that the implementation is simpler. If you think that having a separate transformation for pooling, with common utilities shared with conv, should be reconsider... that should be fine! I personally can't suggest anything more specific without diving more into the code and the structured generator infra.

Hopefully that helps!

I think @rsuderman was suggesting to keep more code in common. Anyway, I still think these two ops are so close to each other than we have to share a lot of code between them. We can go over cleaning it up further in a future patch.

vmurali marked 7 inline comments as done.Jan 4 2023, 2:35 PM
vmurali added inline comments.
mlir/test/Dialect/Linalg/vectorize-convolution.mlir
585

I was basically following the style of the convolution functions. I can change these.

vmurali updated this revision to Diff 486399.Jan 4 2023, 2:36 PM
vmurali marked an inline comment as done.

Addressing Diego's comments

hanchung added inline comments.Jan 5 2023, 11:21 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1863–1865

The comment and the check mismatch. One is RHS and the other is RES.

Should this be ||, not &&?

1881

Declaring a variable for rhsShapedType.getRank() is easier to parse. And we can write if (!A &&!B) instead of if (!(A || B)).

1916–1917

Where is the f dimension? Shouldn't this be bindShapeDims(rhsShapedType, kwSize, cSize, fSize);?

1994–1997

There could have more enum values. The comment might be inaccurate. We should either update the comment or the if-check..

The comment should be combined with the previous comment as well...

2063

nit: s/Perform/perform

2130–2132

please remove braces for single statement...

2388

maybe rename it to OperType or OperKind?

2407–2409

nit: we usually name it with numBlockArguments.

mlir/test/Dialect/Linalg/vectorize-convolution.mlir
586

why not name them to INPUT, FILTER, OUTPUT? It can match the naming in the input MLIR better and help reviewing the test.

I don't know if we have naming style guide for lit test, but we should keep them consistent. All the variables in this file are upper cases with _.

Thanks a lot for addressing all the comments! It looks great to me! Please, wait for others to finish the review before committing. Thanks!

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1874–1876

you can use auto for this type.

2406

We use LogicalResult instead of bool for this.

dcaballe accepted this revision.Jan 5 2023, 11:46 AM
vmurali added inline comments.Jan 5 2023, 2:42 PM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1863–1865

Yes, the code that I changed from my previous version is wrong :).

This is precisely why I think we should just implement what is written in the comments instead of applying demorgan's law and then implementing it, even if it doesn't conform to what exists currently in the LLVM/MLIR code base - it eliminates trivial bugs like this, and is probably easier to optimize for the c++ compiler than for humans.

vmurali marked an inline comment as done.Jan 5 2023, 2:44 PM
vmurali added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1916–1917

f is already bound using outputs, so I don't rebind it now.

vmurali marked 10 inline comments as done.Jan 5 2023, 3:04 PM

Addressed Hanhan's and Diego's comments

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
1881

I have changed (!( A || B || C)) into its distributed equivalent (applying demorgan's law), but as I said in my earlier comment, I think doing this conversion is error prone (at least for me, as demonstrated in my previous revision, when I did this transform in a different location), especially when the conditions become more complicated. It's easier to list all the acceptable alternatives using an OR (where each alternative itself is a complicated expression). Since the semantics of this function is to call return whenever we hit a failure, we should just negate the list.

1994–1997

I am not sure if this code can handle any operation other than Conv and Pooling. I can fix the comment saying this is done only for Conv currently.

2406

I think it just clutters the code (all the returns based on expressions should be converted into LogicalResult using the constructor, example in line 2409). It's probably needed if it's a transformation pass, but this is a helper utility which has nothing to do with LLVM.

mlir/test/Dialect/Linalg/vectorize-convolution.mlir
586

These tests were auto constructed from the output after examining the outputs to be correct, and my script doesn't transform the names other than adding a "V". I can fix it manually, but it's probably not very important, given that these tests will fail appropriately. There are too many variations of convolution to write these tests by hand.

This revision was not accepted when it landed; it landed in state Needs Review.Jan 5 2023, 3:16 PM
This revision was automatically updated to reflect the committed changes.
vmurali marked 3 inline comments as done.
dcaballe added inline comments.Jan 6 2023, 10:40 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
2406

We only have to s/true/success() and s/false/failure(), right?

LogicalResult is used in general all over the place to signal if an something succeeded or failed. It's kind of a convention. It provides more context than just a bool, that could be interpreted in different ways... but most importantly, it enforces the caller to check if the result is success or failure and handle both outcomes properly (if the return value if the function is ignored you will get a warning). In general, there are a few benefits and I think this is a common case where it's used.