shylie/llvm-project
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[mlir][vector] fix: unroll vector.from_elements in gpu pipelines (#154774)

Browse Source 
### Problem

PR #142944 introduced a new canonicalization pattern which caused
failures in the following GPU-related integration tests:

-
mlir/test/Integration/GPU/CUDA/TensorCore/sm80/transform-mma-sync-matmul-f16-f16-accum.mlir
-
mlir/test/Integration/GPU/CUDA/TensorCore/sm80/transform-mma-sync-matmul-f32.mlir

The issue occurs because the new canonicalization pattern can generate
multi-dimensional `vector.from_elements` operations (rank > 1), but the
GPU lowering pipelines were not equipped to handle these during the
conversion to LLVM.

### Fix

This PR adds `vector::populateVectorFromElementsLoweringPatterns` to the
GPU lowering passes that are integrated in `gpu-lower-to-nvvm-pipeline`:

- `GpuToLLVMConversionPass`: the general GPU-to-LLVM conversion pass.
- `LowerGpuOpsToNVVMOpsPass`: the NVVM-specific lowering pass.

Co-authored-by: Yang Bai <yangb@nvidia.com>
This commit is contained in:
Yang Bai 2025-08-22 10:46:06 +08:00 committed by GitHub
parent 418fb50301
commit f1f194bf10
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 7 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlir/lib/Conversion/GPUCommon/GPUToLLVMConversion.cpp
View File

@ -532,6 +532,9 @@ void GpuToLLVMConversionPass::runOnOperation() {
// Vector transfer ops with rank > 1 should be lowered with VectorToSCF. // Vector transfer ops with rank > 1 should be lowered with VectorToSCF.
vector::populateVectorTransferLoweringPatterns(patterns, vector::populateVectorTransferLoweringPatterns(patterns,
/*maxTransferRank=*/1); /*maxTransferRank=*/1);
// Transform N-D vector.from_elements to 1-D vector.from_elements before
// conversion.
vector::populateVectorFromElementsLoweringPatterns(patterns);
if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
return signalPassFailure(); return signalPassFailure();
} }

4
mlir/lib/Conversion/GPUToNVVM/LowerGpuOpsToNVVMOps.cpp
View File

@ -27,6 +27,7 @@
#include "mlir/Dialect/Math/IR/Math.h" #include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h" #include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h" #include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/Transforms/DialectConversion.h" #include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h" #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@ -369,6 +370,9 @@ struct LowerGpuOpsToNVVMOpsPass final
{ {
RewritePatternSet patterns(m.getContext()); RewritePatternSet patterns(m.getContext());
populateGpuRewritePatterns(patterns); populateGpuRewritePatterns(patterns);
// Transform N-D vector.from_elements to 1-D vector.from_elements before
// conversion.
vector::populateVectorFromElementsLoweringPatterns(patterns);
if (failed(applyPatternsGreedily(m, std::move(patterns)))) if (failed(applyPatternsGreedily(m, std::move(patterns))))
return signalPassFailure(); return signalPassFailure();
} }