[mlir][sparse] Removing the DimLvlType and DimLevelFormat types

[wrengr]

Summary

This removes another massive source of redundancy, and instead has the Merger.{h,cpp} reuse the SparseTensorEnums library.

Depends On D136005

[wrengr]

rebase, and addressing the MSVC error

[wrengr]

rebase

[Peiming]

Nice change! looks clean!

LGTM!

[Peiming]

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

147

Yeah, I agree that it should be unreachable here

[wrengr]

Changing more constexpr annotations to inline for MSVC

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

 /// This enum defines all the sparse representations supportable by
 /// the SparseTensor dialect.  We use a lightweight encoding to encode
 /// both the "format" per se (dense, compressed, singleton) as well as
 /// the "properties" (ordered, unique).  The encoding is chosen for
 /// performance of the runtime library, and thus may change in future
 /// versions; consequently, client code should use the predicate functions
 /// defined below, rather than relying on knowledge about the particular
 /// binary encoding.
+///
+/// The `Undef` "format" is a special value used internally for cases
+/// where we need to store an undefined or indeterminate `DimLevelType`.
+/// It should not be used externally, since it does not indicate an
+/// actual/representable format.
 enum class DimLevelType : uint8_t {
+  Undef = 0,           // 0b000_00
   Dense = 4,           // 0b001_00
   Compressed = 8,      // 0b010_00
   CompressedNu = 9,    // 0b010_01
   CompressedNo = 10,   // 0b010_10
   CompressedNuNo = 11, // 0b010_11
   Singleton = 16,      // 0b100_00
   SingletonNu = 17,    // 0b100_01
   SingletonNo = 18,    // 0b100_10
   SingletonNuNo = 19,  // 0b100_11
 };
 
+/// Check that the `DimLevelType` contains a valid (possibly undefined) value.
+constexpr bool isValidDLT(DimLevelType dlt) {
+  const uint8_t formatBits = static_cast<uint8_t>(dlt) >> 2;
+  const uint8_t propertyBits = static_cast<uint8_t>(dlt) & 3;
+  // If undefined or dense, then must be unique and ordered.
+  // Otherwise, the format must be one of the known ones.
+  return (formatBits <= 1) ? (propertyBits == 0)
+                           : (formatBits == 2 || formatBits == 4);
+}
+
+/// Check if the `DimLevelType` is the special undefined value.
+constexpr bool isUndefDLT(DimLevelType dlt) {
+  return dlt == DimLevelType::Undef;
+}
+
 /// Check if the `DimLevelType` is dense.
 constexpr bool isDenseDLT(DimLevelType dlt) {
   return dlt == DimLevelType::Dense;
 }
 
+// We use the idiom `(dlt & ~3) == format` in order to only return true
+// for valid DLTs.  Whereas the `dlt & format` idiom is a bit faster but
+// can return false-positives on invalid DLTs.
+
 /// Check if the `DimLevelType` is compressed (regardless of properties).
 constexpr bool isCompressedDLT(DimLevelType dlt) {
-  return static_cast<uint8_t>(dlt) &
+  return (static_cast<uint8_t>(dlt) & ~3) ==
          static_cast<uint8_t>(DimLevelType::Compressed);
 }
 
 /// Check if the `DimLevelType` is singleton (regardless of properties).
 constexpr bool isSingletonDLT(DimLevelType dlt) {
-  return static_cast<uint8_t>(dlt) &
+  return (static_cast<uint8_t>(dlt) & ~3) ==
          static_cast<uint8_t>(DimLevelType::Singleton);
 }
 
 /// Check if the `DimLevelType` is ordered (regardless of storage format).
 constexpr bool isOrderedDLT(DimLevelType dlt) {
   return !(static_cast<uint8_t>(dlt) & 2);
 }
 
 /// Check if the `DimLevelType` is unique (regardless of storage format).
 constexpr bool isUniqueDLT(DimLevelType dlt) {
   return !(static_cast<uint8_t>(dlt) & 1);
 }
 
 // Ensure the above predicates work as intended.
