[MLIR] Fix use-after-move for DEBUG builds, and broken assert logic.

[noclowns]

These issues affect only Debug builds, and Release builds with asserts enabled.

1. In SparseTensor.h a variable is moved-from within an assert, introducing a side effect that alters its subsequent use, and causes divergence between Debug and Release builds (with asserts disabled).

2. In IterationGraphSorter.cpp, the class constructor arguments are moved-from to initialize class member variables via the initializer list. Because both the arguments and class members are identically named, there's a naming collision where the arguments shadow their identically-named member variables counterparts inside the constructor body. In the original code, unqualified names inside the asserts, referred to the constructor arguments. This is wrong, because these have already been moved-from. It's not just a UB, but is broken. These SmallVector types when moved-from are reset i.e. the size resets to 0. This actually renders the affected asserts ineffective, since the comparisons operate on two hollowed-out objects and always succeed. This name ambiguity is fixed by using 'this->' to correctly refer to the initialized member variables carrying the relevant state.

3. While the fix 2 above made the asserts act as intended, it also unexpectedly broke one mlir test: llvm-lit -v mlir/test/Dialect/SparseTensor/sparse_scalars.mlir This required fixing the assert logic itself, which likely has never worked and went unnoticed all this time due to the bug 2. Specifically, in the failing test that uses mlir/test/Dialect/SparseTensor/sparse_scalars.mlir the '%argq' of 'ins' is defined as 'f32' scalar type, but the original code inside the assert had no support for scalar types as written, and was breaking the test.

Testing:

ninja check-mlir
llvm-lit -v mlir/test/Dialect/SparseTensor/sparse_scalars.mlir

[llvmbot]

@llvm/pr-subscribers-mlir-sparse

@llvm/pr-subscribers-mlir

Author: Slava Gurevich (noclowns)

Changes

Fix use-after-move issues that occur only in DEBUG builds.

In one, constructor arguments are moved in the initializer list to initialize member variables, but the function body mistakenly accesses the moved-from arguments instead of the corresponding class members with identical names. While this is a UB, here it actually renders the affected asserts ineffective, since the comparisons operate on two hollowed-out objects, by implementation.

In a second fix, a variable is moved from within an assert, introducing a side effect that alters its subsequent use and causes divergence between DEBUG and RELEASE builds.

Testing:
ninja check-llvm
ninja check-llvm-unit

---

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

2 Files Affected:

• (modified) mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h (+2-4)
• (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp (+3-3)

diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
index d0a3f01afe871..43e48a6d34026 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h
@@ -158,16 +158,14 @@ namespace sparse_tensor {
 /// Convenience method to abbreviate casting `getType()`.
 template <typename T>
 inline RankedTensorType getRankedTensorType(T &&t) {
-  assert(static_cast<bool>(std::forward<T>(t)) &&
-         "getRankedTensorType got null argument");
+  assert(static_cast<bool>(t) && "getRankedTensorType got null argument");
   return dyn_cast<RankedTensorType>(std::forward<T>(t).getType());
 }
 
 /// Convenience method to abbreviate casting `getType()`.
 template <typename T>
 inline MemRefType getMemRefType(T &&t) {
-  assert(static_cast<bool>(std::forward<T>(t)) &&
-         "getMemRefType got null argument");
+  assert(static_cast<bool>(t) && "getMemRefType got null argument");
   return cast<MemRefType>(std::forward<T>(t).getType());
 }
 
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
index 73e0f3d2891d7..c70fd3b010406 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/IterationGraphSorter.cpp
@@ -159,12 +159,12 @@ IterationGraphSorter::IterationGraphSorter(
       loop2OutLvl(loop2OutLvl), iterTypes(std::move(iterTypes)),
       strategy(strategy) {
   // One map per tensor.
-  assert(loop2InsLvl.size() == ins.size());
+  assert(this->loop2InsLvl.size() == this->ins.size());
   // All the affine maps have the same number of dimensions (loops).
   assert(llvm::all_equal(llvm::map_range(
-      loop2InsLvl, [](AffineMap m) { return m.getNumDims(); })));
+      this->loop2InsLvl, [](AffineMap m) { return m.getNumDims(); })));
   // The number of results of the map should match the rank of the tensor.
-  assert(llvm::all_of(llvm::zip(loop2InsLvl, ins), [](auto mvPair) {
+  assert(llvm::all_of(llvm::zip(this->loop2InsLvl, this->ins), [](auto mvPair) {
     auto [m, v] = mvPair;
     return m.getNumResults() == cast<ShapedType>(v.getType()).getRank();
   }));

[joker-eph]

@noclowns : the change in mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h looks right, why did you close?
(I don't this the change in the other file though)

[noclowns]

@noclowns : the change in mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h looks right, why did you close? (I don't this the change in the other file though)

It's not passing the QA gate, need more time to investigate.

[noclowns]

This could use a good review for part 3 (the functional fix) to make sure the original intent is correctly inferred and implemented. Thanks.

[noclowns]

Updated with suggested changes, and added handing for UnrankedTensorType.

[joker-eph]

Still LGTM