[mlir][sparse] Avoid generating DimOp in conversion passes.

[aartbik]

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

1172

oh wait, remove this comment part.

[Peiming]

Address comments + avoid generating dimOp in codegen pass as well

[Peiming]

remove redundant cast

[aartbik]

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

79 ↗

(On Diff #459126)

Use /// for top level comment

79 ↗

(On Diff #459126)

Use "s" form.

Gets the dimension
or

Returns the dimension....

80 ↗

(On Diff #459126)

not sure if we need this much detail for short a short helper, but okay if you prefer it this way

(but use returns in that case)

381 ↗

(On Diff #459126)

People often prefer the success on main path

if (!sz)

return failure();

rewriter.replaceOp(op, *sz);
return success();

is is also slightly shorter

518 ↗

(On Diff #459126)

I would not use inline comment for this, but attached comment

assert(sz); // this is for sure a sparse tensor

[aartbik]

And I forgot to say: nice work! clean way to avoid the rewriting old ops problem!

[Peiming]

address comments from Aart

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

genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
   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;
-  sizesFromSrc(rewriter, sizes, loc, op.getSrc());
+  sizesFromPtr(rewriter, sizes, loc, encSrc, srcTp, adaptor.getSrc());
   newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, sizes,
             adaptor.getSrc());
   Value iter = genNewCall(rewriter, loc, params);
   // Start a new COO for the destination tensor.
   sizes.clear();
   params.clear();
   // Fills sizes array using the sizes from destination type.
   assert(dstTp.hasStaticShape());
matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
     ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
     Type eltType = srcType.getElementType();
     Type boolType = rewriter.getIntegerType(1);
     Type idxType = rewriter.getIndexType();
     // All initialization should be done on entry of the loop nest.
     rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
     // Determine the size for access expansion (always the innermost stored
     // dimension size, translated back to original dimension). Note that we
     // recursively rewrite the new DimOp on the **original** tensor.
     auto enc = getSparseTensorEncoding(srcType);
     unsigned innerDim = srcType.getRank() - 1;
     if (AffineMap p = enc.getDimOrdering())
       innerDim = p.getDimPosition(innerDim);
-    Value sz = rewriter.create<tensor::DimOp>(loc, op.getTensor(), innerDim);
+    auto sz = sizeFromPtrAtDim(rewriter, loc, enc, srcType, adaptor.getTensor(),
+                               innerDim);
     // Allocate temporary buffers for values, filled-switch, and indices.
     // We do not use stack buffers for this, since the expanded size may
     // be rather large (as it envelops a single expanded dense dimension).
     Value values = genAlloc(rewriter, loc, sz, eltType);
     Value filled = genAlloc(rewriter, loc, sz, boolType);
     Value indices = genAlloc(rewriter, loc, sz, idxType);
     Value zero = constantZero(rewriter, loc, idxType);
     // Reset the values/filled-switch to all-zero/false. Note that this