[Mlir-commits] [mlir] [mlir][LLVM] Prevent vector bit-width overflow (PR #220898)

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/220898

Compute vector bit widths from the rank-one MLIR dimension and saturate the 64-bit product. Otherwise, large dimensions can overflow the multiplication or be truncated by ElementCount and bypass the 2^17-bit verifier limit.

Add regression coverage for both vector.insert and vector.extract. Before the fix, the oversized vector.insert case was accepted and its expected diagnostic was not produced.

Found by Coverity.

Assisted-by: Codex

>From fe83eeb40d255936e185dba4412dadeba70c4bf6 Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 3 Sep 2026 04:24:40 -0700
Subject: [PATCH] [mlir][LLVM] Prevent vector bit-width overflow

Compute vector bit widths from the rank-one MLIR dimension and saturate the
64-bit product. Otherwise, large dimensions can overflow the multiplication or
be truncated by ElementCount and bypass the 2^17-bit verifier limit.

Add regression coverage for both vector.insert and vector.extract. Before the
fix, the oversized vector.insert case was accepted and its expected diagnostic
was not produced.

Found by Coverity.

Assisted-by: Codex
---
 .../mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td   | 20 ++++++---
 mlir/test/Dialect/LLVMIR/invalid.mlir         | 42 +++++++++++++++++++
 2 files changed, 56 insertions(+), 6 deletions(-)

diff --git a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
index f40f3f5ebda93..30a2a1d9ffc20 100644
--- a/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
+++ b/mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td
@@ -1589,9 +1589,13 @@ def LLVM_vector_insert
     "type($srcvec) `into` type($res)";
   let extraClassDeclaration = [{
     uint64_t getVectorBitWidth(Type vector) {
-      return getVectorNumElements(vector).getKnownMinValue() *
-             ::llvm::cast<VectorType>(vector).getElementType()
-                .getIntOrFloatBitWidth();
+      auto vectorType = ::llvm::cast<VectorType>(vector);
+      uint64_t numElements = vectorType.getDimSize(0);
+      uint64_t elementBitWidth =
+          vectorType.getElementType().getIntOrFloatBitWidth();
+      if (numElements > uint64_t(-1) / elementBitWidth)
+        return uint64_t(-1);
+      return numElements * elementBitWidth;
     }
     uint64_t getSrcVectorBitWidth() {
       return getVectorBitWidth(getSrcvec().getType());
@@ -1622,9 +1626,13 @@ def LLVM_vector_extract
     "type($res) `from` type($srcvec)";
   let extraClassDeclaration = [{
     uint64_t getVectorBitWidth(Type vector) {
-      return getVectorNumElements(vector).getKnownMinValue() *
-             ::llvm::cast<VectorType>(vector).getElementType()
-                .getIntOrFloatBitWidth();
+      auto vectorType = ::llvm::cast<VectorType>(vector);
+      uint64_t numElements = vectorType.getDimSize(0);
+      uint64_t elementBitWidth =
+          vectorType.getElementType().getIntOrFloatBitWidth();
+      if (numElements > uint64_t(-1) / elementBitWidth)
+        return uint64_t(-1);
+      return numElements * elementBitWidth;
     }
     uint64_t getSrcVectorBitWidth() {
       return getVectorBitWidth(getSrcvec().getType());
diff --git a/mlir/test/Dialect/LLVMIR/invalid.mlir b/mlir/test/Dialect/LLVMIR/invalid.mlir
index eb0fb9aebb34b..765925568b829 100644
--- a/mlir/test/Dialect/LLVMIR/invalid.mlir
+++ b/mlir/test/Dialect/LLVMIR/invalid.mlir
@@ -1419,6 +1419,27 @@ func.func @insert_vector_invalid_source_vector_size(%arg0 : vector<16385xi8>, %a
 
 // -----
 
+func.func @insert_vector_overflowing_source_vector_size(%arg0 : vector<536870913xi8>, %arg1 : vector<16xi8>) {
+  // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
+  %0 = llvm.intr.vector.insert %arg0, %arg1[0] : vector<536870913xi8> into vector<16xi8>
+}
+
+// -----
+
+func.func @insert_vector_truncating_source_vector_size(%arg0 : vector<4294967296xi8>, %arg1 : vector<16xi8>) {
+  // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
+  %0 = llvm.intr.vector.insert %arg0, %arg1[0] : vector<4294967296xi8> into vector<16xi8>
+}
+
+// -----
+
+func.func @insert_vector_saturating_source_vector_size(%arg0 : vector<2305843009213693953xi8>, %arg1 : vector<16xi8>) {
+  // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
+  %0 = llvm.intr.vector.insert %arg0, %arg1[0] : vector<2305843009213693953xi8> into vector<16xi8>
+}
+
+// -----
+
 func.func @insert_vector_invalid_dest_vector_size(%arg0 : vector<16xi8>, %arg1 : vector<[16385]xi8>) {
   // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
   %0 = llvm.intr.vector.insert %arg0, %arg1[0] : vector<16xi8> into vector<[16385]xi8>
@@ -1440,6 +1461,27 @@ func.func @extract_vector_invalid_source_vector_size(%arg0 : vector<[16385]xi8>)
 
 // -----
 
+func.func @extract_vector_overflowing_source_vector_size(%arg0 : vector<536870913xi8>) {
+  // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
+  %0 = llvm.intr.vector.extract %arg0[0] : vector<16xi8> from vector<536870913xi8>
+}
+
+// -----
+
+func.func @extract_vector_truncating_source_vector_size(%arg0 : vector<4294967296xi8>) {
+  // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
+  %0 = llvm.intr.vector.extract %arg0[0] : vector<16xi8> from vector<4294967296xi8>
+}
+
+// -----
+
+func.func @extract_vector_saturating_source_vector_size(%arg0 : vector<2305843009213693953xi8>) {
+  // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
+  %0 = llvm.intr.vector.extract %arg0[0] : vector<16xi8> from vector<2305843009213693953xi8>
+}
+
+// -----
+
 func.func @extract_vector_invalid_result_vector_size(%arg0 : vector<[16]xi8>) {
   // expected-error at +1 {{op failed to verify that vectors are not bigger than 2^17 bits.}}
   %0 = llvm.intr.vector.extract %arg0[0] : vector<16385xi8> from vector<[16]xi8>