[mlir][sparse] Implement custom reduction

[jim22k]

Summary

• Adds logic within linalg.generic lowering to handle

sparse_tensor.reduce. The identity argument is used
to initialize the loop for lex insert, or as the starting
value during expanded access pattern.

• Adds default identity values for standard reduction operators.
• Tests verify correct behavior.

[jim22k]

@aartbik This is mostly the same code as before the split. I moved getReductionIdentity out of Merger.cpp and it's now fully contained in codegen. I also broke up genInsertionLoad into two functions -- one to handle reductions and one to handle non-reduction cases. This eliminates the duplicate if (codegen.redKind == kNoReduc) checks.

I'm still not sure what all needs to happen for vectorized reductions.

[aartbik]

Sorry for the delay getting back to you on this, I was OOO for 2 weeks. I had a look, and the amount of changes in the merger seem about right. The changes in sparsification, however, feel too elaborate. I will hack a bit around tomorrow with this and let you know if there are simpler solutions, and also get back on the vectorization question. In the meantime, the changes to sparse_out_reduction.mlir seem self-contained, so that could be a quick revision first, once you addressed the memsan issue.

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

866

Note that reductions should be fully "scalarized", so I never expect tensor loads having to deal with reductions

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

1038

why is this commented out? do we need some verification on the value of the third parameter?

1040

This probably should become getRegion() when you rebase with some recent changes in MLIR core

mlir/test/Dialect/SparseTensor/roundtrip.mlir

292

change seems unrelated to this revision

mlir/test/Dialect/SparseTensor/sparse_kernels.mlir

58

why do we have changes in the original reductions?

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir

1

The changes in this file merely seem to increase coverage of existing reductions, but they could be done in their own revision, since it is unrelated to this change.

15

note that this is a DCSC if you want to use the standard naming (you can rename the one above to DCSR then to be consistent)

78

maybe redprod1 and redprod2, or otherwise redprodDCSR and redprocDCSC?

165

you will need to release smr and smc as well!

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir

50

although I of course applaud composing the new feature with other features in our tests, I would expect that for simplicity you also have one integration test that just introduces the new sparse_tensor.reduce, without using the binary or unary ops also

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

265

Should this not be in the binary section?!

I think the confusion is because some other tags are misplaced here ;-)

[aartbik]

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

1116

I think you will only need to change the code slightly in this file at *two* places.

(1) At this point (start of a scalarized reduction), when you are in a custom reduction, don't call genTensorLoad, but initialize the "load" Value with the identity value provided by the sparse_tensor.reduce third parameter.

(2) At genInsertionLoad, we need to deal with regular insertions and access pattern expansion for custom reductions. In the former, we load the identity rather than zero, and are done. In the latter case, we cannot simply load the zero from the expanded access pattern, but need to do

if (!expFilled[i])

value= identity

else

value = expValues[i]

You have some of that logic already, but it happens during a regular tensor load. I think you will just need to do it during genInsertionLoad

Will that provide the semantics you expect?

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

1039

redop -> redOp

1230

This seem small enough to just do

case kReduce: {

  ReduceOp redOp = cast<ReduceOp>(tensorExps[e].op);
  return insertYieldOp(rewriter, loc, redOp.getRegion(), {v0, v1});
}

[aartbik]

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

121

Comment needs to be updated for kReduce

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

395

not needed, see below

446

not needed, see below

476

Note that you do not have to worry about vectorization, since the "block" of code in the semi-ring operation cannot be vectorized anyway. You will have to make sure, however, that vectorization is disabled in the presence of a custom reduction.

798

I would simplify this into

return builder.create<arith::SelectOp>(loc, isFilled, valAtIndex, identity);

866

Let me refine my comment, I would never expect this to happen in pure tensor load, but in an insertion load only. So although I would have done it inside that method, this makes sense, except that I would just test for custom reduction, not all the others, and proceed with the code you had just for custom reduction.

You will have to check if we are inside a custom reduction more widely, since redKind is not necessarily set (when we are in an insertion loop). Having a better detection here, for example by setting a customRed value during the recursion, will avoid the huge block of code later! All logic will be here.

Also, this avoids the changes to existing reductions you had.

1042

this block of code is not needed when the custom reduction is properly detected during insertion load, see my code suggestion above

1117

So, this branch would become, where getCustomRedId fetches the third argument from a custom reduction operation.

if (atStart) {
  Kind kind = merger.exp(last).kind;
  Value load = kind == Kind::kReduce
      ? getCustomRedId(merger.exp(last).op)
      : genTensorLoad(merger, codegen, builder, op, exp);
  codegen.redKind = getReduction(kind);
  codegen.redExp = exp;
  updateReduc(merger, codegen, load, "START");
} else {
  ....

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir

23

unrelated to this change, but I actually noticed that when we use the custom reduction, we don't seem to verify the iterator types as carefully as otherwise is the case. Did you notice that too. If so, we probably need to fix that in the verifier (hopefully we don't need to go into linalg for that).

151

This is actually WAI, sort of. When we do a reduction within the innermost loop, we always come out with a value, even if the loop is actually never taken. This behavior also happens for regular reductions, and is similar to what the TACO compiler generates (at least at the moment). Agreed, we could improve this, but that will take some more work beyond this new operation.

173

Note that recent changes unified this with the bufferization dialect ops:

bufferization.dealloc_tensor

[jim22k]

@aartbik Thank you taking a look at my code and giving me feedback.

I do have a question about the approach you are suggesting. If I check for custom reduction and only use the identity in that case, it will leave the current behavior of standard reductions unchanged, including the behavior that reduction with mulf always returns 0.0. Is that the direction you want to go? My PR tried to address that, but I'm okay leaving it alone and only focusing on custom reduction if you feel that is better.

mlir/test/Dialect/SparseTensor/sparse_kernels.mlir

58

This PR doesn't only implement custom reductions. It also adds the getReductionIdentity which applies to standard reductions like mulf. Without those changes, reduction using mulf will always return 0.0.

If you want to leave the existing behavior of standard reductions unchanged, that is fine with me. It will then require a custom reduction to use mulfwith a specified identity of 1.0.

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir

1

The tests I added here are meant to add coverage of existing (non-custom) reductions, specifically mulf, which I consider as being in a broken state. That is why I included them here.

Per my note above, I'm okay if we want to leave existing behavior alone. In that case, I would remove these tests.

[aartbik]

I propose that we leave current behavior unchanged (although we should fix both issues eventually, i.e. mul-reduction with 1 as well as innermost loop skipped yields no value) and focus on adding new op only.
After that we can address the two issues pointed out.

Sounds good?

[jim22k]

• Updates based on feedback

[jim22k]

@aartbik This is definitely simpler than what I had in Sparsification.cpp. Thanks for pointing me in the right direction.

[aartbik]

Almost there! I believe we can make it even simpler with some recursive set/unset.
But then good to go

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

868

since the if returns, no else, and indent the second return to the left

1037–1053

I believe that this function can be simpler too. Simply remove the "last" and all codegen code you added
Then, as else branch of L985 use

if (merger.exp(exp).kind == Kind::kReduce)

codegen.redCustom = getCustomRedId(merger.exp(exp).op);

and then right before exit at L1291 do

if (merger.exp(exp).kind == Kind::kReduce)

codegen.redCustom = Value()

you can also make redCustom an unsigned and use -1/exp value instead.
This way, the recursive nature will take care of handling the tensor load in the code above, without special handling

Of course, this assume custom reductions are never nested, something that makes sense, but perhaps we should verify

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir

214

make this more precise, it it not about the Inf per se, but about not having full "fill in" due to the reduction loop

[jim22k]

Further refinement based on feedback

[jim22k]

I hope this is what you had in mind. I didn't see how to make redCustom unsigned, as it needs to store a Value.

[aartbik]

In D130573#3727091, @jim22k wrote:

I hope this is what you had in mind. I didn't see how to make redCustom unsigned, as it needs to store a Value.

Oh, I just meant to store the exp index, and then later extract the value using that; but this is fine too

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

1040

we can actually assert here to avoid the nested problem

assert(!codegen.redCustom);
codegen.redCustom = ....

[aartbik]

