[mlir][sparse] Making SortMask into an enum-class

[wrengr]

Summary

This helps to reduce the confusion from using unsigned everywhere.

Depends On D145606

[Peiming]

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

549

I probably need to rethink about how it works with includeDenseInput/Output differently, but this is irrelevant to this patch.

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

 
 //===----------------------------------------------------------------------===//
 // Declarations
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 /// Iteration graph sorting.
-enum SortMask {
+enum class SortMask : unsigned {
   // The individual mask bits.
   kIncludeDenseOutput = 0x1, // b001
   kIncludeDenseInput = 0x2,  // b010
   kIncludeUndef = 0x4,       // b100
   // The subsets of mask bits.
   kIncludeAll = 0x7,   // b111
   kIncludeDense = 0x3, // b011
   kSparseOnly = 0x0,   // b000
 };
 
-/// SortMask tests on individual bits.
-inline static bool includeDenseInput(unsigned mask) {
-  return mask & SortMask::kIncludeDenseInput;
+inline static bool includesAny(SortMask mask1, SortMask mask2) {
+  return static_cast<unsigned>(mask1) & static_cast<unsigned>(mask2);
 }
 
-inline static bool includeDenseOutput(unsigned mask) {
-  return mask & SortMask::kIncludeDenseOutput;
+inline static bool includesDenseInput(SortMask mask) {
+  return includesAny(mask, SortMask::kIncludeDenseInput);
 }
 
-inline static bool includeUndef(unsigned mask) {
-  return mask & SortMask::kIncludeUndef;
+inline static bool includesDenseOutput(SortMask mask) {
+  return includesAny(mask, SortMask::kIncludeDenseOutput);
+}
+
+inline static bool includesDense(SortMask mask) {
+  return includesAny(mask, SortMask::kIncludeDense);
+}
+
+inline static bool includesUndef(SortMask mask) {
+  return includesAny(mask, SortMask::kIncludeUndef);
 }
 
 /// A helper class that visits an affine expression and tries to find an
 /// AffineDimExpr to which the corresponding iterator from a GenericOp matches
 /// the desired iterator type.
 class AffineDimFinder : public AffineExprVisitor<AffineDimFinder> {
 public:
   explicit AffineDimFinder(linalg::GenericOp op)

 }
 
 /// Computes a topologically sorted iteration graph for the linalg operation.
 /// Ensures all tensors are visited in natural coordinate order.  This is
 /// essential for sparse storage formats since these only support access
 /// along fixed levels.  Even for dense storage formats, however, the natural
 /// coordinate order yields innermost unit-stride access with better spatial
 /// locality.
-static bool computeIterationGraph(CodegenEnv &env, unsigned mask,
+static bool computeIterationGraph(CodegenEnv &env, SortMask mask,
                                   OpOperand *skip = nullptr) {
   // Set up an n x n from/to adjacency matrix of the iteration graph
   // for the implicit loop indices i_0 .. i_n-1.
   const unsigned n = env.merger().getNumLoops();
   std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
   std::vector<unsigned> inDegree(n, 0); // in-degree of each node.
   const auto iteratorTypes = env.op().getIteratorTypesArray();
   // Iterate over the indexing maps of every tensor in the tensor expression.
   for (OpOperand &t : env.op()->getOpOperands()) {
     // Get map and encoding.
     const auto map = env.op().getMatchingIndexingMap(&t);
     const auto enc = getSparseTensorEncoding(t.get().getType());
     assert(map.getNumDims() + getNumCompoundAffineOnSparseDims(env.op()) == n);
 
-    bool isDenseInput = !enc && env.op().isDpsInput(&t);
-    bool isDenseOutput = !enc && !isDenseInput;
-
     // Skips dense inputs/outputs when not requested.
-    if ((isDenseInput && !includeDenseInput(mask)) ||
-        (isDenseOutput && !includeDenseOutput(mask)))
+    const bool isDenseInput = !enc && env.op().isDpsInput(&t);
+    const bool isDenseOutput = !enc && !isDenseInput;
+    if ((isDenseInput && !includesDenseInput(mask)) ||
+        (isDenseOutput && !includesDenseOutput(mask)))
       continue;
 
     // Push unrelated loops into sparse iteration space, so these
     // will be skipped more often.
     // TODO: Do we really need this?
-    if (includeUndef(mask)) {
+    if (includesUndef(mask)) {
       unsigned tensor = t.getOperandNumber();
       for (unsigned i = 0; i < n; i++) {
         if (isCompressedDLT(env.dlt(tensor, i)) ||
             isSingletonDLT(env.dlt(tensor, i))) {
           for (unsigned j = 0; j < n; j++)
             if (isUndefDLT(env.dlt(tensor, j))) {
               adjM[i][j] = true;
               inDegree[j]++;
for (Level l = 0; l < lvlRank; l++) {
         std::optional<unsigned> fldx =
             env.merger().getLoopIdx(t.getOperandNumber(), l - 1);
 
         // Applying order constraints on every pair of dimExpr between two
         // compound affine expressions can sometime too strict:
         // E.g, for [dense, dense] -> (d0 + d1, d2 + d3).
         // It is totally fine to have loop sequence d0->d2->d1->d3 instead of
         // requiring d0 < d2, d1 < d2, d0 < d3, d1 < d3.
-        if (!(mask & SortMask::kIncludeDense))
+        if (!includesDense(mask))
           tryLoosenAffineDenseConstraints(env.op(), fldx, fa, tldx, ta);
 
         // (d0 + d1) < (d2 + d3), or
         // filter_loop_d-1 < (d2 + d3), or
         // (d0 + d1) < filter_loop_d, or
         // filter_loop_d-1 < filter_loop_d depending on whether fa/ta is reset
         // above.
         addAffineOrderings(adjM, inDegree, fa, ta, fldx, tldx);