[mlir][sparse] refactor handling of merger leafs and ops

[aartbik]

Summary

Using "default:" in the switch statemements that handle all our
merger ops has become a bit cumbersome since it is easy to overlook
parts of the code that need to handle ops specifically. By enforcing
full switch statements without "default:", we get a compiler warning
when cases are overlooked.

[wrengr]

I like this change a lot, since default cases are too fragile when there are so many cases to consider.

I do wonder though if it might be time to start refactoring the Kind enum to better handle the natural groupings of things. In particular I'm thinking it should be treated like a "bit enum" rather than an "int enum" (e.g., so that the I/F/C/etc suffix is encoded into, say, the low-order bits, whether something is leaf/unary/binary is encoded in the high-order bits, and then the op itself in the middle bits). Of course that should be done in a separate CL

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

1101–1102

Why is the template keyword needed here when we don't need it for any of the other .cast<T>() method calls?

[aartbik]

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

1101–1102

Good question, I am not sure how this was introduced, but we can do without. I removed all occurrences in the code.

[aartbik]

removed template keywords

[aartbik]

In D127263#3567584, @wrengr wrote:

I do wonder though if it might be time to start refactoring the Kind enum to better handle the natural groupings of things.

Agreed. A bit more "horizontal" encoding of type will simplify the long switch statements a bit, and get an intuitive way of getting the type from an op. But yeah, seems follow up CL.

[NathanielMcVicar]

Removing the default: on the switch in isSingleCondition broke the Windows bot, but unfortunately it was down so there was no timely notification. Sorry about that!
https://lab.llvm.org/buildbot/#/builders/13/builds/21881