Assuming you address my last few nits (can't help myself), this is good to ship!
Thanks for your patience during this review.

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

1038

if you move this below the next few if-s that return

you can do this without the else, and without braces

if ...

return

if (merger.exp(exp).kind == Kind::kReduce) {

codegen.redCustom =

t=

[jim22k]

Final updates

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

enum Kind {
   kSubI,
   kAndI,
   kOrI,
   kXorI,
   kShrS, // signed
   kShrU, // unsigned
   kShlI,
   kBinary, // semiring binary op
+  kReduce, // semiring reduction op
 };
 
 /// Children subexpressions of tensor operations.
 struct Children {
   unsigned e0;
   unsigned e1;
 };
 
struct TensorExp {
   /// Direct link to IR for an invariant or the destination value (to
   /// infer destination type) of a cast operation During code generation,
   /// this field may be used to cache "hoisted" loop invariant tensor loads.
   Value val;
 
   /// Code blocks used by semirings. For the case of kUnary and
   /// kBinary, this holds the original operation with all regions. For
   /// kBinaryBranch, this holds the YieldOp for the left or right half
   /// to be merged into a nested scf loop.
   Operation *op;
 };
 
 /// Lattice point. Each lattice point consists of a conjunction of tensor
 /// loop indices (encoded in a bitvector) and the index of the corresponding
 /// tensor expression.
 struct LatPoint {
   LatPoint(unsigned n, unsigned e, unsigned b);

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

 enum SortMask {
   kSparseOnly = 0x0,
   kIncludeDense = 0x1,
   kIncludeUndef = 0x2,
   kIncludeAll = 0x3
 };
 
 // Reduction kinds.
-enum Reduction { kNoReduc, kSum, kProduct, kAnd, kOr, kXor };
+enum Reduction { kNoReduc, kSum, kProduct, kAnd, kOr, kXor, kCustom };
 
 // Code generation.
 struct CodeGen {
   CodeGen(SparsificationOptions o, unsigned numTensors, unsigned numLoops,
           OpOperand *op, unsigned nest)
       : options(o), loops(numLoops), sizes(numLoops), buffers(numTensors),
         pointers(numTensors, std::vector<Value>(numLoops)),
         indices(numTensors, std::vector<Value>(numLoops)),

 //===----------------------------------------------------------------------===//
 // Sparse compiler synthesis methods (reductions).
 //===----------------------------------------------------------------------===//
 
 /// Maps reduction kind to vector::CombiningKind.
 static vector::CombiningKind getCombiningKind(Reduction kind) {
   switch (kind) {
   case kNoReduc:
+  case kCustom:
     break;
   case kSum:
     return vector::CombiningKind::ADD;
   case kProduct:
     return vector::CombiningKind::MUL;
   case kAnd:
     return vector::CombiningKind::AND;
   case kOr:
     return vector::CombiningKind::OR;
   case kXor:
     return vector::CombiningKind::XOR;
   }
   llvm_unreachable("unknown reduction kind");
 }
 
+static bool isValidReduction(Kind kind) {
+  switch (kind) {
+  case Kind::kAddF:
+  case Kind::kAddC:
+  case Kind::kAddI:
+  case Kind::kSubF:
+  case Kind::kSubC:
+  case Kind::kSubI:
+  case Kind::kMulF:
+  case Kind::kMulC:
+  case Kind::kMulI:
+  case Kind::kAndI:
+  case Kind::kOrI:
+  case Kind::kXorI:
+  case Kind::kReduce:
+    return true;
+  default:
+    return false;
+  }
+}
+
 /// Maps operation to reduction.
 static Reduction getReduction(Kind kind) {
   switch (kind) {
   case Kind::kAddF:
   case Kind::kAddC:
   case Kind::kAddI:
   case Kind::kSubF:
   case Kind::kSubC:
   case Kind::kSubI:
     return kSum;
   case Kind::kMulF:
   case Kind::kMulC:
   case Kind::kMulI:
     return kProduct;
   case Kind::kAndI:
     return kAnd;
   case Kind::kOrI:
     return kOr;
   case Kind::kXorI:
     return kXor;
+  case Kind::kReduce:
+    return kCustom;
   default:
     llvm_unreachable("unexpected reduction operator");
   }
 }
 
+/// Generates the reduction identity value based on the kind of operation.
+/// The identity is meant to have no impact on the final reduction value
+/// (i.e. x op identity == x).
+static Value genReductionIdentity(Merger &merger, CodeGen &codegen,
+                                  OpBuilder &builder, linalg::GenericOp op,
+                                  Type tp) {
+  Location loc = op.getLoc();
+  Kind kind = merger.exp(codegen.redExp).kind;
+  switch (kind) {
+  case kMulF:
+    return builder.create<arith::ConstantOp>(loc, tp,
+                                             builder.getFloatAttr(tp, 1.0));
+  case kMulI:
+  case kAndI:
+    return builder.create<arith::ConstantOp>(loc, tp,
+                                             builder.getIntegerAttr(tp, 1));
+  case kReduce:
+    return dyn_cast<sparse_tensor::ReduceOp>(merger.exp(codegen.redExp).op)
+        .identity();
+  default:
+    return builder.create<arith::ConstantOp>(loc, tp, builder.getZeroAttr(tp));
+  }
+}
+
 /// Generates an initial value for a vector reduction, following the scheme
 /// given in Chapter 5 of "The Software Vectorization Handbook", where the
 /// initial scalar value is correctly embedded in the vector reduction value,
 /// and a straightforward horizontal reduction will complete the operation.
 static Value genVectorReducInit(CodeGen &codegen, OpBuilder &builder,
                                 Location loc, VectorType vtp) {
   Value r = codegen.redVal;
   switch (codegen.redKind) {
   case kNoReduc:
+  case kCustom:
     break;
   case kSum:
   case kXor:
     // Initialize reduction vector to: | 0 | .. | 0 | r |
     return builder.create<vector::InsertElementOp>(
         loc, r, constantZero(builder, loc, vtp),
         constantIndex(builder, loc, 0));
   case kProduct:
if (!codegen.expValues) {
     Type tp = getElementTypeOrSelf(t->get().getType());
     return constantZero(builder, loc, tp);
   }
   // Load from expanded access pattern.
   Value index = genIndex(codegen, op, t);
   return builder.create<memref::LoadOp>(loc, codegen.expValues, index);
 }
 
+/// Generates insertion code to implement dynamic tensor load for reduction.
+static Value genInsertionLoadReduce(Merger &merger, CodeGen &codegen,
+                                    OpBuilder &builder, linalg::GenericOp op,
+                                    OpOperand *t) {
+  Location loc = op.getLoc();
+  Type tp = getElementTypeOrSelf(t->get().getType());
+  Value identity = genReductionIdentity(merger, codegen, builder, op, tp);
+  // Direct lexicographic index order, tensor loads as identity.
+  if (!codegen.expValues) {
+    return identity;
+  }
+  // Load from expanded access pattern if filled, identity otherwise.
+  Value index = genIndex(codegen, op, t);
+  Value isFilled =
+      builder.create<memref::LoadOp>(loc, codegen.expFilled, index);
+  scf::IfOp ifIsFilled =
+      builder.create<scf::IfOp>(loc, tp, isFilled, /*else=*/true);
+  // True branch
+  builder.setInsertionPointToStart(ifIsFilled.thenBlock());
+  Value valAtIndex =
+      builder.create<memref::LoadOp>(loc, codegen.expValues, index);
+  builder.create<scf::YieldOp>(loc, valAtIndex);
+  // False branch
+  builder.setInsertionPointToStart(ifIsFilled.elseBlock());
+  builder.create<scf::YieldOp>(loc, identity);
+  builder.setInsertionPointAfter(ifIsFilled);
+  // End if
+  return ifIsFilled.getResult(0);
+}
+
 /// Generates insertion code to implement dynamic tensor store.
 static void genInsertionStore(CodeGen &codegen, OpBuilder &builder,
                               linalg::GenericOp op, OpOperand *t, Value rhs) {
   Location loc = op.getLoc();
   // Direct insertion in lexicographic index order.
   if (!codegen.expValues) {
     builder.create<memref::StoreOp>(loc, rhs, codegen.lexVal);
     builder.create<LexInsertOp>(loc, t->get(), codegen.lexIdx, codegen.lexVal);
static Value genTensorLoad(Merger &merger, CodeGen &codegen, OpBuilder &builder,
   Value val = merger.exp(exp).val;
   if (val) {
     if (codegen.curVecLength > 1 && !val.getType().isa<VectorType>())
       return genVectorInvariantValue(codegen, builder, val);
     return val;
   }
   // Load during insertion.
   OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
-  if (t == codegen.sparseOut)
+  if (t == codegen.sparseOut) {
+    if (codegen.redKind == kNoReduc)
       return genInsertionLoad(codegen, builder, op, t);
+    else
+      return genInsertionLoadReduce(merger, codegen, builder, op, t);
+  }
   // Actual load.
   SmallVector<Value, 4> args;
   Value ptr = genSubscript(codegen, builder, op, t, args);
   if (codegen.curVecLength > 1)
     return genVectorLoad(codegen, builder, ptr, args);
   return builder.create<memref::LoadOp>(op.getLoc(), ptr, args);
 }
 
if (def->getBlock() == block) {
             i, relinkBranch(codegen, rewriter, block, def->getOperand(i), ldx));
     }
   }
   return e;
 }
 
 /// Recursively generates tensor expression.
 static Value genExp(Merger &merger, CodeGen &codegen, RewriterBase &rewriter,
-                    linalg::GenericOp op, unsigned exp, unsigned ldx) {
+                    linalg::GenericOp op, unsigned exp, unsigned ldx,
+                    unsigned last = 0) {
   Location loc = op.getLoc();
   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::kTensor) {
+    // Handle reductions for access pattern expansion. The trigger is when the
+    // output tensor is also an operand parameter, although this can also apply
+    // to inplace updates which are not reductions. To avoid errors, check that
+    // the operation is a permitted reduction.
+    bool validRed = isValidReduction(merger.exp(last).kind);
+    OpOperand *t = op.getInputAndOutputOperands()[merger.exp(exp).tensor];
+    OpOperand *lhs = op.getOutputOperand(0);
+    if (validRed && lhs == t) {
+      codegen.redKind = getReduction(merger.exp(last).kind);
+      codegen.redExp = last; // handling for reduction identity
+    }
+    Value redVal = genTensorLoad(merger, codegen, rewriter, op, exp);
+    if (validRed && lhs == t)
+      codegen.redExp = exp;
+    return redVal;
+  }
   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);
-  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 v0 = genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e0,
+                    ldx, exp);
+  Value v1 = genExp(merger, codegen, rewriter, op, merger.exp(exp).children.e1,
+                    ldx, exp);
   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::kBinaryBranch ||
+             merger.exp(exp).kind == Kind::kReduce))
     ee = relinkBranch(codegen, rewriter, ee.getParentBlock(), ee, ldx);
   return ee;
 }
 
 /// Determines if affine expression is invariant.
 static bool isInvariantAffine(const CodeGen &codegen, AffineExpr a,
                               unsigned ldx, bool &atLevel) {
   switch (a.getKind()) {
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, Kind last = Kind::kTensor) {
+                          bool atStart, unsigned last = 0) {
   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());
     for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
       AffineExpr a = map.getResult(perm(enc, d));
       if (!isInvariantAffine(codegen, a, ldx, atLevel))
         return; // still in play
     }
     // All exhausted at this level (atLevel denotes exactly at this level).
     if (!atLevel)
       return;
     OpOperand *lhs = op.getOutputOperand(0);
     if (lhs == t) {
       // Start or end a scalarized reduction
       if (atStart) {
+        codegen.redKind = getReduction(merger.exp(last).kind);
+        codegen.redExp = last; // handling for reduction identity
         Value load = genTensorLoad(merger, codegen, builder, op, exp);
-        codegen.redKind = getReduction(last);
         codegen.redExp = exp;
         updateReduc(merger, codegen, load);
       } else {
         Value redVal = codegen.redVal;
         updateReduc(merger, codegen, Value());
         codegen.redExp = -1u;
         codegen.redKind = kNoReduc;
         genTensorStore(merger, codegen, builder, op, exp, redVal);
       }
     } else {
       // Start or end loop invariant hoisting of a tensor load.
       merger.exp(exp).val =
           atStart ? genTensorLoad(merger, codegen, builder, op, exp) : Value();
     }
   } else if (merger.exp(exp).kind != Kind::kInvariant &&
              merger.exp(exp).kind != Kind::kIndex) {
     // 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.
-    Kind last = merger.exp(exp).kind;
     unsigned e0 = merger.exp(exp).children.e0;
     unsigned e1 = merger.exp(exp).children.e1;
-    genInvariants(merger, codegen, builder, op, e0, ldx, atStart, last);
-    genInvariants(merger, codegen, builder, op, e1, ldx, atStart, last);
+    genInvariants(merger, codegen, builder, op, e0, ldx, atStart, exp);
+    genInvariants(merger, codegen, builder, op, e1, ldx, atStart, exp);
   }
 }
 
 /// Generates an expanded access pattern in innermost dimension.
 static void genExpansion(Merger &merger, CodeGen &codegen, OpBuilder &builder,
                          linalg::GenericOp op, unsigned at, bool atStart) {
   OpOperand *lhs = codegen.sparseOut;
   if (!lhs || codegen.outerParNest != op.getRank(lhs) - 1 ||

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

TensorExp::TensorExp(Kind k, unsigned x, unsigned y, Value v, Operation *o)
   case kShrS:
   case kShrU:
   case kShlI:
     assert(x != -1u && y != -1u && !v && !o);
     children.e0 = x;
     children.e1 = y;
     break;
   case kBinary:
+  case kReduce:
     assert(x != -1u && y != -1u && !v && o);
     children.e0 = x;
     children.e1 = y;
     break;
   }
 }
 
 LatPoint::LatPoint(unsigned n, unsigned e, unsigned b)
