[mlir][sparse] fix crashes when the tensor that defines the loop bound can not be found

[Peiming]

fix errors.

[Peiming]

add test cases.

[aartbik]

mlir/test/Dialect/SparseTensor/sampled_dense_dense_matmul_with_spy.mlir

82

we should avoid the double reduction and try to get something like this to work (I think bringing in outside values into block should work)

%result = linalg.generic #trait_sampled_dense_dense

  ins(%argA, %argB: tensor<8x8xf64>, tensor<8x8xf64>)
  outs(%argS: tensor<8x8xf64, #SM>) {
    ^bb(%a: f64, %b: f64, %x: f64):
       %sel = sparse_tensor.unary %x : f64 to f64
         present={
            ^bb0(%p: f64):
              sparse_tensor.yield %f1 : f64
         }
         absent={}
       %mul = arith.mulf %a, %b :  f64
       %2 = sparse_tensor.reduce %x, %sel, %f0 : f64 {
          ^bb0(%p: f64, %q: f64):
            %add = arith.addf %p, %mul : f64
            sparse_tensor.yield %add : f64
        }
       linalg.yield %2 : f64
} -> tensor<8x8xf64, #SM>

[Peiming]

update check tests.

[aartbik]

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

548–549

else if () {
}

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h

82

typo: Optional

[Peiming]

address comments.

[Peiming]

cleanup check test.

mlir/lib/Dialect/SparseTensor/Transforms/LoopEmitter.h

 //===----------------------------------------------------------------------===//
 
 class LoopEmitter {
 public:
   /// Optional callback function to setup dense output tensors when
   /// initializing the loop emitter (e.g., to fill a dense output with zeros).
   using OutputUpdater = function_ref<Value(OpBuilder &builder, Location loc,
                                            Value memref, Value tensor)>;
+  using SynTensorBoundSetter =
+      function_ref<Value(OpBuilder &builder, Location loc, Level lvl)>;
+
   // Map from [tid, dim] to a list of dependent [tid, dim] for affine expression
   // index on sparse tensors.
   // E.g., for affine index (d0 + d1), it depends on two [tid, dim] that defines
   // d0 and d1 (for affine expression reduction).
   // If the list is empty, it means that there is no affine expression on the
   // input [tid, dim].
   // NOTE: The caller is responsible to ensure that the order of the returned
   // list to be consistent with the topological order of the iteration graph,
public:
   explicit LoopEmitter(ValueRange tensors, StringAttr loopTag = nullptr,
                        bool hasOutput = false, bool isSparseOut = false,
                        ArrayRef<LoopId> topSort = {},
                        DependentLvlGetter getter = nullptr);
 
   /// Starts a loop emitting session by generating all the buffers needed
   /// for iterating over the tensors.
   void initializeLoopEmit(OpBuilder &builder, Location loc,
-                          OutputUpdater updater = nullptr);
+                          OutputUpdater updater = nullptr,
+                          SynTensorBoundSetter synSetter = nullptr);
 
   /// Generates code to compute an affine expression whose variables are
   /// `LoopId`s (i.e., `a.cast<AffineDimExpr>().getPosition()` is a valid
   /// `LoopId`).
   Value genAffine(OpBuilder &builder, Location loc, AffineExpr a);
 
   /// Enters a new loop sequence, the loops within the same sequence starts
   /// from the break points of previous loop instead of starting over from 0.

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

void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
   // Construct the inverse of the `topSort` from the sparsifier.
   // This is needed to map `AffineDimExpr`s back to the `LoopOrd`
   // used in loop emitter.
   // FIXME: This map should be maintained outside loop emitter.
   for (LoopOrd n = 0; n < numLoops; n++)
     loopIdToOrd[topSort[n]] = n;
 }
 