FAILED: tools/mlir/lib/Dialect/SparseTensor/Utils/CMakeFiles/obj.MLIRSparseTensorUtils.dir/Merger.cpp.obj 
C:\PROGRA~2\MICROS~3\2019\COMMUN~1\VC\Tools\MSVC\1429~1.301\bin\Hostx64\x64\cl.exe  ... /Fotools\mlir\lib\Dialect\SparseTensor\Utils\CMakeFiles\obj.MLIRSparseTensorUtils.dir\Merger.cpp.obj /Fdtools\mlir\lib\Dialect\SparseTensor\Utils\CMakeFiles\obj.MLIRSparseTensorUtils.dir\ /FS -c C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Utils\Merger.cpp
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Utils\Merger.cpp(380) : error C2220: the following warning is treated as an error
C:\buildbot\mlir-x64-windows-ninja\llvm-project\mlir\lib\Dialect\SparseTensor\Utils\Merger.cpp(380) : warning C4715: 'mlir::sparse_tensor::Merger::isSingleCondition': not all control paths return a value

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

 
 //===----------------------------------------------------------------------===//
 // Constructors.
 //===----------------------------------------------------------------------===//
 
 TensorExp::TensorExp(Kind k, unsigned x, unsigned y, Value v, Operation *o)
     : kind(k), val(v), op(o) {
   switch (kind) {
+  // Leaf.
   case kTensor:
     assert(x != -1u && y == -1u && !v && !o);
     tensor = x;
     break;
   case kInvariant:
     assert(x == -1u && y == -1u && v && !o);
     break;
   case kIndex:
     assert(x != -1u && y == -1u && !v && !o);
     index = x;
     break;
+  // Unary operations.
   case kAbsF:
   case kAbsC:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
case kBinaryBranch:
     break;
   case kUnary:
     // No assertion on y can be made, as the branching paths involve both
     // a unary (mapSet) and binary (takeDisj) pathway.
     assert(x != -1u && !v && o);
     children.e0 = x;
     children.e1 = y;
     break;
-  case kBinary:
-    assert(x != -1u && y != -1u && !v && o);
+  // Binary operations.
+  case kMulF:
+  case kMulC:
+  case kMulI:
+  case kDivF:
+  case kDivC:
+  case kDivS:
+  case kDivU:
+  case kAddF:
+  case kAddC:
+  case kAddI:
+  case kSubF:
+  case kSubC:
+  case kSubI:
+  case kAndI:
+  case kOrI:
+  case kXorI:
+  case kShrS:
+  case kShrU:
+  case kShlI:
+    assert(x != -1u && y != -1u && !v && !o);
     children.e0 = x;
     children.e1 = y;
     break;
-  default:
-    assert(x != -1u && y != -1u && !v && !o);
+  case kBinary:
+    assert(x != -1u && y != -1u && !v && o);
     children.e0 = x;
     children.e1 = y;
     break;
   }
 }
 
 LatPoint::LatPoint(unsigned n, unsigned e, unsigned b)
     : bits(n, false), simple(), exp(e) {
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;
 }
 
 bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   switch (tensorExps[e].kind) {
+  // Leaf.
   case kTensor:
     return tensorExps[e].tensor == t;
+  case kInvariant:
+  case kIndex:
+    return false;
+  // Unary operations.
   case kAbsF:
   case kAbsC:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
bool Merger::isSingleCondition(unsigned t, unsigned e) const {
   case kCastS:
   case kCastU:
   case kCastIdx:
   case kTruncI:
   case kCIm:
   case kCRe:
   case kBitCast:
     return isSingleCondition(t, tensorExps[e].children.e0);
+  case kBinaryBranch:
+  case kUnary:
+    return false;
+  // Binary operations.
   case kDivF: // note: x / c only
   case kDivC:
   case kDivS:
   case kDivU:
     assert(!maybeZero(tensorExps[e].children.e1));
     return isSingleCondition(t, tensorExps[e].children.e0);
   case kShrS: // note: x >> inv only
   case kShrU:
case kAndI:
     if (isSingleCondition(t, tensorExps[e].children.e1))
       return isInvariant(tensorExps[e].children.e0);
     return false;
   case kAddF:
   case kAddC:
   case kAddI:
     return isSingleCondition(t, tensorExps[e].children.e0) &&
            isSingleCondition(t, tensorExps[e].children.e1);
-  default:
+  case kSubF:
+  case kSubC:
+  case kSubI:
+  case kOrI:
+  case kXorI:
+  case kBinary:
     return false;
   }
 }
 
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
 // Print methods (for debugging).
 //===----------------------------------------------------------------------===//
 
 static const char *kindToOpSymbol(Kind kind) {
   switch (kind) {
+  // Leaf.
   case kTensor:
     return "tensor";
   case kInvariant:
     return "invariant";
   case kIndex:
     return "index";
+  // Unary operations.
   case kAbsF:
   case kAbsC:
     return "abs";
   case kCeilF:
     return "ceil";
   case kFloorF:
     return "floor";
   case kSqrtF:
static const char *kindToOpSymbol(Kind kind) {
   case kCRe:
     return "complex.re";
   case kBitCast:
     return "cast";
   case kBinaryBranch:
     return "binary_branch";
   case kUnary:
     return "unary";
+  // Binary operations.
   case kMulF:
   case kMulC:
   case kMulI:
     return "*";
   case kDivF:
   case kDivC:
   case kDivS:
   case kDivU:
static const char *kindToOpSymbol(Kind kind) {
   case kBinary:
     return "binary";
   }
   llvm_unreachable("unexpected kind for symbol");
 }
 
 void Merger::dumpExp(unsigned e) const {
   switch (tensorExps[e].kind) {
+  // Leaf.
   case kTensor:
     if (tensorExps[e].tensor == syntheticTensor)
       llvm::dbgs() << "synthetic_";
     else if (tensorExps[e].tensor == outTensor)
       llvm::dbgs() << "output_";
     llvm::dbgs() << "tensor_" << tensorExps[e].tensor;
     break;
   case kInvariant:
     llvm::dbgs() << "invariant";
     break;
   case kIndex:
     llvm::dbgs() << "index_" << tensorExps[e].index;
     break;
+  // Unary operations.
   case kAbsF:
+  case kAbsC:
   case kCeilF:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
   case kLog1pF:
+  case kLog1pC:
   case kSinF:
+  case kSinC:
   case kTanhF:
   case kTanhC:
   case kNegF:
+  case kNegC:
   case kNegI:
   case kTruncF:
   case kExtF:
   case kCastFS:
   case kCastFU:
   case kCastSF:
   case kCastUF:
   case kCastS:
   case kCastU:
   case kCastIdx:
   case kTruncI:
+  case kCIm:
+  case kCRe:
   case kBitCast:
+  case kBinaryBranch:
+  case kUnary:
     llvm::dbgs() << kindToOpSymbol(tensorExps[e].kind) << " ";
     dumpExp(tensorExps[e].children.e0);
     break;
-  default:
+  // Binary operations.
+  case kMulF:
+  case kMulC:
+  case kMulI:
+  case kDivF:
+  case kDivC:
+  case kDivS:
+  case kDivU:
+  case kAddF:
+  case kAddC:
+  case kAddI:
+  case kSubF:
+  case kSubC:
+  case kSubI:
+  case kAndI:
+  case kOrI:
+  case kXorI:
+  case kShrS:
+  case kShrU:
+  case kShlI:
+  case kBinary:
     llvm::dbgs() << "(";
     dumpExp(tensorExps[e].children.e0);
     llvm::dbgs() << " " << kindToOpSymbol(tensorExps[e].kind) << " ";
     dumpExp(tensorExps[e].children.e1);
     llvm::dbgs() << ")";
   }
 }
 

 
 //===----------------------------------------------------------------------===//
 // Builder methods.
 //===----------------------------------------------------------------------===//
 
 unsigned Merger::buildLattices(unsigned e, unsigned i) {
   Kind kind = tensorExps[e].kind;
   switch (kind) {
+  // Leaf.
   case kTensor:
   case kInvariant:
   case kIndex: {
     // Either the index is really used in the tensor expression, or it is
     // set to the undefined index in that dimension. An invariant expression,
     // a proper index value, and a truly dynamic sparse output tensor are set
     // to a synthetic tensor with undefined indices only to ensure the
     // iteration space is not skipped as a result of their contents.
     unsigned s = addSet();
     unsigned t = syntheticTensor;
     if (kind == kTensor) {
       t = tensorExps[e].tensor;
       if (hasSparseOut && t == outTensor)
         t = syntheticTensor;
     }
     latSets[s].push_back(addLat(t, i, e));
     return s;
   }
+  // Unary operations.
   case kAbsF:
   case kAbsC:
   case kCeilF:
-  case kCIm:
-  case kCRe:
   case kFloorF:
   case kSqrtF:
   case kSqrtC:
   case kExpm1F:
   case kExpm1C:
   case kLog1pF:
   case kLog1pC:
   case kSinF:
   case kSinC:
   case kTanhF:
   case kTanhC:
   case kNegF:
   case kNegC:
   case kNegI:
   case kTruncF:
   case kExtF:
   case kCastFS:
   case kCastFU:
   case kCastSF:
   case kCastUF:
   case kCastS:
   case kCastU:
   case kCastIdx:
   case kTruncI:
+  case kCIm:
+  case kCRe:
   case kBitCast:
     // A zero preserving operation (viz. f(0) = 0, [Bik96,Ch5]) maps the
     // lattice set of the operand through the operator into a new set.
     //
     //  -y|!y | y |
     //  --+---+---+
     //    | 0 |-y |
     return mapSet(kind, buildLattices(tensorExps[e].children.e0, i),
//      | absent() | present(y) |
       } // Use a disjunction with `unop` on the left and the absent value as an
       // invariant on the right.
       Block &absentBlock = absentRegion.front();
       YieldOp absentYield = cast<YieldOp>(absentBlock.getTerminator());
       Value absentVal = absentYield.result();
       unsigned rhs = addExp(kInvariant, absentVal);
       return takeDisj(kind, child0, buildLattices(rhs, i), unop);
     }
+  // Binary operations.
   case kMulF:
   case kMulC:
   case kMulI:
   case kAndI:
     // A multiplicative operation only needs to be performed
     // for the conjunction of sparse iteration spaces.
     //
     //  x*y|!y | y |
if (overlapRegion.empty())
     // output.
     return Value();
   return insertYieldOp(rewriter, loc, overlapRegion, {v0, v1});
 }
 
 Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
                        Value v0, Value v1) {
   switch (tensorExps[e].kind) {
+  // Leaf.
   case kTensor:
   case kInvariant:
   case kIndex:
     llvm_unreachable("unexpected non-op");
-  // Unary ops.
+  // Unary operations.
   case kAbsF:
     return rewriter.create<math::AbsOp>(loc, v0);
   case kAbsC: {
-    auto type = v0.getType().template cast<ComplexType>();
-    auto eltType = type.getElementType().template cast<FloatType>();
+    auto type = v0.getType().cast<ComplexType>();
+    auto eltType = type.getElementType().cast<FloatType>();
     return rewriter.create<complex::AbsOp>(loc, eltType, v0);
   }
   case kCeilF:
     return rewriter.create<math::CeilOp>(loc, v0);
   case kFloorF:
     return rewriter.create<math::FloorOp>(loc, v0);
   case kSqrtF:
     return rewriter.create<math::SqrtOp>(loc, v0);
Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
   case kCastS:
     return rewriter.create<arith::ExtSIOp>(loc, inferType(e, v0), v0);
   case kCastU:
     return rewriter.create<arith::ExtUIOp>(loc, inferType(e, v0), v0);
   case kCastIdx:
     return rewriter.create<arith::IndexCastOp>(loc, inferType(e, v0), v0);
   case kTruncI:
     return rewriter.create<arith::TruncIOp>(loc, inferType(e, v0), v0);
-  case kCIm:
-  case kCRe: {
-    auto type = v0.getType().template cast<ComplexType>();
-    auto eltType = type.getElementType().template cast<FloatType>();
-    if (tensorExps[e].kind == kCIm)
-      return rewriter.create<complex::ImOp>(loc, eltType, v0);
-
+  case kCIm: {
+    auto type = v0.getType().cast<ComplexType>();
+    auto eltType = type.getElementType().cast<FloatType>();
+  }
+  case kCRe: {
+    auto type = v0.getType().cast<ComplexType>();
+    auto eltType = type.getElementType().cast<FloatType>();
     return rewriter.create<complex::ReOp>(loc, eltType, v0);
   }
   case kBitCast:
     return rewriter.create<arith::BitcastOp>(loc, inferType(e, v0), v0);
-  // Binary ops.
+  // Binary operations.
   case kMulF:
     return rewriter.create<arith::MulFOp>(loc, v0, v1);
   case kMulC:
     return rewriter.create<complex::MulOp>(loc, v0, v1);
   case kMulI:
     return rewriter.create<arith::MulIOp>(loc, v0, v1);
   case kDivF:
     return rewriter.create<arith::DivFOp>(loc, v0, v1);
Value Merger::buildExp(RewriterBase &rewriter, Location loc, unsigned e,
   case kXorI:
     return rewriter.create<arith::XOrIOp>(loc, v0, v1);
   case kShrS:
     return rewriter.create<arith::ShRSIOp>(loc, v0, v1);
   case kShrU:
     return rewriter.create<arith::ShRUIOp>(loc, v0, v1);
   case kShlI:
     return rewriter.create<arith::ShLIOp>(loc, v0, v1);
-  // Semiring ops with custom logic.
-  case kBinaryBranch:
+  case kBinaryBranch: // semi-ring ops with custom logic.
     return insertYieldOp(rewriter, loc,
                          *tensorExps[e].op->getBlock()->getParent(), {v0});
   case kUnary:
     return buildUnaryPresent(rewriter, loc, tensorExps[e].op, v0);
   case kBinary:
     return buildBinaryOverlap(rewriter, loc, tensorExps[e].op, v0, v1);
   }
   llvm_unreachable("unexpected expression kind in build");
 }
 
 } // namespace sparse_tensor
 } // namespace mlir

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

protected:
   }
 
   /// Check that there are n lattice points in set s.
   void expectNumLatPoints(unsigned s, unsigned n) {
     EXPECT_THAT(merger.set(s).size(), n);
   }
 
   /// Compares expressions for equality. Equality is defined recursively as:
