[Mlir-commits] [mlir] b082ea1 - [MLIR][LLVMIR] Fix llvm.extractvalue folder (#201838)

[llvmlistbot at llvm.org]

Author: Vadim Curcă
Date: 2026-06-05T16:01:55+02:00
New Revision: b082ea1f961b22e02a870c0cb6f3f8bcd031b2a8

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

LOG: [MLIR][LLVMIR] Fix llvm.extractvalue folder (#201838)

Fix the llvm.extractvalue folder when trying to extract a value from a
multidimensional constant. Add lit tests that would crash without the
fix.

Added: 
    

Modified: 
    mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
    mlir/test/Dialect/LLVMIR/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
index 4f21db172d813..ac6cf03e057dc 100644
--- a/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
+++ b/mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp
@@ -1943,12 +1943,16 @@ static Type getInsertExtractValueElementType(Type llvmType,
 }
 
 /// Extracts the element at the given index from an attribute. For
-/// `ElementsAttr` and `ArrayAttr`, returns the element at the specified index.
-/// For `ZeroAttr`, `UndefAttr`, and `PoisonAttr`, returns the attribute itself
-/// unchanged. Returns `nullptr` if the attribute is not one of these types or
-/// if the index is out of bounds.
+/// `ElementsAttr`, returns the element at the specified index, or `nullptr` if
+/// the shaped type does not have rank 1. For `ArrayAttr`, returns the element
+/// at the specified index. For `ZeroAttr`, `UndefAttr`, and `PoisonAttr`,
+/// returns the attribute itself unchanged. Returns `nullptr` if the attribute
+/// is not one of these types or if the index is out of bounds.
 static Attribute extractElementAt(Attribute attr, size_t index) {
   if (auto elementsAttr = dyn_cast<ElementsAttr>(attr)) {
+    ShapedType shapedType = elementsAttr.getShapedType();
+    if (!shapedType.hasRank() || shapedType.getRank() != 1)
+      return nullptr;
     if (index < static_cast<size_t>(elementsAttr.getNumElements()))
       return elementsAttr.getValues<Attribute>()[index];
     return nullptr;

diff  --git a/mlir/test/Dialect/LLVMIR/canonicalize.mlir b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
index 544666fb30099..11e28d52355cd 100644
--- a/mlir/test/Dialect/LLVMIR/canonicalize.mlir
+++ b/mlir/test/Dialect/LLVMIR/canonicalize.mlir
@@ -185,6 +185,30 @@ llvm.func @fold_extract_sparse() -> f32 {
 
 // -----
 
+// CHECK-LABEL: no_fold_extract_splat_rank_mismatch
+llvm.func @no_fold_extract_splat_rank_mismatch() -> vector<2xi32> {
+  %0 = llvm.mlir.constant(dense<12> : vector<2xi32>) : vector<2xi32>
+  %1 = llvm.mlir.constant(dense<23> : vector<4x2xi32>) : !llvm.array<4 x vector<2xi32>>
+  // CHECK: extractvalue
+  %2 = llvm.extractvalue %1[0] : !llvm.array<4 x vector<2xi32>>
+  %3 = llvm.shl %0, %2 : vector<2xi32>
+  llvm.return %3 : vector<2xi32>
+}
+
+// -----
+
+// CHECK-LABEL: no_fold_extract_sparse_rank_mismatch
+llvm.func @no_fold_extract_sparse_rank_mismatch() -> vector<2xi32> {
+  %0 = llvm.mlir.constant(dense<12> : vector<2xi32>) : vector<2xi32>
+  %1 = llvm.mlir.constant(sparse<[[0, 0]], [23]> : vector<4x2xi32>) : !llvm.array<4 x vector<2xi32>>
+  // CHECK: extractvalue
+  %2 = llvm.extractvalue %1[0] : !llvm.array<4 x vector<2xi32>>
+  %3 = llvm.shl %0, %2 : vector<2xi32>
+  llvm.return %3 : vector<2xi32>
+}
+
+// -----
+
 // CHECK-LABEL: fold_zero
 llvm.func @fold_zero() -> i32 {
   // CHECK-NOT: insertvalue