[Mlir-commits] [mlir] c42ecce - [mlir][sparse] optimizing permutation validity check in toMLIRSparseTensor

[wren romano]

Author: wren romano
Date: 2022-09-29T15:08:46-07:00
New Revision: c42ecce7b92239085c4955f92100b2e961638d8c

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

LOG: [mlir][sparse] optimizing permutation validity check in toMLIRSparseTensor

The previous sorting-based check was O(n*log n). The new implementation is O(n).

This is a followup to the refactoring of D133462, D133830, D133831, and D133833.

Depends On D133833

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D133838

Added: 
    

Modified: 
    mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
index 58830af6ab37..992c85ac827c 100644
--- a/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseTensorUtils.cpp
@@ -78,11 +78,13 @@ toMLIRSparseTensor(uint64_t rank, uint64_t nse, const uint64_t *shape,
                    const uint64_t *perm, const DimLevelType *sparsity) {
 #ifndef NDEBUG
   // Verify that perm is a permutation of 0..(rank-1).
-  std::vector<uint64_t> order(perm, perm + rank);
-  std::sort(order.begin(), order.end());
-  for (uint64_t i = 0; i < rank; ++i)
-    if (i != order[i])
+  std::vector<bool> seen(rank, false);
+  for (uint64_t i = 0; i < rank; ++i) {
+    const uint64_t j = perm[i];
+    if (j >= rank || seen[j])
       MLIR_SPARSETENSOR_FATAL("Not a permutation of 0..%" PRIu64 "\n", rank);
+    seen[j] = true;
+  }
 
   // Verify that the sparsity values are supported.
   for (uint64_t i = 0; i < rank; ++i)