diff options
| author | Guray Ozen <guray.ozen@gmail.com> | 2024-02-13 09:17:09 +0100 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-02-13 09:17:09 +0100 |
| commit | e892f323ea234a997888b0fa44356b0678eda5f0 (patch) | |
| tree | 0d71fd5d336c4a1de2a3d925de9e8ba660611dc6 | |
| parent | 87d771193490e6604f132752438bd8404c83498c (diff) | |
| download | llvm-e892f323ea234a997888b0fa44356b0678eda5f0.zip llvm-e892f323ea234a997888b0fa44356b0678eda5f0.tar.gz llvm-e892f323ea234a997888b0fa44356b0678eda5f0.tar.bz2 | |
[mlir][nvgpu] Allow TMA's last dim to be non-128B without swizzling (#81499)
Allow TMA's last dimension to be non-128B when swizzling mode is not
set.
Test `tma_load_64x8_8x128_noswizzle.mlir` is failing due to the
verifier. This PR will fix that
| -rw-r--r-- | mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp | 3 |
1 files changed, 2 insertions, 1 deletions
diff --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp index 4b632747..26f831f 100644 --- a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp +++ b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp @@ -362,7 +362,8 @@ std::optional<InFlightDiagnostic> verifyTmaDescriptorWithMemref( << kMaxTMADimension << " but it is " << dim; } } - if (descMemref.getRank() > 1) { + if (descMemref.getRank() > 1 && + descType.getSwizzle() != TensorMapSwizzleKind::SWIZZLE_NONE) { unsigned lastDimensionByte = descMemref.getElementTypeBitWidth() * descMemref.getShape().back() / 8; if (lastDimensionByte != kMaxTMALastdimByte) |