[mlir][nvgpu] Fix a division by zero crash in OptimizeSharedMemoryPass (#174931)

Fixes #173553.
This commit is contained in:
Longsheng Mou 2026-01-16 10:42:08 +08:00 committed by GitHub
parent 02e615d654
commit ab95de00b3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 14 additions and 0 deletions

View File

@ -166,6 +166,9 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
// Check if this is necessary given the assumption of 128b accesses:
// If dim[rank-1] is small enough to fit 8 rows in a 128B line.
const int64_t rowSize = memRefType.getDimSize(memRefType.getRank() - 1);
if (ShapedType::isDynamic(rowSize) || rowSize == 0)
return failure();
const int64_t rowsPerLine =
(8 * kSharedMemoryLineSizeBytes / memRefType.getElementTypeBitWidth()) /
rowSize;

View File

@ -248,3 +248,14 @@ func.func @test_0_d() -> memref<i32, #gpu.address_space<workgroup>> {
%alloc = memref.alloc() : memref<i32, #gpu.address_space<workgroup>>
return %alloc : memref<i32, #gpu.address_space<workgroup>>
}
// -----
// Ensure the case with zero or dynamic dim not crash.
// CHECK-LABEL: func @test_dynamic_and_zero_dim
func.func @test_dynamic_and_zero_dim(%arg0 : index) {
%alloc = memref.alloc() : memref<0xf32, 3>
%alloc_1 = memref.alloc(%arg0) : memref<?xf32, 3>
return
}