[mlir][sparse] fix crash when using sparse_tensor::UnaryOp and ReduceOp.

[Peiming]

revert unintended changes.

[Peiming]

revert unintended changes.

[aartbik]

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

202–203

just an idea for later, but do we want to increase # tensors by one also?

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

1496

typo: synthetic

1593

yeah, this is good logic

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

144

repeating Wren's feedback on the typo originating from the other revision: "particular"

[Peiming]

address comments.

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

202–203

Yeah! We should! Also it now matches the number of tensor in Merger. But I will do it in a separate patch...

[wrengr]

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

202–203

I think it'd be good to have the getNumTensors method give the +1 result; so that it stays consistent with the Merger method of the same name, as well as avoiding needing to spell out all these off-by-one cases

420

Moving this assertion defeats the purpose. We need to check the validity of tid before evaluating collapseReassoc[tid]

[Peiming]

address comments.

[Peiming]

address comments.

[wrengr]

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

235–238

Nit: I think it'd be better to avoid saying "real" here, since that evokes the mathematical meaning of "real" meaning that the elements are real-numbers.

Often the antonym of "synthetic" is "analytic", though I don't think that's appropriate either (since it evokes other logical/mathematical things). Hmm... I would suggest "numOperands", but this includes the output tensor too... Perhaps, "numManifestTensors"?

[Peiming]

renaming varaible.

[Peiming]

fix typo.

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

public:
   /// Gets loop induction variable for the given `LoopOrd`.
   Value getLoopIV(LoopOrd n) const {
     return n < getCurrentDepth() ? loopStack[n].iv : Value();
   }
 
   /// Gets the total number of tensors that loopEmitter is operating on.
   unsigned getNumTensors() const { return tensors.size(); }
 
