[Mlir-commits] [mlir] e892f32 - [mlir][nvgpu] Allow TMA's last dim to be non-128B without swizzling (#81499)

[llvmlistbot at llvm.org]

Author: Guray Ozen
Date: 2024-02-13T09:17:09+01:00
New Revision: e892f323ea234a997888b0fa44356b0678eda5f0

URL: https://github.com/llvm/llvm-project/commit/e892f323ea234a997888b0fa44356b0678eda5f0
DIFF: https://github.com/llvm/llvm-project/commit/e892f323ea234a997888b0fa44356b0678eda5f0.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp b/mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
index 4b6327479a219c..26f831f10a4e40 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)