[Mlir-commits] [mlir] d8ae3a9 - [mlir][sparse] Fix strict weak ordering in sorting

[Aart Bik]

Author: Aart Bik
Date: 2023-08-22T17:15:10-07:00
New Revision: d8ae3a9b2533fbd11c0be6aaf66c16fb2b918f2a

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

LOG: [mlir][sparse] Fix strict weak ordering in sorting

These checks have been fired when comparing same elements like comp(a, a). Let's fix it. I don't have commit rights.

 Danila Kutenin
 kutdanila at yandex.ru

Reviewed By: aartbik

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

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
index 23b278317f589d..854b6f5587073b 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenEnv.cpp
@@ -31,20 +31,20 @@ static bool isMaterializing(Value val) {
 /// Makes target array's elements sorted according to the `order` array.
 static void sortArrayBasedOnOrder(std::vector<LoopId> &target,
                                   ArrayRef<LoopId> order) {
-  std::sort(target.begin(), target.end(), [&order](LoopId l, LoopId r) {
-    assert(l != r);
-    int idxL = -1, idxR = -1;
-    for (int i = 0, e = order.size(); i < e; i++) {
-      if (order[i] == l)
-        idxL = i;
-      if (order[i] == r)
-        idxR = i;
-    }
-    assert(idxL >= 0 && idxR >= 0);
-    return idxL < idxR;
-  });
+  std::sort(target.begin(), target.end(),
+            [&order](const LoopId &l, const LoopId &r) {
+              assert(std::addressof(l) == std::addressof(r) || l != r);
+              int idxL = -1, idxR = -1;
+              for (int i = 0, e = order.size(); i < e; i++) {
+                if (order[i] == l)
+                  idxL = i;
+                if (order[i] == r)
+                  idxR = i;
+              }
+              assert(idxL >= 0 && idxR >= 0);
+              return idxL < idxR;
+            });
 }
-
 //===----------------------------------------------------------------------===//
 // Code generation environment constructor and general methods
 //===----------------------------------------------------------------------===//

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
index 06d101886ae726..ce77d7a519877d 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
@@ -566,6 +566,8 @@ void sparse_tensor::foreachInSparseConstant(
                   continue;
                 return lhsCoords[l].getInt() < rhsCoords[l].getInt();
               }
+              if (std::addressof(lhs) == std::addressof(rhs))
+                return false;
               llvm_unreachable("no equal coordinate in sparse element attr");
             });