[mlir][sparse] Factoring out allocaIndices()

[wrengr]

Summary

This is preliminary work towards D110790. Depends On D110882.

[aartbik]

LG, but two nits before accepting

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

374

Please move this into the "Helper methods" section (L32-297), don't intersperse with conversion methods.

Also, please add a one line /// Doc comment on method.

[wrengr]

Addressing comments

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

for (unsigned i = 0; i < rank; i++) {
     Value val = rewriter.create<tensor::ExtractOp>(loc, indices,
                                                    ValueRange{ivs[0], idx});
     val = rewriter.create<IndexCastOp>(loc, val, rewriter.getIndexType());
     rewriter.create<memref::StoreOp>(loc, val, ind, idx);
   }
   return rewriter.create<tensor::ExtractOp>(loc, values, ivs[0]);
 }
 
+/// Generates code to stack-allocate a `memref<?xindex>` where the `?`
+/// is the given `rank`.  This array is intended to serve as a reusable
+/// buffer for storing the indices of a single tensor element, to avoid
+/// allocation in the body of loops.
+static Value allocaIndices(ConversionPatternRewriter &rewriter, Location loc,
+                           int64_t rank) {
+  auto indexTp = rewriter.getIndexType();
+  auto memTp = MemRefType::get({ShapedType::kDynamicSize}, indexTp);
+  Value arg = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(rank));
+  return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
+}
+
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<ReturnOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
if (!enc)
       return failure();
     Value perm;
     rewriter.replaceOp(
         op, genNewCall(rewriter, op, enc, 0, perm, adaptor.getOperands()[0]));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto encDst = getSparseTensorEncoding(resType);
     auto encSrc = getSparseTensorEncoding(op.source().getType());
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     // Note that the dense tensor traversal code is actually implemented
     // using MLIR IR to avoid having to expose too much low-level
     // memref traversal details to the runtime support library.
     // Also note that the code below only generates the "new" ops and
     // the loop-nest per se; whereas the entire body of the innermost
     // loop is generated by genAddElt().
     Location loc = op->getLoc();
     ShapedType shape = resType.cast<ShapedType>();
-    auto memTp =
-        MemRefType::get({ShapedType::kDynamicSize}, rewriter.getIndexType());
     Value perm;
     Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
-    Value arg = rewriter.create<ConstantOp>(
-        loc, rewriter.getIndexAttr(shape.getRank()));
-    Value ind = rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
+    Value ind = allocaIndices(rewriter, loc, shape.getRank());
     SmallVector<Value> lo;
     SmallVector<Value> hi;
     SmallVector<Value> st;
     Value zero = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(0));
     Value one = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(1));
     Value tensor = adaptor.getOperands()[0];
     auto indicesValues = genSplitSparseConstant(rewriter, op, tensor);
     bool isCOOConstant = indicesValues.hasValue();