Refactor TensorExp parameters into a union

[gussmith23]

Summary

To make TensorExp clearer, this change refactors the e0/e1 fields into a union: e0/e1 for a binary op tensor expression, and tensor_num for a tensor-kinded tensor expression.

[aartbik]

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

55–61

can we make this an anonymous struct, defined inline, so that at least we don't need the xxx.params.yyy everywhere?

[aartbik]

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

32

also, just call this tensor

(the _ in the name is against our style, but the use of just tensor is also consistent with comments and code elsewhere)

[gussmith23]

Make union anonymous; rename tensor_num->tensor

[gussmith23]

Make struct anonymous

[aartbik]

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

55–61

I would remove this comment completely, since that is true for all fields of this struct

You have adequate comments inside the union

[gussmith23]

Remove comment

[gussmith23]

Rebase

[gussmith23]

Fix clang-tidy

[gussmith23]

Again, attempt to fix clang-tidy

[aartbik]

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

55–61

assuming clang tidy accepts this, this will need to move up in the struct (or perhaps even outside for readability), since within classes/structs, typedefs should go first

[gussmith23]

Move struct outside

[gussmith23]

Rebase

[gussmith23]

Fix after rebase

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

 namespace sparse_tensor {
 
 /// Tensor expression kind.
 enum class Kind { kTensor, kInvariant, kMulF, kMulI, kAddF, kAddI };
 
 /// Dimension level type for a tensor (undef means index does not appear).
 enum class Dim { kSparse, kDense, kSingle, kUndef };
 
+/// Children expressions of a binary TensorExp.
+struct Children {
+  unsigned e0;
+  unsigned e1;
+};
+
 /// Tensor expression. Represents a MLIR expression in tensor index notation.
 /// For tensors, e0 denotes the tensor index. For invariants, the IR value is
 /// stored directly. For binary operations, e0 and e1 denote the index of the
 /// children tensor expressions.
 struct TensorExp {
-  TensorExp(Kind k, unsigned x, unsigned y, Value v)
-      : kind(k), e0(x), e1(y), val(v) {
-    assert((kind == Kind::kTensor && e0 != -1u && e1 == -1u && !val) ||
-           (kind == Kind::kInvariant && e0 == -1u && e1 == -1u && val) ||
-           (kind >= Kind::kMulF && e0 != -1u && e1 != -1u && !val));
+  TensorExp(Kind k, unsigned x, unsigned y, Value v) : kind(k), val(v) {
+    assert((kind == Kind::kTensor && x != -1u && y == -1u && !val) ||
+           (kind == Kind::kInvariant && x == -1u && y == -1u && val) ||
+           (kind >= Kind::kMulF && x != -1u && y != -1u && !val));
+    if (kind == Kind::kTensor) {
+      tensor = x;
+    } else if (kind >= Kind::kMulF) {
+      children.e0 = x;
+      children.e1 = y;
+    }
   }
 
   /// Tensor expression kind.
   Kind kind;
 
-  /// Indices of children expression(s).
-  unsigned e0;
-  unsigned e1;
+  union {
+    /// Expressions representing tensors simply have a tensor number.
+    unsigned tensor;
+
+    /// Binary operations hold the indices of their child expressions.
+    Children children;
+  };
 
   /// Direct link to IR for an invariant. During code generation,
   /// field is used to cache "hoisted" loop invariant tensor loads.
   Value val;
 };
 
 /// Lattice point. Each lattice point consists of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding

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

 }
 
 /// Returns true if given tensor co-iterates with conjunction only.
 /// For the output tensor, this defines a "simply dynamic" operation.
 /// For instance: A(I) = A(I) * B(I) * C(I)
 static unsigned isConjunction(Merger &merger, unsigned tensor, unsigned exp) {
   switch (merger.exp(exp).kind) {
   case Kind::kTensor:
-    return merger.exp(exp).e0 == tensor;
+    return merger.exp(exp).tensor == tensor;
   case Kind::kMulF:
   case Kind::kMulI:
-    return isConjunction(merger, tensor, merger.exp(exp).e0) ||
-           isConjunction(merger, tensor, merger.exp(exp).e1);
+    return isConjunction(merger, tensor, merger.exp(exp).children.e0) ||
+           isConjunction(merger, tensor, merger.exp(exp).children.e1);
   default:
     return false;
   }
 }
 
 /// Returns true when the tensor expression is admissable for codegen.
 /// Since all sparse input tensors are admissable, we just need to check
 /// whether the output tensor in the tensor expression codegen is admissable.
static Value genTensorLoad(Merger &merger, CodeGen &codegen,
   Value val = merger.exp(exp).val;
   if (val) {
     if (codegen.curVecLength > 1 && !val.getType().isa<VectorType>())
       return genVectorInvariantValue(codegen, rewriter, val);
     return val;
   }
   // Actual load.
   SmallVector<Value, 4> args;
-  OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).e0];
+  OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
   unsigned tensor = t->getOperandNumber();
   auto map = op.getTiedIndexingMap(t);
   auto enc = getSparseTensorEncoding(t->get().getType());
   unsigned rank = map.getNumResults();
   if (enc) {
     unsigned idx = map.getDimPosition(perm(enc, rank - 1));
     assert(codegen.pidxs[tensor][idx] != nullptr);
     args.push_back(codegen.pidxs[tensor][idx]); // position index

 
 /// Recursively generates tensor expression.
 static Value genExp(Merger &merger, CodeGen &codegen, PatternRewriter &rewriter,
                     linalg::GenericOp op, unsigned exp) {
   if (merger.exp(exp).kind == Kind::kTensor)
     return genTensorLoad(merger, codegen, rewriter, op, exp);
   else if (merger.exp(exp).kind == Kind::kInvariant)
     return genInvariantValue(merger, codegen, rewriter, exp);
-  Value v0 = genExp(merger, codegen, rewriter, op, merger.exp(exp).e0);
-  Value v1 = genExp(merger, codegen, rewriter, op, merger.exp(exp).e1);
+  Value v0 = genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e0);
+  Value v1 = genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e1);
   switch (merger.exp(exp).kind) {
   case Kind::kTensor:
   case Kind::kInvariant:
     llvm_unreachable("handled above");
   case Kind::kMulF:
     return rewriter.create<MulFOp>(op.getLoc(), v0, v1);
   case Kind::kMulI:
     return rewriter.create<MulIOp>(op.getLoc(), v0, v1);
   case Kind::kAddF:
     return rewriter.create<AddFOp>(op.getLoc(), v0, v1);
   case Kind::kAddI:
     return rewriter.create<AddIOp>(op.getLoc(), v0, v1);
   }
   llvm_unreachable("unexpected expression kind");
 }
 
 /// Hoists loop invariant tensor loads for which indices have been exhausted.
 static void genInvariants(Merger &merger, CodeGen &codegen,
                           PatternRewriter &rewriter, linalg::GenericOp op,
                           unsigned exp, unsigned ldx, bool hoist) {
   if (merger.exp(exp).kind == Kind::kTensor) {
     // Inspect tensor indices.
     bool atLevel = ldx == -1u;
-    OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).e0];
+    OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
     auto map = op.getTiedIndexingMap(t);
     auto enc = getSparseTensorEncoding(t->get().getType());
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       unsigned idx = map.getDimPosition(perm(enc, d));
       if (!codegen.loops[idx])
         return; // still in play
       else if (idx == ldx)
         atLevel = true;
     }
     // All exhausted at this level (atLevel denotes exactly at this level).
     OpOperand *lhs = op.getOutputOperand(0);
     if (lhs == t) {
       codegen.redExp = hoist ? exp : -1u;
     } else if (atLevel) {
       merger.exp(exp).val =
           hoist ? genTensorLoad(merger, codegen, rewriter, op, exp) : Value();
     }
   } else if (merger.exp(exp).kind != Kind::kInvariant) {
     // Traverse into the binary operations. Note that we only hoist
     // tensor loads, since subsequent MLIR/LLVM passes know how to
     // deal with all other kinds of derived loop invariants.
-    unsigned e0 = merger.exp(exp).e0;
-    unsigned e1 = merger.exp(exp).e1;
+    unsigned e0 = merger.exp(exp).children.e0;
+    unsigned e1 = merger.exp(exp).children.e1;
     genInvariants(merger, codegen, rewriter, op, e0, ldx, hoist);
     genInvariants(merger, codegen, rewriter, op, e1, ldx, hoist);
   }
 }
 
 /// Generates initialization code for the subsequent loop sequence at
 /// current index level. Returns true if the loop sequence needs to
 /// maintain the universal index.

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