case kAddI:
     return isSingleCondition(t, tensorExps[e].children.e0) &&
            isSingleCondition(t, tensorExps[e].children.e1);
   case kSubF:
   case kSubC:
   case kSubI:
   case kOrI:
   case kXorI:
   case kBinary:
+  case kReduce:
     return false;
   }
   llvm_unreachable("unexpected kind");
 }
 
 #ifndef NDEBUG
 
 //===----------------------------------------------------------------------===//
static const char *kindToOpSymbol(Kind kind) {
   case kShrS:
     return "a>>";
   case kShrU:
     return ">>";
   case kShlI:
     return "<<";
   case kBinary:
     return "binary";
+  case kReduce:
+    return "reduce";
   }
   llvm_unreachable("unexpected kind for symbol");
 }
 
 void Merger::dumpExp(unsigned e) const {
   switch (tensorExps[e].kind) {
   // Leaf.
   case kTensor:
void Merger::dumpExp(unsigned e) const {
   case kSubI:
   case kAndI:
   case kOrI:
   case kXorI:
   case kShrS:
   case kShrU:
   case kShlI:
   case kBinary:
+  case kReduce:
     llvm::dbgs() << "(";
     dumpExp(tensorExps[e].children.e0);
     llvm::dbgs() << " " << kindToOpSymbol(tensorExps[e].kind) << " ";
     dumpExp(tensorExps[e].children.e1);
     llvm::dbgs() << ")";
   }
 }
 
//     x  | left(x) | overlap(x,y) |
         rightYield = rightBlock.getTerminator();
       }
       bool includeLeft = binop.getLeftIdentity() || !leftRegion.empty();
       bool includeRight = binop.getRightIdentity() || !rightRegion.empty();
       return takeCombi(kBinary, child0, child1, binop, includeLeft,
                        kBinaryBranch, leftYield, includeRight, kBinaryBranch,
                        rightYield);
     }
+  case kReduce:
+    // A custom reduce operation.
+    return takeConj(kind, buildLattices(tensorExps[e].children.e0, i),
+                    buildLattices(tensorExps[e].children.e1, i),
+                    tensorExps[e].op);
   }
   llvm_unreachable("unexpected expression kind");
 }
 
 Optional<unsigned> Merger::buildTensorExpFromLinalg(linalg::GenericOp op) {
   // Build the linalg semantics backward from yield.
   Operation *yield = op.region().front().getTerminator();
   assert(isa<linalg::YieldOp>(yield));
if (x.has_value()) {
         if (isAdmissableBranch(unop, unop.getPresentRegion()) &&
             isAdmissableBranch(unop, unop.getAbsentRegion()))
           return addExp(kUnary, e, Value(), def);
       }
     }
   }
   // Construct binary operations if subexpressions can be built.
   // See buildLattices() for an explanation of rejecting certain
-  // division and shift operations
+  // division and shift operations.
   if (def->getNumOperands() == 2) {
     auto x = buildTensorExp(op, def->getOperand(0));
     auto y = buildTensorExp(op, def->getOperand(1));
     if (x.has_value() && y.has_value()) {
       unsigned e0 = x.value();
       unsigned e1 = y.value();
       if (isa<arith::MulFOp>(def))
         return addExp(kMulF, e0, e1);
if (x.has_value() && y.has_value()) {
             (binop.getLeftIdentity() ||
              isAdmissableBranch(binop, binop.getLeftRegion())) &&
             (binop.getRightIdentity() ||
              isAdmissableBranch(binop, binop.getRightRegion())))
           return addExp(kBinary, e0, e1, Value(), def);
       }
     }
   }
+  // 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.hasValue() && y.hasValue() && z.hasValue()) {
+      unsigned e0 = x.getValue();
+      unsigned e1 = y.getValue();
+      // unsigned e2 = z.getValue();
+      if (auto redop = dyn_cast<sparse_tensor::ReduceOp>(def)) {
+        if (isAdmissableBranch(redop, redop.region()))
+          return addExp(kReduce, e0, e1, Value(), def);
+      }
+    }
+  }
   // Cannot build.
   return None;
 }
 
 static Value insertYieldOp(RewriterBase &rewriter, Location loc, Region &region,
                            ValueRange vals) {
   // Make a clone of overlap region.
   Region tmpRegion;
static Value buildBinaryOverlap(RewriterBase &rewriter, Location loc,
   Region &overlapRegion = binop.getOverlapRegion();
   if (overlapRegion.empty())
     // Uninitialized Value() will be interpreted as missing data in the
     // output.
     return Value();
   return insertYieldOp(rewriter, loc, overlapRegion, {v0, v1});
 }
 
+static Value buildReduce(RewriterBase &rewriter, Location loc, Operation *op,
+                         Value v0, Value v1) {
+  ReduceOp redop = cast<ReduceOp>(op);
+  Region &formula = redop.region();
+  return insertYieldOp(rewriter, loc, formula, {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");
case kShlI:
     return rewriter.create<arith::ShLIOp>(loc, v0, v1);
   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);
+  case kReduce:
+    return buildReduce(rewriter, loc, tensorExps[e].op, v0, v1);
   }
   llvm_unreachable("unexpected expression kind in build");
 }
 
 } // namespace sparse_tensor
 } // namespace mlir

mlir/test/Dialect/SparseTensor/roundtrip.mlir

 //       CHECK: }
 func.func @sparse_reduce_2d_to_1d(%arg0: f64, %arg1: f64) -> f64 {
   %cf0 = arith.constant 0.0 : f64
   %r = sparse_tensor.reduce %arg0, %arg1, %cf0 : f64 {
       ^bb0(%x: f64, %y: f64):
         sparse_tensor.yield %x : f64
     }
   return %r : f64
-}
+}
 

mlir/test/Dialect/SparseTensor/sparse_kernels.mlir

%0 = linalg.matmul
     ins(%a, %b: tensor<10x20xf32, #DCSR>, tensor<20x30xf32>)
     outs(%c: tensor<10x30xf32>) -> tensor<10x30xf32>
   return %0 : tensor<10x30xf32>
 }
 
 //
 // Computes C = A x B with all matrices sparse (SpMSpM) in DCSR.
 //
-// CHECK-LABEL:   func @matmul2(
+// CHECK-LABEL:   func.func @matmul2(
 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>>,
-// CHECK-SAME:      %[[VAL_1:.*]]: tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> {
-// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 0 : index
-// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 1 : index
-// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 2 : index
+// CHECK-SAME:      %[[VAL_1:.*]]: tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>> {
+// CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 2 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0.000000e+00 : f64
 // CHECK-DAG:       %[[VAL_6:.*]] = arith.constant false
 // CHECK-DAG:       %[[VAL_7:.*]] = arith.constant true
 // CHECK:           %[[VAL_8:.*]] = bufferization.alloc_tensor() : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_4]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_4]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_9:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_2]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_10:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_2]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
 // CHECK:           %[[VAL_13:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<4x8xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf64>
-// CHECK:           %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_3]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_3]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_4]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_17:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_4]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_14:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_2]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_15:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_2]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_16:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_3]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_17:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_3]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
 // CHECK:           %[[VAL_18:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<8x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf64>
-// CHECK:           %[[VAL_19:.*]] = memref.alloca(%[[VAL_5]]) : memref<?xindex>
-// CHECK:           %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_4]]] : memref<?xindex>
-// CHECK:           scf.for %[[VAL_22:.*]] = %[[VAL_20]] to %[[VAL_21]] step %[[VAL_4]] {
+// CHECK:           %[[VAL_19:.*]] = memref.alloca(%[[VAL_4]]) : memref<?xindex>
+// CHECK:           %[[VAL_20:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:           %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:           scf.for %[[VAL_22:.*]] = %[[VAL_20]] to %[[VAL_21]] step %[[VAL_3]] {
 // CHECK:             %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_22]]] : memref<?xindex>