-  /// - Two expressions can only be equal if they have the same Kind.
-  /// - Two binary expressions are equal if they have the same Kind and their
-  ///     children are equal.
-  /// - Expressions with Kind invariant or tensor are equal if they have the
-  ///     same expression id.
+  /// - Operations are equal if they have the same kind and children.
+  /// - Leaf tensors are equal if they refer to the same tensor.
   bool compareExpression(unsigned e, const std::shared_ptr<Pattern> &pattern) {
     auto tensorExp = merger.exp(e);
     if (tensorExp.kind != pattern->kind)
       return false;
-    assert(tensorExp.kind != Kind::kInvariant &&
-           "Invariant comparison not yet supported");
     switch (tensorExp.kind) {
-    case Kind::kTensor:
+    // Leaf.
+    case kTensor:
       return tensorExp.tensor == pattern->tensorNum;
-    case Kind::kAbsF:
-    case Kind::kCeilF:
-    case Kind::kFloorF:
-    case Kind::kNegF:
-    case Kind::kNegI:
+    case kInvariant:
+    case kIndex:
+      llvm_unreachable("invariant not handled yet");
+    // Unary operations.
+    case kAbsF:
+    case kAbsC:
+    case kCeilF:
+    case kFloorF:
+    case kSqrtF:
+    case kSqrtC:
+    case kExpm1F:
+    case kExpm1C:
+    case kLog1pF:
+    case kLog1pC:
+    case kSinF:
+    case kSinC:
+    case kTanhF:
+    case kTanhC:
+    case kNegF:
+    case kNegC:
+    case kNegI:
+    case kTruncF:
+    case kExtF:
+    case kCastFS:
+    case kCastFU:
+    case kCastSF:
+    case kCastUF:
+    case kCastS:
+    case kCastU:
+    case kCastIdx:
+    case kTruncI:
+    case kCIm:
+    case kCRe:
+    case kBitCast:
+    case kBinaryBranch:
+    case kUnary:
+    case kShlI:
+    case kBinary:
       return compareExpression(tensorExp.children.e0, pattern->e0);
-    case Kind::kMulF:
-    case Kind::kMulI:
-    case Kind::kDivF:
-    case Kind::kDivS:
-    case Kind::kDivU:
-    case Kind::kAddF:
-    case Kind::kAddI:
-    case Kind::kSubF:
-    case Kind::kSubI:
-    case Kind::kAndI:
-    case Kind::kOrI:
-    case Kind::kXorI:
+    // Binary operations.
+    case kMulF:
+    case kMulC:
+    case kMulI:
+    case kDivF:
+    case kDivC:
+    case kDivS:
+    case kDivU:
+    case kAddF:
+    case kAddC:
+    case kAddI:
+    case kSubF:
+    case kSubC:
+    case kSubI:
+    case kAndI:
+    case kOrI:
+    case kXorI:
+    case kShrS:
+    case kShrU:
       return compareExpression(tensorExp.children.e0, pattern->e0) &&
              compareExpression(tensorExp.children.e1, pattern->e1);
-    default:
-      llvm_unreachable("Unhandled Kind");
     }
+    llvm_unreachable("unexpected kind");
   }
 
   unsigned numTensors;
   unsigned numLoops;
   Merger merger;
 };
 
 class MergerTest3T1L : public MergerTestBase {