[mlir][sparse] Supporting (non)uniqueness in `SparseTensorStorage::lexDiff`

[wrengr]

Summary

This change corrects a crash condition

[wrengr]

adding regression test

[wrengr]

fixing NBSP characters

[wrengr]

updating commentary in the test file to indicate that it's difficult to recreate the crash conditions.

[wrengr]

rebase

[wrengr]

Updated the test to reliably trigger the crash. Unfortunately, the changes to lexDiff do not in fact correct this crash.

[aartbik]

https://reviews.llvm.org/D158578 with different test

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

void insPath(const uint64_t *lvlCoords, uint64_t diffLvl, uint64_t full,
     }
     values.push_back(val);
   }
 
   /// Finds the lexicographically first level where the level-coordinates
   /// in the argument differ from those in the current cursor.
   uint64_t lexDiff(const uint64_t *lvlCoords) const {
     const uint64_t lvlRank = getLvlRank();
-    for (uint64_t l = 0; l < lvlRank; ++l)
-      if (lvlCoords[l] > lvlCursor[l])
+    for (uint64_t l = 0; l < lvlRank; ++l) {
+      const auto crd = lvlCoords[l];
+      const auto cur = lvlCursor[l];
+      if (crd > cur || (crd == cur && !isUniqueLvl(l)))
         return l;
-      else
-        assert(lvlCoords[l] == lvlCursor[l] && "non-lexicographic insertion");
-    assert(0 && "duplicate insertion");
+      if (crd < cur) {
+        assert(false && "non-lexicographic insertion");
+        return -1u;
+      }
+    }
+    assert(false && "duplicate insertion");
     return -1u;
   }
 
   // Allow `SparseTensorEnumerator` to access the data-members (to avoid
   // the cost of virtual-function dispatch in inner loops), without
   // making them public to other client code.
   friend class SparseTensorEnumerator<P, C, V>;
 

mlir/test/Integration/Dialect/SparseTensor/CPU/nonunique_lexdiff.mlir

+// This test ensures proper handling of non-unique levels.
+// In particular, without the D156946 changes to `SparseTensorStorage::lexDiff`
+// this code causes the size-assertion in `SparseTensorStorage::toCOO` to fail.
+
+// RUN: mlir-opt %s \
+// RUN:   --sparse-compiler=enable-runtime-library=true \
+// RUN: | mlir-cpu-runner \
+// RUN:   -e entry -entry-point-result=void \
+// RUN:   -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils \
+// RUN: | FileCheck %s
+
+!Dense = tensor<2x2xi32>
+!COO = tensor<2x2xi32, #sparse_tensor.encoding<{lvlTypes = ["compressed-nu", "singleton"]}>>
+!Vec = vector<2x2xi32>
+module {
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %i0 = arith.constant 0 : i32
+    %i2 = arith.constant 2 : i32
+
+    %t0 = arith.constant dense<[[1, 0], [3, 4]]> : !Dense
+    // This conversion works just fine.
+    %t1 = sparse_tensor.convert %t0 : !Dense to !COO
+    // This insertion is needed to trigger the crash;
+    // without it, the conversion back to dense will work just fine.
+    %t2 = sparse_tensor.insert %i2 into %t1[%c0, %c1] : !COO
+    // This conversion is the thing that crashes.
+    %t3 = sparse_tensor.convert %t2 : !COO to !Dense
+
+    // CHECK: ( ( 1, 2 ), ( 3, 4 ) )
+    %v = vector.transfer_read %t3[%c0, %c0], %i0 : !Dense, !Vec
+    vector.print %v : !Vec
+
+    // Release the resources.
+    bufferization.dealloc_tensor %t1 : !COO
+    bufferization.dealloc_tensor %t2 : !COO
+
+    return
+  }
+}