[mlir][sparse] fix a bug in sparse2sparse reshape.

[Peiming]

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

567

The reason why I change it to a template function is that I can use more meaningful op.getSrc() instead of op.getOperand()[0], but let me know if you prefer the previous one.

[aartbik]

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

567

I am fine with the template, looks cleaner. But as discussed offline, the fact you can change this without hitting a test change probably indicates we are missing a test case ;-) Please see if you can find one.

[Peiming]

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

567

I can't find one, it would be triggered when the dst tensor are always static shaped.

567

I can't find one, it would be triggered when the dst tensor are always static shaped.

Sorry, would NOT be triggered

[aartbik]

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

575

note that this one sets the size to 0 for dynamic (used for new to simply verify that the static sizes are as expected), but for EmptyCOO that is not allowed

so should we assert on the dynamic size then?

766

yeah, much cleaner like this!

[Peiming]

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

575

Or assert on the dstTp before calling it? I am not sure whether 0 is chose for some special purposes in other converters.

[aartbik]

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

575

yeah, you don't have to change the util, but just check for static shape prior to using it
(I think we will anyway assert in support lib for zero sizes other than when reading a tensor, but always good to be extra safe)

[Peiming]

Address comments from Aart

[aartbik]

Ship it! Thanks for carefully looking at this.

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

 ///       https://github.com/llvm/llvm-project/issues/56477
 ///
 ///   iter = src->toCOO();
 ///   coo = newSparseCOO()
 ///   while (elem = iter->getNext()) {
 ///     coo->add(reshape(elem.indices), elem.value)
 ///   }
 ///   s = newSparseTensor(coo)
+template <typename ReshapeOp>
 static LogicalResult
-genSparse2SparseReshape(Operation *op, ConversionPatternRewriter &rewriter,
-                        ArrayRef<ReassociationIndices> reassociation, Value src,
-                        RankedTensorType dstTp, RankedTensorType srcTp) {
-  Location loc = op->getLoc();
-  auto encDst = getSparseTensorEncoding(dstTp);
+genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
+                        ConversionPatternRewriter &rewriter) {
+  Location loc = op.getLoc();
+  auto srcTp = op.getSrc().getType().template cast<RankedTensorType>();
+  auto dstTp = op.getResult().getType().template cast<RankedTensorType>();
   auto encSrc = getSparseTensorEncoding(srcTp);
-  assert(encDst && encSrc);
+  auto encDst = getSparseTensorEncoding(dstTp);
+  if (!encDst || !encSrc)
+    return failure();
+
   unsigned srcRank = srcTp.getRank();
   unsigned dstRank = dstTp.getRank();
   Type elemTp = srcTp.getElementType();
   assert(elemTp == dstTp.getElementType() &&
          "reshape should not change element type");
   // Start an iterator over the source tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
       op->getContext(), encSrc.getDimLevelType(), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
   SmallVector<Value, 4> sizes;
   SmallVector<Value, 8> params;
-  sizesFromPtr(rewriter, sizes, loc, noPerm, srcTp, src);
+  sizesFromSrc(rewriter, sizes, loc, op.getSrc());
   newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, sizes,
-            src);
+            adaptor.getSrc());
   Value iter = genNewCall(rewriter, loc, params);
   // Start a new COO for the destination tensor.
   sizes.clear();
   params.clear();
-  sizesFromPtr(rewriter, sizes, loc, encDst, dstTp, src);
+  // Fills sizes array using the sizes from destination type.
+  assert(dstTp.hasStaticShape());
+  sizesFromType(rewriter, sizes, loc, dstTp);
   newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, sizes);
   Value coo = genNewCall(rewriter, loc, params);
   Value dstPerm = params[2];
   // Construct a while loop over the iterator.
   Value srcIdx = genAlloca(rewriter, loc, srcRank, rewriter.getIndexType());
   Value dstIdx = genAlloca(rewriter, loc, dstRank, rewriter.getIndexType());
   Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
   SmallVector<Value> noArgs;
   SmallVector<Type> noTypes;
   auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
   Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
   rewriter.setInsertionPointToEnd(before);
   Value cond = genGetNextCall(rewriter, loc, iter, srcIdx, elemPtr);
   rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
   // Translate indices from source to target and insert. Note that we do
   // not need to store the value in elemPtr, as the value is still there.
   Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
   rewriter.setInsertionPointToStart(after);
-  translateIndices(loc, rewriter, reassociation, dstTp, srcTp, dstIdx, srcIdx);
+  translateIndices(loc, rewriter, op.getReassociationIndices(), dstTp, srcTp,
+                   dstIdx, srcIdx);
   genAddEltCall(rewriter, loc, elemTp, coo, elemPtr, dstIdx, dstPerm);
   rewriter.create<scf::YieldOp>(loc);
   // Final call to construct sparse tensor storage and free temporary resources.
   rewriter.setInsertionPointAfter(whileOp);
   params[6] = constantAction(rewriter, loc, Action::kFromCOO);
   params[7] = coo;
   Value dst = genNewCall(rewriter, loc, params);
   genDelCOOCall(rewriter, loc, elemTp, coo);

 template <typename ReshapeOp>
 class SparseReshapeConverter : public OpConversionPattern<ReshapeOp> {
 public:
   using OpAdaptor = typename OpConversionPattern<ReshapeOp>::OpAdaptor;
   using OpConversionPattern<ReshapeOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ReshapeOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Type dstType = op.getResult().getType();
-    Type srcType = op.getSrc().getType();
-    auto encDst = getSparseTensorEncoding(dstType);
-    auto encSrc = getSparseTensorEncoding(srcType);
-    if (encDst && encSrc)
-      return genSparse2SparseReshape(
-          op, rewriter, op.getReassociationIndices(), adaptor.getOperands()[0],
-          dstType.cast<RankedTensorType>(), srcType.cast<RankedTensorType>());
-    return failure(); // handled elsewhere
+    return genSparse2SparseReshape(op, adaptor, rewriter);
   }
 };
 
 /// Sparse conversion rule for the new operator.
 class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult