[mlir][nvgpu] Fix a division by zero crash in OptimizeSharedMemoryPass (#174931)
Fixes #173553.
This commit is contained in:
parent
02e615d654
commit
ab95de00b3
@ -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;
|
||||
|
||||
@ -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
|
||||
}
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user