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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Mehdi Amini (joker-eph)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/220898.diff


2 Files Affected:

- (modified) mlir/include/mlir/Dialect/LLVMIR/LLVMIntrinsicOps.td (+14-6) 
- (modified) mlir/test/Dialect/LLVMIR/invalid.mlir (+42) 


``````````diff
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>

``````````

</details>


https://github.com/llvm/llvm-project/pull/220898