[mlir][sparse] Marking off todos and updating commentary

[wrengr]

Summary

Depends On D155999

[wrengr]

Moved null-expr handling from Ranks::isValid(DimLvlExpr) to DimSpec::isValid(Ranks), and updated comments/documentation accordingly.

[wrengr]

Removing todo regarding ASan, since running things under ASan doesn't find any problems

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.h

public:
   void setExpr(DimExpr newExpr) {
     assert(!hasExpr());
     expr = newExpr;
   }
   constexpr bool canElideExpr() const { return elideExpr; }
   void setElideExpr(bool b) { elideExpr = b; }
   constexpr SparseTensorDimSliceAttr getSlice() const { return slice; }
 
-  /// Checks whether the variables bound/used by this spec are valid
-  /// with respect to the given ranks.
+  /// Checks whether the variables bound/used by this spec are valid with
+  /// respect to the given ranks.  Note that null `DimExpr` is considered
+  /// to be vacuously valid, and therefore calling `setExpr` invalidates
+  /// the result of this predicate.
   bool isValid(Ranks const &ranks) const;
 
   // TODO(wrengr): Use it or loose it.
   bool isFunctionOf(Var var) const;
   bool isFunctionOf(VarSet const &vars) const;
   void getFreeVars(VarSet &vars) const;
 
   void print(llvm::raw_ostream &os, bool wantElision = true) const;

mlir/lib/Dialect/SparseTensor/IR/Detail/DimLvlMap.cpp

void DimLvlExpr::printAffineExprInternal(
   } else if (op != AffineExprKind::Add) {
     // Default rule for tightly binding binary operators.
     // (Including `Mul` that didn't match the previous rule.)
     lhs.printStrong(os);
     os << binopSpelling;
     rhs.printStrong(os);
   } else {
     // Combination of all the special rules for addition/subtraction.
-    // TODO(wrengr): despite being succinct, this is prolly too confusing for
-    // readers.
     lhs.printWeak(os);
     const auto rx = matchNeg(rhs);
     os << (rx ? " - " : " + ");
     const auto &rlhs = rx ? rx->getLHS() : rhs;
     const auto rrhs = rx ? rx->getRHS() : -1; // value irrelevant when `!rx`
     const bool nonunit = rrhs != -1;          // value irrelevant when `!rx`
     const bool isStrong =
         rx && rlhs && (nonunit || rlhs.getAffineKind() == AffineExprKind::Add);

 //===----------------------------------------------------------------------===//
 // `DimSpec` implementation.
 //===----------------------------------------------------------------------===//
 
 DimSpec::DimSpec(DimVar var, DimExpr expr, SparseTensorDimSliceAttr slice)
     : var(var), expr(expr), slice(slice) {}
 
 bool DimSpec::isValid(Ranks const &ranks) const {
-  return ranks.isValid(var) && ranks.isValid(expr);
-  // TODO(wrengr): is there anything in `slice` that needs validation?
+  // Nothing in `slice` needs additional validation.
+  // We explicitly consider null-expr to be vacuously valid.
+  return ranks.isValid(var) && (!expr || ranks.isValid(expr));
 }
 
 bool DimSpec::isFunctionOf(VarSet const &vars) const {
   return vars.occursIn(expr);
 }
 
 void DimSpec::getFreeVars(VarSet &vars) const { vars.add(expr); }
 

 
 LvlSpec::LvlSpec(LvlVar var, LvlExpr expr, DimLevelType type)
     : var(var), expr(expr), type(type) {
   assert(expr);
   assert(isValidDLT(type) && !isUndefDLT(type));
 }
 
 bool LvlSpec::isValid(Ranks const &ranks) const {
+  // Nothing in `type` needs additional validation.
   return ranks.isValid(var) && ranks.isValid(expr);
-  // TODO(wrengr): is there anything in `type` that needs validation?
 }
 
 bool LvlSpec::isFunctionOf(VarSet const &vars) const {
   return vars.occursIn(expr);
 }
 
 void LvlSpec::getFreeVars(VarSet &vars) const { vars.add(expr); }
 

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.h

 
 protected:
   /// Protected ctor for the RTTI methods to use.
   constexpr explicit Var(Impl impl) : impl(impl) {}
 
 public:
   constexpr Var(VarKind vk, Num n) : impl(Impl(vk, n)) {}
   Var(AffineSymbolExpr sym) : Var(VarKind::Symbol, sym.getPosition()) {}
-  Var(VarKind vk, AffineDimExpr var) : Var(vk, var.getPosition()) {}
+  // TODO(wrengr): Should make the first argument an `ExprKind` instead...?
+  Var(VarKind vk, AffineDimExpr var) : Var(vk, var.getPosition()) {
+    assert(vk != VarKind::Symbol);
+  }
 
   constexpr bool operator==(Var other) const { return impl == other.impl; }
   constexpr bool operator!=(Var other) const { return !(*this == other); }
 
   constexpr VarKind getKind() const { return impl.getKind(); }
   constexpr Num getNum() const { return impl.getNum(); }
 
   template <typename U>

 /// associated variable.
 class VarInfo final {
 public:
   /// Newtype for unique identifiers of `VarInfo` records, to ensure
   /// they aren't confused with `Var::Num`.
   enum class ID : unsigned {};
 
 private:
-  // FUTURE_CL(wrengr): We could use the high-bit of `Var::Impl` to
-  // store the `std::optional` bit, therefore allowing us to bitbash the
-  // `num` and `kind` fields together.
-  //
   StringRef name;              // The bare-id used in the MLIR source.
   llvm::SMLoc loc;             // The location of the first occurence.
-                               // TODO(wrengr): See the above `LocatedVar` note.
   ID id;                       // The unique `VarInfo`-identifier.
   std::optional<Var::Num> num; // The unique `Var`-identifier (if resolved).
   VarKind kind;                // The kind of variable.
 
 public:
   constexpr VarInfo(ID id, StringRef name, llvm::SMLoc loc, VarKind vk,
                     std::optional<Var::Num> n = {})
       : name(name), loc(loc), id(id), num(n), kind(vk) {

mlir/lib/Dialect/SparseTensor/IR/Detail/Var.cpp

void Var::dump() const {
   llvm::errs() << "\n";
 }
 
 //===----------------------------------------------------------------------===//
 // `Ranks` implementation.
 //===----------------------------------------------------------------------===//
 
 bool Ranks::isValid(DimLvlExpr expr) const {
-  // FIXME(wrengr): we have cases without affine expr at an early point
-  if (!expr.getAffineExpr())
-    return true;
-  // Each `DimLvlExpr` only allows one kind of non-symbol variable.
+  assert(expr);
+  // Compute the maximum identifiers for symbol-vars and dim/lvl-vars
+  // (each `DimLvlExpr` only allows one kind of non-symbol variable).
   int64_t maxSym = -1, maxVar = -1;
-  // TODO(wrengr): If we run into ASan issues, that may be due to the
-  // "`{{...}}`" syntax; so we may want to try using local-variables instead.
   mlir::getMaxDimAndSymbol<ArrayRef<AffineExpr>>({{expr.getAffineExpr()}},
                                                  maxVar, maxSym);
   // TODO(wrengr): We may want to add a call to `LLVM_DEBUG` like
   // `willBeValidAffineMap` does.  And/or should return `InFlightDiagnostic`
   // instead of bool.
   return maxSym < getSymRank() && maxVar < getRank(expr.getAllowedVarKind());
 }