-// CHECK:             memref.store %[[VAL_23]], %[[VAL_19]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:             memref.store %[[VAL_23]], %[[VAL_19]]{{\[}}%[[VAL_2]]] : memref<?xindex>
 // CHECK:             %[[VAL_24:.*]], %[[VAL_25:.*]], %[[VAL_26:.*]], %[[VAL_27:.*]] = sparse_tensor.expand %[[VAL_8]] : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf64>, memref<?xi1>, memref<?xindex>, index
 // CHECK:             %[[VAL_28:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_22]]] : memref<?xindex>
-// CHECK:             %[[VAL_29:.*]] = arith.addi %[[VAL_22]], %[[VAL_4]] : index
+// CHECK:             %[[VAL_29:.*]] = arith.addi %[[VAL_22]], %[[VAL_3]] : index
 // CHECK:             %[[VAL_30:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_29]]] : memref<?xindex>
-// CHECK:             %[[VAL_31:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_4]]] : memref<?xindex>
+// CHECK:             %[[VAL_31:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_3]]] : memref<?xindex>
 // CHECK:             %[[VAL_33:.*]]:3 = scf.while (%[[VAL_34:.*]] = %[[VAL_28]], %[[VAL_35:.*]] = %[[VAL_31]], %[[VAL_36:.*]] = %[[VAL_27]]) : (index, index, index) -> (index, index, index) {
 // CHECK:               %[[VAL_37:.*]] = arith.cmpi ult, %[[VAL_34]], %[[VAL_30]] : index
 // CHECK:               %[[VAL_38:.*]] = arith.cmpi ult, %[[VAL_35]], %[[VAL_32]] : index
 // CHECK:               %[[VAL_39:.*]] = arith.andi %[[VAL_37]], %[[VAL_38]] : i1
 // CHECK:               scf.condition(%[[VAL_39]]) %[[VAL_34]], %[[VAL_35]], %[[VAL_36]] : index, index, index
 // CHECK:             } do {
 // CHECK:             ^bb0(%[[VAL_40:.*]]: index, %[[VAL_41:.*]]: index, %[[VAL_42:.*]]: index):
 // CHECK:               %[[VAL_43:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_40]]] : memref<?xindex>
 // CHECK:               %[[VAL_44:.*]] = memref.load %[[VAL_15]]{{\[}}%[[VAL_41]]] : memref<?xindex>
 // CHECK:               %[[VAL_45:.*]] = arith.cmpi ult, %[[VAL_44]], %[[VAL_43]] : index
 // CHECK:               %[[VAL_46:.*]] = arith.select %[[VAL_45]], %[[VAL_44]], %[[VAL_43]] : index
 // CHECK:               %[[VAL_47:.*]] = arith.cmpi eq, %[[VAL_43]], %[[VAL_46]] : index
 // CHECK:               %[[VAL_48:.*]] = arith.cmpi eq, %[[VAL_44]], %[[VAL_46]] : index
 // CHECK:               %[[VAL_49:.*]] = arith.andi %[[VAL_47]], %[[VAL_48]] : i1
 // CHECK:               %[[VAL_50:.*]] = scf.if %[[VAL_49]] -> (index) {
 // CHECK:                 %[[VAL_51:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_40]]] : memref<?xf64>
 // CHECK:                 %[[VAL_52:.*]] = memref.load %[[VAL_16]]{{\[}}%[[VAL_41]]] : memref<?xindex>
-// CHECK:                 %[[VAL_53:.*]] = arith.addi %[[VAL_41]], %[[VAL_4]] : index
+// CHECK:                 %[[VAL_53:.*]] = arith.addi %[[VAL_41]], %[[VAL_3]] : index
 // CHECK:                 %[[VAL_54:.*]] = memref.load %[[VAL_16]]{{\[}}%[[VAL_53]]] : memref<?xindex>
-// CHECK:                 %[[VAL_55:.*]] = scf.for %[[VAL_56:.*]] = %[[VAL_52]] to %[[VAL_54]] step %[[VAL_4]] iter_args(%[[VAL_57:.*]] = %[[VAL_42]]) -> (index) {
+// CHECK:                 %[[VAL_55:.*]] = scf.for %[[VAL_56:.*]] = %[[VAL_52]] to %[[VAL_54]] step %[[VAL_3]] iter_args(%[[VAL_57:.*]] = %[[VAL_42]]) -> (index) {
 // CHECK:                   %[[VAL_58:.*]] = memref.load %[[VAL_17]]{{\[}}%[[VAL_56]]] : memref<?xindex>
-// CHECK:                   %[[VAL_59:.*]] = memref.load %[[VAL_24]]{{\[}}%[[VAL_58]]] : memref<?xf64>
-// CHECK:                   %[[VAL_60:.*]] = memref.load %[[VAL_18]]{{\[}}%[[VAL_56]]] : memref<?xf64>
-// CHECK:                   %[[VAL_61:.*]] = arith.mulf %[[VAL_51]], %[[VAL_60]] : f64
-// CHECK:                   %[[VAL_62:.*]] = arith.addf %[[VAL_59]], %[[VAL_61]] : f64
-// CHECK:                   %[[VAL_63:.*]] = memref.load %[[VAL_25]]{{\[}}%[[VAL_58]]] : memref<?xi1>
-// CHECK:                   %[[VAL_64:.*]] = arith.cmpi eq, %[[VAL_63]], %[[VAL_6]] : i1
-// CHECK:                   %[[VAL_65:.*]] = scf.if %[[VAL_64]] -> (index) {
+// CHECK:                   %[[VAL_59:.*]] = memref.load %[[VAL_25]]{{\[}}%[[VAL_58]]] : memref<?xi1>
+// CHECK:                   %[[VAL_60:.*]] = scf.if %[[VAL_59]] -> (f64) {
+// CHECK:                     %[[VAL_61:.*]] = memref.load %[[VAL_24]]{{\[}}%[[VAL_58]]] : memref<?xf64>
+// CHECK:                     scf.yield %[[VAL_61]] : f64
+// CHECK:                   } else {
+// CHECK:                     scf.yield %[[VAL_5]] : f64
+// CHECK:                   }
+// CHECK:                   %[[VAL_62:.*]] = memref.load %[[VAL_18]]{{\[}}%[[VAL_56]]] : memref<?xf64>
+// CHECK:                   %[[VAL_63:.*]] = arith.mulf %[[VAL_51]], %[[VAL_62]] : f64
+// CHECK:                   %[[VAL_64:.*]] = arith.addf %[[VAL_65:.*]], %[[VAL_63]] : f64
+// CHECK:                   %[[VAL_66:.*]] = memref.load %[[VAL_25]]{{\[}}%[[VAL_58]]] : memref<?xi1>
+// CHECK:                   %[[VAL_67:.*]] = arith.cmpi eq, %[[VAL_66]], %[[VAL_6]] : i1
+// CHECK:                   %[[VAL_68:.*]] = scf.if %[[VAL_67]] -> (index) {
 // CHECK:                     memref.store %[[VAL_7]], %[[VAL_25]]{{\[}}%[[VAL_58]]] : memref<?xi1>
 // CHECK:                     memref.store %[[VAL_58]], %[[VAL_26]]{{\[}}%[[VAL_57]]] : memref<?xindex>
-// CHECK:                     %[[VAL_66:.*]] = arith.addi %[[VAL_57]], %[[VAL_4]] : index
-// CHECK:                     scf.yield %[[VAL_66]] : index
+// CHECK:                     %[[VAL_69:.*]] = arith.addi %[[VAL_57]], %[[VAL_3]] : index
+// CHECK:                     scf.yield %[[VAL_69]] : index
 // CHECK:                   } else {
 // CHECK:                     scf.yield %[[VAL_57]] : index
 // CHECK:                   }
-// CHECK:                   memref.store %[[VAL_62]], %[[VAL_24]]{{\[}}%[[VAL_58]]] : memref<?xf64>
-// CHECK:                   scf.yield %[[VAL_67:.*]] : index
+// CHECK:                   memref.store %[[VAL_64]], %[[VAL_24]]{{\[}}%[[VAL_58]]] : memref<?xf64>
+// CHECK:                   scf.yield %[[VAL_70:.*]] : index
 // CHECK:                 }
-// CHECK:                 scf.yield %[[VAL_68:.*]] : index
+// CHECK:                 scf.yield %[[VAL_71:.*]] : index
 // CHECK:               } else {
 // CHECK:                 scf.yield %[[VAL_42]] : index
 // CHECK:               }
-// CHECK:               %[[VAL_69:.*]] = arith.cmpi eq, %[[VAL_43]], %[[VAL_46]] : index
-// CHECK:               %[[VAL_70:.*]] = arith.addi %[[VAL_40]], %[[VAL_4]] : index
-// CHECK:               %[[VAL_71:.*]] = arith.select %[[VAL_69]], %[[VAL_70]], %[[VAL_40]] : index
-// CHECK:               %[[VAL_72:.*]] = arith.cmpi eq, %[[VAL_44]], %[[VAL_46]] : index
-// CHECK:               %[[VAL_73:.*]] = arith.addi %[[VAL_41]], %[[VAL_4]] : index
-// CHECK:               %[[VAL_74:.*]] = arith.select %[[VAL_72]], %[[VAL_73]], %[[VAL_41]] : index
-// CHECK:               scf.yield %[[VAL_71]], %[[VAL_74]], %[[VAL_75:.*]] : index, index, index
+// CHECK:               %[[VAL_72:.*]] = arith.cmpi eq, %[[VAL_43]], %[[VAL_46]] : index
+// CHECK:               %[[VAL_73:.*]] = arith.addi %[[VAL_40]], %[[VAL_3]] : index
+// CHECK:               %[[VAL_74:.*]] = arith.select %[[VAL_72]], %[[VAL_73]], %[[VAL_40]] : index
+// CHECK:               %[[VAL_75:.*]] = arith.cmpi eq, %[[VAL_44]], %[[VAL_46]] : index
+// CHECK:               %[[VAL_76:.*]] = arith.addi %[[VAL_41]], %[[VAL_3]] : index
+// CHECK:               %[[VAL_77:.*]] = arith.select %[[VAL_75]], %[[VAL_76]], %[[VAL_41]] : index
+// CHECK:               scf.yield %[[VAL_74]], %[[VAL_77]], %[[VAL_78:.*]] : index, index, index
 // CHECK:             }
