[mlir][sparse] fix crash when sparsifying broadcast operations.

[aartbik]

can you please add a CHECK test (that exhibits a fail before/pass after behavior)

[Peiming]

add test cases.

[Peiming]

minor fix

[Peiming]

fix typo.

[Peiming]

rebase

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

Value genAddress(OpBuilder &builder, Location loc, size_t tid, size_t dim,
     Value add = builder.create<arith::AddIOp>(loc, mul, iv);
     return add;
   }
 
   bool isOutputTensor(size_t tid) {
     return hasOutput && tid == tensors.size() - 1;
   }
 
+  bool isSparseOutput(size_t tid) { return isOutputTensor(tid) && isSparseOut; }
+
   /// Setups [lo, hi] for iterating tensor[dim], it assumes that tensor[0
   /// ...dims-1] has already been setup.
   void prepareLoopOverTensorAtDim(OpBuilder &builder, Location loc, size_t tid,
                                   size_t dim);
 
   /// Emits extra locals, since the locals might not be in simplified lattices
   /// point used to generate the loops, but are still required to generates
   /// expressions.
private:
 
   /// Exits a while loop, returns the reduction results.
   SmallVector<Value, 2> exitCoiterationLoop(OpBuilder &builder, Location loc,
                                             ArrayRef<Value> reduc);
 
   // Whether the loop emitter needs to treat the last tensor as the output
   // tensor.
   bool hasOutput;
+  bool isSparseOut;
   /// Input and (optional) output tensors.
   std::vector<Value> tensors;
   /// The dim type array for each tensor.
   std::vector<std::vector<DimLevelType>> dimTypes;
   /// Sparse iteration information (by tensor and dim). These arrays
   /// are updated to remain current within the current loop.
   std::vector<std::vector<Value>> pidxs;
   std::vector<std::vector<Value>> coord;

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

 
 //===----------------------------------------------------------------------===//
 // Sparse tensor loop emitter class implementations
 //===----------------------------------------------------------------------===//
 
 SparseTensorLoopEmitter::SparseTensorLoopEmitter(ValueRange tensors,
                                                  bool hasOutput,
                                                  bool isSparseOut)
-    : hasOutput(hasOutput), tensors(tensors.begin(), tensors.end()),
-      dimTypes(tensors.size()), pidxs(tensors.size()), coord(tensors.size()),
-      highs(tensors.size()), ptrBuffer(tensors.size()),
-      idxBuffer(tensors.size()), valBuffer(tensors.size()), loopStack() {
+    : hasOutput(hasOutput), isSparseOut(isSparseOut),
+      tensors(tensors.begin(), tensors.end()), dimTypes(tensors.size()),
+      pidxs(tensors.size()), coord(tensors.size()), highs(tensors.size()),
+      ptrBuffer(tensors.size()), idxBuffer(tensors.size()),
+      valBuffer(tensors.size()), loopStack() {
   for (size_t tid = 0, e = tensors.size(); tid < e; tid++) {
     auto t = tensors[tid];
     // a scalar or 0-dimension tensors
     if (isZeroRankedTensorOrScalar(t.getType()))
       continue;
     auto rtp = t.getType().cast<RankedTensorType>();
     auto rank = static_cast<size_t>(rtp.getRank());
     auto enc = getSparseTensorEncoding(rtp);
if (isSparseInput) {
     // Generating a load on the indices array yields the coordinate.
     Value ptr = idxBuffer[tid][dim];
     coord[tid][dim] = genIndexLoad(builder, loc, ptr, iv);
   } else {
     // Dense tensor, the coordinates is the inducation variable.
     coord[tid][dim] = iv;
     // generate pidx for dense dim (pidx = i * sz + j)
     auto enc = getSparseTensorEncoding(tensors[tid].getType());
-    if (enc)
+    if (enc && !isSparseOutput(tid))
       pidxs[tid][dim] = genAddress(builder, loc, tid, dim, iv);
   }
 
   // NOTE: we can also prepares for next dim here in advance
   // Push the loop into stack
   loopStack.emplace_back(ArrayRef<size_t>(tid), ArrayRef<size_t>(dim), forOp,
                          coord[tid][dim]);
   // Emit extra locals.
Operation *SparseTensorLoopEmitter::enterCoIterationOverTensorsAtDims(
 
   for (auto [tid, dim] : llvm::zip(tids, dims)) {
     // All dense dim (as well as sparse output tensor) shared the same pidx in
     // the while loop.
     if (isDenseDLT(dimTypes[tid][dim])) {
       pidxs[tid][dim] = min;
       // generate pidx for dense dim (pidx = i * sz + j)
       auto enc = getSparseTensorEncoding(tensors[tid].getType());
-      if (enc)
+      if (enc && !isSparseOutput(tid))
         pidxs[tid][dim] = genAddress(builder, loc, tid, dim, min);
     }
     // NOTE: we can also prepares for next dim here in advance
   }
   // Sets up the loop stack.
   loopStack.emplace_back(tids, dims, whileOp, min);
   assert(loopStack.size() == loopSeqStack.size());
 
void SparseTensorLoopEmitter::emitExtraLocalsForTensorsAtDenseDims(
     OpBuilder &builder, Location loc, ArrayRef<size_t> tids,
     ArrayRef<size_t> dims) {
   // Initialize dense positions. 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 codegen.
   for (auto [tid, dim] : llvm::zip(tids, dims)) {
     assert(isDenseDLT(dimTypes[tid][dim]));
     auto enc = getSparseTensorEncoding(tensors[tid].getType());
-    if (enc) {
+    if (enc && !isSparseOutput(tid)) {
       bool validPidx = dim == 0 || pidxs[tid][dim - 1];
       if (!validPidx) {
         // We might not find the pidx for the sparse output tensor as it is
         // unconditionally required by the sparsification.
         assert(isOutputTensor(tid));
         continue;
       }
       pidxs[tid][dim] = genAddress(builder, loc, tid, dim, loopStack.back().iv);

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

static void translateBitsToTidDimPairs(Merger &merger, CodeGen &codegen,
   // First converts bits to array + dim pair
   for (unsigned b = 0, e = simple.size(); b < e; b++) {
     size_t tid = merger.tensor(b);
     if (simple.test(b)) {
       // the simplified condition must be a subset of the original condition.
       assert(all.test(b));
       assert(merger.index(b) == idx);
       if (isUndefDLT(merger.getDimLevelType(b))) {
-        // This could be a synthetic tensor (for invariants and sparse output
-        // tensor).
-        // In both cases, we mean to generate loops over output tensor.
-        // e.g.,
-        // out[i][j] = invariant;
-        if (merger.getSynTensorID() == tid)
-          tid = merger.getOutTensorID();
+        // An undefined dlt in the lattices, we probably mean to iterate based
+        // on the dim of output tensor.
+        // E.g., this could be a synthetic tensor (for invariants and sparse
+        // output tensor).
+        // out[i][j] = invariant; or a broadcast
+        // out[i][j] = int[i]
       }
       auto dim = codegen.loopIdxToDim[tid][idx];
       if (dim != INVALID_ID) {
         // dim could be invalid if this is a zero ranked tensor
         condTids.push_back(tid);
         condDims.push_back(dim);
       }
     } else if ((all.test(b) || merger.isOutTensor(b, idx)) &&