[Mlir-commits] [mlir] [mlir][sparse] Reject tensor files with unsupported ranks (PR #220709)

[Mehdi Amini]

https://github.com/joker-eph updated https://github.com/llvm/llvm-project/pull/220709

>From 77bd56fe20791e0345dc8b245be246fbf336b91c Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Wed, 2 Sep 2026 07:34:16 -0700
Subject: [PATCH] [mlir][sparse] Reject tensor files with unsupported ranks

The sparse tensor reader stores the rank, nonzero count, and dimension
sizes in a fixed 512-element array. Reject ranks above 510 before reading
the dimension sizes so an untrusted FROSTT header cannot write past it.

The pre-fix test failed under Valgrind with:

  ==623833== Invalid write of size 8
  ==623833==    by SparseTensorReader::readExtFROSTTHeader()
  ==623833== Address ... is 0 bytes after a block of size 5,152 alloc'd

The post-fix run rejects rank 700 before reading dimensions and reports
zero Valgrind errors.

This addresses CID 1499743.

Found by Coverity.

Assisted-by: Codex
---
 .../mlir/ExecutionEngine/SparseTensor/File.h  |  3 +-
 .../lib/ExecutionEngine/SparseTensor/File.cpp |  7 +++
 .../SparseTensor/CPU/sparse_file_rank.mlir    | 44 +++++++++++++++++++
 mlir/test/Integration/data/rank_511.tns       |  3 ++
 4 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_rank.mlir
 create mode 100644 mlir/test/Integration/data/rank_511.tns

diff --git a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
index a6f199c0a0c1e..3ce4159e2e398 100644
--- a/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
+++ b/mlir/include/mlir/ExecutionEngine/SparseTensor/File.h
@@ -267,12 +267,13 @@ class SparseTensorReader final {
   /// nonzeros, and the dimensions sizes (one per rank) of the sparse tensor.
   void readExtFROSTTHeader();
 
+  static constexpr uint64_t kMaxRank = 510;
   static constexpr int kColWidth = 1025;
   const char *const filename;
   FILE *file = nullptr;
   ValueKind valueKind_ = ValueKind::kInvalid;
   bool isSymmetric_ = false;
-  uint64_t idata[512];
+  uint64_t idata[kMaxRank + 2];
   char line[kColWidth];
 };
 
diff --git a/mlir/lib/ExecutionEngine/SparseTensor/File.cpp b/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
index 63b29ba34ba02..224dea1f4efe9 100644
--- a/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensor/File.cpp
@@ -187,6 +187,13 @@ void SparseTensorReader::readExtFROSTTHeader() {
     fprintf(stderr, "Cannot find metadata in %s\n", filename);
     exit(1);
   }
+  if (idata[0] > kMaxRank) {
+    fprintf(stderr,
+            "Tensor rank %" PRIu64 " exceeds maximum supported rank %" PRIu64
+            " in %s\n",
+            idata[0], kMaxRank, filename);
+    exit(1);
+  }
   // Followed by a line with the dimension sizes (one per rank).
   for (uint64_t r = 0; r < idata[0]; r++) {
     if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1) {
diff --git a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_rank.mlir b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_rank.mlir
new file mode 100644
index 0000000000000..5cb3898e3de6f
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_rank.mlir
@@ -0,0 +1,44 @@
+//--------------------------------------------------------------------------------------------------
+// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
+//
+// Set-up that's shared across all tests in this directory. In principle, this
+// config could be moved to lit.local.cfg. However, there are downstream users that
+//  do not use these LIT config files. Hence why this is kept inline.
+//
+// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
+// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
+// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
+// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
+// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
+// DEFINE: %{run_libs_sve} = -shared-libs=%native_mlir_runner_utils,%native_mlir_c_runner_utils
+// DEFINE: %{run_opts} = -e main -entry-point-result=void
+// DEFINE: %{run} = mlir-runner %{run_opts} %{run_libs}
+// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve" %{run_opts} %{run_libs_sve}
+//
+// DEFINE: %{env} =
+//--------------------------------------------------------------------------------------------------
+
+// REDEFINE: %{env} = TENSOR0="%mlir_src_dir/test/Integration/data/rank_511.tns"
+// RUN: %{compile} | not env %{env} %{run} 2>&1 | FileCheck %s
+
+// CHECK: Tensor rank 511 exceeds maximum supported rank 510
+
+!Filename = !llvm.ptr
+
+#SparseTensor = #sparse_tensor.encoding<{
+  map = (d0) -> (d0 : compressed)
+}>
+
+module {
+  func.func private @getTensorFilename(index) -> !Filename
+
+  func.func @main() {
+    %c0 = arith.constant 0 : index
+    %fileName = call @getTensorFilename(%c0) : (index) -> !Filename
+    %tensor = sparse_tensor.new %fileName
+      : !Filename to tensor<1xf64, #SparseTensor>
+    bufferization.dealloc_tensor %tensor
+      : tensor<1xf64, #SparseTensor>
+    return
+  }
+}
diff --git a/mlir/test/Integration/data/rank_511.tns b/mlir/test/Integration/data/rank_511.tns
new file mode 100644
index 0000000000000..6d7be368f0df3
--- /dev/null
+++ b/mlir/test/Integration/data/rank_511.tns
@@ -0,0 +1,3 @@
+# Extended FROSTT tensor with an unsupported rank.
+511 0
+1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1 1