Skip to content

Conversation

@rupprecht
Copy link
Collaborator

When we create a SparseIterator, we sometimes wrap it in a FilterIterator, which delegates some calls to the underlying SparseIterator.

After construction, e.g. in makeNonEmptySubSectIterator(), we call setSparseEmitStrategy(). This sets the strategy only in one of the filters -- if we call setSparseEmitStrategy() immediately after creating the SparseIterator, then the wrapped SparseIterator will have the right strategy, and the FilterIterator strategy will be unintialized; if we call setSparseEmitStrategy() after wrapping the iterator in FilterIterator, then the opposite happens.

If we make setSparseEmitStrategy() a virtual method so that it's included in the FilterIterator pattern, and then do all reads of emitStrategy via a virtual method as well, it's pretty simple to ensure that the value of strategy is being set consistently and correctly.

Without this, the UB of strategy being uninitialized manifests as a sporadic test failure in mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir, when run downstream with the right flags (e.g. asan + assertions off). The test sometimes fails with ne_sub<trivial<dense[0,1]>>.begin' op created with unregistered dialect. It can also be directly observed w/ msan that this uninitialized read is the cause of that issue, but msan causes other problems w/ this test.

@llvmbot
Copy link
Member

llvmbot commented Oct 29, 2025

@llvm/pr-subscribers-mlir

Author: Jordan Rupprecht (rupprecht)

Changes

When we create a SparseIterator, we sometimes wrap it in a FilterIterator, which delegates some calls to the underlying SparseIterator.

After construction, e.g. in makeNonEmptySubSectIterator(), we call setSparseEmitStrategy(). This sets the strategy only in one of the filters -- if we call setSparseEmitStrategy() immediately after creating the SparseIterator, then the wrapped SparseIterator will have the right strategy, and the FilterIterator strategy will be unintialized; if we call setSparseEmitStrategy() after wrapping the iterator in FilterIterator, then the opposite happens.

If we make setSparseEmitStrategy() a virtual method so that it's included in the FilterIterator pattern, and then do all reads of emitStrategy via a virtual method as well, it's pretty simple to ensure that the value of strategy is being set consistently and correctly.

Without this, the UB of strategy being uninitialized manifests as a sporadic test failure in mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir, when run downstream with the right flags (e.g. asan + assertions off). The test sometimes fails with ne_sub&lt;trivial&lt;dense[0,1]&gt;&gt;.begin' op created with unregistered dialect. It can also be directly observed w/ msan that this uninitialized read is the cause of that issue, but msan causes other problems w/ this test.


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