-void LoopEmitter::initializeLoopEmit(OpBuilder &builder, Location loc,
-                                     LoopEmitter::OutputUpdater updater) {
+void LoopEmitter::initializeLoopEmit(
+    OpBuilder &builder, Location loc, LoopEmitter::OutputUpdater updater,
+    LoopEmitter::SynTensorBoundSetter synSetter) {
+
+  // For every synthetic tensor, set the high bound by calling the callback.
+  if (synSetter)
+    for (unsigned i = 0, e = highs[getSynTensorId()].size(); i < e; i++)
+      highs[getSynTensorId()][i] = synSetter(builder, loc, i);
+
   // For every manifest tensor:
   // * get the values buffer.
   // * For every level:
   //   * get the positions and coordinates buffers
   //   * get/compute the level-size, which is also used as the upper-bound
   //     on positions.
   for (TensorId t = 0, numTensors = getNumManifestTensors(); t < numTensors;
        t++) {

 
 void LoopEmitter::enterNewLoopSeq(OpBuilder &builder, Location loc,
                                   ArrayRef<TensorLevel> tidLvls) {
   // TODO: sort
   assert(loopSeqStack.size() == loopStack.size());
   // Prepares for all the tensors used in the current loop sequence.
   std::vector<std::tuple<TensorId, Level, bool>> slicedTids;
 
-  bool hasSynTensor = false;
-  std::optional<std::pair<TensorId, Level>> loopBoundDefLevel = std::nullopt;
   for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
     if (!dependentLvlMap[tid][lvl].empty()) {
       bool fullyRed = genSliceBegin(builder, loc, tid, lvl);
       slicedTids.emplace_back(tid, lvl, fullyRed);
     } else {
-      if (isSynTensor(tid)) {
-        hasSynTensor = true;
-      } else {
-        loopBoundDefLevel = std::make_pair(tid, lvl);
+      if (!isSynTensor(tid)) {
         prepareLoopOverTensorAtLvl(builder, loc, tid, lvl);
       }
     }
   }
 
-  if (hasSynTensor && loopBoundDefLevel.has_value()) {
-    // TODO: compute the loopBound for index reduction by d - sum(unres_lvls).
-    highs[getSynTensorId()][getCurrentDepth()] =
-        lvlSizes[loopBoundDefLevel->first][loopBoundDefLevel->second];
-  }
   // Universal Index starts from 0.
   loopSeqStack.emplace_back(C_IDX(0), std::move(slicedTids));
 }
 
 void LoopEmitter::exitCurrentLoopSeq(OpBuilder &builder, Location loc) {
   assert(loopSeqStack.size() == loopStack.size() + 1);
 
   const auto &slicedTids = loopSeqStack.back().second;

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

 //===----------------------------------------------------------------------===//
 
 /// Local bufferization of all dense and sparse data structures.
 static void genBuffers(CodegenEnv &env, OpBuilder &builder) {
   linalg::GenericOp op = env.op();
   Location loc = op.getLoc();
   assert(op.getNumOperands() == op.getNumDpsInputs() + 1);
 
+  SmallVector<Range, 4> loopRange =
+      llvm::cast<linalg::LinalgOp>(op.getOperation())
+          .createLoopRanges(builder, loc);
+
+  assert(loopRange.size() == env.merger().getStartingFilterLoopId());
+  SmallVector<Range, 4> sortedRange;
+  for (unsigned i = 0, e = env.topSortSize(); i < e; i++) {
+    LoopId ldx = env.topSortAt(i);
+    // FIXME: Gets rid of filter loops since we have a better algorithm to deal
+    // with affine index expression.
+    if (ldx < env.merger().getStartingFilterLoopId()) {
+      sortedRange.push_back(loopRange[ldx]);
+    }
+  }
+
   env.emitter().initializeLoopEmit(
       builder, loc,
       /// Generates buffer for the output tensor.
       /// Note that all sparse kernels assume that when all elements are written
       /// to (viz. x(i) = y(i) * z(i)), the output buffer is already initialized
       /// to all zeroes and only nonzeroes values are computed and written out.
       /// For updates (viz. x(i) += y(i) * z(i)), only nonzeroes values are used
       /// for the updates and no assumption on the original contents of the
env.emitter().initializeLoopEmit(
         Value init = memref;
         if (!isInit) {
           Value zero = constantZero(builder, loc,
                                     getElementTypeOrSelf(tensor.getType()));
           builder.create<linalg::FillOp>(loc, ValueRange{zero},
                                          ValueRange{init});
         }
         return init;
+      },
+      [&sortedRange, &env](OpBuilder &b, Location loc, Level l) {
+        assert(l < env.topSortSize());
+        // FIXME: Remove filter loop since we have a better algorithm to
+        // deal with affine index expression.
+        if (l >= env.merger().getStartingFilterLoopId())
+          return Value();
+
+        return mlir::getValueOrCreateConstantIndexOp(b, loc,
+                                                     sortedRange[l].size);
       });
 }
 
 /// Generates index for load/store on sparse tensor.
 // FIXME: It's not entirely clear what "index" means here (i.e., is it
 // a "coordinate", or "Ldx", or what).  So the function should be renamed
 // and/or the documentation expanded in order to clarify.
 static Value genIndex(CodegenEnv &env, OpOperand *t) {
env.merger().foreachTensorLoopId(
             numloopCond++;
             return;
           }
           if (isUndefDLT(dlt)) {
             // An undefined dlt in the lattices, we probably mean to
             // iterate based on the level of output tensor.  E.g., this
             // could be a synthetic tensor (for invariants and sparse
             // output tensor).
-            if (env.isReduc() && env.merger().getSynTensorID() == tid) {
+            auto itType = env.op().getIteratorTypesArray()[ldx];
+            if (linalg::isReductionIterator(itType) &&
+                env.merger().getSynTensorID() == tid) {
               // Coiterating with an invariant, and this is a reduction loop
               // e.g., out = prod(in[i][j] op invariant);
               // In this case, we can not infer the loop bound from output
               // (whose level is reduced). Instead we use the synthetic tensor
               // to infer the bound.
               // The level of the synthetic tensor is the current loop depth;
               // the rank of the synthetic tensor equals to number of loops.
               lvl = env.emitter().getCurrentDepth();
static bool translateBitsToTidLvlPairs(
 
   if (isDenseDLT(env.dlt(outTid, ldx))) {
     // Note that we generate dense indices of the output tensor
     // unconditionally, since they may not appear in the lattice, but may be
     // needed for linearized env.
     tidLvls.push_back(env.makeTensorLevel(outTid, *outLvl));
   }
 
-  assert(numloopCond > 0);
+  if (numloopCond == 0) {
+    // Corner cases where the loop bound is defined by a *unused* operand, in
+    // this case, we just generate a dense "fake" loop by iterating over the
+    // synthetic tensor.
+    tidLvls.push_back(env.makeTensorLevel(env.merger().getSynTensorID(),
+                                          env.emitter().getCurrentDepth()));
+    numloopCond++;
+  }
   // If we just need to one loop conditions and the conditions is not imposed on
   // non-unique level, the loop can be generated by a for loop.
   return numloopCond == 1 && !hasNonUnique;
 }
 
 /// Starts a single loop in current sequence.
 static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
                                               OpBuilder &builder, LoopOrd at,

mlir/test/Dialect/SparseTensor/sampled_dense_dense_matmul_with_spy.mlir

+// RUN: mlir-opt %s -sparsification --canonicalize --cse | FileCheck %s
+
+#SM = #sparse_tensor.encoding<{ lvlTypes = [ "dense", "compressed" ] }>
+#trait_sampled_dense_dense = {
+  indexing_maps = [
+    affine_map<(i,j,k) -> (i,j)>,  // S
+    affine_map<(i,j,k) -> (i,k)>,  // A
+    affine_map<(i,j,k) -> (k,j)>,  // B
+    affine_map<(i,j,k) -> (i,j)>   // S_out
+  ],
+  iterator_types = ["parallel", "parallel", "reduction"],
+  doc = "S_out(i,j) = spy[S(i,j)] x SUM_k A(i,k) B(k,j)"
+}
+
+// CHECK-LABEL:   func.func @sparse_sampled_dd(
+// CHECK-SAME:      %[[VAL_0:.*]]: tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>,
+// CHECK-SAME:      %[[VAL_1:.*]]: tensor<8x8xf64>,
+// CHECK-SAME:      %[[VAL_2:.*]]: tensor<8x8xf64>) -> tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> {
+// CHECK-DAG:       %[[VAL_3:.*]] = arith.constant 8 : index
+// CHECK-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
+// CHECK-DAG:       %[[VAL_5:.*]] = arith.constant 1 : index
+// CHECK-DAG:       %[[VAL_6:.*]] = arith.constant false
+// CHECK-DAG:       %[[VAL_7:.*]] = arith.constant true
+// CHECK-DAG:       %[[VAL_8:.*]] = bufferization.alloc_tensor() : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK-DAG:       %[[VAL_9:.*]] = sparse_tensor.positions %[[VAL_0]] {level = 1 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK-DAG:       %[[VAL_10:.*]] = sparse_tensor.coordinates %[[VAL_0]] {level = 1 : index} : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
+// CHECK:           %[[VAL_11:.*]] = scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] iter_args(%[[VAL_13:.*]] = %[[VAL_8]]) -> (tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>) {
+// CHECK:             %[[VAL_14:.*]], %[[VAL_15:.*]], %[[VAL_16:.*]], %[[VAL_17:.*]] = sparse_tensor.expand %[[VAL_8]] : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>, memref<?xi1>, memref<?xindex>
+// CHECK:             %[[VAL_18:.*]] = scf.for %[[VAL_19:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] iter_args(%[[VAL_20:.*]] = %[[VAL_17]]) -> (index) {
+// CHECK:               %[[VAL_21:.*]] = tensor.extract %[[VAL_1]]{{\[}}%[[VAL_12]], %[[VAL_19]]] : tensor<8x8xf64>
+// CHECK:               %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<?xindex>
+// CHECK:               %[[VAL_23:.*]] = arith.addi %[[VAL_12]], %[[VAL_5]] : index
+// CHECK:               %[[VAL_24:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_23]]] : memref<?xindex>
+// CHECK:               %[[VAL_25:.*]] = scf.for %[[VAL_26:.*]] = %[[VAL_22]] to %[[VAL_24]] step %[[VAL_5]] iter_args(%[[VAL_27:.*]] = %[[VAL_20]]) -> (index) {
+// CHECK:                 %[[VAL_28:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_26]]] : memref<?xindex>
+// CHECK:                 %[[VAL_29:.*]] = tensor.extract %[[VAL_2]]{{\[}}%[[VAL_19]], %[[VAL_28]]] : tensor<8x8xf64>
+// CHECK:                 %[[VAL_30:.*]] = arith.mulf %[[VAL_21]], %[[VAL_29]] : f64
+// CHECK:                 %[[VAL_31:.*]] = memref.load %[[VAL_15]]{{\[}}%[[VAL_28]]] : memref<?xi1>
+// CHECK:                 %[[VAL_32:.*]] = arith.cmpi eq, %[[VAL_31]], %[[VAL_6]] : i1
+// CHECK:                 %[[VAL_33:.*]] = scf.if %[[VAL_32]] -> (index) {
+// CHECK:                   memref.store %[[VAL_7]], %[[VAL_15]]{{\[}}%[[VAL_28]]] : memref<?xi1>
+// CHECK:                   memref.store %[[VAL_28]], %[[VAL_16]]{{\[}}%[[VAL_27]]] : memref<?xindex>
+// CHECK:                   %[[VAL_34:.*]] = arith.addi %[[VAL_27]], %[[VAL_5]] : index
+// CHECK:                   scf.yield %[[VAL_34]] : index
+// CHECK:                 } else {
+// CHECK:                   scf.yield %[[VAL_27]] : index
+// CHECK:                 }
+// CHECK:                 memref.store %[[VAL_30]], %[[VAL_14]]{{\[}}%[[VAL_28]]] : memref<?xf64>
+// CHECK:                 scf.yield %[[VAL_35:.*]] : index
+// CHECK:               } {"Emitted from" = "linalg.generic"}
+// CHECK:               scf.yield %[[VAL_36:.*]] : index
+// CHECK:             } {"Emitted from" = "linalg.generic"}
+// CHECK:             %[[VAL_37:.*]] = sparse_tensor.compress %[[VAL_14]], %[[VAL_15]], %[[VAL_16]], %[[VAL_38:.*]] into %[[VAL_13]]{{\[}}%[[VAL_12]]] : memref<?xf64>, memref<?xi1>, memref<?xindex>, tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:             scf.yield %[[VAL_37]] : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:           } {"Emitted from" = "linalg.generic"}
+// CHECK:           %[[VAL_39:.*]] = sparse_tensor.load %[[VAL_40:.*]] hasInserts : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:           return %[[VAL_39]] : tensor<8x8xf64, #sparse_tensor.encoding<{{.*}}>>
+// CHECK:         }
+func.func @sparse_sampled_dd(%argS: tensor<8x8xf64, #SM>,
+                             %argA: tensor<8x8xf64>,
+                             %argB: tensor<8x8xf64>) -> tensor<8x8xf64, #SM> {
+  %f0 = arith.constant 0.0 : f64
+  %f1 = arith.constant 1.0 : f64
+  %init = bufferization.alloc_tensor() : tensor<8x8xf64, #SM>
+  %result = linalg.generic #trait_sampled_dense_dense
+    ins(%argS, %argA, %argB: tensor<8x8xf64, #SM>, tensor<8x8xf64>, tensor<8x8xf64>)
+    outs(%init: tensor<8x8xf64, #SM>) {
+    ^bb(%s: f64, %a: f64, %b: f64, %x: f64):
+       // We only care whether %s is present or not, but not the actually value in it.
+       %exist = sparse_tensor.unary %s : f64 to f64
+         present={
+         ^bb0(%p: f64):
+             sparse_tensor.yield %f1 : f64
+         }
+         absent={}
+       %mul = arith.mulf %a, %b :  f64
+       %2 = sparse_tensor.reduce %mul, %exist, %f0 : f64 {
+         ^bb0(%m: f64, %e: f64):
+           %sel = arith.mulf %m, %e : f64 // should fold
+           sparse_tensor.yield %sel : f64
+       }
+       %add = arith.addf %x, %2 : f64
+       linalg.yield %2 : f64
+    } -> tensor<8x8xf64, #SM>
+  return %result : tensor<8x8xf64, #SM>
+}