+static_assert((isValidDLT(DimLevelType::Undef) &&
+               isValidDLT(DimLevelType::Dense) &&
+               isValidDLT(DimLevelType::Compressed) &&
+               isValidDLT(DimLevelType::CompressedNu) &&
+               isValidDLT(DimLevelType::CompressedNo) &&
+               isValidDLT(DimLevelType::CompressedNuNo) &&
+               isValidDLT(DimLevelType::Singleton) &&
+               isValidDLT(DimLevelType::SingletonNu) &&
+               isValidDLT(DimLevelType::SingletonNo) &&
+               isValidDLT(DimLevelType::SingletonNuNo)),
+              "isValidDLT definition is broken");
+
 static_assert((!isCompressedDLT(DimLevelType::Dense) &&
                isCompressedDLT(DimLevelType::Compressed) &&
                isCompressedDLT(DimLevelType::CompressedNu) &&
                isCompressedDLT(DimLevelType::CompressedNo) &&
                isCompressedDLT(DimLevelType::CompressedNuNo) &&
                !isCompressedDLT(DimLevelType::Singleton) &&
                !isCompressedDLT(DimLevelType::SingletonNu) &&
                !isCompressedDLT(DimLevelType::SingletonNo) &&

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 /// Convenience method to get a sparse encoding attribute from a type.
 /// Returns null-attribute for any type without an encoding.
 SparseTensorEncodingAttr getSparseTensorEncoding(Type type);
 
 //
 // Dimension level types.
 //
 
-// Cannot be constexpr, because `getRank` isn't constexpr.  However,
-// for some strange reason, the wrapper functions below don't trigger
-// the same [-Winvalid-constexpr] warning (despite this function not
-// being constexpr).
-inline DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {
-  assert(d < static_cast<uint64_t>(type.getRank()));
-  if (auto enc = getSparseTensorEncoding(type))
-    return enc.getDimLevelType()[d];
+// MSVC does not allow this to be constexpr, because
+// `SparseTensorEncodingAttr::operator bool` isn't declared constexpr.
+inline DimLevelType getDimLevelType(const SparseTensorEncodingAttr &enc,
+                                    uint64_t d) {
+  if (enc) {
+    auto types = enc.getDimLevelType();
+    assert(d < types.size() && "Dimension out of bounds");
+    return types[d];
+  }
   return DimLevelType::Dense; // unannotated tensor is dense
 }
 
+constexpr DimLevelType getDimLevelType(RankedTensorType type, uint64_t d) {
+  return getDimLevelType(getSparseTensorEncoding(type), d);
+}
+
 /// Convenience function to test for dense dimension (0 <= d < rank).
 constexpr bool isDenseDim(RankedTensorType type, uint64_t d) {
   return isDenseDLT(getDimLevelType(type, d));
 }
 
 /// Convenience function to test for compressed dimension (0 <= d < rank).
 constexpr bool isCompressedDim(RankedTensorType type, uint64_t d) {
   return isCompressedDLT(getDimLevelType(type, d));

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/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SparseTensor/IR/Enums.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 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.
   kAbsF,
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<DimLevelFormat>(
-                        l, DimLevelFormat(DimLvlType::kUndef))) {}
+        dimTypes(t + 1, std::vector<DimLevelType>(l, DimLevelType::Undef)) {}
 
   /// 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 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;
 
-  /// 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 sparse dimension level type.
   bool hasAnySparse(const BitVector &bits) const;
 
-  /// Dimension level format getter.
-  DimLevelFormat getDimLevelFormat(unsigned t, unsigned i) const {
+  /// Gets the dimension level type of the `i`th loop of the `t`th tensor.
+  DimLevelType getDimLevelType(unsigned t, unsigned i) const {
     assert(t < numTensors && i < numLoops);
-    return dims[t][i];
+    return dimTypes[t][i];
+  }
+
+  /// Gets the dimension level type of `b`.
+  DimLevelType getDimLevelType(unsigned b) const {
+    return getDimLevelType(tensor(b), index(b));
   }
 
-  /// Dimension level format setter.
-  void setDimLevelFormat(unsigned t, unsigned i, DimLevelFormat d) {
-    dims[t][i] = d;
+  /// Sets the dimension level type of the `i`th loop of the `t`th tensor.
+  void setDimLevelType(unsigned t, unsigned i, DimLevelType d) {
+    assert(isValidDLT(d));
+    dimTypes[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
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<DimLevelFormat>> dims;
+  std::vector<std::vector<DimLevelType>> dimTypes;
   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/IR/SparseTensorDialect.cpp

return parser.getChecked<SparseTensorEncodingAttr>(
       parser.getContext(), dlt, dimOrd, higherOrd, ptr, ind);
 }
 
 void SparseTensorEncodingAttr::print(AsmPrinter &printer) const {
   // Print the struct-like storage in dictionary fashion.
   printer << "<{ dimLevelType = [ ";
   for (unsigned i = 0, e = getDimLevelType().size(); i < e; i++) {
     switch (getDimLevelType()[i]) {
+    case DimLevelType::Undef:
+      // TODO: should probably raise an error instead of printing it...
+      printer << "\"undef\"";
+      break;
     case DimLevelType::Dense:
       printer << "\"dense\"";
       break;
     case DimLevelType::Compressed:
       printer << "\"compressed\"";
       break;
     case DimLevelType::CompressedNu:
       printer << "\"compressed-nu\"";

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

for (unsigned i = 0; i < sz; i++)
     inv[i] = m.getDimPosition(i);
   // Construct the permutation.
   SmallVector<unsigned, 4> perm(sz);
   for (unsigned i = 0; i < sz; i++)
     perm[i] = inv[topSort[i]];
   return AffineMap::getPermutationMap(perm, context);
 }
 
-/// 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) {
-    switch (enc.getDimLevelType()[d]) {
-    case DimLevelType::Dense:
-      return DimLevelFormat(DimLvlType::kDense);
-    case DimLevelType::Compressed:
-      return DimLevelFormat(DimLvlType::kCompressed);
-    case DimLevelType::CompressedNu:
-      return DimLevelFormat(DimLvlType::kCompressed, true, false);
-    case DimLevelType::CompressedNo:
-      return DimLevelFormat(DimLvlType::kCompressed, false, true);
-    case DimLevelType::CompressedNuNo:
-      return DimLevelFormat(DimLvlType::kCompressed, false, false);
-    case DimLevelType::Singleton:
-      return DimLevelFormat(DimLvlType::kSingleton);
-    case DimLevelType::SingletonNu:
-      return DimLevelFormat(DimLvlType::kSingleton, true, false);
-    case DimLevelType::SingletonNo:
-      return DimLevelFormat(DimLvlType::kSingleton, false, true);
-    case DimLevelType::SingletonNuNo:
-      return DimLevelFormat(DimLvlType::kSingleton, false, false);
-    }
-  }
-  return DimLevelFormat(DimLvlType::kDense);
-}
-
 /// Helper method to inspect affine expressions. Rejects cases where the
 /// 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) {
+                       DimLevelType dim) {
   switch (a.getKind()) {
   case AffineExprKind::DimId: {
     unsigned idx = a.cast<AffineDimExpr>().getPosition();
-    if (!merger.isDimLevelType(tensor, idx, DimLvlType::kUndef))
+    if (!isUndefDLT(merger.getDimLevelType(tensor, idx)))
       return false; // used more than once
-    merger.setDimLevelFormat(tensor, idx, dim);
+    merger.setDimLevelType(tensor, idx, dim);
     return true;
   }
   case AffineExprKind::Add:
   case AffineExprKind::Mul: {
-    if (dim.levelType != DimLvlType::kDense)
+    if (!isDenseDLT(dim))
       return false; // compound only in dense dim
     auto binOp = a.cast<AffineBinaryOpExpr>();
     return findAffine(merger, tensor, binOp.getLHS(), dim) &&
            findAffine(merger, tensor, binOp.getRHS(), dim);
   }
   case AffineExprKind::Constant:
-    return dim.levelType == DimLvlType::kDense; // const only in dense dim
+    return isDenseDLT(dim); // 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 admissible. Returns false if
 /// no annotations are found or inadmissible constructs occur.
 static bool findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
   bool annotated = false;
   for (OpOperand &t : op->getOpOperands()) {
     auto map = op.getMatchingIndexingMap(&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(toOrigDim(enc, d));
-      if (!findAffine(merger, tensor, a, toDimLevelFormat(enc, d)))
+      if (!findAffine(merger, tensor, a, getDimLevelType(enc, d)))
         return false; // inadmissible affine expression
     }
   }
   return annotated;
 }
 
 /// A helper to compute a topological sort. O(n^2) time complexity
 /// as we use adj matrix for the graph.
for (unsigned d = 1, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr t = map.getResult(toOrigDim(enc, d));
       addAffineOrderings(adjM, inDegree, 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.isDimLevelType(tensor, i, DimLvlType::kCompressed) ||
-            merger.isDimLevelType(tensor, i, DimLvlType::kSingleton)) {
+        if (isCompressedDLT(merger.getDimLevelType(tensor, i)) ||
+            isSingletonDLT(merger.getDimLevelType(tensor, i))) {
           for (unsigned j = 0; j < n; j++)
-            if (merger.isDimLevelType(tensor, j, DimLvlType::kUndef)) {
+            if (isUndefDLT(merger.getDimLevelType(tensor, j))) {
               adjM[i][j] = true;
               inDegree[j]++;
             }
         } else {
-          assert(merger.isDimLevelType(tensor, i, DimLvlType::kDense) ||
-                 merger.isDimLevelType(tensor, i, DimLvlType::kUndef));
+          assert(isDenseDLT(merger.getDimLevelType(tensor, i)) ||
+                 isUndefDLT(merger.getDimLevelType(tensor, i)));
         }
     }
   }
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
   topSort.clear();
   topSort.reserve(n);
   return topSortOptimal(n, iteratorTypes, topSort, inDegree, adjM);
static bool isAdmissibleTensorExp(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 admissible since insertions cannot occur.
   bool allDense = true;
   auto iteratorTypes = op.getIteratorTypesArray();
   unsigned numLoops = iteratorTypes.size();
   for (unsigned i = 0; i < numLoops; i++)
-    if (merger.isDimLevelType(tensor, i, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(tensor, i, DimLvlType::kSingleton)) {
+    if (isCompressedDLT(merger.getDimLevelType(tensor, i)) ||
+        isSingletonDLT(merger.getDimLevelType(tensor, i))) {
       allDense = false;
       break;
     } else {
-      assert(merger.isDimLevelType(tensor, i, DimLvlType::kDense) ||
-             merger.isDimLevelType(tensor, i, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(tensor, i)) ||
+             isUndefDLT(merger.getDimLevelType(tensor, i)));
     }
   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 admissible without special codegen.
   if (merger.isSingleCondition(tensor, exp))
     return true;
for (OpOperand &t : op->getOpOperands()) {
     // Scan all dimensions of current tensor.
     args.clear();
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr a = map.getResult(toOrigDim(enc, d));
       if (a.getKind() != AffineExprKind::DimId)
         continue; // compound
       unsigned idx = a.cast<AffineDimExpr>().getPosition();
       // Handle the different storage schemes.
-      if (merger.isDimLevelType(tensor, idx, DimLvlType::kCompressed)) {
+      if (isCompressedDLT(merger.getDimLevelType(tensor, idx))) {
         // Compressed dimension, fetch pointer and indices.
         auto ptrTp =
             MemRefType::get(dynShape, getPointerOverheadType(builder, enc));
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         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)) {
+      } else if (isSingletonDLT(merger.getDimLevelType(tensor, idx))) {
         // Singleton dimension, fetch indices.
         auto indTp =
             MemRefType::get(dynShape, getIndexOverheadType(builder, enc));
         auto dim = builder.getIndexAttr(d);
         codegen.indices[tensor][idx] =
             builder.create<ToIndicesOp>(loc, indTp, t.get(), dim);
       } else {
         // Dense dimension, nothing to fetch.
-        assert(merger.isDimLevelType(tensor, idx, DimLvlType::kDense));
+        assert(isDenseDLT(merger.getDimLevelType(tensor, idx)));
       }
       // Find upper bound in current dimension.
       unsigned p = toOrigDim(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;
static bool genInit(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   unsigned idx = topSort[at];
 
   // Initialize sparse positions.
   for (unsigned b = 0, be = inits.size(); b < be; b++) {
     if (!inits[b])
       continue;
     unsigned tensor = merger.tensor(b);
     assert(idx == merger.index(b));
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed)) {
+    if (isCompressedDLT(merger.getDimLevelType(b))) {
       // Initialize sparse index that will implement the iteration:
       //   for pidx_idx = pointers(pidx_idx-1), pointers(1+pidx_idx-1)
       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)) {
+    } else if (isSingletonDLT(merger.getDimLevelType(b))) {
       // Initialize sparse index that will implement the "iteration":
       //   for pidx_idx = pidx_idx-1, 1+pidx_idx-1
       // We rely on subsequent loop unrolling to get rid of the loop
       // if it is not involved in co-iteration with anything else.
       unsigned pat = at;
       for (; pat != 0; pat--) {
         if (codegen.pidxs[tensor][topSort[pat - 1]])
           break;
       }
       Value one = constantIndex(builder, loc, 1);
       Value p0 = (pat == 0) ? constantIndex(builder, loc, 0)
                             : codegen.pidxs[tensor][topSort[pat - 1]];
       codegen.pidxs[tensor][idx] = p0;
       codegen.highs[tensor][idx] = builder.create<arith::AddIOp>(loc, p0, one);
     } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(b)) ||
+             isUndefDLT(merger.getDimLevelType(b)));
       // Dense index still in play.
       needsUniv = true;
     }
   }
 
   // Initialize the universal dense index.
   codegen.loops[idx] = constantIndex(builder, loc, 0);
   return needsUniv;

 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.getIteratorTypesArray();
   bool isReduction = linalg::isReductionIterator(iteratorTypes[idx]);
-  bool isSparse = merger.isDimLevelType(fb, DimLvlType::kCompressed) ||
-                  merger.isDimLevelType(fb, DimLvlType::kSingleton);
+  bool isSparse = isCompressedDLT(merger.getDimLevelType(fb)) ||
+                  isSingletonDLT(merger.getDimLevelType(fb));
   bool isVector = isVectorFor(codegen, isInner, isReduction, isSparse) &&
                   denseUnitStrides(merger, op, idx);
   bool isParallel =
       isParallelFor(codegen, isOuter, isReduction, isSparse, isVector);
 
   // Prepare vector length.
   if (isVector)
     codegen.curVecLength = codegen.options.vectorLength;
static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                            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])
       continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+    if (isCompressedDLT(merger.getDimLevelType(b)) ||
+        isSingletonDLT(merger.getDimLevelType(b))) {
       unsigned tensor = merger.tensor(b);
       assert(idx == merger.index(b));
       types.push_back(indexType);
       operands.push_back(codegen.pidxs[tensor][idx]);
     } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(b)) ||
+             isUndefDLT(merger.getDimLevelType(b)));
     }
   }
   if (codegen.redVal) {
     types.push_back(codegen.redVal.getType());
     operands.push_back(codegen.redVal);
   }
   if (codegen.expValues) {
     types.push_back(indexType);
static Operation *genWhile(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   // 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])
       continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+    if (isCompressedDLT(merger.getDimLevelType(b)) ||
+        isSingletonDLT(merger.getDimLevelType(b))) {
       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++);
     } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(b)) ||
+             isUndefDLT(merger.getDimLevelType(b)));
     }
   }
   if (codegen.redVal)
     updateReduc(merger, codegen, after->getArgument(o++));
   if (codegen.expValues)
     codegen.expCount = after->getArgument(o++);
   if (needsUniv)
     codegen.loops[idx] = after->getArgument(o++);