unsigned Merger::optimizeSet(unsigned s0) {
   unsigned s = addSet();
   assert(latSets[s0].size() != 0);
   unsigned p0 = latSets[s0][0];
   for (unsigned p1 : latSets[s0]) {
     bool add = true;
     if (p0 != p1) {
       // Is this a straightforward copy?
       unsigned e = latPoints[p1].exp;
-      if (tensorExps[e].kind == Kind::kTensor && tensorExps[e].e0 == outTensor)
+      if (tensorExps[e].kind == Kind::kTensor &&
+          tensorExps[e].tensor == outTensor)
         continue;
       // Conjunction already covered?
       for (unsigned p2 : latSets[s]) {
         assert(!latGT(p1, p2)); // Lj => Li would be bad
         if (onlyDenseDiff(p2, p1)) {
           add = false;
           break;
         }

 
 //
 // Print methods (for debugging).
 //
 
 void Merger::dumpExp(unsigned e) const {
   switch (tensorExps[e].kind) {
   case Kind::kTensor:
-    if (tensorExps[e].e0 == syntheticTensor)
+    if (tensorExps[e].tensor == syntheticTensor)
       llvm::dbgs() << "synthetic_";
-    else if (tensorExps[e].e0 == outTensor)
+    else if (tensorExps[e].tensor == outTensor)
       llvm::dbgs() << "output_";
-    llvm::dbgs() << "tensor_" << tensorExps[e].e0;
+    llvm::dbgs() << "tensor_" << tensorExps[e].tensor;
     break;
   case Kind::kInvariant:
     llvm::dbgs() << "invariant";
     break;
   default:
   case Kind::kMulI:
     llvm::dbgs() << "(";
-    dumpExp(tensorExps[e].e0);
+    dumpExp(tensorExps[e].children.e0);
     llvm::dbgs() << " * ";
-    dumpExp(tensorExps[e].e1);
+    dumpExp(tensorExps[e].children.e1);
     llvm::dbgs() << ")";
     break;
   case Kind::kAddF:
   case Kind::kAddI:
     llvm::dbgs() << "(";
-    dumpExp(tensorExps[e].e0);
+    dumpExp(tensorExps[e].children.e0);
     llvm::dbgs() << " + ";
-    dumpExp(tensorExps[e].e1);
+    dumpExp(tensorExps[e].children.e1);
     llvm::dbgs() << ")";
     break;
   }
 }
 
 void Merger::dumpLat(unsigned p) const {
   llvm::dbgs() << "lat(";
   dumpBits(latPoints[p].bits);

 
 unsigned Merger::buildLattices(unsigned e, unsigned idx) {
   Kind kind = tensorExps[e].kind;
   if (kind == Kind::kTensor || kind == Kind::kInvariant) {
     // Either the index is really used in the tensor expression, or it is
     // set to the undefined index in that dimension. An invariant expression
     // is set to a synthetic tensor with undefined indices only.
     unsigned s = addSet();
-    unsigned t = kind == Kind::kTensor ? tensorExps[e].e0 : syntheticTensor;
+    unsigned t =
+        kind == Kind::kTensor ? tensorExps[e].children.e0 : syntheticTensor;
     latSets[s].push_back(addLat(t, idx, e));
     return s;
   }
-  unsigned s0 = buildLattices(tensorExps[e].e0, idx);
-  unsigned s1 = buildLattices(tensorExps[e].e1, idx);
+  unsigned s0 = buildLattices(tensorExps[e].children.e0, idx);
+  unsigned s1 = buildLattices(tensorExps[e].children.e1, idx);
   switch (kind) {
   case Kind::kTensor:
   case Kind::kInvariant:
     llvm_unreachable("handled above");
   case Kind::kMulF:
   case Kind::kMulI:
     return takeConj(kind, s0, s1);
   case Kind::kAddF: