[mlir][vector][nvgpu] Move MMA contraction preparation to VectorUtils
authorJakub Kuderski <kubak@google.com>
Thu, 9 Mar 2023 19:56:20 +0000 (14:56 -0500)
committerJakub Kuderski <kubak@google.com>
Thu, 9 Mar 2023 19:56:21 +0000 (14:56 -0500)
commitfb7ef637a84652dbd3d973a1ba7db9470181b5aa
tree7a594f47e8ff3a36ff6da3a6ede3b1380b4020f7
parente99b5ad38381ab263820b23a184d217a4112519c
[mlir][vector][nvgpu] Move MMA contraction preparation to VectorUtils

This pattern is not specific to nvgpu; I intend to use in SPIR-V codegen. `VectorTransforms` seems like a more generally useful place.

In addition:
-  Fix a bug in the second condition (the dimensions were swapped for RHS).
-  Add tests.
-  Add support for externally provided filter functions, similar to other vector transforms.
-  Prefer to transpose before zero/sign-extending inputs.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D145638
mlir/include/mlir/Dialect/NVGPU/Utils/MMAUtils.h
mlir/include/mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h
mlir/lib/Conversion/VectorToGPU/VectorToGPU.cpp
mlir/lib/Dialect/NVGPU/Utils/MMAUtils.cpp
mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
mlir/test/Dialect/Vector/vector-contract-matmul-transforms.mlir [new file with mode: 0644]
mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp