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

[mlir][nvgpu] Add tma last dim bytes check (#153451)

Browse Source 
Add the check the number of bytes in the last dimension of Tma must be a
multiple of 16.
This commit is contained in:
lonely eagle 2025-08-14 20:14:20 +08:00 committed by GitHub
parent 87de48d11f
commit 6d08a39eeb
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 17 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mlir/lib/Dialect/NVGPU/IR/NVGPUDialect.cpp
View File

@ -422,6 +422,12 @@ std::optional<InFlightDiagnostic> verifyTmaDescriptorWithMemref(
<< descMemref << " != " << dstMemref;
}
int lastDimBytes =
descMemref.getShape().back() * descMemref.getElementTypeBitWidth() / 8;
if (lastDimBytes % 16 != 0) {
return op->emitError() << "the bytes in the last dimension of the tensor "
"map must be a multiple of 16";
}
return std::nullopt;
}

11
mlir/test/Dialect/NVGPU/invalid.mlir
View File

@ -378,3 +378,14 @@ func.func @check_matrixC_dim(%arg0: vector<4x4xf16>, %arg1: vector<2x2xf16>, %ar
%d = nvgpu.mma.sync (%arg0, %arg1, %arg2) {mmaShape = [16, 8, 16]} : (vector<4x4xf16>, vector<2x2xf16>, vector<4xf16>) -> vector<2x2xf16>
return %d : vector<2x2xf16>
}
// -----
!desc = !nvgpu.tensormap.descriptor<tensor = memref<32x8xi8,3>, swizzle=none, l2promo = none, oob = zero, interleave = none>
!mbarrier = !nvgpu.mbarrier.group<memorySpace = #gpu.address_space<workgroup>>
func.func @tma_last_dim_bytes(%desc: !desc, %buffer: memref<32x8xi8,3>, %mbarrier: !mbarrier) {
%c0 = arith.constant 0 : index
// expected-error @+1 {{the bytes in the last dimension of the tensor map must be a multiple of 16}}
nvgpu.tma.async.load %desc[%c0, %c0], %mbarrier[%c0] to %buffer : !desc, !mbarrier -> memref<32x8xi8,3>
return
}