-// CHECK:             sparse_tensor.compress %[[VAL_8]], %[[VAL_19]], %[[VAL_24]], %[[VAL_25]], %[[VAL_26]], %[[VAL_76:.*]]#2 : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
+// CHECK:             sparse_tensor.compress %[[VAL_8]], %[[VAL_19]], %[[VAL_24]], %[[VAL_25]], %[[VAL_26]], %[[VAL_79:.*]]#2 : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>, memref<?xindex>, memref<?xf64>, memref<?xi1>, memref<?xindex>, index
 // CHECK:           }
-// CHECK:           %[[VAL_77:.*]] = sparse_tensor.load %[[VAL_8]] hasInserts : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           return %[[VAL_77]] : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           %[[VAL_80:.*]] = sparse_tensor.load %[[VAL_8]] hasInserts : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           return %[[VAL_80]] : tensor<4x4xf64, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:         }
 func.func @matmul2(%A: tensor<4x8xf64, #DCSR>,
               %B: tensor<8x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR> {
   %c4 = arith.constant 4 : index
   %C = bufferization.alloc_tensor() : tensor<4x4xf64, #DCSR>
   %D = linalg.matmul
     ins(%A, %B: tensor<4x8xf64, #DCSR>, tensor<8x4xf64, #DCSR>)
        outs(%C: tensor<4x4xf64, #DCSR>) -> tensor<4x4xf64, #DCSR>

mlir/test/Dialect/SparseTensor/sparse_out.mlir

indexing_maps = [
     affine_map<(i,j,k) -> (i,k)>, // A
     affine_map<(i,j,k) -> (k,j)>, // B
     affine_map<(i,j,k) -> (i,j)>  // C (out)
   ],
   iterator_types = ["parallel", "parallel", "reduction"],
   doc = "C(i,j) = SUM_k A(i,k) * B(k,j)"
 }
 
-// CHECK-LABEL:   func @matmat(
+// CHECK-LABEL:   func.func @matmat(
 // CHECK-SAME:      %[[VAL_0:.*]]: tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>,
-// CHECK-SAME:      %[[VAL_1:.*]]: tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> {
+// CHECK-SAME:      %[[VAL_1:.*]]: tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>) -> tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> {
 // CHECK-DAG:       %[[VAL_2:.*]] = arith.constant 0 : index
 // CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
 // CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 2 : index
-// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant false
-// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant true
-// CHECK:           %[[VAL_7:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_8:.*]] = tensor.dim %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_9:.*]] = bufferization.alloc_tensor(%[[VAL_7]], %[[VAL_8]]) : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           %[[VAL_10:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_11:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_12:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_13:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_14:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf32>
-// CHECK:           %[[VAL_15:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_16:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_17:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_18:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
-// CHECK:           %[[VAL_19:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf32>
-// CHECK:           %[[VAL_20:.*]] = memref.alloca(%[[VAL_4]]) : memref<?xindex>
-// CHECK:           %[[VAL_21:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_2]]] : memref<?xindex>
-// CHECK:           %[[VAL_22:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:           scf.for %[[VAL_23:.*]] = %[[VAL_21]] to %[[VAL_22]] step %[[VAL_3]] {
-// CHECK:             %[[VAL_24:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_23]]] : memref<?xindex>
-// CHECK:             memref.store %[[VAL_24]], %[[VAL_20]]{{\[}}%[[VAL_2]]] : memref<?xindex>
-// CHECK:             %[[VAL_25:.*]], %[[VAL_26:.*]], %[[VAL_27:.*]], %[[VAL_28:.*]] = sparse_tensor.expand %[[VAL_9]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf32>, memref<?xi1>, memref<?xindex>, index
-// CHECK:             %[[VAL_29:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_23]]] : memref<?xindex>
-// CHECK:             %[[VAL_30:.*]] = arith.addi %[[VAL_23]], %[[VAL_3]] : index
-// CHECK:             %[[VAL_31:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_30]]] : memref<?xindex>
-// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_15]]{{\[}}%[[VAL_2]]] : memref<?xindex>
-// CHECK:             %[[VAL_33:.*]] = memref.load %[[VAL_15]]{{\[}}%[[VAL_3]]] : memref<?xindex>
-// CHECK:             %[[VAL_34:.*]]:3 = scf.while (%[[VAL_35:.*]] = %[[VAL_29]], %[[VAL_36:.*]] = %[[VAL_32]], %[[VAL_37:.*]] = %[[VAL_28]]) : (index, index, index) -> (index, index, index) {
-// CHECK:               %[[VAL_38:.*]] = arith.cmpi ult, %[[VAL_35]], %[[VAL_31]] : index
-// CHECK:               %[[VAL_39:.*]] = arith.cmpi ult, %[[VAL_36]], %[[VAL_33]] : index
-// CHECK:               %[[VAL_40:.*]] = arith.andi %[[VAL_38]], %[[VAL_39]] : i1
-// CHECK:               scf.condition(%[[VAL_40]]) %[[VAL_35]], %[[VAL_36]], %[[VAL_37]] : index, index, index
+// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 0.000000e+00 : f32
+// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant false
+// CHECK-DAG:       %[[VAL_7:.*]] = arith.constant true
+// CHECK:           %[[VAL_8:.*]] = tensor.dim %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           %[[VAL_9:.*]] = tensor.dim %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           %[[VAL_10:.*]] = bufferization.alloc_tensor(%[[VAL_8]], %[[VAL_9]]) : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           %[[VAL_11:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_12:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_13:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_14:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_15:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf32>
+// CHECK:           %[[VAL_16:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_17:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_2]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_18:.*]] = sparse_tensor.pointers %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_19:.*]] = sparse_tensor.indices %[[VAL_1]], %[[VAL_3]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xindex>
+// CHECK:           %[[VAL_20:.*]] = sparse_tensor.values %[[VAL_1]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf32>
+// CHECK:           %[[VAL_21:.*]] = memref.alloca(%[[VAL_4]]) : memref<?xindex>
+// CHECK:           %[[VAL_22:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:           %[[VAL_23:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:           scf.for %[[VAL_24:.*]] = %[[VAL_22]] to %[[VAL_23]] step %[[VAL_3]] {
+// CHECK:             %[[VAL_25:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_24]]] : memref<?xindex>
+// CHECK:             memref.store %[[VAL_25]], %[[VAL_21]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:             %[[VAL_26:.*]], %[[VAL_27:.*]], %[[VAL_28:.*]], %[[VAL_29:.*]] = sparse_tensor.expand %[[VAL_10]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>> to memref<?xf32>, memref<?xi1>, memref<?xindex>, index
+// CHECK:             %[[VAL_30:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_24]]] : memref<?xindex>
+// CHECK:             %[[VAL_31:.*]] = arith.addi %[[VAL_24]], %[[VAL_3]] : index
+// CHECK:             %[[VAL_32:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_31]]] : memref<?xindex>
+// CHECK:             %[[VAL_33:.*]] = memref.load %[[VAL_16]]{{\[}}%[[VAL_2]]] : memref<?xindex>
+// CHECK:             %[[VAL_34:.*]] = memref.load %[[VAL_16]]{{\[}}%[[VAL_3]]] : memref<?xindex>
+// CHECK:             %[[VAL_35:.*]]:3 = scf.while (%[[VAL_36:.*]] = %[[VAL_30]], %[[VAL_37:.*]] = %[[VAL_33]], %[[VAL_38:.*]] = %[[VAL_29]]) : (index, index, index) -> (index, index, index) {
+// CHECK:               %[[VAL_39:.*]] = arith.cmpi ult, %[[VAL_36]], %[[VAL_32]] : index
+// CHECK:               %[[VAL_40:.*]] = arith.cmpi ult, %[[VAL_37]], %[[VAL_34]] : index
+// CHECK:               %[[VAL_41:.*]] = arith.andi %[[VAL_39]], %[[VAL_40]] : i1
+// CHECK:               scf.condition(%[[VAL_41]]) %[[VAL_36]], %[[VAL_37]], %[[VAL_38]] : index, index, index
 // CHECK:             } do {
-// CHECK:             ^bb0(%[[VAL_41:.*]]: index, %[[VAL_42:.*]]: index, %[[VAL_43:.*]]: index):
-// CHECK:               %[[VAL_44:.*]] = memref.load %[[VAL_13]]{{\[}}%[[VAL_41]]] : memref<?xindex>
-// CHECK:               %[[VAL_45:.*]] = memref.load %[[VAL_16]]{{\[}}%[[VAL_42]]] : memref<?xindex>
-// CHECK:               %[[VAL_46:.*]] = arith.cmpi ult, %[[VAL_45]], %[[VAL_44]] : index
-// CHECK:               %[[VAL_47:.*]] = arith.select %[[VAL_46]], %[[VAL_45]], %[[VAL_44]] : index
-// CHECK:               %[[VAL_48:.*]] = arith.cmpi eq, %[[VAL_44]], %[[VAL_47]] : index
-// CHECK:               %[[VAL_49:.*]] = arith.cmpi eq, %[[VAL_45]], %[[VAL_47]] : index
-// CHECK:               %[[VAL_50:.*]] = arith.andi %[[VAL_48]], %[[VAL_49]] : i1
-// CHECK:               %[[VAL_51:.*]] = scf.if %[[VAL_50]] -> (index) {
-// CHECK:                 %[[VAL_52:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_41]]] : memref<?xf32>
-// CHECK:                 %[[VAL_53:.*]] = memref.load %[[VAL_17]]{{\[}}%[[VAL_42]]] : memref<?xindex>
-// CHECK:                 %[[VAL_54:.*]] = arith.addi %[[VAL_42]], %[[VAL_3]] : index
-// CHECK:                 %[[VAL_55:.*]] = memref.load %[[VAL_17]]{{\[}}%[[VAL_54]]] : memref<?xindex>
-// CHECK:                 %[[VAL_56:.*]] = scf.for %[[VAL_57:.*]] = %[[VAL_53]] to %[[VAL_55]] step %[[VAL_3]] iter_args(%[[VAL_58:.*]] = %[[VAL_43]]) -> (index) {
-// CHECK:                   %[[VAL_59:.*]] = memref.load %[[VAL_18]]{{\[}}%[[VAL_57]]] : memref<?xindex>
-// CHECK:                   %[[VAL_60:.*]] = memref.load %[[VAL_25]]{{\[}}%[[VAL_59]]] : memref<?xf32>
-// CHECK:                   %[[VAL_61:.*]] = memref.load %[[VAL_19]]{{\[}}%[[VAL_57]]] : memref<?xf32>
-// CHECK:                   %[[VAL_62:.*]] = arith.mulf %[[VAL_52]], %[[VAL_61]] : f32
-// CHECK:                   %[[VAL_63:.*]] = arith.addf %[[VAL_60]], %[[VAL_62]] : f32
-// CHECK:                   %[[VAL_64:.*]] = memref.load %[[VAL_26]]{{\[}}%[[VAL_59]]] : memref<?xi1>
-// CHECK:                   %[[VAL_65:.*]] = arith.cmpi eq, %[[VAL_64]], %[[VAL_5]] : i1
-// CHECK:                   %[[VAL_66:.*]] = scf.if %[[VAL_65]] -> (index) {
-// CHECK:                     memref.store %[[VAL_6]], %[[VAL_26]]{{\[}}%[[VAL_59]]] : memref<?xi1>
-// CHECK:                     memref.store %[[VAL_59]], %[[VAL_27]]{{\[}}%[[VAL_58]]] : memref<?xindex>
-// CHECK:                     %[[VAL_67:.*]] = arith.addi %[[VAL_58]], %[[VAL_3]] : index
-// CHECK:                     scf.yield %[[VAL_67]] : index
+// CHECK:             ^bb0(%[[VAL_42:.*]]: index, %[[VAL_43:.*]]: index, %[[VAL_44:.*]]: index):
+// CHECK:               %[[VAL_45:.*]] = memref.load %[[VAL_14]]{{\[}}%[[VAL_42]]] : memref<?xindex>
+// CHECK:               %[[VAL_46:.*]] = memref.load %[[VAL_17]]{{\[}}%[[VAL_43]]] : memref<?xindex>
+// CHECK:               %[[VAL_47:.*]] = arith.cmpi ult, %[[VAL_46]], %[[VAL_45]] : index
+// CHECK:               %[[VAL_48:.*]] = arith.select %[[VAL_47]], %[[VAL_46]], %[[VAL_45]] : index
+// CHECK:               %[[VAL_49:.*]] = arith.cmpi eq, %[[VAL_45]], %[[VAL_48]] : index
+// CHECK:               %[[VAL_50:.*]] = arith.cmpi eq, %[[VAL_46]], %[[VAL_48]] : index
+// CHECK:               %[[VAL_51:.*]] = arith.andi %[[VAL_49]], %[[VAL_50]] : i1
+// CHECK:               %[[VAL_52:.*]] = scf.if %[[VAL_51]] -> (index) {
+// CHECK:                 %[[VAL_53:.*]] = memref.load %[[VAL_15]]{{\[}}%[[VAL_42]]] : memref<?xf32>
+// CHECK:                 %[[VAL_54:.*]] = memref.load %[[VAL_18]]{{\[}}%[[VAL_43]]] : memref<?xindex>
+// CHECK:                 %[[VAL_55:.*]] = arith.addi %[[VAL_43]], %[[VAL_3]] : index
+// CHECK:                 %[[VAL_56:.*]] = memref.load %[[VAL_18]]{{\[}}%[[VAL_55]]] : memref<?xindex>
+// CHECK:                 %[[VAL_57:.*]] = scf.for %[[VAL_58:.*]] = %[[VAL_54]] to %[[VAL_56]] step %[[VAL_3]] iter_args(%[[VAL_59:.*]] = %[[VAL_44]]) -> (index) {
+// CHECK:                   %[[VAL_60:.*]] = memref.load %[[VAL_19]]{{\[}}%[[VAL_58]]] : memref<?xindex>
+// CHECK:                   %[[VAL_61:.*]] = memref.load %[[VAL_27]]{{\[}}%[[VAL_60]]] : memref<?xi1>
+// CHECK:                   %[[VAL_62:.*]] = scf.if %[[VAL_61]] -> (f32) {
+// CHECK:                     %[[VAL_63:.*]] = memref.load %[[VAL_26]]{{\[}}%[[VAL_60]]] : memref<?xf32>
+// CHECK:                     scf.yield %[[VAL_63]] : f32
 // CHECK:                   } else {
-// CHECK:                     scf.yield %[[VAL_58]] : index
+// CHECK:                     scf.yield %[[VAL_5]] : f32
 // CHECK:                   }
-// CHECK:                   memref.store %[[VAL_63]], %[[VAL_25]]{{\[}}%[[VAL_59]]] : memref<?xf32>
-// CHECK:                   scf.yield %[[VAL_68:.*]] : index
+// CHECK:                   %[[VAL_64:.*]] = memref.load %[[VAL_20]]{{\[}}%[[VAL_58]]] : memref<?xf32>
+// CHECK:                   %[[VAL_65:.*]] = arith.mulf %[[VAL_53]], %[[VAL_64]] : f32
+// CHECK:                   %[[VAL_66:.*]] = arith.addf %[[VAL_67:.*]], %[[VAL_65]] : f32
+// CHECK:                   %[[VAL_68:.*]] = memref.load %[[VAL_27]]{{\[}}%[[VAL_60]]] : memref<?xi1>
+// CHECK:                   %[[VAL_69:.*]] = arith.cmpi eq, %[[VAL_68]], %[[VAL_6]] : i1
+// CHECK:                   %[[VAL_70:.*]] = scf.if %[[VAL_69]] -> (index) {
+// CHECK:                     memref.store %[[VAL_7]], %[[VAL_27]]{{\[}}%[[VAL_60]]] : memref<?xi1>
+// CHECK:                     memref.store %[[VAL_60]], %[[VAL_28]]{{\[}}%[[VAL_59]]] : memref<?xindex>
+// CHECK:                     %[[VAL_71:.*]] = arith.addi %[[VAL_59]], %[[VAL_3]] : index
+// CHECK:                     scf.yield %[[VAL_71]] : index
+// CHECK:                   } else {
+// CHECK:                     scf.yield %[[VAL_59]] : index
+// CHECK:                   }
+// CHECK:                   memref.store %[[VAL_66]], %[[VAL_26]]{{\[}}%[[VAL_60]]] : memref<?xf32>
+// CHECK:                   scf.yield %[[VAL_72:.*]] : index
 // CHECK:                 }
-// CHECK:                 scf.yield %[[VAL_69:.*]] : index
+// CHECK:                 scf.yield %[[VAL_73:.*]] : index
 // CHECK:               } else {
-// CHECK:                 scf.yield %[[VAL_43]] : index
+// CHECK:                 scf.yield %[[VAL_44]] : index
 // CHECK:               }
-// CHECK:               %[[VAL_70:.*]] = arith.cmpi eq, %[[VAL_44]], %[[VAL_47]] : index
-// CHECK:               %[[VAL_71:.*]] = arith.addi %[[VAL_41]], %[[VAL_3]] : index
-// CHECK:               %[[VAL_72:.*]] = arith.select %[[VAL_70]], %[[VAL_71]], %[[VAL_41]] : index
-// CHECK:               %[[VAL_73:.*]] = arith.cmpi eq, %[[VAL_45]], %[[VAL_47]] : index
-// CHECK:               %[[VAL_74:.*]] = arith.addi %[[VAL_42]], %[[VAL_3]] : index
-// CHECK:               %[[VAL_75:.*]] = arith.select %[[VAL_73]], %[[VAL_74]], %[[VAL_42]] : index
-// CHECK:               scf.yield %[[VAL_72]], %[[VAL_75]], %[[VAL_76:.*]] : index, index, index
+// CHECK:               %[[VAL_74:.*]] = arith.cmpi eq, %[[VAL_45]], %[[VAL_48]] : index
+// CHECK:               %[[VAL_75:.*]] = arith.addi %[[VAL_42]], %[[VAL_3]] : index
+// CHECK:               %[[VAL_76:.*]] = arith.select %[[VAL_74]], %[[VAL_75]], %[[VAL_42]] : index
+// CHECK:               %[[VAL_77:.*]] = arith.cmpi eq, %[[VAL_46]], %[[VAL_48]] : index
+// CHECK:               %[[VAL_78:.*]] = arith.addi %[[VAL_43]], %[[VAL_3]] : index
+// CHECK:               %[[VAL_79:.*]] = arith.select %[[VAL_77]], %[[VAL_78]], %[[VAL_43]] : index
+// CHECK:               scf.yield %[[VAL_76]], %[[VAL_79]], %[[VAL_80:.*]] : index, index, index
 // CHECK:             }
-// CHECK:             sparse_tensor.compress %[[VAL_9]], %[[VAL_20]], %[[VAL_25]], %[[VAL_26]], %[[VAL_27]], %[[VAL_77:.*]]#2 : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>, memref<?xindex>, memref<?xf32>, memref<?xi1>, memref<?xindex>, index
+// CHECK:             sparse_tensor.compress %[[VAL_10]], %[[VAL_21]], %[[VAL_26]], %[[VAL_27]], %[[VAL_28]], %[[VAL_81:.*]]#2 : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>, memref<?xindex>, memref<?xf32>, memref<?xi1>, memref<?xindex>, index
 // CHECK:           }
-// CHECK:           %[[VAL_78:.*]] = sparse_tensor.load %[[VAL_9]] hasInserts : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
-// CHECK:           return %[[VAL_78]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           %[[VAL_82:.*]] = sparse_tensor.load %[[VAL_10]] hasInserts : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
+// CHECK:           return %[[VAL_82]] : tensor<?x?xf32, #sparse_tensor.encoding<{{{.*}}}>>
 // CHECK:         }
 func.func @matmat(%arga: tensor<?x?xf32, #DCSR>,
              %argb: tensor<?x?xf32, #DCSR>) -> tensor<?x?xf32, #DCSR> {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %d0 = tensor.dim %arga, %c0 : tensor<?x?xf32, #DCSR>
   %d1 = tensor.dim %argb, %c1 : tensor<?x?xf32, #DCSR>
   %cinit = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf32, #DCSR>

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_out_reduction.mlir

 // RUN: mlir-opt %s --sparse-compiler | \
 // RUN: mlir-cpu-runner \
 // RUN:  -e entry -entry-point-result=void  \
 // RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
+#SparseVector = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed" ]
+}>
+
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed" ]
 }>
 
+#SparseCSCMatrix = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
 #SparseTensor = #sparse_tensor.encoding<{
   dimLevelType = [ "compressed", "compressed", "compressed" ]
 }>
 
 #redsum = {
   indexing_maps = [
     affine_map<(i,j,k) -> (i,j,k)>, // A
     affine_map<(i,j,k) -> (i,j,k)>, // B
     affine_map<(i,j,k) -> (i,j)>    // X (out)
   ],
   iterator_types = ["parallel", "parallel", "reduction"],
   doc = "X(i,j) = SUM_k A(i,j,k) * B(i,j,k)"
 }
 
