[mlir][sparse] moving kInvalidId into "detail" namespace

[wrengr]

Summary

In the next few commits I will be converting the various Merger identifier typedefs into newtypes; and once that's done, the kInvalidId constant will only be used internally and therefore does not need to be part of the public mlir::sparse_tensor namespace.

Depends On D146673

[aartbik]

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

1114

are you deliberately enforcing the ::mlir top level instead of relying on the using mlir::sparse_tensor with the shorter version of this?
(I assume yes, but just checking)

[wrengr]

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

1114

yeah it's deliberate, just for paranoia's sake really. I'm pretty sure just using sparse_tensor::detail::kInvalidId would work too. But in the next few CLs this'll be replaced with if (!e.isValid()) so I'm not woo worried about the verbiage in the interim

[wrengr]

rebase

[wrengr]

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

[wrengr]

rebase

[wrengr]

rebase

mlir/include/mlir/Dialect/SparseTensor/Utils/Merger.h

 /// and serve as unique identifiers for the corresponding `LatPoint` object.
 using LatPointId = unsigned;
 
 /// `LatSet` identifiers.  These are allocated by `Merger::addSet` (and
 /// by other methods calling that one), and serve as unique identifiers
 /// for the corresponding `SmallVector<LatPointId>` object.
 using LatSetId = unsigned;
 