static void genLocals(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   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])
       continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+    if (isCompressedDLT(merger.getDimLevelType(b)) ||
+        isSingletonDLT(merger.getDimLevelType(b))) {
       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) {
           Value cmp = builder.create<arith::CmpIOp>(
               loc, arith::CmpIPredicate::ult, load, min);
           min = builder.create<arith::SelectOp>(loc, cmp, load, min);
         } else {
           min = load;
         }
       }
     } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(b)) ||
+             isUndefDLT(merger.getDimLevelType(b)));
     }
   }
 
   // Merge dense universal index over minimum.
   if (min) {
     assert(!needsUniv);
     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.isDimLevelType(b, DimLvlType::kDense)) {
+        isDenseDLT(merger.getDimLevelType(b))) {
       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,
   // 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])
       continue;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+    if (isCompressedDLT(merger.getDimLevelType(b)) ||
+        isSingletonDLT(merger.getDimLevelType(b))) {
       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);
       operands.push_back(builder.create<arith::SelectOp>(loc, cmp, add, op3));
       codegen.pidxs[tensor][idx] = whileOp->getResult(o++);
     } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(b)) ||
+             isUndefDLT(merger.getDimLevelType(b)));
     }
   }
   if (codegen.redVal) {
     operands.push_back(codegen.redVal);
     updateReduc(merger, codegen, whileOp->getResult(o++));
   }
   if (codegen.expValues) {
     operands.push_back(codegen.expCount);
static scf::IfOp genIf(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   SmallVector<Type, 4> types;
   Value cond;
   for (unsigned b = 0, be = conditions.size(); b < be; b++) {
     if (!conditions[b])
       continue;
     unsigned tensor = merger.tensor(b);
     assert(idx == merger.index(b));
     Value clause;
-    if (merger.isDimLevelType(b, DimLvlType::kCompressed) ||
-        merger.isDimLevelType(b, DimLvlType::kSingleton)) {
+    if (isCompressedDLT(merger.getDimLevelType(b)) ||
+        isSingletonDLT(merger.getDimLevelType(b))) {
       Value op1 = codegen.idxs[tensor][idx];
       Value op2 = codegen.loops[idx];
       clause = builder.create<arith::CmpIOp>(loc, arith::CmpIPredicate::eq, op1,
                                              op2);
     } else {
-      assert(merger.isDimLevelType(b, DimLvlType::kDense) ||
-             merger.isDimLevelType(b, DimLvlType::kUndef));
+      assert(isDenseDLT(merger.getDimLevelType(b)) ||
+             isUndefDLT(merger.getDimLevelType(b)));
       clause = constantI1(builder, loc, true);
     }
     cond = cond ? builder.create<arith::AndIOp>(loc, cond, clause) : clause;
   }
   if (codegen.redVal)
     types.push_back(codegen.redVal.getType());
   if (codegen.expValues)
     types.push_back(builder.getIndexType());

mlir/lib/Dialect/SparseTensor/Utils/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorUtils
   Merger.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/SparseTensor
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRComplexDialect
   MLIRIR
   MLIRLinalgDialect
+  MLIRSparseTensorEnums
 )

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

BitVector Merger::simplifyCond(unsigned s0, unsigned p0) {
   BitVector simple = latPoints[p0].bits;
   bool reset = isSingleton && hasAnySparse(simple);
   unsigned be = simple.size();
   unsigned offset = 0; // relative to the end
   if (!reset)
     // Starts resetting from a dense dimension, so that the first bit (if kept)
     // is not undefined dimension type.
     for (unsigned b = 0; b < be; b++) {
-      if (simple[b] && isDimLevelType(b, DimLvlType::kDense)) {
+      if (simple[b] && isDenseDLT(getDimLevelType(b))) {
         offset = be - b - 1; // relative to the end
         break;
       }
     }
 
   // Now apply the two basic rules. We also iterate the bits reversely to always
   // keep the rightmost bit (which could possibly be a synthetic tensor).
   for (unsigned b = be - 1 - offset, i = 0; i < be;
        b = b == 0 ? be - 1 : b - 1, i++) {
-    if (simple[b] && (!isDimLevelType(b, DimLvlType::kCompressed) &&
-                      !isDimLevelType(b, DimLvlType::kSingleton))) {
+    if (simple[b] && (!isCompressedDLT(getDimLevelType(b)) &&
+                      !isSingletonDLT(getDimLevelType(b)))) {
       if (reset)
         simple.reset(b);
       reset = true;
     }
   }
   return simple;
 }
 
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   case kReduce:
     return false;
   }
   llvm_unreachable("unexpected kind");
 }
 
 bool Merger::hasAnySparse(const BitVector &bits) const {
   for (unsigned b = 0, be = bits.size(); b < be; b++)
-    if (bits[b] && (isDimLevelType(b, DimLvlType::kCompressed) ||
-                    isDimLevelType(b, DimLvlType::kSingleton)))
+    if (bits[b] && (isCompressedDLT(getDimLevelType(b)) ||
+                    isSingletonDLT(getDimLevelType(b))))
       return true;
   return false;
 }
 
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Print methods (for debugging).
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];
+      DimLevelType dlt = dimTypes[t][i];
       llvm::dbgs() << " i_" << t << "_" << i << "_";
-      switch (f.levelType) {
-      case DimLvlType::kDense:
+      if (isDenseDLT(dlt))
         llvm::dbgs() << "D";
-        break;
-      case DimLvlType::kCompressed:
+      else if (isCompressedDLT(dlt))
         llvm::dbgs() << "C";
-        break;
-      case DimLvlType::kSingleton:
+      else if (isSingletonDLT(dlt))
         llvm::dbgs() << "S";
-        break;
-      case DimLvlType::kUndef:
+      else if (isUndefDLT(dlt))
         llvm::dbgs() << "U";
-        break;
-      }
-      llvm::dbgs() << "[O=" << f.isOrdered << ",U=" << f.isUnique << "]";
+      llvm::dbgs() << "[O=" << isOrderedDLT(dlt) << ",U=" << isUniqueDLT(dlt)
+                   << "]";
     }
   }
 }
 
 #endif // NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Builder methods.

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.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t0, l0, DimLevelType::Compressed);
 
     // Tensor 1: sparse input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t1, l0, DimLevelType::Compressed);
 
     // Tensor 2: dense output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kDense));
