[mlir][sparse] reject kernels with non-sparsfiable reduction expression.

[Peiming]

Summary

To address https://github.com/llvm/llvm-project/issues/59394.

Reduction on negation of the output tensor is a non-sparsifiable kernel, it creates cyclic dependency.

This patch reject those cases instead of crashing.

[Peiming]

small fix

[aartbik]

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

588

It would then require ... (no s)

604

remove empty line

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

456

This should be an OR, right, not an AND?

out = (-out) + a

460

could we write this more concisely by testing that for every sub, check that LHS operand does not contain tensor, either direct, or deeper nested.

[qcolombet]

Thanks for the quick fix!
Could you add a test case as well? (E.g., the one from the bug.)

[Peiming]

address comments.

[Peiming]

In D139659#3982710, @qcolombet wrote:

Thanks for the quick fix!
Could you add a test case as well? (E.g., the one from the bug.)

The patch simply does not crash, your test file will still not be compiled ;-(

[aartbik]

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

275

nit, just t is the style in this file

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

99

move typedef and static helper *before* the Constructors section

(and leave empty vspace)

388

super nit: but we use unsigned t in this file for tensor (tid is your loop generator convention)

[Peiming]

address comments.

[wrengr]

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

27

+1 :)

[wrengr]

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

389

nit: since you only use this variable as the scrutinee of the switch statement, it'd be nicer to inline it there

423

ditto

[aartbik]

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

274

last nit: Returns

277

last nit: Returns

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

389

yeah, agreed, it will easily fit inside the switch (..)

[qcolombet]

In D139659#3982793, @Peiming wrote:

In D139659#3982710, @qcolombet wrote:

Thanks for the quick fix!
Could you add a test case as well? (E.g., the one from the bug.)