+namespace detail {
 /// A constant serving as the canonically invalid identifier, regardless
 /// of the identifier type.
 static constexpr unsigned kInvalidId = -1u;
+} // namespace detail
 
 /// Tensor expression. Represents an MLIR expression in tensor index notation.
 struct TensorExp final {
   enum class Kind;
 
   /// Child subexpressions for non-leaf expressions.
   struct Children final {
     ExprId e0;
public:
   // gave the number of input tensors, instead of the current number of
   // input+output tensors.
   Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
          unsigned numFilterLoops, unsigned maxLvlRank);
 
   /// Constructs a new tensor expression, and returns its identifier.
   /// The type of the `e0` argument varies according to the value of the
   /// `k` argument, as described by the `TensorExp` ctor.
-  ExprId addExp(TensorExp::Kind k, unsigned e0, ExprId e1 = kInvalidId,
+  ExprId addExp(TensorExp::Kind k, unsigned e0, ExprId e1 = detail::kInvalidId,
                 Value v = Value(), Operation *op = nullptr);
   ExprId addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op = nullptr) {
-    return addExp(k, e, kInvalidId, v, op);
+    return addExp(k, e, detail::kInvalidId, v, op);
   }
   ExprId addExp(TensorExp::Kind k, Value v, Operation *op = nullptr) {
-    return addExp(k, kInvalidId, kInvalidId, v, op);
+    return addExp(k, detail::kInvalidId, detail::kInvalidId, v, op);
   }
 
   /// Constructs a new iteration lattice point, and returns its identifier.
   LatPointId addLat(TensorId t, LoopId i, ExprId e);
   LatPointId addLat(const BitVector &bits, ExprId e);
 
   /// Constructs a new (initially empty) set, and returns its identifier.
   LatSetId addSet();

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

public:
   Value getExpandCount() const { return expCount; }
   void endExpand();
 
   //
   // Reduction methods.
   //
 
   void startReduc(ExprId exp, Value val);
-  bool isReduc() const { return redExp != kInvalidId; }
+  bool isReduc() const { return redExp != detail::kInvalidId; }
   void updateReduc(Value val);
   Value getReduc() const { return redVal; }
   Value endReduc();
   void setValidLexInsert(Value val);
   void clearValidLexInsert();
   Value getValidLexInsert() const { return redValidLexInsert; }
 
   void startCustomReduc(ExprId exp);
-  bool isCustomReduc() const { return redCustom != kInvalidId; }
+  bool isCustomReduc() const { return redCustom != detail::kInvalidId; }
   Value getCustomRedId();
   void endCustomReduc();
 
 private:
   // Linalg operation.
   linalg::GenericOp linalgOp;
 
   // Sparsification options.

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

 
 CodegenEnv::CodegenEnv(linalg::GenericOp linop, SparsificationOptions opts,
                        unsigned numTensors, unsigned numLoops,
                        unsigned numFilterLoops, unsigned maxRank)
     : linalgOp(linop), sparseOptions(opts),
       latticeMerger(numTensors, numLoops, numFilterLoops, maxRank),
       loopEmitter(), topSort(), sparseOut(nullptr), outerParNest(-1u),
       insChain(), expValues(), expFilled(), expAdded(), expCount(), redVal(),
-      redExp(kInvalidId), redCustom(kInvalidId), redValidLexInsert() {}
+      redExp(detail::kInvalidId), redCustom(detail::kInvalidId),
+      redValidLexInsert() {}
 
 LogicalResult CodegenEnv::initTensorExp() {
   // Builds the tensor expression for the Linalg operation in SSA form.
   std::optional<ExprId> optExp = latticeMerger.buildTensorExpFromLinalg(op());
   if (!optExp || !isAdmissibleTensorExp(*optExp))
     return failure();
 
   tensorExp = *optExp;
void CodegenEnv::endExpand() {
   expValues = expFilled = expAdded = expCount = Value();
 }
 
 //===----------------------------------------------------------------------===//
 // Code generation environment reduction methods
 //===----------------------------------------------------------------------===//
 
 void CodegenEnv::startReduc(ExprId exp, Value val) {
-  assert(!isReduc() && exp != kInvalidId);
+  assert(!isReduc() && exp != detail::kInvalidId);
   redExp = exp;
   updateReduc(val);
 }
 
 void CodegenEnv::updateReduc(Value val) {
   assert(isReduc());
   redVal = val;
   // NOTE: `genLoopBoundary` requires that this performs a unilateral
   // update without checking for a previous value first.  (It's not
   // clear whether any other callsites also require that.)
   latticeMerger.updateExprValue(redExp, val);
 }
 
 Value CodegenEnv::endReduc() {
   assert(isReduc());
   Value val = redVal;
   redVal = val;
   latticeMerger.clearExprValue(redExp);
-  redExp = kInvalidId;
+  redExp = detail::kInvalidId;
   return val;
 }
 
 void CodegenEnv::setValidLexInsert(Value val) {
   assert(isReduc() && val);
   redValidLexInsert = val;
 }
 
 void CodegenEnv::clearValidLexInsert() {
   assert(!isReduc());
   redValidLexInsert = Value();
 }
 
 void CodegenEnv::startCustomReduc(ExprId exp) {
-  assert(!isCustomReduc() && exp != kInvalidId);
+  assert(!isCustomReduc() && exp != detail::kInvalidId);
   redCustom = exp;
 }
 
 Value CodegenEnv::getCustomRedId() {
   assert(isCustomReduc());
   return dyn_cast<sparse_tensor::ReduceOp>(exp(redCustom).op).getIdentity();
 }
 
 void CodegenEnv::endCustomReduc() {
   assert(isCustomReduc());
-  redCustom = kInvalidId;
+  redCustom = detail::kInvalidId;
 }

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

 }
 
 /// Recursively generates tensor expression.
 static Value genExp(CodegenEnv &env, RewriterBase &rewriter, ExprId e,
                     LoopId ldx) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
 
-  if (e == kInvalidId)
+  if (e == ::mlir::sparse_tensor::detail::kInvalidId)
     return Value();
   const TensorExp &exp = env.exp(e);
   const auto kind = exp.kind;
   if (kind == TensorExp::Kind::kTensor)
     return genTensorLoad(env, rewriter, e);
   if (kind == TensorExp::Kind::kInvariant)
     return genInvariantValue(env, e);
   if (kind == TensorExp::Kind::kLoopVar)
if (kind == TensorExp::Kind::kSelect)
     env.merger().setExprValue(e, v0); // Preserve value for later use.
 
   return ee;
 }
 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
 static void genInvariants(CodegenEnv &env, OpBuilder &builder, ExprId exp,
                           LoopId ldx, bool atStart) {
-  if (exp == kInvalidId)
+  if (exp == ::mlir::sparse_tensor::detail::kInvalidId)
     return;
   if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
     // Inspect tensor indices.
-    bool isAtLoop = ldx == kInvalidId;
+    bool isAtLoop = ldx == ::mlir::sparse_tensor::detail::kInvalidId;
     linalg::GenericOp op = env.op();
     OpOperand &t = op->getOpOperand(env.exp(exp).tensor);
     auto map = op.getMatchingIndexingMap(&t);
     const auto stt = getSparseTensorType(t.get());
     const Level lvlRank = stt.getLvlRank();
     assert(static_cast<Level>(map.getNumResults()) == lvlRank);
     for (Level l = 0; l < lvlRank; l++) {
       // FIXME: `toOrigDim` is deprecated.
if (at == env.topSortSize()) {
     const LoopId ldx = env.topSortAt(at - 1);
     Value rhs = genExp(env, rewriter, exp, ldx);
     genTensorStore(env, rewriter, exp, rhs);
     return;
   }
 
   // Construct iteration lattices for current loop index, with L0 at top.
   const LoopId idx = env.topSortAt(at);
-  const LoopId ldx = at == 0 ? kInvalidId : env.topSortAt(at - 1);
+  const LoopId ldx = at == 0 ? ::mlir::sparse_tensor::detail::kInvalidId
+                             : env.topSortAt(at - 1);
   const LatSetId lts =
       env.merger().optimizeSet(env.merger().buildLattices(exp, idx));
 
   // Start a loop sequence.
   bool needsUniv = startLoopSeq(env, rewriter, exp, at, idx, ldx, lts);
 
   // Emit a loop for every lattice point L0 >= Li in this loop sequence.
   //

mlir/lib/Dialect/SparseTensor/Utils/Merger.cpp

 // Constructors.
 //===----------------------------------------------------------------------===//
 
 TensorExp::TensorExp(Kind k, unsigned x, ExprId y, Value v, Operation *o)
     : kind(k), val(v), op(o) {
   switch (kind) {
   // Leaf.
   case TensorExp::Kind::kTensor:
-    assert(x != kInvalidId && y == kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     tensor = x;
     return;
   case TensorExp::Kind::kInvariant:
-    assert(x == kInvalidId && y == kInvalidId && v && !o);
+    assert(x == detail::kInvalidId && y == detail::kInvalidId && v && !o);
     return;
   case TensorExp::Kind::kLoopVar:
-    assert(x != kInvalidId && y == kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     loop = x;
     return;
   // Unary operations.
   case TensorExp::Kind::kAbsF:
   case TensorExp::Kind::kAbsC:
   case TensorExp::Kind::kAbsI:
   case TensorExp::Kind::kCeilF:
   case TensorExp::Kind::kFloorF:
   case TensorExp::Kind::kSqrtF:
   case TensorExp::Kind::kSqrtC:
   case TensorExp::Kind::kExpm1F:
   case TensorExp::Kind::kExpm1C:
   case TensorExp::Kind::kLog1pF:
   case TensorExp::Kind::kLog1pC:
   case TensorExp::Kind::kSinF:
   case TensorExp::Kind::kSinC:
   case TensorExp::Kind::kTanhF:
   case TensorExp::Kind::kTanhC:
   case TensorExp::Kind::kNegF:
   case TensorExp::Kind::kNegC:
   case TensorExp::Kind::kNegI:
   case TensorExp::Kind::kCIm:
   case TensorExp::Kind::kCRe:
-    assert(x != kInvalidId && y == kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
     children.e0 = x;
     children.e1 = y;
     return;
   case TensorExp::Kind::kTruncF:
   case TensorExp::Kind::kExtF:
   case TensorExp::Kind::kCastFS:
   case TensorExp::Kind::kCastFU:
   case TensorExp::Kind::kCastSF:
   case TensorExp::Kind::kCastUF:
   case TensorExp::Kind::kCastS:
   case TensorExp::Kind::kCastU:
   case TensorExp::Kind::kCastIdx:
   case TensorExp::Kind::kTruncI:
   case TensorExp::Kind::kBitCast:
-    assert(x != kInvalidId && y == kInvalidId && v && !o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && v && !o);
     children.e0 = x;
     children.e1 = y;
     return;
   case TensorExp::Kind::kBinaryBranch:
   case TensorExp::Kind::kSelect:
-    assert(x != kInvalidId && y == kInvalidId && !v && o);
+    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
     return;
   case TensorExp::Kind::kUnary:
     // No assertion on y can be made, as the branching paths involve both
     // a unary (`mapSet`) and binary (`disjSet`) pathway.
-    assert(x != kInvalidId && !v && o);
+    assert(x != detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
     return;
   // Binary operations.
   case TensorExp::Kind::kMulF:
   case TensorExp::Kind::kMulC:
   case TensorExp::Kind::kMulI:
   case TensorExp::Kind::kDivF:
   case TensorExp::Kind::kDivC:
   case TensorExp::Kind::kDivS:
   case TensorExp::Kind::kDivU:
   case TensorExp::Kind::kAddF:
   case TensorExp::Kind::kAddC:
   case TensorExp::Kind::kAddI:
   case TensorExp::Kind::kSubF:
   case TensorExp::Kind::kSubC:
   case TensorExp::Kind::kSubI:
   case TensorExp::Kind::kAndI:
   case TensorExp::Kind::kOrI:
   case TensorExp::Kind::kXorI:
   case TensorExp::Kind::kShrS:
   case TensorExp::Kind::kShrU:
   case TensorExp::Kind::kShlI:
-    assert(x != kInvalidId && y != kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && y != detail::kInvalidId && !v && !o);
     children.e0 = x;
     children.e1 = y;
     return;
   case TensorExp::Kind::kBinary:
   case TensorExp::Kind::kReduce:
-    assert(x != kInvalidId && y != kInvalidId && !v && o);
+    assert(x != detail::kInvalidId && y != detail::kInvalidId && !v && o);
     children.e0 = x;
     children.e1 = y;
     return;
   }
   llvm_unreachable("unexpected kind");
 }
 
 LatPoint::LatPoint(const BitVector &bits, ExprId e) : bits(bits), exp(e) {}