+    merger.setDimLevelType(t2, l0, DimLevelType::Dense);
   }
 };
 
 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.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t0, l0, DimLevelType::Compressed);
 
     // Tensor 1: sparse input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t1, l0, DimLevelType::Compressed);
 
     // Tensor 2: sparse input vector
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t2, l0, DimLevelType::Compressed);
 
     // Tensor 3: dense output vector
     merger.addExp(Kind::kTensor, t3, -1u);
-    merger.setDimLevelFormat(t3, l0, DimLevelFormat(DimLvlType::kDense));
+    merger.setDimLevelType(t3, l0, DimLevelType::Dense);
   }
 };
 
 ///
 /// 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.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t0, l0, DimLevelType::Compressed);
 
     // Tensor 1: dense input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kDense));
+    merger.setDimLevelType(t1, l0, DimLevelType::Dense);
 
     // Tensor 2: dense output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kDense));
+    merger.setDimLevelType(t2, l0, DimLevelType::Dense);
   }
 };
 
 ///
 /// Tests with both undef and dense input.
 ///
 
 class MergerTest4T1LU : public MergerTestBase {
 protected:
   // Our three tensors (two inputs, one output).
   const unsigned t0 = 0, t1 = 1, t2 = 2, t3 = 3;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest4T1LU() : MergerTestBase(4, 1) {
     // Tensor 0: undef input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
-    merger.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kUndef));
+    merger.setDimLevelType(t0, l0, DimLevelType::Undef);
 
     // Tensor 1: dense input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kDense));