The patch simply does not crash, your test file will still not be compiled ;-(

You could still check that the IR is left untouched.
I assume that’s what you mean by not compiled.

[Peiming]

add rejected test case

[qcolombet]

You could still check that the IR is left untouched.

And you could also do that on more complicated expressions to test hasNegateOnOut more thoroughly (like out = add in, (mul out, -1)).

[Peiming]

address comments

[Peiming]

add some comments in rejected.mlir

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

public:
   /// expressions).
   unsigned getSynTensorID() const { return syntheticTensor; }
 
   bool isFilterLoop(unsigned ldx) const {
     assert(ldx < numLoops);
     return ldx >= numNativeLoops;
   }
 
+  /// Returns true if the expression contains the `t` as an operand.
+  bool expContainsTensor(unsigned e, unsigned t) const;
+
+  /// Returns true if the expression contains a negation on output tensor.
+  /// I.e., `- outTensor` or `exp - outputTensor`
+  /// NOTE: this is an trivial tests in that it does not handle recursive
+  /// negation, i.e., it returns true when the expression is `-(-tensor)`.
+  bool hasNegateOnOut(unsigned e) const;
+
   /// Returns true if given tensor iterates *only* in the given tensor
   /// expression. For the output tensor, this defines a "simply dynamic"
   /// operation [Bik96]. For instance: a(i) *= 2.0 or a(i) += a(i) for
   /// sparse vector a.
   bool isSingleCondition(unsigned t, unsigned e) const;
 
   /// Returns true if any set bit corresponds to sparse dimension level type.
   bool hasAnySparse(const BitVector &bits) const;

 #ifndef NDEBUG
   /// Print methods (for debugging).
   void dumpExp(unsigned e) const;
   void dumpLat(unsigned p) const;
   void dumpSet(unsigned s) const;
   void dumpBits(const BitVector &bits) const;
 #endif
 
-  /// Builds the iteration lattices in a bottom-up traversal given the remaining
-  /// tensor (sub)expression and the next loop index in the iteration graph.
-  /// Returns index of the root expression.
+  /// Builds the iteration lattices in a bottom-up traversal given the
+  /// remaining tensor (sub)expression and the next loop index in the
+  /// iteration graph. Returns index of the root expression.
   unsigned buildLattices(unsigned e, unsigned i);
 
   /// Builds a tensor expression from the given Linalg operation.
   /// Returns index of the root expression on success.
   Optional<unsigned> buildTensorExpFromLinalg(linalg::GenericOp op);
 
   /// Rebuilds SSA format from a tensor expression.
   Value buildExp(RewriterBase &rewriter, Location loc, unsigned e, Value v0,
private:
   const unsigned syntheticTensor;
   const unsigned numTensors;
   const unsigned numNativeLoops;
   const unsigned numLoops;
   bool hasSparseOut;
   // Map that converts pair<tensor id, loop id> to the corresponding dimension
   // level type.
   std::vector<std::vector<DimLevelType>> dimTypes;
-  // Map that converts pair<tensor id, loop id> to the corresponding dimension.
+  // Map that converts pair<tensor id, loop id> to the corresponding
+  // dimension.
   std::vector<std::vector<Optional<unsigned>>> loopIdxToDim;
   // Map that converts pair<tensor id, dim> to the corresponding loop id.
   std::vector<std::vector<Optional<unsigned>>> dimToLoopIdx;
   llvm::SmallVector<TensorExp> tensorExps;
   llvm::SmallVector<LatPoint> latPoints;
   llvm::SmallVector<SmallVector<unsigned>> latSets;
 };
 
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_UTILS_MERGER_H_

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

 /// Since all sparse input tensors are admissible, we just need to check
 /// whether the out tensor in the tensor expression codegen is admissible.
 /// Sets `sparseOut` to the tensor and `outerParNest` to the outer injective
 /// nesting depth when a "truly dynamic" sparse tensor output occurs.
 static bool isAdmissibleTensorExp(Merger &merger, linalg::GenericOp op,
                                   std::vector<unsigned> &topSort, unsigned exp,
                                   OpOperand **sparseOut,
                                   unsigned &outerParNest) {
+  // We reject any expression that makes a reduction from `-outTensor`, as those
+  // expression create dependency between the current iteration (i) and the
+  // previous iteration (i-1). It would then require iterating over the whole
+  // coordinate space, which prevent us from exploiting sparsity for faster
+  // code.
+  for (utils::IteratorType it : op.getIteratorTypesArray()) {
+    if (it == utils::IteratorType::reduction) {
+      if (merger.hasNegateOnOut(exp))
+        return false;
+      break;
+    }
+  }
+
   OpOperand *lhs = op.getDpsInitOperand(0);
   unsigned tensor = lhs->getOperandNumber();
   auto enc = getSparseTensorEncoding(lhs->get().getType());
   // An non-annotated output tensor is assumed dense, and becomes a random
   // access n-dim memref. Admissible since insertions cannot occur.
   if (!enc)
     return true;
   // An all-dense annotated "sparse" output tensor becomes a linearized random
   // access 1-dim memref. Also admissible since insertions cannot occur.
   bool allDense = true;
   unsigned numLoops = merger.getNumLoops(); // numNativeLoops + numFilterLoops
   for (unsigned i = 0; i < merger.getNumLoops(); i++)
     if (isCompressedDLT(merger.getDimLevelType(tensor, i)) ||
         isSingletonDLT(merger.getDimLevelType(tensor, i))) {

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

 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 
 #include "mlir/IR/Operation.h"
 #include "llvm/Support/Debug.h"
 
 namespace mlir {
 namespace sparse_tensor {
 
+enum class ExpArity {
+  kNullary,
+  kUnary,
+  kBinary,
+};
+
+static ExpArity getExpArity(Kind k) {
+  switch (k) {
+  // Leaf.
+  case kTensor:
+  case kInvariant:
+  case kIndex:
+    return ExpArity::kNullary;
+  case kAbsF:
+  case kAbsC:
+  case kAbsI:
+  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 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 kSelect:
+  case kNegF:
+  case kNegC:
+  case kNegI:
+    return ExpArity::kUnary;
+  // Binary operations.
+  case kDivF:
+  case kDivC:
+  case kDivS:
+  case kDivU:
+  case kShrS:
+  case kShrU:
+  case kShlI:
+  case kMulF:
+  case kMulC:
+  case kMulI:
+  case kAndI:
+  case kAddF:
+  case kAddC:
+  case kAddI:
+  case kOrI:
+  case kXorI:
+  case kBinary:
+  case kReduce:
+  case kSubF:
+  case kSubC:
+  case kSubI:
+    return ExpArity::kBinary;
+  }
+  llvm_unreachable("unexpected kind");
+}
+
 //===----------------------------------------------------------------------===//
 // 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;

 }
 
 bool Merger::onlyDenseDiff(unsigned i, unsigned j) {
   BitVector tmp = latPoints[j].bits;
   tmp ^= latPoints[i].bits;
   return !hasAnySparse(tmp);
 }
 
+bool Merger::expContainsTensor(unsigned e, unsigned t) const {
+  if (tensorExps[e].kind == kTensor)
+    return tensorExps[e].tensor == t;
+
+  switch (getExpArity(tensorExps[e].kind)) {
+  case ExpArity::kNullary:
+    return false;
+  case ExpArity::kUnary: {
+    unsigned op = tensorExps[e].children.e0;
+    if (tensorExps[op].kind == kTensor && tensorExps[op].tensor == t)
+      return true;
+    return expContainsTensor(op, t);
+  }
+  case ExpArity::kBinary: {
+    unsigned op1 = tensorExps[e].children.e0;
+    unsigned op2 = tensorExps[e].children.e1;
+    if ((tensorExps[op1].kind == kTensor && tensorExps[op1].tensor == t) ||
+        (tensorExps[op2].kind == kTensor && tensorExps[op2].tensor == t))
+      return true;
+    return expContainsTensor(op1, t) || expContainsTensor(op2, t);
+  }
+  }
+  llvm_unreachable("unexpected arity");
+}
+
+bool Merger::hasNegateOnOut(unsigned e) const {
+  switch (tensorExps[e].kind) {
+  case kNegF:
+  case kNegC:
+  case kNegI:
+    return expContainsTensor(tensorExps[e].children.e0, outTensor);
+  case kSubF:
+  case kSubC:
+  case kSubI:
+    return expContainsTensor(tensorExps[e].children.e1, outTensor) ||
+           hasNegateOnOut(tensorExps[e].children.e0);
+  default: {
+    switch (getExpArity(tensorExps[e].kind)) {
+    case ExpArity::kNullary:
+      return false;
+    case ExpArity::kUnary:
+      return hasNegateOnOut(tensorExps[e].children.e0);
+    case ExpArity::kBinary:
+      return hasNegateOnOut(tensorExps[e].children.e0) ||
+             hasNegateOnOut(tensorExps[e].children.e1);
+    }
+  }
+  }
+  llvm_unreachable("unexpected kind");
+}
+
 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 kAbsI:
   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:

mlir/test/Dialect/SparseTensor/rejected.mlir

+// RUN: mlir-opt %s -sparsification | FileCheck %s
+
+
+// The file contains examples that will be rejected by sparse compiler
+// (we expect the linalg.generic unchanged).
+#SparseVector = #sparse_tensor.encoding<{dimLevelType = ["compressed"]}>
+
+#trait = {
+  indexing_maps = [ 
+    affine_map<(i) -> (i)>,  // a (in)
+    affine_map<(i) -> ()>    // x (out)
+  ],  
+  iterator_types = ["reduction"]
+}
+
+// CHECK-LABEL:   func.func @sparse_reduction_subi(
+// CHECK-SAME:      %[[VAL_0:.*]]: tensor<i32>,
+// CHECK-SAME:      %[[VAL_1:.*]]: tensor<?xi32, #sparse_tensor.encoding<{ dimLevelType = [ "compressed" ] }>>) -> tensor<i32> {
+// CHECK:           %[[VAL_2:.*]] = linalg.generic
+// CHECK:           ^bb0(%[[VAL_3:.*]]: i32, %[[VAL_4:.*]]: i32):
+// CHECK:             %[[VAL_5:.*]] = arith.subi %[[VAL_3]], %[[VAL_4]] : i32
+// CHECK:             linalg.yield %[[VAL_5]] : i32
+// CHECK:           } -> tensor<i32>
+// CHECK:           return %[[VAL_6:.*]] : tensor<i32>
+func.func @sparse_reduction_subi(%argx: tensor<i32>,
+                             %arga: tensor<?xi32, #SparseVector>)
+ -> tensor<i32> {
+  %0 = linalg.generic #trait
+     ins(%arga: tensor<?xi32, #SparseVector>)
+      outs(%argx: tensor<i32>) {
+      ^bb(%a: i32, %x: i32):
+        // NOTE: `subi %a, %x` is the reason why the program is rejected by the sparse compiler.
+        // It is because we do not allow `-outTensor` in reduction loops as it creates cyclic
+        // dependences.
+        %t = arith.subi %a, %x: i32 
+        linalg.yield %t : i32 
+  } -> tensor<i32>
+  return %0 : tensor<i32>
+}