[mlir][sparse] Factoring out LoopEmitter::isValidLevel

[wrengr]

Summary

Depends On D146674

[wrengr]

rebase

[wrengr]

arc diff $(git merge-base HEAD origin) --update D146676

[wrengr]

rebase

[wrengr]

rebase

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h

private:
   bool isOutputTensor(TensorId tid) const {
     return hasOutput && tid == static_cast<TensorId>(getNumTensors() - 1);
   }
 
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 
+  bool isValidLevel(TensorId tid, Level lvl) const {
+    return tid < lvlTypes.size() && lvl < lvlTypes[tid].size();
+  }
+
   /// Prepares loop for iterating over `tensor[lvl]`, under the assumption
   /// that `tensor[0...lvl-1]` loops have already been set up.
   void prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                   TensorId tid, Level lvl);
 
   /// Emits extra locals, since the locals might not be in simplified lattices
   /// point used to generate the loops, but are still required to generate
   /// expressions.

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.cpp

Operation *LoopEmitter::enterLoopOverTensorAtLvl(
   assert(!isParallel || reduc.size() <= 1);
   bool isSparseInput = false;
   TensorId tid = tids.front();
   Level dstLvl = lvls.front();
   assert(tids.size() == lvls.size());
   for (auto [t, l] : llvm::zip(tids, lvls)) {
     // TODO: this check for validity of the (t,l) pairs should be
     // checked/enforced at the callsites, if possible.
-    assert(t < lvlTypes.size() && l < lvlTypes[t].size());
+    assert(isValidLevel(t, l));
     assert(!coords[t][l]); // We cannot re-enter the same level
     const auto lvlTp = lvlTypes[t][l];
     const bool isSparse = isCompressedDLT(lvlTp) || isSingletonDLT(lvlTp);
     // Must be a recognizable level-type.
     assert(isSparse || isDenseDLT(lvlTp));
     // We can at most have one sparse input, otherwise, a while loop is required
     // to co-iterate multiple sparse tensors.
     assert(!isSparseInput || !isSparse);
Operation *LoopEmitter::enterLoopOverTensorAtLvl(
   emitExtraLocalsForTensorsAtDenseLvls(builder, loc, tids, lvls);
 
   return loop;
 }
 
 Operation *LoopEmitter::enterFilterLoopOverTensorAtLvl(
     OpBuilder &builder, Location loc, TensorId tid, Level lvl,
     AffineExpr affine, MutableArrayRef<Value> reduc) {
-  assert(tid < lvlTypes.size() && lvl < lvlTypes[tid].size());
+  assert(isValidLevel(tid, lvl));
   assert(!affine.isa<AffineDimExpr>() && !isDenseDLT(lvlTypes[tid][lvl]));
   // We can not re-enter the same level.
   assert(!coords[tid][lvl]);
 
   // TODO: We should instead use a whileOp for filter loop to allow early
   // break when exceeding (for ordered levels).
   // TODO: There are many other potiential opportunities that we might apply in
   // the future. E.g., we could use binary search to locate positions.
Operation *LoopEmitter::enterCoIterationOverTensorsAtLvls(
   for (unsigned i = 0, e = reduc.size(); i < e; i++)
     reduc[i] = after->getArgument(o + i);
 
   return whileOp;
 }
 
 void LoopEmitter::prepareLoopOverTensorAtLvl(OpBuilder &builder, Location loc,
                                              TensorId tid, Level dstLvl) {
-  assert(tid < lvlTypes.size() && dstLvl < lvlTypes[tid].size());
+  assert(isValidLevel(tid, dstLvl));
   const auto lvlTp = lvlTypes[tid][dstLvl];
 
   if (isDenseDLT(lvlTp))
     return;
 
   const Value c0 = constantIndex(builder, loc, 0);
   const Value c1 = constantIndex(builder, loc, 1);
   for (const Level srcLvl : getCollapseReassociation(tid, dstLvl)) {