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

[Jueon Park]

https://github.com/JueonPark updated https://github.com/llvm/llvm-project/pull/179111

>From f4aac04a3d329189f36dbb035599f0f1dad75dd5 Mon Sep 17 00:00:00 2001
From: rebel-jueonpark <jueonpark at rebellions.ai>
Date: Sun, 1 Feb 2026 23:21:32 +0900
Subject: [PATCH] [mlir][nvgpu] Fix crash in optimize-shared-memory pass with
 vector element types

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.

Add an early-exit guard to return failure() when the memref's element
type is not a scalar int or float.

Fixes #177823
---
 .../Dialect/NVGPU/Transforms/OptimizeSharedMemory.cpp |  5 +++++
 mlir/test/Dialect/NVGPU/optimize-shared-memory.mlir   | 11 +++++++++++
 2 files changed, 16 insertions(+)

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
+}