+#redprod = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>, // A
+    affine_map<(i,j) -> (i)>    // X (out)
+  ],
+  iterator_types = ["parallel", "reduction"],
+  doc = "X(i) = PROD_j A(i,j)"
+}
+
 module {
   func.func @redsum(%arga: tensor<?x?x?xi32, #SparseTensor>,
                %argb: tensor<?x?x?xi32, #SparseTensor>)
              -> tensor<?x?xi32, #SparseMatrix> {
     %c0 = arith.constant 0 : index
     %c1 = arith.constant 1 : index
     %d0 = tensor.dim %arga, %c0 : tensor<?x?x?xi32, #SparseTensor>
     %d1 = tensor.dim %arga, %c1 : tensor<?x?x?xi32, #SparseTensor>
     %xinit = bufferization.alloc_tensor(%d0, %d1): tensor<?x?xi32, #SparseMatrix>
     %0 = linalg.generic #redsum
       ins(%arga, %argb: tensor<?x?x?xi32, #SparseTensor>,
                         tensor<?x?x?xi32, #SparseTensor>)
       outs(%xinit: tensor<?x?xi32, #SparseMatrix>) {
         ^bb(%a: i32, %b: i32, %x: i32):
           %0 = arith.muli %a, %b : i32
           %1 = arith.addi %x, %0 : i32
           linalg.yield %1 : i32
     } -> tensor<?x?xi32, #SparseMatrix>
     return %0 : tensor<?x?xi32, #SparseMatrix>
   }
 
+  func.func @redprod(%arga: tensor<?x?xi32, #SparseMatrix>) -> tensor<?xi32, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d0 = tensor.dim %arga, %c0 : tensor<?x?xi32, #SparseMatrix>
+    %xinit = bufferization.alloc_tensor(%d0): tensor<?xi32, #SparseVector>
+    %0 = linalg.generic #redprod
+      ins(%arga: tensor<?x?xi32, #SparseMatrix>)
+      outs(%xinit: tensor<?xi32, #SparseVector>) {
+        ^bb(%a: i32, %x: i32):
+          %0 = arith.muli %x, %a : i32
+          linalg.yield %0 : i32
+    } -> tensor<?xi32, #SparseVector>
+    return %0 : tensor<?xi32, #SparseVector>
+  }
+
+  func.func @redprod2(%arga: tensor<?x?xi32, #SparseCSCMatrix>) -> tensor<?xi32, #SparseVector> {
+    %c0 = arith.constant 0 : index
+    %d0 = tensor.dim %arga, %c0 : tensor<?x?xi32, #SparseCSCMatrix>
+    %xinit = bufferization.alloc_tensor(%d0): tensor<?xi32, #SparseVector>
+    %0 = linalg.generic #redprod
+      ins(%arga: tensor<?x?xi32, #SparseCSCMatrix>)
+      outs(%xinit: tensor<?xi32, #SparseVector>) {
+        ^bb(%a: i32, %x: i32):
+          %0 = arith.muli %x, %a : i32
+          linalg.yield %0 : i32
+    } -> tensor<?xi32, #SparseVector>
+    return %0 : tensor<?xi32, #SparseVector>
+  }
+
+  // Dumps a sparse vector.
+  func.func @dumpvec(%arg0: tensor<?xi32, #SparseVector>) {
+    // Dump the values array to verify only sparse contents are stored.
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1 : i32
+    %0 = sparse_tensor.values %arg0 : tensor<?xi32, #SparseVector> to memref<?xi32>
+    %1 = vector.transfer_read %0[%c0], %d0: memref<?xi32>, vector<4xi32>
+    vector.print %1 : vector<4xi32>
+    // Dump the dense vector to verify structure is correct.
+    %dv = sparse_tensor.convert %arg0 : tensor<?xi32, #SparseVector> to tensor<?xi32>
+    %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xi32>, vector<4xi32>
+    vector.print %2 : vector<4xi32>
+    return
+  }
+
   // Driver method to call and verify tensor kernel.
   func.func @entry() {
     %c0 = arith.constant 0 : index
     %i0 = arith.constant -1 : i32
 
     // Setup very sparse 3-d tensors.
     %t1 = arith.constant sparse<
        [ [1,1,3], [2,0,0], [2,2,1], [2,2,2], [2,2,3] ], [ 1, 2, 3, 4, 5 ]
     > : tensor<3x3x4xi32>
     %t2 = arith.constant sparse<
        [ [1,0,0], [1,1,3], [2,2,1], [2,2,3] ], [ 6, 7, 8, 9 ]
     > : tensor<3x3x4xi32>
     %st1 = sparse_tensor.convert %t1
       : tensor<3x3x4xi32> to tensor<?x?x?xi32, #SparseTensor>
     %st2 = sparse_tensor.convert %t2
       : tensor<3x3x4xi32> to tensor<?x?x?xi32, #SparseTensor>
 
+    // Setup sparse 2-d tensors.
+    %m1 = arith.constant sparse<
+       [ [0, 3], [0, 4], [2, 3], [3, 0], [3, 2], [3, 4] ], [ 1, 2, 3, 4, 5, 6 ]
+    >: tensor<4x5xi32>
+    %smr = sparse_tensor.convert %m1
+      : tensor<4x5xi32> to tensor<?x?xi32, #SparseMatrix>
+    %smc = sparse_tensor.convert %m1
+      : tensor<4x5xi32> to tensor<?x?xi32, #SparseCSCMatrix>
+
     // Call kernel.
     %0 = call @redsum(%st1, %st2)
       : (tensor<?x?x?xi32, #SparseTensor>,
          tensor<?x?x?xi32, #SparseTensor>) -> tensor<?x?xi32, #SparseMatrix>
+    %1 = call @redprod(%smr)
+      : (tensor<?x?xi32, #SparseMatrix>) -> tensor<?xi32, #SparseVector>
+    %2 = call @redprod2(%smc)
+      : (tensor<?x?xi32, #SparseCSCMatrix>) -> tensor<?xi32, #SparseVector>
 
     //
     // Verify results. Only two entries stored in result. Correct structure.
     //
     // CHECK: ( 7, 69, -1, -1 )
     // CHECK-NEXT: ( ( 0, 0, 0 ), ( 0, 7, 0 ), ( 0, 0, 69 ) )
+    // CHECK-NEXT: ( 2, 3, 120, -1 )
+    // CHECK-NEXT: ( 2, 0, 3, 120 )
+    // CHECK-NEXT: ( 2, 3, 120, -1 )
+    // CHECK-NEXT: ( 2, 0, 3, 120 )
     //
     %val = sparse_tensor.values %0
       : tensor<?x?xi32, #SparseMatrix> to memref<?xi32>
     %vv = vector.transfer_read %val[%c0], %i0: memref<?xi32>, vector<4xi32>
     vector.print %vv : vector<4xi32>
     %dm = sparse_tensor.convert %0
       : tensor<?x?xi32, #SparseMatrix> to tensor<?x?xi32>
     %vm = vector.transfer_read %dm[%c0, %c0], %i0: tensor<?x?xi32>, vector<3x3xi32>
     vector.print %vm : vector<3x3xi32>
+    call @dumpvec(%1) : (tensor<?xi32, #SparseVector>) -> ()
+    call @dumpvec(%2) : (tensor<?xi32, #SparseVector>) -> ()
 
     // Release the resources.
     sparse_tensor.release %st1 : tensor<?x?x?xi32, #SparseTensor>
     sparse_tensor.release %st2 : tensor<?x?x?xi32, #SparseTensor>
     sparse_tensor.release %0 : tensor<?x?xi32, #SparseMatrix>
+    sparse_tensor.release %1 : tensor<?xi32, #SparseVector>
+    sparse_tensor.release %2 : tensor<?xi32, #SparseVector>
     return
   }
 }

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_reduce_custom.mlir

+// RUN: mlir-opt %s --sparse-compiler | \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+#CSR = #sparse_tensor.encoding<{dimLevelType = ["dense", "compressed"]}>
+#CSC = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+//
+// Traits for tensor operations.
+//
+#trait_matmul = {
+  indexing_maps = [
+    affine_map<(i,j,k) -> (i,k)>, // A
+    affine_map<(i,j,k) -> (k,j)>, // B
+    affine_map<(i,j,k) -> (i,j)>  // C (out)
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"],
+  doc = "C(i,j) = SUM_k A(i,k) * B(k,j)"
+}
+
+#trait_mat_reduce = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A (in)
+    affine_map<(i,j) -> (i)>   // X (out)
+  ],
+  iterator_types = ["parallel", "reduce"]
+}
+
+module {
+  // Creates a new sparse vector using the minimum values from two input sparse vectors.
+  // When there is no overlap, include the present value in the output.
+  func.func @min_plus_csrcsr(%arga: tensor<?x?xf64, #CSR>,
+                             %argb: tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR> {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %maxf = arith.constant 1.0e999 : f64
+    %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
+    %d1 = tensor.dim %argb, %c1 : tensor<?x?xf64, #CSR>
+    %xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #CSR>
+    %0 = linalg.generic #trait_matmul
+       ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSR>)
+        outs(%xm: tensor<?x?xf64, #CSR>) {
+        ^bb(%a: f64, %b: f64, %output: f64):
+          %1 = sparse_tensor.binary %a, %b : f64, f64 to f64
+            overlap = {
+              ^bb0(%x: f64, %y: f64):
+                %3 = arith.addf %x, %y : f64
+                sparse_tensor.yield %3 : f64
+            }
+            left={}
+            right={}
+          %2 = sparse_tensor.reduce %1, %output, %maxf : f64 {
+              ^bb0(%x: f64, %y: f64):
+                %cmp = arith.cmpf "olt", %x, %y : f64
+                %3 = arith.select %cmp, %x, %y : f64
+                sparse_tensor.yield %3 : f64
+            }
+          linalg.yield %2 : f64
+    } -> tensor<?x?xf64, #CSR>
+    return %0 : tensor<?x?xf64, #CSR>
+  }
+
+  func.func @min_plus_csrcsc(%arga: tensor<?x?xf64, #CSR>,
+                             %argb: tensor<?x?xf64, #CSC>) -> tensor<?x?xf64, #CSR> {
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %maxf = arith.constant 1.0e999 : f64
+    %d0 = tensor.dim %arga, %c0 : tensor<?x?xf64, #CSR>
+    %d1 = tensor.dim %argb, %c1 : tensor<?x?xf64, #CSC>
+    %xm = bufferization.alloc_tensor(%d0, %d1) : tensor<?x?xf64, #CSR>
+    %0 = linalg.generic #trait_matmul
+       ins(%arga, %argb: tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSC>)
+        outs(%xm: tensor<?x?xf64, #CSR>) {
+        ^bb(%a: f64, %b: f64, %output: f64):
+          %1 = sparse_tensor.binary %a, %b : f64, f64 to f64
+            overlap = {
+              ^bb0(%x: f64, %y: f64):
+                %3 = arith.addf %x, %y : f64
+                sparse_tensor.yield %3 : f64
+            }
+            left={}
+            right={}
+          %2 = sparse_tensor.reduce %1, %output, %maxf : f64 {
+              ^bb0(%x: f64, %y: f64):
+                %cmp = arith.cmpf "olt", %x, %y : f64
+                %3 = arith.select %cmp, %x, %y : f64
+                sparse_tensor.yield %3 : f64
+            }
+          linalg.yield %2 : f64
+    } -> tensor<?x?xf64, #CSR>
+    return %0 : tensor<?x?xf64, #CSR>
+  }
+
+  // Dumps a sparse vector of type f64.
+  func.func @dump_vec(%arg0: tensor<?xf64, #SparseVector>) {
+    // Dump the values array to verify only sparse contents are stored.
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f64
+    %0 = sparse_tensor.values %arg0 : tensor<?xf64, #SparseVector> to memref<?xf64>
+    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64>
+    vector.print %1 : vector<16xf64>
+    // Dump the dense vector to verify structure is correct.
+    %dv = sparse_tensor.convert %arg0 : tensor<?xf64, #SparseVector> to tensor<?xf64>
+    %2 = vector.transfer_read %dv[%c0], %d0: tensor<?xf64>, vector<32xf64>
+    vector.print %2 : vector<32xf64>
+    return
+  }
+
+  // Dump a sparse matrix.
+  func.func @dump_mat(%arg0: tensor<?x?xf64, #CSR>) {
+    // Dump the values array to verify only sparse contents are stored.
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f64
+    %0 = sparse_tensor.values %arg0 : tensor<?x?xf64, #CSR> to memref<?xf64>
+    %1 = vector.transfer_read %0[%c0], %d0: memref<?xf64>, vector<16xf64>
+    vector.print %1 : vector<16xf64>
+    %dm = sparse_tensor.convert %arg0 : tensor<?x?xf64, #CSR> to tensor<?x?xf64>
+    %2 = vector.transfer_read %dm[%c0, %c0], %d0: tensor<?x?xf64>, vector<5x5xf64>
+    vector.print %2 : vector<5x5xf64>
+    return
+  }
+
+  // Driver method to call and verify vector kernels.
+  func.func @entry() {
+    %c0 = arith.constant 0 : index
+
+    // Setup sparse matrices.
+    %m1 = arith.constant sparse<
+       [ [0,0], [0,1], [1,0], [2,2], [2,3], [2,4], [3,0], [3,2], [3,3] ],
+         [ 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0 ]
+    > : tensor<4x5xf64>
+    %m2 = arith.constant sparse<
+       [ [0,0], [1,3], [2,0], [2,3], [3,1], [4,1] ],
+         [6.0, 5.0, 4.0, 3.0, 2.0, 11.0 ]
+    > : tensor<5x4xf64>
+    %sm1 = sparse_tensor.convert %m1 : tensor<4x5xf64> to tensor<?x?xf64, #CSR>
+    %sm2r = sparse_tensor.convert %m2 : tensor<5x4xf64> to tensor<?x?xf64, #CSR>
+    %sm2c = sparse_tensor.convert %m2 : tensor<5x4xf64> to tensor<?x?xf64, #CSC>
+
+    // Call sparse matrix kernels.
+    %5 = call @min_plus_csrcsr(%sm1, %sm2r)
+      : (tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSR>) -> tensor<?x?xf64, #CSR>
+    // COM: This is broken because the lex-insert version of matmul always inserts the
+    // COM:   identity value, even if there is nothing to accumulate, resulting in a
+    // COM:   dense output.
+    // COM: %6 = call @min_plus_csrcsc(%sm1, %sm2c)
+    // COM:   : (tensor<?x?xf64, #CSR>, tensor<?x?xf64, #CSC>) -> tensor<?x?xf64, #CSR>
+
+    //
+    // Verify the results.
+    //
+    // CHECK:      ( 1, 2, 3, 4, 5, 6, 7, 8, 9, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( ( 1, 2, 0, 0, 0 ), ( 3, 0, 0, 0, 0 ), ( 0, 0, 4, 5, 6 ), ( 7, 0, 8, 9, 0 ), ( -1, -1, -1, -1, -1 ) )
+    // CHECK-NEXT: ( 6, 5, 4, 3, 2, 11, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( ( 6, 0, 0, 0, -1 ), ( 0, 0, 0, 5, -1 ), ( 4, 0, 0, 3, -1 ), ( 0, 2, 0, 0, -1 ), ( 0, 11, 0, 0, -1 ) )
+    // CHECK-NEXT: ( 7, 7, 9, 8, 7, 7, 12, 11, 11, -1, -1, -1, -1, -1, -1, -1 )
+    // CHECK-NEXT: ( ( 7, 0, 0, 7, -1 ), ( 9, 0, 0, 0, -1 ), ( 8, 7, 0, 7, -1 ), ( 12, 11, 0, 11, -1 ), ( -1, -1, -1, -1, -1 ) )
+    // COM: CHECK-NEXT: ( 7, 7, 9, 8, 7, 7, 12, 11, 11, -1, -1, -1, -1, -1, -1, -1 )
+    // COM: CHECK-NEXT: ( ( 7, 0, 0, 7, -1 ), ( 9, 0, 0, 0, -1 ), ( 8, 7, 0, 7, -1 ), ( 12, 11, 0, 11, -1 ), ( -1, -1, -1, -1, -1 ) )
+    //
+    call @dump_mat(%sm1) : (tensor<?x?xf64, #CSR>) -> ()
+    call @dump_mat(%sm2r) : (tensor<?x?xf64, #CSR>) -> ()
+    call @dump_mat(%5) : (tensor<?x?xf64, #CSR>) -> ()
+    // COM: call @dump_mat(%6) : (tensor<?x?xf64, #CSR>) -> ()
+
+    // Release the resources.
+    sparse_tensor.release %sm1 : tensor<?x?xf64, #CSR>
+    sparse_tensor.release %sm2r : tensor<?x?xf64, #CSR>
+    sparse_tensor.release %sm2c : tensor<?x?xf64, #CSC>
+    sparse_tensor.release %5 : tensor<?x?xf64, #CSR>
+    // COM: sparse_tensor.release %6 : tensor<?x?xf64, #CSR>
+    return
+  }
+}

mlir/unittests/Dialect/SparseTensor/MergerTest.cpp

bool compareExpression(unsigned e, const std::shared_ptr<Pattern> &pattern) {
     case kTruncI:
     case kCIm:
     case kCRe:
     case kBitCast:
     case kBinaryBranch:
     case kUnary:
     case kShlI:
     case kBinary:
+    case kReduce:
       return compareExpression(tensorExp.children.e0, pattern->e0);
     // Binary operations.
     case kMulF:
     case kMulC:
     case kMulI:
     case kDivF:
     case kDivC:
     case kDivS: