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

grypp · web-flow · commit e892f323ea23 · 2024-02-13T09:17:09.000+01:00
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
diff --git 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)

Original file line number	Diff line number	Diff line change
`@@ -362,7 +362,8 @@ std::optional<InFlightDiagnostic> verifyTmaDescriptorWithMemref(`
`362`	`362`	`<< kMaxTMADimension << " but it is " << dim;`
`363`	`363`	`}`
`364`	`364`	`}`
`365`		`- if (descMemref.getRank() > 1) {`
	`365`	`+ if (descMemref.getRank() > 1 &&`
	`366`	`+ descType.getSwizzle() != TensorMapSwizzleKind::SWIZZLE_NONE) {`
`366`	`367`	`unsigned lastDimensionByte =`
`367`	`368`	`descMemref.getElementTypeBitWidth() * descMemref.getShape().back() / 8;`
`368`	`369`	`if (lastDimensionByte != kMaxTMALastdimByte)`