[mlir][sparse] Introducing newtypes for ExprId, LatPointId, LatSetId

[wrengr]

Summary

This commit performs the actual change from using typedefs to using newtypes. The ExprId, LatPointId, LatSetId types are especially simple, because: (1) the Merger has compete control over allocating them, and (2) client code only ever needs the equality operator (if it even needs that).

Depends On D146684

[wrengr]

rebase

[wrengr]

Addressing the -Wmismatched-tags warnings on Debian.

[Peiming]

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

590

Why not supporting an implicit type conversion from the new types to unsigned (but not vice versa) so that they can be used directly for vector/array indices?

[wrengr]

Postponing the migration of TensorId, LoopId, TensorLoopId into MergerNewtypes.h until D146691 (and D146693)

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

 //===- Merger.h - Utilities for defining lattices ---------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This header file defines utilities for dealing with iteration lattices.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_
 #define MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_
 
+#include "mlir/Dialect/SparseTensor/Utils/MergerNewtypes.h"
+
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/BitVector.h"
 #include <optional>
 
 namespace mlir {
 namespace sparse_tensor {
 
 // TODO: These type aliases currently only serve to make the code more
 // self-documenting, however because they are not type-checked they can
-// do nothing to prevent mixups.  We should really change them from mere
-// aliases to actual struct definitions, so that we can type-check them.
+// do nothing to prevent mixups.  They should be made into newtypes and
+// moved into MergerNewtypes.h with the other newtypes.
 
 /// Tensor identifiers.  The valid set of identifiers is defined by the
 /// first argument passed to the `Merger` ctor.
 using TensorId = unsigned;
 
 /// Loop identifiers.  The valid set of identifiers is defined by the
 /// second two arguments to the `Merger` ctor.
 ///

 using LoopId = unsigned;
 
 /// A compressed representation of `std::pair<TensorId, LoopId>`.
 /// The compression scheme is such that this also serves as an index
 /// into the bitvector stored in `LatPoint` (since that bitvector is
 /// just the implementation for a set of `TensorLoopId` values).
 using TensorLoopId = unsigned;
 
-/// `TensorExp` identifiers.  These are allocated by `Merger::addExp`,
-/// and serve as unique identifiers for the corresponding `TensorExp` object.
-using ExprId = unsigned;
-
-/// `LatPoint` identifiers.  These are allocated by `Merger::addLat`,
-/// 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;
     ExprId e1;
public:
 
   /// Constructs a new tensor expression, and returns its identifier.
   ExprId addTensorExp(TensorId t);
   /// Constructs a new loop-variable expression, and returns its identifier.
   ExprId addLoopVarExp(LoopId i);
   /// Constructs a new invariant expression, and returns its identifier.
   ExprId addInvariantExp(Value v);
   /// Constructs a new unary or binary expression, and returns its identifier.
-  ExprId addExp(TensorExp::Kind k, ExprId e0, ExprId e1 = detail::kInvalidId,
+  ExprId addExp(TensorExp::Kind k, ExprId e0, ExprId e1 = ExprId(),
                 Operation *op = nullptr);
   /// Constructs a new sesquinary expression, and returns its identifier.
   /// Currently no sesquinary `Kind` allows specifying the `op`, but we
   /// allow it anyways because `mapSet` is designed to allow it.
   ExprId addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op = nullptr);
 
   /// Constructs a new iteration lattice point, and returns its identifier.
   LatPointId addLat(TensorId t, LoopId i, ExprId e);
public:
   /// time (e.g., as in `TensorExpr &te = merger.exp(e)`), since insertions
   /// into the merger can cause data movement which will invalidate the
   /// underlying memory address.  This isn't just a problem with the `&`
   /// references, but also applies to the `ArrayRef`.  In particular,
   /// using `for (LatPointId p : merger.set(s))` will run into the same
   /// dangling-reference problems if the loop body inserts new sets.
   const TensorExp &exp(ExprId e) const {
     assert(isValidExprId(e));
-    return tensorExps[e];
+    return tensorExps[e.value];
   }
   const LatPoint &lat(LatPointId p) const {
     assert(isValidLatPointId(p));
-    return latPoints[p];
+    return latPoints[p.value];
   }
   ArrayRef<LatPointId> set(LatSetId s) const {
     assert(isValidLatSetId(s));
-    return latSets[s];
+    return latSets[s.value];
   }
 
   /// Checks whether the given expression has an associated value.
   bool hasExprValue(ExprId e) const { return static_cast<bool>(exp(e).val); }
 
   /// Sets the expression to have the associated value.  Asserts that
   /// the new value is defined, and that the expression does not already
   /// have a value.  If you want to overwrite a previous associated value,
   /// use `updateExprValue` instead.
   void setExprValue(ExprId e, Value v) {
     assert(isValidExprId(e));
     assert(v && "Got an undefined value");
-    auto &val = tensorExps[e].val;
+    auto &val = tensorExps[e.value].val;
     assert(!val && "Expression already has an associated value");
     val = v;
   }
 
   /// Clears the value associated with the expression.  Asserts that the
   /// expression does indeed have an associated value before clearing it.
   /// If you don't want to check for a previous associated value first,
   /// then use `updateExprValue` instead.
   void clearExprValue(ExprId e) {
     assert(isValidExprId(e));
-    auto &val = tensorExps[e].val;
+    auto &val = tensorExps[e.value].val;
     assert(val && "Expression does not have an associated value to clear");
     val = Value();
   }
 
   /// Unilaterally updates the expression to have the associated value.
   /// That is, unlike `setExprValue` and `clearExprValue`, this method
   /// does not perform any checks on whether the expression had a
   /// previously associated value nor whether the new value is defined.
   //
   // TODO: The unilateral update semantics are required by the
   // current implementation of `CodegenEnv::genLoopBoundary`; however,
   // that implementation seems a bit dubious.  We would much rather have
   // the semantics `{ clearExprValue(e); setExprValue(e, v); }` or
   // `{ clearExprValue(e); if (v) setExprValue(e, v); }` since those
   // provide better invariants.
   void updateExprValue(ExprId e, Value v) {
     assert(isValidExprId(e));
-    tensorExps[e].val = v;
+    tensorExps[e.value].val = v;
   }
 
 #ifndef NDEBUG
   /// Print methods (for debugging).
   void dumpExp(ExprId e) const;
   void dumpLat(LatPointId p) const;
   void dumpSet(LatSetId s) const;
   void dumpBits(const BitVector &bits) const;
private:
   constexpr bool isValidTensorId(TensorId t) const { return t < numTensors; }
   constexpr bool isValidLoopId(LoopId i) const {
     return i != detail::kInvalidId && i < numLoops;
   }
   bool isValidLevel(TensorId t, Level lvl) const {
     return isValidTensorId(t) && lvl < lvlToLoop[t].size();
   }
   bool isValidExprId(ExprId e) const {
-    return e != detail::kInvalidId && e < tensorExps.size();
+    return e.isValid() && e.value < tensorExps.size();
   }
   bool isValidLatPointId(LatPointId p) const {
-    return p != detail::kInvalidId && p < latPoints.size();
+    return p.isValid() && p.value < latPoints.size();
   }
   bool isValidLatSetId(LatSetId s) const {
-    return s != detail::kInvalidId && s < latSets.size();
+    return s.isValid() && s.value < latSets.size();
   }
   bool maybeZero(ExprId e) const;
   bool isInvariant(ExprId e) const {
     return exp(e).kind == TensorExp::Kind::kInvariant;
   }
   Type inferType(ExprId e, Value src) const;
 
   /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.

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

+//===- MergerNewtypes.h - Newtypes for the `Merger` class -------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This header defines a number of "newtypes" (i.e., data types which are
+// zero-cost abstractions for wrapping some underlying type while ensuring
+// that the compiler keeps the new type distinct from the old type),
+// along with related classes for iterating over them, etc.
+//
+// To guarantee that we achieve the goal of being zero-cost abstractions,
+// we assert `std::is_trivially_copyable` and `std::is_trivially_destructible`
+// of all these newtypes.  These predicates license the compiler to make
+// several optimizations; some of which are explicitly documented by the
+// C++ standard:
+// <https://en.cppreference.com/w/cpp/types/is_trivially_copyable#Notes>
+// <https://en.cppreference.com/w/cpp/types/is_trivially_destructible#Notes>
+// However, some key optimizations aren't mentioned by the standard; e.g.,
+// that trivially-copyable enables passing-by-value, and the conjunction
+// of trivially-copyable and trivially-destructible enables passing those
+// values in registers rather than on the stack (cf.,
+// <https://www.agner.org/optimize/calling_conventions.pdf>).
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_UTILS_MERGERNEWTYPES_H_
+#define MLIR_DIALECT_SPARSETENSOR_UTILS_MERGERNEWTYPES_H_
+
+#include <cassert>
+#include <type_traits>
+
+namespace mlir {
+namespace sparse_tensor {
+
+namespace detail {
+/// A constant serving as the canonically invalid identifier,
+/// regardless of the identifier type.
+static constexpr unsigned kInvalidId = -1u;
+} // namespace detail
+
+//===----------------------------------------------------------------------===//
+/// `TensorExp` identifiers.  These are allocated by `Merger::addExp`,
+/// and serve as unique identifiers for the corresponding `TensorExp` object.
+class ExprId final {
+  friend class Merger;
+  friend struct TensorExp; // Must be "struct" for [-Wmismatched-tags].
+  explicit constexpr ExprId(unsigned value) : value(value) {}
+
+public:
+  /// Constructs a new expression identifier with a dedicated
+  /// known-invalid value.
+  explicit constexpr ExprId() : value(detail::kInvalidId) {}
+  /// Checks whether the expression identifier has the dedicated
+  /// known-invalid value.
+  constexpr bool isValid() const { return value != detail::kInvalidId; }
+  constexpr bool operator==(ExprId rhs) const { return value == rhs.value; }
+  constexpr bool operator!=(ExprId rhs) const { return value != rhs.value; }
+
+private:
+  unsigned value;
+};
+static_assert(std::is_trivially_copyable_v<ExprId> &&
+              std::is_trivially_destructible_v<ExprId>);
+
+//===----------------------------------------------------------------------===//
+/// `LatPoint` identifiers.  These are allocated by `Merger::addLat`,
+/// and serve as unique identifiers for the corresponding `LatPoint` object.
+class LatPointId final {
+  friend class Merger;
+  explicit constexpr LatPointId(unsigned value) : value(value) {}
+
+public:
+  /// Constructs a new lattice-point identifier with a dedicated
+  /// known-invalid value.
+  explicit constexpr LatPointId() : value(detail::kInvalidId) {}
+  /// Checks whether the lattice-point identifier has the dedicated
+  /// known-invalid value.
+  constexpr bool isValid() const { return value != detail::kInvalidId; }
+  constexpr bool operator==(LatPointId rhs) const { return value == rhs.value; }
+  constexpr bool operator!=(LatPointId rhs) const { return value != rhs.value; }
+
+private:
+  unsigned value;
+};
+static_assert(std::is_trivially_copyable_v<LatPointId> &&
+              std::is_trivially_destructible_v<LatPointId>);
+
+//===----------------------------------------------------------------------===//
+/// `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.
+class LatSetId final {
+  friend class Merger;
+  explicit constexpr LatSetId(unsigned value) : value(value) {}
+
+public:
+  /// Constructs a new lattice-set identifier with a dedicated
+  /// known-invalid value.
+  explicit constexpr LatSetId() : value(detail::kInvalidId) {}
+  /// Checks whether the lattice-set identifier has the dedicated
+  /// known-invalid value.
+  constexpr bool isValid() const { return value != detail::kInvalidId; }
+  constexpr bool operator==(LatSetId rhs) const { return value == rhs.value; }
+  constexpr bool operator!=(LatSetId rhs) const { return value != rhs.value; }
+
+private:
+  unsigned value;
+};
+static_assert(std::is_trivially_copyable_v<LatSetId> &&
+              std::is_trivially_destructible_v<LatSetId>);
+
+} // namespace sparse_tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGERNEWTYPES_H_

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 != detail::kInvalidId; }
+  bool isReduc() const { return redExp.isValid(); }
   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 != detail::kInvalidId; }
+  bool isCustomReduc() const { return redCustom.isValid(); }
   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(detail::kInvalidId), redCustom(detail::kInvalidId),
-      redValidLexInsert() {}
+      redExp(), redCustom(), 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 != detail::kInvalidId);
+  assert(!isReduc() && exp.isValid());
   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 = detail::kInvalidId;
+  redExp = ExprId();
   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 != detail::kInvalidId);
+  assert(!isCustomReduc() && exp.isValid());
   redCustom = exp;
 }
 
 Value CodegenEnv::getCustomRedId() {
   assert(isCustomReduc());
   return dyn_cast<sparse_tensor::ReduceOp>(exp(redCustom).op).getIdentity();
 }
 
 void CodegenEnv::endCustomReduc() {
   assert(isCustomReduc());
-  redCustom = detail::kInvalidId;
+  redCustom = ExprId();
 }

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 == ::mlir::sparse_tensor::detail::kInvalidId)
+  if (!e.isValid())
     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 == ::mlir::sparse_tensor::detail::kInvalidId)
+  if (!exp.isValid())
     return;
   if (env.exp(exp).kind == TensorExp::Kind::kTensor) {
     // Inspect tensor indices.
     bool isAtLoop = ldx == ::mlir::sparse_tensor::detail::kInvalidId;
     linalg::GenericOp op = env.op();
     OpOperand &t = op->getOpOperand(env.exp(exp).tensor);
     const auto map = op.getMatchingIndexingMap(&t);
     const auto stt = getSparseTensorType(t.get());

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

 //===----------------------------------------------------------------------===//
 
 TensorExp::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 != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && !y.isValid() && !v && !o);
     tensor = x;
     return;
   case TensorExp::Kind::kInvariant:
-    assert(x == detail::kInvalidId && y == detail::kInvalidId && v && !o);
+    assert(x == detail::kInvalidId && !y.isValid() && v && !o);
     return;
   case TensorExp::Kind::kLoopVar:
-    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
+    assert(x != detail::kInvalidId && !y.isValid() && !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 != detail::kInvalidId && y == detail::kInvalidId && !v && !o);
-    children.e0 = x;
+    assert(x != detail::kInvalidId && !y.isValid() && !v && !o);
+    children.e0 = ExprId{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 != detail::kInvalidId && y == detail::kInvalidId && v && !o);
-    children.e0 = x;
+    assert(x != detail::kInvalidId && !y.isValid() && v && !o);
+    children.e0 = ExprId{x};
     children.e1 = y;
     return;
   case TensorExp::Kind::kBinaryBranch:
   case TensorExp::Kind::kSelect:
-    assert(x != detail::kInvalidId && y == detail::kInvalidId && !v && o);
-    children.e0 = x;
+    assert(x != detail::kInvalidId && !y.isValid() && !v && o);
+    children.e0 = ExprId{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 != detail::kInvalidId && !v && o);
-    children.e0 = x;
+    children.e0 = ExprId{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 != detail::kInvalidId && y != detail::kInvalidId && !v && !o);
-    children.e0 = x;
+    assert(x != detail::kInvalidId && y.isValid() && !v && !o);
+    children.e0 = ExprId{x};
     children.e1 = y;
     return;
   case TensorExp::Kind::kBinary:
   case TensorExp::Kind::kReduce:
-    assert(x != detail::kInvalidId && y != detail::kInvalidId && !v && o);
-    children.e0 = x;
+    assert(x != detail::kInvalidId && y.isValid() && !v && o);
+    children.e0 = ExprId{x};
     children.e1 = y;
     return;
   }
   llvm_unreachable("unexpected kind");
 }
 
 Merger::Merger(unsigned numInputOutputTensors, unsigned numNativeLoops,
                unsigned numFilterLoops, unsigned maxLvlRank)

 
 //===----------------------------------------------------------------------===//
 // Lattice methods.
 //===----------------------------------------------------------------------===//
 
 ExprId Merger::addTensorExp(TensorId t) {
   assert(isValidTensorId(t));
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(TensorExp::Kind::kTensor, t, detail::kInvalidId,
-                          Value(), nullptr);
+  tensorExps.emplace_back(TensorExp::Kind::kTensor, t, ExprId(), Value(),
+                          nullptr);
   return eNew;
 }
 
 ExprId Merger::addLoopVarExp(LoopId i) {
   assert(isValidLoopId(i));
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(TensorExp::Kind::kLoopVar, i, detail::kInvalidId,
-                          Value(), nullptr);
+  tensorExps.emplace_back(TensorExp::Kind::kLoopVar, i, ExprId(), Value(),
+                          nullptr);
   return eNew;
 }
 
 ExprId Merger::addInvariantExp(Value v) {
   const ExprId eNew(tensorExps.size());
   tensorExps.emplace_back(TensorExp::Kind::kInvariant, detail::kInvalidId,
-                          detail::kInvalidId, v, nullptr);
+                          ExprId(), v, nullptr);
   return eNew;
 }
 
 ExprId Merger::addExp(TensorExp::Kind k, ExprId e0, ExprId e1, Operation *op) {
   assert(k > TensorExp::Kind::kLoopVar);
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(k, e0, e1, Value(), op);
+  tensorExps.emplace_back(k, e0.value, e1, Value(), op);
   return eNew;
 }
 
 ExprId Merger::addExp(TensorExp::Kind k, ExprId e, Value v, Operation *op) {
   assert(k > TensorExp::Kind::kLoopVar);
   const ExprId eNew(tensorExps.size());
-  tensorExps.emplace_back(k, e, detail::kInvalidId, v, op);
+  tensorExps.emplace_back(k, e.value, ExprId(), v, op);
   return eNew;
 }
 
 LatPointId Merger::addLat(TensorId t, LoopId i, ExprId e) {
   const LatPointId pNew(latPoints.size());
   const unsigned size = numLoops * numTensors;
   const TensorLoopId b = makeTensorLoopId(t, i);
   latPoints.emplace_back(size, e);
-  latPoints[pNew].bits.set(b);
+  latPoints[pNew.value].bits.set(b);
   return pNew;
 }
 
 LatPointId Merger::addLat(const BitVector &bits, ExprId e) {
   assert(bits.size() == numLoops * numTensors);
   const LatPointId pNew(latPoints.size());
   latPoints.emplace_back(bits, e);
   return pNew;
LatPointId Merger::conjLat(TensorExp::Kind kind, LatPointId p0, LatPointId p1,
   const ExprId e = addExp(kind, point0.exp, point1.exp, op);
   latPoints.emplace_back(bits, e);
   return pNew;
 }
 
 LatSetId Merger::conjSet(TensorExp::Kind kind, LatSetId s0, LatSetId s1,
                          Operation *op) {
   const LatSetId sNew = addSet();
-  auto &setNew = latSets[sNew];
+  auto &setNew = latSets[sNew.value];
   for (const LatPointId p0 : set(s0))
     for (const LatPointId p1 : set(s1))
       setNew.push_back(conjLat(kind, p0, p1, op));
   return sNew;
 }
 
 LatSetId Merger::disjSet(TensorExp::Kind kind, LatSetId s0, LatSetId s1,
                          Operation *op) {
   const LatSetId sNew = conjSet(kind, s0, s1, op);
   // Followed by all in s0.
-  latSets[sNew].append(latSets[s0]);
+  latSets[sNew.value].append(latSets[s0.value]);
   // Map binary 0-y to unary -y.
   // TODO: move this if-else logic into buildLattices
   if (kind == TensorExp::Kind::kSubF)
     s1 = mapSet(TensorExp::Kind::kNegF, s1);
   else if (kind == TensorExp::Kind::kSubC)
     s1 = mapSet(TensorExp::Kind::kNegC, s1);
   else if (kind == TensorExp::Kind::kSubI)
     s1 = mapSet(TensorExp::Kind::kNegI, s1);
   // Followed by all in s1.
-  latSets[sNew].append(latSets[s1]);
+  latSets[sNew.value].append(latSets[s1.value]);
   return sNew;
 }
 
 LatSetId Merger::combiSet(TensorExp::Kind kind, LatSetId s0, LatSetId s1,
                           Operation *orig, bool includeLeft,
                           TensorExp::Kind ltrans, Operation *opleft,
                           bool includeRight, TensorExp::Kind rtrans,
                           Operation *opright) {
   const LatSetId sNew = conjSet(kind, s0, s1, orig);
   // Left Region.
   if (includeLeft) {
     if (opleft)
       s0 = mapSet(ltrans, s0, Value(), opleft);
-    latSets[sNew].append(latSets[s0]);
+    latSets[sNew.value].append(latSets[s0.value]);
   }
   // Right Region.
   if (includeRight) {
     if (opright)
       s1 = mapSet(rtrans, s1, Value(), opright);
-    latSets[sNew].append(latSets[s1]);
+    latSets[sNew.value].append(latSets[s1.value]);
   }
   return sNew;
 }
 
 LatSetId Merger::mapSet(TensorExp::Kind kind, LatSetId s0, Value v,
                         Operation *op) {
   assert(TensorExp::Kind::kAbsF <= kind && kind <= TensorExp::Kind::kSelect);
   const LatSetId sNew = addSet();
-  auto &setNew = latSets[sNew];
+  auto &setNew = latSets[sNew.value];
   for (const LatPointId p : set(s0)) {
-    const auto &point = latPoints[p];
+    const auto &point = latPoints[p.value];
     setNew.push_back(addLat(point.bits, addExp(kind, point.exp, v, op)));
   }
   return sNew;
 }
 
 LatSetId Merger::optimizeSet(LatSetId s0) {
   const LatSetId sNew = addSet();
-  auto &setNew = latSets[sNew];
+  auto &setNew = latSets[sNew.value];
   const auto &set0 = set(s0);
   assert(!set0.empty());
   const LatPointId p0 = set0[0];
   for (const LatPointId p1 : set0) {
     bool add = true;
     if (p0 != p1) {
       // Check whether this is a straightforward copy.
-      if (expIsTensor(latPoints[p1].exp, outTensor))
+      if (expIsTensor(latPoints[p1.value].exp, outTensor))
         continue;
       // Check whether this conjunction is already covered.
       for (const LatPointId p2 : setNew) {
         assert(!latGT(p1, p2)); // Lj => Li would be bad
         if (onlyDenseDiff(p2, p1)) {
           add = false;
           break;
         }
       }
       assert(!add || latGT(p0, p1));
     }
     if (add)
       setNew.push_back(p1);
   }
   for (const LatPointId p : setNew)
-    latPoints[p].simple = simplifyCond(sNew, p);
+    latPoints[p.value].simple = simplifyCond(sNew, p);
   return sNew;
 }
 
 BitVector Merger::simplifyCond(LatSetId s0, LatPointId p0) {
   // First determine if this lattice point is a *singleton*, i.e.,
   // the last point in a lattice, no other is less than this one.
   bool isSingleton = true;
   for (const LatPointId p1 : set(s0)) {
case TensorExp::Kind::kLoopVar: {
     // iteration space is not skipped as a result of their contents.
     const LatSetId s = addSet();
     TensorId t = syntheticTensor;
     if (kind == TensorExp::Kind::kTensor) {
       t = expr.tensor;
       if (hasSparseOut && t == outTensor)
         t = syntheticTensor;
     }
-    latSets[s].push_back(addLat(t, i, e));
+    latSets[s.value].push_back(addLat(t, i, e));
     return s;
   }
   // Unary operations.
   case TensorExp::Kind::kAbsF:
   case TensorExp::Kind::kAbsC:
   case TensorExp::Kind::kAbsI:
   case TensorExp::Kind::kCeilF:
   case TensorExp::Kind::kFloorF: