Page MenuHomePhabricator

[mlir][Linalg] Convolution tiling added to ConvOp vectorization pass
ClosedPublic

Authored by limo1996 on Sep 15 2020, 2:23 AM.

Details

Summary

ConvOp vectorization supports now only convolutions of static shapes with dimensions
of size either 3(vectorized) or 1(not) as underlying vectors have to be of static
shape as well. In this commit we add support for convolutions of any size as well as
dynamic shapes by leveraging existing matmul infrastructure for tiling of both input
and kernel to sizes accepted by the previous version of ConvOp vectorization.
In the future this pass can be extended to take "tiling mask" as a user input which
will enable vectorization of user specified dimensions.

Diff Detail

Event Timeline

limo1996 created this revision.Sep 15 2020, 2:23 AM
Herald added a project: Restricted Project. · View Herald TranscriptSep 15 2020, 2:23 AM
limo1996 requested review of this revision.Sep 15 2020, 2:23 AM
nicolasvasilache accepted this revision.Sep 15 2020, 5:29 AM

Looks good!

This revision is now accepted and ready to land.Sep 15 2020, 5:29 AM
ftynse added inline comments.Sep 15 2020, 6:03 AM
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
35

SmallVectorImpl

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

Nit: ForOp reads like scf::ForOp, I'd just drop it and expand "Vect" to "Vectorization" instead

445

Could we factor out constant strings into named variables (constexpr static StringRef kTiledMarker = "TILED"; works)

452

No need to prefix SmallVector with llvm::

469

Would it make sense to declare const int64_t noTile = 1 and const int64_t tile = 3 and use those instead of magic numbers below?

471

Probably not for this commit, but it should be possible to somehow emit the dimensionality of the convolution op in the op class itself. So one would write populateVectPatterns<ConvNWCOp, ConvNWCOp::NumDims> (or actually query this inside the function instead of leaking magic numbers here).

mlir/test/lib/Transforms/TestConvVectorization.cpp
81

The line number will get out-of-sync reeeally fast. Point to the function or provide some more meaningful description on how to find it.

limo1996 updated this revision to Diff 291895.Sep 15 2020, 6:45 AM
limo1996 marked 5 inline comments as done.

Comments of Alex addressed

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

I think it will make code definitely more readable but for the longer arrays it might look annoying.. So I don't have a preference. Also in the future I would let user to pass the tiling mask as a pass argument so it will disappear.

471

I see. Well I can do that together with adding it as a pass argument. I would keep those as default values in case user does not want to specify them..

ftynse added inline comments.Sep 15 2020, 6:54 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
469

There are other ways of reducing verbosity without relying to magic numbers. There is currently nothing in this file that explains what "1" and "3" mean. I could have guessed today, but two weeks from now I and likely you will have forgotten what they mean.

Something like

const int64_t noTile = 1;
const int64_t tile = 3;
auto makeTileSizes = [](unsigned numNoTile, unsigned numTile) {
  SmallVector<int64_t, 10> result(numNoTile, noTile);
  result.append(numTile, tile);
  return result;
};

with appropriate comments used as populate<...>(..., makeTileSizes(/*numNoTile=*/3, /*numTile*/=2) goes a long way towards maintainability of the code.

471

Maybe I am misunderstanding the magic values that are passed as template parameters. I assumed they were numbers of loops in convolutions (NCW has 3 and NCDHW has 5, for example). These are properties of the _op_, I would strongly object to passing them as pass arguments and being anything else than what the op expects, unless the op was specifically designed for this purpose.

limo1996 marked 2 inline comments as done.Sep 15 2020, 7:23 AM
limo1996 added inline comments.
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
471

Ahh sorry I meant the array not template argument. You are right. I will do it in separate commit.

limo1996 updated this revision to Diff 291927.Sep 15 2020, 7:50 AM
limo1996 marked an inline comment as done.

Refactored "Magic numbers" away by introducing NoTile and TileSize constants in ConvOpVectorization
rewrite pattern.

limo1996 added inline comments.Sep 15 2020, 7:52 AM
mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp
469

Done. But I moved noTile and tile into ConvOpVectorization as those "magic numbers" are used elsewhere as well.

limo1996 updated this revision to Diff 291953.Sep 15 2020, 9:06 AM

NoTile -> noTile, TileSize -> tileSize

ftynse accepted this revision.Sep 17 2020, 1:45 AM
ftynse added inline comments.
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
593

Let's have a TODO that this should be parameterizable somehow.

limo1996 marked 2 inline comments as done.Sep 17 2020, 1:50 AM
limo1996 added inline comments.
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h
593

As a comment in the code?

limo1996 marked an inline comment as done.Sep 17 2020, 2:27 AM
This revision was landed with ongoing or failed builds.Sep 17 2020, 2:40 AM
This revision was automatically updated to reflect the committed changes.