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.
Details
Diff Detail
- Repository
- rG LLVM Github Monorepo
Event Timeline
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. |
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.. |
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. |
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. |
Refactored "Magic numbers" away by introducing NoTile and TileSize constants in ConvOpVectorization
rewrite pattern.
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. |
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h | ||
---|---|---|
593 | Let's have a TODO that this should be parameterizable somehow. |
mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h | ||
---|---|---|
593 | As a comment in the code? |
SmallVectorImpl