2 Files Affected:

  • (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp (+13-5)
  • (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h (+5-1)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
index 46d0baac58f06..61b5ad600a16e 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
@@ -504,6 +504,14 @@ class SimpleWrapIterator : public SparseIterator {
                      unsigned extraCursorVal = 0)
       : SparseIterator(kind, *wrap, extraCursorVal), wrap(std::move(wrap)) {}
 
+  void setSparseEmitStrategy(SparseEmitStrategy strategy) override {
+    wrap->setSparseEmitStrategy(strategy);
+  }
+
+  SparseEmitStrategy getSparseEmitStrategy() const override {
+    return wrap->getSparseEmitStrategy();
+  }
+
   SmallVector<Type> getCursorValTypes(OpBuilder &b) const override {
     return wrap->getCursorValTypes(b);
   }
@@ -979,7 +987,7 @@ class SubSectIterator : public SparseIterator {
 
 void SparseIterator::genInit(OpBuilder &b, Location l,
                              const SparseIterator *p) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *begin = b.create(l, b.getStringAttr(prefix + ".begin"), {},
                                 getCursorValTypes(b));
@@ -994,7 +1002,7 @@ void SparseIterator::genInit(OpBuilder &b, Location l,
 }
 
 Value SparseIterator::genNotEnd(OpBuilder &b, Location l) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *notEnd = b.create(l, b.getStringAttr(prefix + ".not_end"),
                                  getCursor(), b.getI1Type());
@@ -1005,7 +1013,7 @@ Value SparseIterator::genNotEnd(OpBuilder &b, Location l) {
 }
 
 void SparseIterator::locate(OpBuilder &b, Location l, Value crd) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     SmallVector<Value> args = getCursor();
     args.push_back(crd);
@@ -1019,7 +1027,7 @@ void SparseIterator::locate(OpBuilder &b, Location l, Value crd) {
 }
 
 Value SparseIterator::deref(OpBuilder &b, Location l) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     SmallVector<Value> args = getCursor();
     Operation *deref = b.create(l, b.getStringAttr(prefix + ".deref"),
@@ -1032,7 +1040,7 @@ Value SparseIterator::deref(OpBuilder &b, Location l) {
 
 ValueRange SparseIterator::forward(OpBuilder &b, Location l) {
   assert(!randomAccessible());
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *next = b.create(l, b.getStringAttr(prefix + ".next"),
                                getCursor(), getCursorValTypes(b));
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h
index 642cb1afa156b..3636f3f01adb5 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h
@@ -177,10 +177,14 @@ class SparseIterator {
 public:
   virtual ~SparseIterator() = default;
 
-  void setSparseEmitStrategy(SparseEmitStrategy strategy) {
+  virtual void setSparseEmitStrategy(SparseEmitStrategy strategy) {
     emitStrategy = strategy;
   }
 
+  virtual SparseEmitStrategy getSparseEmitStrategy() const {
+    return emitStrategy;
+  }
+
   virtual std::string getDebugInterfacePrefix() const = 0;
   virtual SmallVector<Type> getCursorValTypes(OpBuilder &b) const = 0;
 

@llvmbot
Copy link
Member

llvmbot commented Oct 29, 2025

@llvm/pr-subscribers-mlir-sparse

Author: Jordan Rupprecht (rupprecht)

Changes

When we create a SparseIterator, we sometimes wrap it in a FilterIterator, which delegates some calls to the underlying SparseIterator.

After construction, e.g. in makeNonEmptySubSectIterator(), we call setSparseEmitStrategy(). This sets the strategy only in one of the filters -- if we call setSparseEmitStrategy() immediately after creating the SparseIterator, then the wrapped SparseIterator will have the right strategy, and the FilterIterator strategy will be unintialized; if we call setSparseEmitStrategy() after wrapping the iterator in FilterIterator, then the opposite happens.

If we make setSparseEmitStrategy() a virtual method so that it's included in the FilterIterator pattern, and then do all reads of emitStrategy via a virtual method as well, it's pretty simple to ensure that the value of strategy is being set consistently and correctly.

Without this, the UB of strategy being uninitialized manifests as a sporadic test failure in mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir, when run downstream with the right flags (e.g. asan + assertions off). The test sometimes fails with ne_sub&lt;trivial&lt;dense[0,1]&gt;&gt;.begin' op created with unregistered dialect. It can also be directly observed w/ msan that this uninitialized read is the cause of that issue, but msan causes other problems w/ this test.


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

2 Files Affected:

  • (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp (+13-5)
  • (modified) mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h (+5-1)
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
index 46d0baac58f06..61b5ad600a16e 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.cpp
@@ -504,6 +504,14 @@ class SimpleWrapIterator : public SparseIterator {
                      unsigned extraCursorVal = 0)
       : SparseIterator(kind, *wrap, extraCursorVal), wrap(std::move(wrap)) {}
 
+  void setSparseEmitStrategy(SparseEmitStrategy strategy) override {
+    wrap->setSparseEmitStrategy(strategy);
+  }
+
+  SparseEmitStrategy getSparseEmitStrategy() const override {
+    return wrap->getSparseEmitStrategy();
+  }
+
   SmallVector<Type> getCursorValTypes(OpBuilder &b) const override {
     return wrap->getCursorValTypes(b);
   }
@@ -979,7 +987,7 @@ class SubSectIterator : public SparseIterator {
 
 void SparseIterator::genInit(OpBuilder &b, Location l,
                              const SparseIterator *p) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *begin = b.create(l, b.getStringAttr(prefix + ".begin"), {},
                                 getCursorValTypes(b));
@@ -994,7 +1002,7 @@ void SparseIterator::genInit(OpBuilder &b, Location l,
 }
 
 Value SparseIterator::genNotEnd(OpBuilder &b, Location l) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *notEnd = b.create(l, b.getStringAttr(prefix + ".not_end"),
                                  getCursor(), b.getI1Type());
@@ -1005,7 +1013,7 @@ Value SparseIterator::genNotEnd(OpBuilder &b, Location l) {
 }
 
 void SparseIterator::locate(OpBuilder &b, Location l, Value crd) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     SmallVector<Value> args = getCursor();
     args.push_back(crd);
@@ -1019,7 +1027,7 @@ void SparseIterator::locate(OpBuilder &b, Location l, Value crd) {
 }
 
 Value SparseIterator::deref(OpBuilder &b, Location l) {
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     SmallVector<Value> args = getCursor();
     Operation *deref = b.create(l, b.getStringAttr(prefix + ".deref"),
@@ -1032,7 +1040,7 @@ Value SparseIterator::deref(OpBuilder &b, Location l) {
 
 ValueRange SparseIterator::forward(OpBuilder &b, Location l) {
   assert(!randomAccessible());
-  if (emitStrategy == SparseEmitStrategy::kDebugInterface) {
+  if (getSparseEmitStrategy() == SparseEmitStrategy::kDebugInterface) {
     std::string prefix = getDebugInterfacePrefix();
     Operation *next = b.create(l, b.getStringAttr(prefix + ".next"),
                                getCursor(), getCursorValTypes(b));
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h
index 642cb1afa156b..3636f3f01adb5 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/Utils/SparseTensorIterator.h
@@ -177,10 +177,14 @@ class SparseIterator {
 public:
   virtual ~SparseIterator() = default;
 
-  void setSparseEmitStrategy(SparseEmitStrategy strategy) {
+  virtual void setSparseEmitStrategy(SparseEmitStrategy strategy) {
     emitStrategy = strategy;
   }
 
+  virtual SparseEmitStrategy getSparseEmitStrategy() const {
+    return emitStrategy;
+  }
+
   virtual std::string getDebugInterfacePrefix() const = 0;
   virtual SmallVector<Type> getCursorValTypes(OpBuilder &b) const = 0;
 

@aartbik
Copy link
Contributor

aartbik commented Oct 29, 2025

If this fixes an UB, then it is arguably no longer an NFC, right? [I am not 100% sure about our terminology].
In any case, thanks for fixing.

@rupprecht rupprecht changed the title [mlir][sparse][NFC] Include sparse emit strategy in wrapping iterator [mlir][sparse] Include sparse emit strategy in wrapping iterator Oct 29, 2025
@rupprecht
Copy link
Collaborator Author

Thanks!

If this fixes an UB, then it is arguably no longer an NFC, right? [I am not 100% sure about our terminology]. In any case, thanks for fixing.

I've heard a wide range of opinions for what's considered NFC. This patch feels NFC to me because, outside of the UB issue, nothing should be changing. But I can totally see how removing UB is not NFC. I'll err on the side of not claiming this is NFC.

@rupprecht rupprecht merged commit 34c58c8 into llvm:main Oct 29, 2025
13 checks passed
@rupprecht rupprecht deleted the sparse-emit-strategy-msan branch October 29, 2025 20:40
aokblast pushed a commit to aokblast/llvm-project that referenced this pull request Oct 30, 2025
…m#165611)

When we create a `SparseIterator`, we sometimes wrap it in a
`FilterIterator`, which delegates _some_ calls to the underlying
`SparseIterator`.

After construction, e.g. in `makeNonEmptySubSectIterator()`, we call
`setSparseEmitStrategy()`. This sets the strategy only in one of the
filters -- if we call `setSparseEmitStrategy()` immediately after
creating the `SparseIterator`, then the wrapped `SparseIterator` will
have the right strategy, and the `FilterIterator` strategy will be
unintialized; if we call `setSparseEmitStrategy()` after wrapping the
iterator in `FilterIterator`, then the opposite happens.

If we make `setSparseEmitStrategy()` a virtual method so that it's
included in the `FilterIterator` pattern, and then do all reads of
`emitStrategy` via a virtual method as well, it's pretty simple to
ensure that the value of `strategy` is being set consistently and
correctly.

Without this, the UB of strategy being uninitialized manifests as a
sporadic test failure in
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir,
when run downstream with the right flags (e.g. asan + assertions off).
The test sometimes fails with `ne_sub<trivial<dense[0,1]>>.begin' op
created with unregistered dialect`. It can also be directly observed w/
msan that this uninitialized read is the cause of that issue, but msan
causes other problems w/ this test.
DEBADRIBASAK pushed a commit to DEBADRIBASAK/llvm-project that referenced this pull request Nov 3, 2025
…m#165611)

When we create a `SparseIterator`, we sometimes wrap it in a
`FilterIterator`, which delegates _some_ calls to the underlying
`SparseIterator`.

After construction, e.g. in `makeNonEmptySubSectIterator()`, we call
`setSparseEmitStrategy()`. This sets the strategy only in one of the
filters -- if we call `setSparseEmitStrategy()` immediately after
creating the `SparseIterator`, then the wrapped `SparseIterator` will
have the right strategy, and the `FilterIterator` strategy will be
unintialized; if we call `setSparseEmitStrategy()` after wrapping the
iterator in `FilterIterator`, then the opposite happens.

If we make `setSparseEmitStrategy()` a virtual method so that it's
included in the `FilterIterator` pattern, and then do all reads of
`emitStrategy` via a virtual method as well, it's pretty simple to
ensure that the value of `strategy` is being set consistently and
correctly.

Without this, the UB of strategy being uninitialized manifests as a
sporadic test failure in
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_strided_conv_2d_nhwc_hwcf.mlir,
when run downstream with the right flags (e.g. asan + assertions off).
The test sometimes fails with `ne_sub<trivial<dense[0,1]>>.begin' op
created with unregistered dialect`. It can also be directly observed w/
msan that this uninitialized read is the cause of that issue, but msan
causes other problems w/ this test.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

mlir:sparse Sparse compiler in MLIR mlir

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants