[mlir][sparse] properly record dimension level type and properties

[aartbik]

Summary

A next step towards supporting the new dimension level types and
properties. This changes properly records the properties in the
Merger, so that subsequent computations (lattice optimizations)
and code generation (during sparsification) can do the right thing.

https://github.com/llvm/llvm-project/issues/51658

[Peiming]

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

268

Should be a cheap copy for now. If we are supporting more property in the future, maybe a bitvector for property is better?

But I am okay with it now

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

309–311

Why not just !isDimLevelType(tensor, i, kDense)

357–359

Same here?

1162

I actually prefer an enum class over enum for dimLvlType

[aartbik]

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

268

Yeah, we can make that const & but that has all sort of complications with inserting in the vector while "holding" on to the reference

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

309–311

Because there is also undef that we want to move in the other direction

1162

Okay, also note DimLvl instead of DimLevel since we had another enum of that name already ;-)

[aartbik]

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

357–359

It may be possible here. Will clean that up when implementing.

[aartbik]

addressed comments, also fixed a has_value/value warning that was lingering in this file

[aartbik]

rebased with main

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

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/BitVector.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
 /// Dimension level type for a tensor (undef means index does not appear).
-enum Dim { kSparse, kDense, kUndef };
+enum class DimLvlType { kDense, kCompressed, kSingleton, kUndef };
+
+/// Per-dimension level format (type and properties). Dense and undefined
+/// level types should always be marked ordered and unique.
+struct DimLevelFormat {
+  DimLevelFormat(DimLvlType tp, bool o = true, bool u = true)
+      : levelType(tp), isOrdered(o), isUnique(u) {
+    assert((tp == DimLvlType::kCompressed || tp == DimLvlType::kSingleton) ||
+           (o && u));
+  }
+  DimLvlType levelType;
+  bool isOrdered;
+  bool isUnique;
+};
 
 /// Tensor expression kind.
 enum Kind {
   // Leaf.
   kTensor = 0,
   kInvariant,
   kIndex,
   // Unary operations.

 public:
   /// Constructs a merger for the given number of tensors and loops. The
   /// user supplies the number of tensors involved in the kernel, with the
   /// last tensor in this set denoting the output tensor. The merger adds an
   /// additional synthetic tensor at the end of this set to represent all
   /// invariant expressions in the kernel.
   Merger(unsigned t, unsigned l)
       : outTensor(t - 1), syntheticTensor(t), numTensors(t + 1), numLoops(l),
-        hasSparseOut(false), dims(t + 1, std::vector<Dim>(l, Dim::kUndef)) {}
+        hasSparseOut(false),
+        dims(t + 1, std::vector<DimLevelFormat>(
+                        l, DimLevelFormat(DimLvlType::kUndef))) {}
 
   /// Adds a tensor expression. Returns its index.
   unsigned addExp(Kind k, unsigned e0, unsigned e1 = -1u, Value v = Value(),
                   Operation *op = nullptr);
   unsigned addExp(Kind k, unsigned e, Value v, Operation *op = nullptr) {
     return addExp(k, e, -1u, v, op);
   }
   unsigned addExp(Kind k, Value v, Operation *op = nullptr) {
public:
 
   /// Returns true if Li and Lj only differ in dense.
   bool onlyDenseDiff(unsigned i, unsigned j);
 
   /// Bit translation.
   unsigned tensor(unsigned b) const { return b % numTensors; }
   unsigned index(unsigned b) const { return b / numTensors; }
 
-  /// Returns true if bit corresponds to queried dim.
-  bool isDim(unsigned b, Dim d) const { return isDim(tensor(b), index(b), d); }
-
   /// Returns true if bit corresponds to index of output tensor.
   bool isOutTensor(unsigned b, unsigned i) const {
     return tensor(b) == outTensor && index(b) == i;
   }
 
-  /// Returns true if tensor access at given index has queried dim.
-  bool isDim(unsigned t, unsigned i, Dim d) const {
-    assert(t < numTensors && i < numLoops);
-    return dims[t][i] == d;
-  }
-
-  /// Returns true if any set bit corresponds to queried dim.
-  bool hasAnyDimOf(const BitVector &bits, Dim d) const;
-
   /// Returns true if given tensor iterates *only* in the given tensor
   /// expression. For the output tensor, this defines a "simply dynamic"
   /// operation [Bik96]. For instance: a(i) *= 2.0 or a(i) += a(i) for
   /// sparse vector a.
   bool isSingleCondition(unsigned t, unsigned e) const;
 
-  /// Dimension setter.
-  void setDim(unsigned t, unsigned i, Dim d) { dims[t][i] = d; }
+  /// Returns true if bit corresponds to given dimension level type.
+  bool isDimLevelType(unsigned b, DimLvlType tp) const {
+    return isDimLevelType(tensor(b), index(b), tp);
+  }
+
+  /// Returns true if tensor access at index has given dimension level type.
+  bool isDimLevelType(unsigned t, unsigned i, DimLvlType tp) const {
+    return getDimLevelFormat(t, i).levelType == tp;
+  }
+
+  /// Returns true if any set bit corresponds to given dimension level type.
+  bool hasAnyDimLevelTypeOf(const BitVector &bits, DimLvlType tp) const;
+
+  /// Dimension level format getter.
+  DimLevelFormat getDimLevelFormat(unsigned t, unsigned i) const {
+    assert(t < numTensors && i < numLoops);
+    return dims[t][i];
+  }
+
+  /// Dimension level format setter.
+  void setDimLevelFormat(unsigned t, unsigned i, DimLevelFormat d) {
+    dims[t][i] = d;
+  }
 
   // Has sparse output tensor setter.
   void setHasSparseOut(bool s) { hasSparseOut = s; }
 
   /// Convenience getters to immediately access the stored nodes.
   /// Typically it is inadvisible to keep the reference around, as in
   /// "TensorExpr &te = merger.exp(e))", since insertions into the merger
   /// may cause data movement and invalidate the underlying memory address.
private:
   Optional<unsigned> buildTensorExp(linalg::GenericOp op, Value v);
 
   /// Merger data structures.
   const unsigned outTensor;
   const unsigned syntheticTensor;
   const unsigned numTensors;
   const unsigned numLoops;
   bool hasSparseOut;
-  std::vector<std::vector<Dim>> dims;
+  std::vector<std::vector<DimLevelFormat>> dims;
   llvm::SmallVector<TensorExp, 32> tensorExps;
   llvm::SmallVector<LatPoint, 16> latPoints;
   llvm::SmallVector<SmallVector<unsigned, 16>, 8> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

if (enc) {
     if (order) {
       assert(order.isPermutation());
       return order.getDimPosition(d);
     }
   }
   return d;
 }
 
-/// Helper method to translate dim level type to internal representation.
-static Dim toDim(const SparseTensorEncodingAttr &enc, unsigned d) {
+/// Helper method to obtain the dimension level format from the encoding.
+//
+//  TODO: note that we store, but currently completely *ignore* the properties
+//
+static DimLevelFormat toDimLevelFormat(const SparseTensorEncodingAttr &enc,
+                                       unsigned d) {
   if (enc) {
-    SparseTensorEncodingAttr::DimLevelType tp = enc.getDimLevelType()[d];
-    if (tp == SparseTensorEncodingAttr::DimLevelType::Compressed)
-      return Dim::kSparse;
+    switch (enc.getDimLevelType()[d]) {
+    case SparseTensorEncodingAttr::DimLevelType::Dense:
+      return DimLevelFormat(DimLvlType::kDense);
+    case SparseTensorEncodingAttr::DimLevelType::Compressed:
+      return DimLevelFormat(DimLvlType::kCompressed);
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
+      return DimLevelFormat(DimLvlType::kCompressed, true, false);
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
+      return DimLevelFormat(DimLvlType::kCompressed, false, true);
+    case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
+      return DimLevelFormat(DimLvlType::kCompressed, false, false);
+    case SparseTensorEncodingAttr::DimLevelType::Singleton:
+      return DimLevelFormat(DimLvlType::kSingleton);
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
+      return DimLevelFormat(DimLvlType::kSingleton, true, false);
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
+      return DimLevelFormat(DimLvlType::kSingleton, false, true);
+    case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
+      return DimLevelFormat(DimLvlType::kSingleton, false, false);
+    }
   }
-  return Dim::kDense;
+  return DimLevelFormat(DimLvlType::kDense);
 }
 
 /// Helper method to inspect affine expressions. Rejects cases where the
-/// same index is used more than once. Also rejects affine expressions
-/// that are not a direct index for annotated tensors.
-// TODO: accept more affine cases for sparse tensors
-static bool findAffine(Merger &merger, unsigned tensor, AffineExpr a, Dim dim,
-                       bool isDense) {
+/// same index is used more than once. Also rejects compound affine
+/// expressions in sparse dimensions.
+static bool findAffine(Merger &merger, unsigned tensor, AffineExpr a,
+                       DimLevelFormat dim) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
-    if (!merger.isDim(tensor, idx, Dim::kUndef))
+    if (!merger.isDimLevelType(tensor, idx, DimLvlType::kUndef))
       return false; // used more than once
-    merger.setDim(tensor, idx, dim);
+    merger.setDimLevelFormat(tensor, idx, dim);
     return true;
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul: {
-    if (!isDense)
-      return false;
+    if (dim.levelType != DimLvlType::kDense)
+      return false; // compound only in dense dim
     auto binOp = a.cast<AffineBinaryOpExpr>();
-    return findAffine(merger, tensor, binOp.getLHS(), dim, isDense) &&
-           findAffine(merger, tensor, binOp.getRHS(), dim, isDense);
+    return findAffine(merger, tensor, binOp.getLHS(), dim) &&
+           findAffine(merger, tensor, binOp.getRHS(), dim);
   }
   case AffineExprKind::Constant:
-    return isDense;
+    return dim.levelType == DimLvlType::kDense; // const only in dense dim
   default:
     return false;
   }
 }
 
 /// Helper method to inspect sparse encodings in the tensor types.
 /// Fills the per-dimension sparsity information for all tensors.
 /// Returns true if the sparse annotations and affine subscript
 /// expressions of all tensors are admissable. Returns false if
 /// no annotations are found or inadmissable constructs occur.
 static bool findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
   bool annotated = false;
   for (OpOperand *t : op.getInputAndOutputOperands()) {
     auto map = op.getTiedIndexingMap(t);
     auto enc = getSparseTensorEncoding(t->get().getType());
     if (enc)
       annotated = true;
     assert(map.getNumResults() == op.getRank(t));
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       unsigned tensor = t->getOperandNumber();
       AffineExpr a = map.getResult(perm(enc, d));
-      if (!findAffine(merger, tensor, a, toDim(enc, d), !enc))
+      if (!findAffine(merger, tensor, a, toDimLevelFormat(enc, d)))
         return false; // inadmissable affine expression
     }
   }
   return annotated;
 }
 
 /// A DFS helper to compute a topological sort. Note that recursion is
 /// bounded by the number of implicit loops, which is always small.
for (unsigned d = 1, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr t = map.getResult(perm(enc, d));
       addAffineOrderings(adjM, f, t, 0);
     }
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     if (mask & SortMask::kIncludeUndef) {
       unsigned tensor = t->getOperandNumber();
       for (unsigned i = 0; i < n; i++)
-        if (merger.isDim(tensor, i, Dim::kSparse))
+        if (merger.isDimLevelType(tensor, i, DimLvlType::kCompressed) ||
+            merger.isDimLevelType(tensor, i, DimLvlType::kSingleton))
           for (unsigned j = 0; j < n; j++)
-            if (merger.isDim(tensor, j, Dim::kUndef))
+            if (merger.isDimLevelType(tensor, j, DimLvlType::kUndef))
               adjM[i][j] = true;
     }
   }
-
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
   topSort.clear();
   topSort.reserve(n);
   std::vector<unsigned> visit(n, 0);
   for (unsigned i = 0; i < n; i++)
     if (visit[i] == 0)
       if (!topSortDFS(i, visit, topSort, adjM))
static bool isAdmissableTensorExp(Merger &merger, linalg::GenericOp op,
   if (!enc)
     return true;
   // An all-dense annotated "sparse" output tensor becomes a linearized random
   // access 1-dim memref. Also admissable since insertions cannot occur.
   bool allDense = true;
   auto iteratorTypes = op.iterator_types().getValue();
   unsigned numLoops = iteratorTypes.size();
   for (unsigned i = 0; i < numLoops; i++)
-    if (merger.isDim(tensor, i, Dim::kSparse)) {
+    if (merger.isDimLevelType(tensor, i, DimLvlType::kCompressed) ||
+        merger.isDimLevelType(tensor, i, DimLvlType::kSingleton)) {
       allDense = false;
       break;
     }
   if (allDense)
     return true;
   // A tensor expression with a sparse output tensor that changes its values
   // but not its nonzero structure, an operation called "simply dynamic" in
   // [Bik96,Ch9], is also admissable without special codegen.
   if (merger.isSingleCondition(tensor, exp))
for (OpOperand *t : op.getInputAndOutputOperands()) {
     // Scan all dimensions of current tensor.
     args.clear();
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
       if (a.getKind() != AffineExprKind::DimId)
         continue; // compound
       unsigned idx = a.cast<AffineDimExpr>().getPosition();
       // Handle sparse storage schemes.
-      if (merger.isDim(tensor, idx, Dim::kSparse)) {
+      if (merger.isDimLevelType(tensor, idx, DimLvlType::kCompressed)) {
         auto dynShape = {ShapedType::kDynamicSize};
         auto ptrTp =
             MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         // Generate sparse primitives to obtains pointer and indices.
         codegen.pointers[tensor][idx] =
             builder.create<ToPointersOp>(loc, ptrTp, t->get(), dim);
         codegen.indices[tensor][idx] =
             builder.create<ToIndicesOp>(loc, indTp, t->get(), dim);
+      } else if (merger.isDimLevelType(tensor, idx, DimLvlType::kSingleton)) {
+        llvm_unreachable("TODO: not implemented yet");
       }
       // Find upper bound in current dimension.
       unsigned p = perm(enc, d);
       Value up = linalg::createOrFoldDimOp(builder, loc, t->get(), p);
       if (ShapedType::isDynamic(shape[p]))
         args.push_back(up);
       assert(codegen.highs[tensor][idx] == nullptr);
       codegen.sizes[idx] = codegen.highs[tensor][idx] = up;
     }
     // Perform the required bufferization. Dense inputs materialize
     // from the input tensors. Dense outputs need special handling.
     // Sparse inputs use sparse primitives to obtain the values.
-    // We also accept in-place all-dense annotated "sparse" outputs.
     Type elementType = getElementTypeOrSelf(t->get().getType());
     if (!enc) {
       // Non-annotated dense tensors.
       auto denseTp = MemRefType::get(shape, elementType);
       if (tensor < op.getNumInputs())
         codegen.buffers[tensor] =
             builder.create<bufferization::ToMemrefOp>(loc, denseTp, t->get());
       else
static Value genExp(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
   if (exp == -1u)
     return Value();
   if (merger.exp(exp).kind == Kind::kTensor)
     return genTensorLoad(merger, codegen, rewriter, op, exp);
   if (merger.exp(exp).kind == Kind::kInvariant)
     return genInvariantValue(merger, codegen, rewriter, exp);
   if (merger.exp(exp).kind == Kind::kIndex)
     return genIndexValue(codegen, rewriter, merger.exp(exp).index, ldx);
+
   if (merger.exp(exp).kind == Kind::kReduce) {
     // Make custom reduction identity accessible for expanded access pattern.
     assert(codegen.redCustom == -1u);
     codegen.redCustom = exp;
   }
+
   Value v0 =
       genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e0, ldx);
   Value v1 =
       genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e1, ldx);
   Value ee = merger.buildExp(rewriter, loc, exp, v0, v1);
   if (ee && (merger.exp(exp).kind == Kind::kUnary ||
              merger.exp(exp).kind == Kind::kBinary ||
              merger.exp(exp).kind == Kind::kBinaryBranch ||
              merger.exp(exp).kind == Kind::kReduce))
     ee = relinkBranch(codegen, rewriter, ee.getParentBlock(), ee, ldx);
-  if (merger.exp(exp).kind == Kind::kReduce)
+
+  if (merger.exp(exp).kind == Kind::kReduce) {
+    assert(codegen.redCustom != -1u);
     codegen.redCustom = -1u;
+  }
+
   return ee;
 }
 
 /// Determines if affine expression is invariant.
 static bool isInvariantAffine(const CodeGen &codegen, AffineExpr a,
                               unsigned ldx, bool &atLevel) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
static bool isInvariantAffine(const CodeGen &codegen, AffineExpr a,
   default:
     return true;
   }
 }
 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
 static void genInvariants(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                           linalg::GenericOp op, unsigned exp, unsigned ldx,
-                          bool atStart, unsigned last = 0) {
+                          bool atStart, unsigned last = -1u) {
   if (exp == -1u)
     return;
   if (merger.exp(exp).kind == Kind::kTensor) {
     // Inspect tensor indices.
     bool atLevel = ldx == -1u;
     OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
     auto map = op.getTiedIndexingMap(t);
     auto enc = getSparseTensorEncoding(t->get().getType());
static bool genInit(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse positions.
   for (unsigned b = 0, be = inits.size(); b < be; b++) {
     if (inits[b]) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
-      if (merger.isDim(b, Dim::kSparse)) {
+      if (merger.isDimLevelType(b, DimLvlType::kCompressed)) {
         // Initialize sparse index.
         unsigned pat = at;
         for (; pat != 0; pat--) {
           if (codegen.pidxs[tensor][topSort[pat - 1]])
             break;
         }
         Value ptr = codegen.pointers[tensor][idx];
         Value one = constantIndex(builder, loc, 1);
         Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
                               : codegen.pidxs[tensor][topSort[pat - 1]];
         codegen.pidxs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p0);
         Value p1 = builder.create<arith::AddIOp>(loc, p0, one);
         codegen.highs[tensor][idx] = genLoad(codegen, builder, loc, ptr, p1);
+      } else if (merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+        llvm_unreachable("TODO: not implemented yet");
       } else {
         // Dense index still in play.
         needsUniv = true;
       }
     }
   }
 
   // Initialize the universal dense index.

 static Operation *genFor(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, bool isOuter, bool isInner,
                          unsigned idx, BitVector &indices) {
   unsigned fb = indices.find_first();
   unsigned tensor = merger.tensor(fb);
   assert(idx == merger.index(fb));
   auto iteratorTypes = op.iterator_types().getValue();
   bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
-  bool isSparse = merger.isDim(fb, Dim::kSparse);
+  bool isSparse = merger.isDimLevelType(fb, DimLvlType::kCompressed);
   bool isVector = isVectorFor(codegen, isInner, isReduction, isSparse) &&
                   denseUnitStrides(merger, op, idx);
   bool isParallel =
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
+  assert(!merger.isDimLevelType(fb, DimLvlType::kSingleton) && "TODO: implement");
+
   // Prepare vector length.
   if (isVector)
     codegen.curVecLength = codegen.options.vectorLength;
 
   // Loop bounds and increment.
   Location loc = op.getLoc();
   Value lo = isSparse ? codegen.pidxs[tensor][idx] : codegen.loops[idx];
   Value hi = isSparse ? codegen.highs[tensor][idx] : codegen.sizes[idx];

 static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                            linalg::GenericOp op, unsigned idx, bool needsUniv,
                            BitVector &indices) {
   SmallVector<Type, 4> types;
   SmallVector<Value, 4> operands;
   // Construct the while-loop with a parameter for each index.
   Type indexType = builder.getIndexType();
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
-    if (indices[b] && merger.isDim(b, Dim::kSparse)) {
+    if (indices[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       types.push_back(indexType);
       operands.push_back(codegen.pidxs[tensor][idx]);
     }
   }
   if (codegen.redVal) {
     types.push_back(codegen.redVal.getType());
static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   Block *after = builder.createBlock(&whileOp.getAfter(), {}, types, locs);
 
   // Build the "before" region, which effectively consists
   // of a conjunction of "i < upper" tests on all induction.
   builder.setInsertionPointToStart(&whileOp.getBefore().front());
   Value cond;
   unsigned o = 0;
   for (unsigned b = 0, be = indices.size(); b < be; b++) {
-    if (indices[b] && merger.isDim(b, Dim::kSparse)) {
+    // TODO: singleton
+    if (indices[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = before->getArgument(o);
       Value op2 = codegen.highs[tensor][idx];
       Value opc = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::ult,
                                                 op1, op2);
       cond = cond ? builder.create<arith::AndIOp>(loc, cond, opc) : opc;
       codegen.pidxs[tensor][idx] = after->getArgument(o++);
static void genLocals(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                       linalg::GenericOp op, std::vector<unsigned> &topSort,
                       unsigned at, bool needsUniv, BitVector &locals) {
   Location loc = op.getLoc();
   unsigned idx = topSort[at];
 
   // Initialize sparse indices.
   Value min;
   for (unsigned b = 0, be = locals.size(); b < be; b++) {
-    if (locals[b] && merger.isDim(b, Dim::kSparse)) {
+    // TODO: singleton
+    if (locals[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value ptr = codegen.indices[tensor][idx];
       Value s = codegen.pidxs[tensor][idx];
       Value load = genLoad(codegen, builder, loc, ptr, s);
       codegen.idxs[tensor][idx] = load;
       if (!needsUniv) {
         if (min) {
if (min) {
     codegen.loops[idx] = min;
   }
 
   // Initialize dense positions. Note that we generate dense indices of the
   // output tensor unconditionally, since they may not appear in the lattice,
   // but may be needed for linearized codegen.
   for (unsigned b = 0, be = locals.size(); b < be; b++) {
     if ((locals[b] || merger.isOutTensor(b, idx)) &&
-        merger.isDim(b, Dim::kDense)) {
+        merger.isDimLevelType(b, DimLvlType::kDense)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       unsigned pat = at;
       for (; pat != 0; pat--)
         if (codegen.pidxs[tensor][topSort[pat - 1]])
           break;
       Value p = (pat == 0) ? constantIndex(builder, loc, 0)
                            : codegen.pidxs[tensor][topSort[pat - 1]];
static void genWhileInduction(Merger &merger, CodeGen &codegen,
   // in the individual if-branches to avoid re-evaluating the conditions.
   // However, that would result in a rather elaborate forest of yield
   // instructions during code generation. Moreover, performing the induction
   // after the if-statements more closely resembles code generated by TACO.
   unsigned o = 0;
   SmallVector<Value, 4> operands;
   Value one = constantIndex(builder, loc, 1);
   for (unsigned b = 0, be = induction.size(); b < be; b++) {
-    if (induction[b] && merger.isDim(b, Dim::kSparse)) {
+    // TODO: singleton
+    if (induction[b] && merger.isDimLevelType(b, DimLvlType::kCompressed)) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value op1 = codegen.idxs[tensor][idx];
       Value op2 = codegen.loops[idx];
       Value op3 = codegen.pidxs[tensor][idx];
       Value cmp = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                 op1, op2);
       Value add = builder.create<arith::AddIOp>(loc, op3, one);
static scf::IfOp genIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   Location loc = op.getLoc();
   SmallVector<Type, 4> types;
   Value cond;
   for (unsigned b = 0, be = conditions.size(); b < be; b++) {
     if (conditions[b]) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       Value clause;
-      if (merger.isDim(b, Dim::kSparse)) {
+      // TODO: singleton
+      if (merger.isDimLevelType(b, DimLvlType::kCompressed)) {
         Value op1 = codegen.idxs[tensor][idx];
         Value op2 = codegen.loops[idx];
         clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq,
                                                op1, op2);
       } else {
         clause = constantI1(builder, loc, true);
       }
       cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
static bool startLoopSeq(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   bool needsUniv =
       genInit(merger, codegen, builder, op, topSort, at, merger.lat(l0).bits);
   // Maintain the universal index only if it is actually
   // consumed by a subsequent lattice point.
   if (needsUniv) {
     unsigned lsize = merger.set(lts).size();
     for (unsigned i = 1; i < lsize; i++) {
       unsigned li = merger.set(lts)[i];
-      if (!merger.hasAnyDimOf(merger.lat(li).simple, Dim::kSparse))
+      if (!merger.hasAnyDimLevelTypeOf(merger.lat(li).simple, DimLvlType::kCompressed) &&
+          !merger.hasAnyDimLevelTypeOf(merger.lat(li).simple, DimLvlType::kSingleton))
         return true;
     }
   }
   return false;
 }
 
 /// Starts a single loop in current sequence.
 static Operation *startLoop(Merger &merger, CodeGen &codegen,

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

BitVector Merger::simplifyCond(unsigned s0, unsigned p0) {
   bool isSingleton = true;
   for (unsigned p1 : latSets[s0]) {
     if (p0 != p1 && latGT(p0, p1)) {
       isSingleton = false;
       break;
     }
   }
   // Now apply the two basic rules.
+  //
+  // TODO: improve for singleton and properties
+  //
   BitVector simple = latPoints[p0].bits;
-  bool reset = isSingleton && hasAnyDimOf(simple, kSparse);
+  bool reset = isSingleton &&
+      (hasAnyDimLevelTypeOf(simple, DimLvlType::kCompressed) ||
+       hasAnyDimLevelTypeOf(simple, DimLvlType::kSingleton));
   for (unsigned b = 0, be = simple.size(); b < be; b++) {
-    if (simple[b] && !isDim(b, kSparse)) {
+    if (simple[b] &&
+        (!isDimLevelType(b, DimLvlType::kCompressed) &&
+         !isDimLevelType(b, DimLvlType::kSingleton))) {
       if (reset)
         simple.reset(b);
       reset = true;
     }
   }
   return simple;
 }
 
 bool Merger::latGT(unsigned i, unsigned j) const {
   const BitVector &bitsi = latPoints[i].bits;
   const BitVector &bitsj = latPoints[j].bits;
   assert(bitsi.size() == bitsj.size());
   if (bitsi.count() > bitsj.count()) {
     for (unsigned b = 0, be = bitsj.size(); b < be; b++)
       if (bitsj[b] && !bitsi[b])
         return false;
     return true;
   }
   return false;
 }
 
 bool Merger::onlyDenseDiff(unsigned i, unsigned j) {
   BitVector tmp = latPoints[j].bits;
   tmp ^= latPoints[i].bits;
-  return !hasAnyDimOf(tmp, kSparse);
-}
-
-bool Merger::hasAnyDimOf(const BitVector &bits, Dim d) const {
-  for (unsigned b = 0, be = bits.size(); b < be; b++)
-    if (bits[b] && isDim(b, d))
-      return true;
-  return false;
+  return !hasAnyDimLevelTypeOf(tmp, DimLvlType::kCompressed) &&
+         !hasAnyDimLevelTypeOf(tmp, DimLvlType::kSingleton);
 }
 
 bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   switch (tensorExps[e].kind) {
   // Leaf.
   case kTensor:
     return tensorExps[e].tensor == t;
   case kInvariant:
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   case kXorI:
   case kBinary:
   case kReduce:
     return false;
   }
   llvm_unreachable("unexpected kind");
 }
 
+bool Merger::hasAnyDimLevelTypeOf(const BitVector &bits, DimLvlType tp) const {
+  for (unsigned b = 0, be = bits.size(); b < be; b++)
+    if (bits[b] && isDimLevelType(b, tp))
+      return true;
+  return false;
+}
+
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Print methods (for debugging).
 //===----------------------------------------------------------------------===//
 
 static const char *kindToOpSymbol(Kind kind) {
   switch (kind) {
void Merger::dumpSet(unsigned s) const {
   llvm::dbgs() << "}\n";
 }
 
 void Merger::dumpBits(const BitVector &bits) const {
   for (unsigned b = 0, be = bits.size(); b < be; b++) {
     if (bits[b]) {
       unsigned t = tensor(b);
       unsigned i = index(b);
+      DimLevelFormat f = dims[t][i];
       llvm::dbgs() << " i_" << t << "_" << i << "_";
-      switch (dims[t][i]) {
-      case kSparse:
-        llvm::dbgs() << "S";
-        break;
-      case kDense:
+      switch (f.levelType) {
+      case DimLvlType::kDense:
         llvm::dbgs() << "D";
         break;
-      case kUndef:
+      case DimLvlType::kCompressed:
+        llvm::dbgs() << "C";
+        break;
+      case DimLvlType::kSingleton:
+        llvm::dbgs() << "S";
+        break;
+      case DimLvlType::kUndef:
         llvm::dbgs() << "U";
         break;
       }
+      llvm::dbgs() << "[O=" << f.isOrdered << ",U=" << f.isUnique << "]";
     }
   }
 }
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.
static bool isAdmissableBranchExp(Operation *op, Block *block, Value v) {
   // Arguments are always admissable.
   if (auto arg = v.dyn_cast<BlockArgument>())
     return true;
   // Accept index anywhere.
   Operation *def = v.getDefiningOp();
   if (isa<linalg::IndexOp>(def))
     return true;
   // Operation defined outside branch.
-  if (def->getBlock() != block) {
+  if (def->getBlock() != block)
     return def->getBlock() != op->getBlock(); // invariant?
-  }
   // Operation defined within branch. Anything is accepted,
   // as long as all subexpressions are admissable.
   for (unsigned i = 0, n = def->getNumOperands(); i < n; i++)
     if (!isAdmissableBranchExp(op, block, def->getOperand(i)))
       return false;
   return true;
 }
 
Optional<unsigned> Merger::buildTensorExp(linalg::GenericOp op, Value v) {
   // Construct ternary operations if subexpressions can be built.
   if (def->getNumOperands() == 3) {
     auto x = buildTensorExp(op, def->getOperand(0));
     auto y = buildTensorExp(op, def->getOperand(1));
     auto z = buildTensorExp(op, def->getOperand(2));
     if (x.has_value() && y.has_value() && z.has_value()) {
       unsigned e0 = x.value();
       unsigned e1 = y.value();
-      // unsigned e2 = z.getValue();
       if (auto redop = dyn_cast<sparse_tensor::ReduceOp>(def)) {
         if (isAdmissableBranch(redop, redop.getRegion()))
           return addExp(kReduce, e0, e1, Value(), def);
       }
     }
   }
   // Cannot build.
   return None;

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

protected:
   const unsigned t0 = 0, t1 = 1, t2 = 2;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest3T1L() : MergerTestBase(3, 1) {
     // Tensor 0: sparse input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
-    merger.setDim(t0, l0, Dim::kSparse);
+    merger.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kCompressed));
 
     // Tensor 1: sparse input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDim(t1, l0, Dim::kSparse);
+    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kCompressed));
 
     // Tensor 2: dense output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDim(t2, l0, Dim::kDense);
+    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kDense));
   }
 };
 
 class MergerTest4T1L : public MergerTestBase {
 protected:
   // Our four tensors (three inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2, t3 = 3;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest4T1L() : MergerTestBase(4, 1) {
     // Tensor 0: sparse input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
-    merger.setDim(t0, l0, Dim::kSparse);
+    merger.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kCompressed));
 
     // Tensor 1: sparse input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDim(t1, l0, Dim::kSparse);
+    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kCompressed));
 
     // Tensor 2: sparse input vector
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDim(t2, l0, Dim::kSparse);
+    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kCompressed));
 
     // Tensor 3: dense output vector
     merger.addExp(Kind::kTensor, t3, -1u);
-    merger.setDim(t3, l0, Dim::kDense);
+    merger.setDimLevelFormat(t3, l0, DimLevelFormat(DimLvlType::kDense));
   }
 };
 
 ///
 /// Tests with both sparse and dense input.
 ///
 
 class MergerTest3T1LD : public MergerTestBase {
 protected:
   // Our three tensors (two inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest3T1LD() : MergerTestBase(3, 1) {
     // Tensor 0: sparse input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
-    merger.setDim(t0, l0, Dim::kSparse);
+    merger.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kCompressed));
 
     // Tensor 1: dense input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDim(t1, l0, Dim::kDense);
+    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kDense));
 
     // Tensor 2: dense output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDim(t2, l0, Dim::kDense);
+    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kDense));
   }
 };
 
 } // namespace
 
 /// Vector addition (disjunction) of 2 vectors. i.e.;
 ///   a(i) = b(i) + c(i)
 /// which should form the 3 lattice points