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

[aartbik]

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

1171

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

Use /// for top level comment

79

Use "s" form.

Gets the dimension
or

Returns the dimension....

80

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)

378

People often prefer the success on main path

if (!sz)

return failure();

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

is is also slightly shorter

516

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/SparseTensorCodegen.cpp

if (auto cast =
       // fields. In this case, take the operands in the cast and replace the
       // sparse tensor output with the flattened type array.
       flattened.append(cast.getOperands().begin(), cast.getOperands().end());
     else
       flattened.push_back(operand);
   }
 }
 
+/// Gets the dimension size for the given sparse tensor at the given dim.
+/// Returns None if no sparse encoding is attached to the tensor type.
+static Optional<Value> sizeFromTensorAtDim(OpBuilder &rewriter, Location loc,
+                                           ShapedType tensorTp,
+                                           Value adaptedValue, unsigned dim) {
+  auto enc = getSparseTensorEncoding(tensorTp);
+  if (!enc)
+    return llvm::None;
+
+  // Access into static dimension can query original type directly.
+  // Note that this is typically already done by DimOp's folding.
+  auto shape = tensorTp.getShape();
+  if (!ShapedType::isDynamic(shape[dim]))
+    return constantIndex(rewriter, loc, shape[dim]);
+
+  // Any other query can consult the dimSizes array at field 0 using,
+  // accounting for the reordering applied to the sparse storage.
+  auto tuple =
+      llvm::cast<UnrealizedConversionCastOp>(adaptedValue.getDefiningOp());
+  return rewriter
+      .create<memref::LoadOp>(loc, tuple.getInputs().front(),
+                              constantIndex(rewriter, loc, toStored(enc, dim)))
+      .getResult();
+}
+
 /// Maps a sparse tensor type to the appropriate compounded buffers.
 static Optional<LogicalResult>
 convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
   auto enc = getSparseTensorEncoding(type);
   if (!enc)
     return llvm::None;
   // Construct the basic types.
   auto context = type.getContext();

 
 /// Sparse codegen rule for dimension accesses.
 class SparseDimOpConverter : public OpConversionPattern<tensor::DimOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    // Only rewrite annotated DimOp with constant index.
-    auto enc = getSparseTensorEncoding(op.getSource().getType());
-    if (!enc)
-      return failure();
     Optional<int64_t> index = op.getConstantIndex();
     if (!index)
       return failure();
-    // Access into static dimension can query original type directly.
-    // Note that this is typically already done by DimOp's folding.
-    Location loc = op->getLoc();
-    auto shape = op.getSource().getType().cast<RankedTensorType>().getShape();
-    if (!ShapedType::isDynamic(shape[*index])) {
-      rewriter.replaceOp(op, constantIndex(rewriter, loc, shape[*index]));
-      return success();
-    }
-    // Any other query can consult the dimSizes array at field 0 using,
-    // accounting for the reordering applied to the sparse storage.
-    auto tuple = llvm::cast<UnrealizedConversionCastOp>(
-        adaptor.getSource().getDefiningOp());
-    rewriter.replaceOpWithNewOp<memref::LoadOp>(
-        op, tuple.getInputs().front(),
-        constantIndex(rewriter, loc, toStored(enc, *index)));
+    auto sz =
+        sizeFromTensorAtDim(rewriter, op.getLoc(),
+                            op.getSource().getType().cast<RankedTensorType>(),
+                            adaptor.getSource(), *index);
+    if (!sz)
+      return failure();
+
+    rewriter.replaceOp(op, *sz);
     return success();
   }
 };
 
 /// Sparse codegen rule for trivial tensor casts.
 class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
     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 = sizeFromTensorAtDim(rewriter, loc, srcType, adaptor.getTensor(),
+                                  innerDim);
+    assert(sz); // This for sure is a sparse tensor
     // Generate a memref for `sz` elements of type `t`.
     auto genAlloc = [&](Type t) {
       auto memTp = MemRefType::get({ShapedType::kDynamicSize}, t);
-      return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{sz});
+      return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{*sz});
     };
     // 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(eltType);
     Value filled = genAlloc(boolType);
     Value indices = genAlloc(idxType);
     Value zero = constantZero(rewriter, loc, idxType);

 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorCodegenPatterns(TypeConverter &typeConverter,
                                                RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseCallConverter, SparseDimOpConverter,
                SparseCastConverter, SparseExpandConverter,
                SparseTensorAllocConverter, SparseTensorDeallocConverter,
                SparseToPointersConverter, SparseToIndicesConverter,
                SparseToValuesConverter, SparseTensorLoadConverter>(
-               typeConverter, patterns.getContext());
 }

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,
     Location loc = op->getLoc();
     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.
+    // dimension size, translated back to original dimension).
     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