[Mlir-commits] [mlir] ee92ac2 - [mlir][nvgpu] Fix crash in optimize-shared-memory pass with vector element types (#179111)

[llvmlistbot at llvm.org]

Author: Jueon Park
Date: 2026-03-04T12:33:40+01:00
New Revision: ee92ac2343f6e0ff787bbe660c9a9293caa5269c

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

LOG: [mlir][nvgpu] Fix crash in optimize-shared-memory pass with vector element types (#179111)

The `--nvgpu-optimize-shared-memory` pass crashed when processing
memrefs with vector element types (e.g., `memref<16x1xvector<16xf16>,
3>`). This occurred because getElementTypeBitWidth() calls
getIntOrFloatBitWidth(), which asserts the element type must be an
integer or float.
Thus, this PR adds an early-exit guard to return failure() when the
memref's element type is not a scalar int or float.

I wasn't sure if we should support vector types (by multiplying element
bit width by vector length) or just reject them. For now, I've
implemented it to return failure on non-scalar types.

Fixes #177823

Co-authored-by: rebel-jueonpark <jueonpark at rebellions.ai>

Added: 
    

Modified: 
    mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
    mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
index c4d00b48a49d3..e06b15febf145 100644
--- a/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
+++ b/mlir/lib/Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp
@@ -156,6 +156,11 @@ mlir::nvgpu::optimizeSharedMemoryReadsAndWrites(Operation *parentOp,
   if (memRefType.getRank() == 0)
     return failure();
 
+  // Only support memrefs with scalar element types (i.e., int or float).
+  // Memrefs with vector element types are not supported.
+  if (!memRefType.getElementType().isIntOrFloat())
+    return failure();
+
   // Abort if the given value has any sub-views; we do not do any alias
   // analysis.
   bool hasSubView = false;

diff  --git a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
index 596d24b94811e..128c4c5d918d9 100644
--- a/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
+++ b/mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir
@@ -259,3 +259,14 @@ func.func @test_dynamic_and_zero_dim(%arg0 : index) {
   %alloc_1 = memref.alloc(%arg0) : memref<?xf32, 3>
   return
 }
+
+// -----
+
+// Ensure memrefs with vector element types do not crash (issue #177823).
+
+// CHECK-LABEL: func @test_vector_element_type
+// CHECK: memref.alloc() : memref<16x1xvector<16xf16>, 3>
+func.func @test_vector_element_type() {
+  %alloc = memref.alloc() : memref<16x1xvector<16xf16>, 3>
+  return
+}