+    merger.setDimLevelType(t1, l0, DimLevelType::Dense);
 
     // Tensor 2: undef input vector.
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kUndef));
+    merger.setDimLevelType(t2, l0, DimLevelType::Undef);
 
     // Tensor 3: dense output vector.
     merger.addExp(Kind::kTensor, t3, -1u);
-    merger.setDimLevelFormat(t3, l0, DimLevelFormat(DimLvlType::kDense));
+    merger.setDimLevelType(t3, l0, DimLevelType::Dense);
   }
 };
 
 ///
 /// Tests with operation on sparse output.
 ///
 
 class MergerTest3T1L_SO : public MergerTestBase {
 protected:
   // Our three tensors (two inputs, one output, one synthetic).
   const unsigned t0 = 0, t1 = 1, t2 = 2, t3 = 3;
 
   // Our single loop.
   const unsigned l0 = 0;
 
   MergerTest3T1L_SO() : MergerTestBase(3, 1) {
     merger.setHasSparseOut(true);
 
     // Tensor 0: undef input vector.
     merger.addExp(Kind::kTensor, t0, -1u);
-    merger.setDimLevelFormat(t0, l0, DimLevelFormat(DimLvlType::kUndef));
+    merger.setDimLevelType(t0, l0, DimLevelType::Undef);
 
     // Tensor 1: undef input vector.
     merger.addExp(Kind::kTensor, t1, -1u);
-    merger.setDimLevelFormat(t1, l0, DimLevelFormat(DimLvlType::kUndef));
+    merger.setDimLevelType(t1, l0, DimLevelType::Undef);
 
     // Tensor 2: sparse output vector.
     merger.addExp(Kind::kTensor, t2, -1u);
-    merger.setDimLevelFormat(t2, l0, DimLevelFormat(DimLvlType::kCompressed));
+    merger.setDimLevelType(t2, l0, DimLevelType::Compressed);
   }
 };
 
 } // namespace
 
 /// Vector multiplication (conjunction) of 3 vectors, i.e.;
 ///   a(i) = b(i) * c(i) * d(i)
 /// which should form the single lattice point

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

cc_library(
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":ComplexDialect",
         ":IR",
         ":LinalgDialect",
         ":MathDialect",
         ":SparseTensorDialect",
+        ":SparseTensorEnums",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "SparseTensorTransforms",
     srcs = glob([
         "lib/Dialect/SparseTensor/Transforms/*.cpp",