+  /// Gets the TensorId for synthetic tensor.
+  TensorId getSynTensorId() const { return tensors.size(); }
+
   /// Compresses a TensorId and Level into a TensorLevel.
   TensorLevel makeTensorLevel(TensorId t, Level l) const {
-    return l * getNumTensors() + t;
+    // TODO: getNumTensor() should include synthetic tensor.
+    return l * (getNumTensors() + 1) + t;
   }
 
   /// De-compresses a TensorLevel back to a pair of TensorId and Level.
   std::pair<TensorId, Level> unpackTensorLevel(TensorLevel tidLvl) const {
-    unsigned nt = getNumTensors();
+    unsigned nt = getNumTensors() + 1;
     return std::make_pair(tidLvl % nt, tidLvl / nt);
   }
 
   /// Converts a range of TensorLevel to a range of std::pair<TensorId, Level>
   template <class ContainerTy>
   auto unpackTensorLevelRange(ContainerTy &&c) const {
     using EltTy = decltype(*c.begin());
     static_assert(std::is_same_v<llvm::remove_cvref_t<EltTy>, TensorLevel>,
private:
 
   /// Generates a predicate to determine whether the tranformed coordinates are
   /// in the given slice.
   /// Returns std::pair<Transformed coordinates, Predicate>
   std::pair<Value, Value> genSliceLegitPredicate(OpBuilder &builder,
                                                  Location loc, Value crd,
                                                  TensorId tid, Level lvl);
 
+  bool isSynTensor(TensorId tid) const { return tid == getNumTensors(); }
+
   bool isOutputTensor(TensorId tid) const {
     return hasOutput && tid == getNumTensors() - 1;
   }
 
   bool isSparseOutput(TensorId tid) const {
     return isOutputTensor(tid) && isSparseOut;
   }
 
private:
   ///
   /// NOTE: the result uses `Level` rather than the `int64_t` of
   /// `ReassociationIndices`, since the former gives clarity to what
   /// the values actually mean.
   ///
   /// TODO: why not do this computation when we first store the reassoc,
   /// instead of doing it every time we look it up?
   SmallVector<Level, 2> getCollapseReassociation(TensorId tid, Level dstLvl) {
-    assert(tid < getNumTensors() && "Invalid TensorId");
-    assert(collapseReassoc.size() == getNumTensors());
+    assert(tid < getNumTensors() + 1 && "Invalid TensorId");
+    assert(collapseReassoc.size() == getNumTensors() + 1);
     if (const auto reassoc = collapseReassoc[tid]) {
+      assert(!isSynTensor(tid) && !isOutputTensor(tid) &&
+             "Output/Synthetic tensor should not have reassociation");
       // TODO: store the dstLvlRank in the LoopEmitter so that we can
       // check `dstLvl < dstLvlRank` at the top; and only here need to
       // assert that `reassoc.size() == dstLvlRank`.
       assert(dstLvl < reassoc.size() && "Level is out-of-bounds");
       const auto srcLvls = cast<ArrayAttr>(reassoc[dstLvl]);
       return llvm::to_vector<2>(
           llvm::map_range(srcLvls, [&](Attribute srcLvl) -> Level {
             // TODO: replace this with the converter for `LevelAttr`.

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

 void LoopEmitter::initialize(ValueRange ts, StringAttr loopTag, bool hasOutput,
                              bool isSparseOut, ArrayRef<LoopId> topSort,
                              DependentLvlGetter dimGetter) {
   // First initialize the top-level type of the fields.
   this->loopTag = loopTag;
   this->hasOutput = hasOutput;
   this->isSparseOut = isSparseOut;
 
-  const unsigned numTensors = ts.size();
+  const unsigned numManifestTensors = ts.size();
+  const unsigned synTensorId = numManifestTensors;
+  const unsigned numTensors = numManifestTensors + 1;
+
   this->tensors.assign(ts.begin(), ts.end());
   this->lvlTypes.assign(numTensors, std::vector<DimLevelType>());
   this->lvlSizes.assign(numTensors, std::vector<Value>());
   this->highs.assign(numTensors, std::vector<Value>());
   this->segHi.assign(numTensors, std::vector<Value>());
   this->posits.assign(numTensors, std::vector<Value>());
   this->coords.assign(numTensors, std::vector<Value>());
   this->positionsBuffers.assign(numTensors, std::vector<Value>());
this->dependentLvlMap.assign(
       numTensors, std::vector<std::vector<std::pair<TensorId, Level>>>());
   this->slicePosBuffer.assign(numTensors, std::vector<std::vector<Value>>());
   this->sliceSizes.assign(numTensors, std::vector<std::vector<Value>>());
   this->sliceStack.assign(numTensors, std::vector<SliceInfo>());
   this->levelReducedDep.assign(numTensors, std::vector<unsigned>());
 
   // Initialize nested types of `TensorId`-indexed fields.
   for (TensorId tid = 0; tid < numTensors; tid++) {
+    Level lvlRank;
+    if (tid == synTensorId) {
+      // Synthetic tensor (conceptually) is an all-dense tensor with rank equal
+      // to the total number of loops (each level can potentially be mapped to
+      // one of the loop being generated).
+      lvlRank = numLoops;
+      lvlTypes[tid].assign(lvlRank, DimLevelType::Dense);
+    } else {
       const Value t = tensors[tid];
       // a scalar or 0-dimension tensors
       if (isZeroRankedTensorOrScalar(t.getType()))
         continue;
 
       auto rtp = getRankedTensorType(t);
       if (auto reshape = t.getDefiningOp<tensor::CollapseShapeOp>();
           isUniqueCOOType(rtp) && reshape) {
         // TODO: Supports more kinds of sparse tensors.
-      // FIXME: We should instead lower reshape operations on sparse tensors to
-      // view change.
+        // FIXME: We should instead lower reshape operations on sparse tensors
+        // to view change.
         collapseReassoc[tid] = reshape.getReassociation();
         rtp = reshape.getSrcType();
         // Overwrites the tensor to the source tensor of reshape operations.
         tensors[tid] = reshape.getSrc();
       }
       const SparseTensorType stt(rtp);
-    const Level lvlRank = stt.getLvlRank();
+      lvlRank = stt.getLvlRank();
+
       // We always treat sparse output tensor as dense so that we always iterate
       // it based on lvl size.
       if (stt.hasEncoding() && !(isOutputTensor(tid) && isSparseOut)) {
         const auto enc = stt.getEncoding();
         isSparseSlices[tid] = enc.isSlice();
         for (auto lvlTp : enc.getLvlTypes())
           lvlTypes[tid].push_back(lvlTp);
       } else {
         lvlTypes[tid].assign(lvlRank, DimLevelType::Dense);
       }
+    }
 
     // Initialize using empty value.
     lvlSizes[tid].assign(lvlRank, Value());
     highs[tid].assign(lvlRank, Value());
     segHi[tid].assign(lvlRank, Value());
     posits[tid].assign(lvlRank, Value());
     coords[tid].assign(lvlRank, Value());
     positionsBuffers[tid].assign(lvlRank, Value());
     coordinatesBuffers[tid].assign(lvlRank, Value());
     sliceOffsets[tid].assign(lvlRank, Value());
     sliceStrides[tid].assign(lvlRank, Value());
 
     // Slice-driven loops related initialization.
     levelReducedDep[tid].assign(lvlRank, 0);
     dependentLvlMap[tid].assign(lvlRank,
                                 std::vector<std::pair<TensorId, Level>>());
     slicePosBuffer[tid].assign(lvlRank, std::vector<Value>());
     sliceSizes[tid].assign(lvlRank, std::vector<Value>());
     sliceStack[tid].emplace_back(/*minCrd=*/Value(),
                                  /*offset=*/Value(), /*isNonEmpty*/ Value(),
                                  std::nullopt, 0);
-    if (dimGetter) {
+    if (dimGetter && !isSynTensor(tid)) {
       auto reassoc = collapseReassoc[tid];
       Level dstRank = reassoc ? reassoc.size() : lvlRank;
       for (Level l = 0; l < dstRank; l++) {
         dependentLvlMap[tid][l] = dimGetter(tid, l);
         unsigned depends = dependentLvlMap[tid][l].size();
         if (depends == 0)
           continue;
         // TODO: View-base collapse and dependent index reduction are not

 }
 
 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);
         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;

 }
 
 void LoopEmitter::emitExtraLocalsForTensorsAtDenseLvls(
     OpBuilder &builder, Location loc, ArrayRef<TensorLevel> tidLvls) {
   // Initialize dense positions. Note that we generate dense coordinates of the
   // output tensor unconditionally, since they may not appear in the lattice,
   // but may be needed for linearized codegen.
   for (auto [tid, lvl] : unpackTensorLevelRange(tidLvls)) {
+    if (isSynTensor(tid))
+      continue;
+
     if (isDenseDLT(lvlTypes[tid][lvl])) {
       // Slice-driven dense level should have be handled already.
       if (!dependentLvlMap[tid][lvl].empty())
         continue;
 
       auto enc = getSparseTensorEncoding(tensors[tid].getType());
       if (enc && !isSparseOutput(tid)) {
         bool validPos = lvl == 0 || posits[tid][lvl - 1];

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

static bool startLoopSeq(CodegenEnv &env, OpBuilder &builder, ExprId exp,
   const LatPointId l0 = env.set(lts)[0];
   bool needsUniv = false;
 
   SmallVector<TensorLevel> tidLvls;
   env.merger().foreachTensorLoopId(l0, [&](TensorLoopId b, TensorId tid,
                                            std::optional<Level> lvl,
                                            DimLevelType dlt, bool isIdxReduc) {
     assert(env.merger().loop(b) == idx);
-    if (isDenseDLT(dlt) || isUndefDLT(dlt))
+    if (isDenseDLT(dlt) || isUndefDLT(dlt)) {
+      if (tid == env.merger().getSynTensorID()) {
+        // Needs loop emitter to set up loop bounds for synthetic tensor too if
+        // there is a loop condition imposed on the synthetic tensor.
+        tidLvls.push_back(
+            env.makeTensorLevel(tid, env.emitter().getCurrentDepth()));
+      }
       needsUniv = true;
+    }
     if (isCompressedDLT(dlt) || isSingletonDLT(dlt) ||
         isCompressedWithHiDLT(dlt) || isIdxReduc) {
       // Only when this is a index reduction loop, can the dlt be undefined.
       assert(!isUndefDLT(dlt) || isIdxReduc);
       // sparse/singleton levels, or a dense/sparse index reduction loop.
       tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
     }
   });
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).
-            // out[i][j] = invariant; or a broadcast
+            if (env.isReduc() && 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();
+            } else {
+              // or a broadcast
               // out[i][j] = in[i] (j is undef for input)
               tid = outTid;
               lvl = outLvl;
               // Skips invalid lvl (e.g., when this is a zero ranked tensor).
               if (!lvl)
                 return;
             }
+          }
           hasNonUnique = !isUniqueDLT(dlt) || hasNonUnique;
           tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
           numloopCond++;
         } else if (isDenseDLT(dlt) || isIdxReduc) {
           tidLvls.push_back(env.makeTensorLevel(tid, *lvl));
         } else {
           assert(isUndefDLT(dlt));
           linalg::GenericOp op = env.op();
static std::pair<Operation *, bool> startLoop(CodegenEnv &env,
   }
 
   // Until now, we have entered every <tid, lvl> pair in {cond, extra,
   // affine}Tids/Lvls. The addresses of the upcoming levels which are dependent
   // on constant affines expression may now be determined.
   auto allTidLvls =
       llvm::concat<TensorLevel>(tidLvls, llvm::make_first_range(affineTidLvls));
   for (auto [tid, lvl] : env.unpackTensorLevelRange(allTidLvls)) {
-    if (tid != env.merger().getOutTensorID())
+    if (tid != env.merger().getOutTensorID() &&
+        tid != env.merger().getSynTensorID())
       genConstantDenseAddressFromLevel(env, builder, tid, lvl + 1);
   }
 
   return std::make_pair(loop, isSingleCond);
 }
 
 /// Ends a single loop in current sequence. Returns new values for needsUniv.
 static bool endLoop(CodegenEnv &env, RewriterBase &rewriter, Operation *loop,
if (getSparseTensorEncoding(resType)) {
     if (Value chain = env.getInsertionChain()) {
       hasInserts = true;
       tensor = chain;
     }
     rewriter.replaceOpWithNewOp<LoadOp>(op, resType, tensor, hasInserts);
   } else {
     // To rematerialize an non-annotated tensor, simply load it
     // from the bufferized value.
-    Value val = env.emitter().getValBuffer().back(); // value array
+    Value val = env.emitter().getValBuffer()[env.merger().getOutTensorID()];
     rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, val);
   }
 }
 
 //===----------------------------------------------------------------------===//
 // Sparse compiler rewriting methods.
 //===----------------------------------------------------------------------===//
 

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

func.func @entry() {
 
     %c_1_f32 = arith.constant dense<[
       1.0, 1.0, 1.0, 3.5, 1.0, 1.0, 1.0, 1.0,
       1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 1.0,
       1.0, 1.0, 1.0, 1.0, 3.0, 1.0, 1.0, 1.0,
       1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 4.0
     ]> : tensor<32xf32>
 
-    // Convert constants to annotated tensors.
+    // Convert constants to annotated tensors. Note that this
+    // particular conversion only stores nonzero elements,
+    // so we will have no explicit zeros, only implicit zeros.
     %d0_i32 = sparse_tensor.convert %c_0_i32
       : tensor<32xi32> to tensor<32xi32, #DV>
     %d0_f32 = sparse_tensor.convert %c_0_f32
       : tensor<32xf32> to tensor<32xf32, #DV>
     %s0_i32 = sparse_tensor.convert %c_0_i32
       : tensor<32xi32> to tensor<32xi32, #SV>
     %s0_f32 = sparse_tensor.convert %c_0_f32
       : tensor<32xf32> to tensor<32xf32, #SV>
     %d1_i32 = sparse_tensor.convert %c_1_i32
       : tensor<32xi32> to tensor<32xi32, #DV>
     %d1_f32 = sparse_tensor.convert %c_1_f32
       : tensor<32xf32> to tensor<32xf32, #DV>
     %s1_i32 = sparse_tensor.convert %c_1_i32
       : tensor<32xi32> to tensor<32xi32, #SV>
     %s1_f32 = sparse_tensor.convert %c_1_f32
       : tensor<32xf32> to tensor<32xf32, #SV>
 
+    // Special case, construct a sparse vector with an explicit zero.
+    %v0 = arith.constant sparse< [ [1] ], [ 0 ] > : tensor<32xi32>
+    %s0 = sparse_tensor.convert %v0: tensor<32xi32> to tensor<32xi32, #SV>
+
     // Call the kernels.
     %0 = call @prod_dreduction_i32(%d0_i32, %ri) : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
     %1 = call @prod_dreduction_f32(%d0_f32, %rf) : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
     %2 = call @prod_sreduction_i32(%s0_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
     %3 = call @prod_sreduction_f32(%s0_f32, %rf) : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
     %4 = call @prod_dreduction_i32(%d1_i32, %ri) : (tensor<32xi32, #DV>, tensor<i32>) -> tensor<i32>
     %5 = call @prod_dreduction_f32(%d1_f32, %rf) : (tensor<32xf32, #DV>, tensor<f32>) -> tensor<f32>
     %6 = call @prod_sreduction_i32(%s1_i32, %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
     %7 = call @prod_sreduction_f32(%s1_f32, %rf) : (tensor<32xf32, #SV>, tensor<f32>) -> tensor<f32>
+    %8 = call @prod_sreduction_i32(%s0,     %ri) : (tensor<32xi32, #SV>, tensor<i32>) -> tensor<i32>
 
     // Verify results. Note that the custom reduction gave permission
     // to treat an explicit vs implicit zero differently to compute the
-    // full product reduction. A "standard" product reduction would
-    // have to return 0 for any implicit zero occurrence too.
+    // full product reduction over stored elements. A "standard" product
+    // reduction would have to return 0 for any implicit zero occurrence
+    // too. An explicit zero nullifies the product, though, as requested.
     //
     // CHECK: 0
+    // CHECK: 0
     // CHECK: 3087
     // CHECK: 14
     // CHECK: 3087
     // CHECK: 168
     // CHECK: 3087
     // CHECK: 168
+    // CHECK: 0
     //
     call @dump_i32(%0) : (tensor<i32>) -> ()
     call @dump_f32(%1) : (tensor<f32>) -> ()
     call @dump_i32(%2) : (tensor<i32>) -> ()
     call @dump_f32(%3) : (tensor<f32>) -> ()
     call @dump_i32(%4) : (tensor<i32>) -> ()
     call @dump_f32(%5) : (tensor<f32>) -> ()
     call @dump_i32(%6) : (tensor<i32>) -> ()
     call @dump_f32(%7) : (tensor<f32>) -> ()
+    call @dump_i32(%8) : (tensor<i32>) -> ()
 
     // Release the resources.
     bufferization.dealloc_tensor %d0_i32 : tensor<32xi32, #DV>
     bufferization.dealloc_tensor %d0_f32 : tensor<32xf32, #DV>
     bufferization.dealloc_tensor %s0_i32 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %s0_f32 : tensor<32xf32, #SV>
     bufferization.dealloc_tensor %d1_i32 : tensor<32xi32, #DV>
     bufferization.dealloc_tensor %d1_f32 : tensor<32xf32, #DV>
     bufferization.dealloc_tensor %s1_i32 : tensor<32xi32, #SV>
     bufferization.dealloc_tensor %s1_f32 : tensor<32xf32, #SV>
+    bufferization.dealloc_tensor %s0     : tensor<32xi32, #SV>
 
     return
   }
 }