[mlir][sparse] Refactor the code that reshapes the values buffer for annotated all dense tensors.

[bixia]

Summary

Move the functionality to codegen utils for sharing with the codegen path.

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

 ///
 /// %c3 = arith.constant 31
 /// %v3 = complex.constant (5.0, 6.0)
 /// callback({%c3}, %v3)
 void foreachInSparseConstant(
     Location loc, RewriterBase &rewriter, SparseElementsAttr attr,
     function_ref<void(ArrayRef<Value>, Value)> callback);
 
+/// Converts the vector indices and store it into the memory pointed by
+/// `ind`, apply (optional) `offset` on `offsetDim`.
+void storeIndices(OpBuilder &builder, Location loc, unsigned rank, Value ind,
+                  ValueRange ivs, unsigned offsetDim = 0,
+                  Value offset = Value());
+
+/// Reshapes the linear values buffer for an annotated all dense sparse tensor
+/// to match the shape of the corresponding dense tensor to support direct
+/// access of the buffer through indices.
+Value reshapeValuesToLevels(OpBuilder &builder, Location loc,
+                            SparseTensorEncodingAttr enc,
+                            const SmallVectorImpl<Value> &dimSizes,
+                            Value valuesBuffer, Value idxBuffer);
+
 //===----------------------------------------------------------------------===//
 // Inlined constant generators.
 //
 // All these functions are just wrappers to improve code legibility;
 // therefore, we mark them as `inline` to avoid introducing any additional
 // overhead due to the legibility.
 //
 // TODO: Ideally these should move upstream, so that we don't

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

if (attr.getElementType().isa<ComplexType>()) {
       // Remaps value.
       val = rewriter.create<arith::ConstantOp>(loc, valAttr);
     }
     assert(val);
     callback(coords, val);
   }
 }
 
+void sparse_tensor::storeIndices(OpBuilder &builder, Location loc,
+                                 unsigned rank, Value ind, ValueRange ivs,
+                                 unsigned offsetDim, Value offset) {
+  for (unsigned i = 0; i < rank; i++) {
+    Value idx = ivs[i];
+    if (offsetDim == i && offset)
+      idx = builder.create<arith::AddIOp>(loc, idx, offset);
+    builder.create<memref::StoreOp>(loc, idx, ind,
+                                    constantIndex(builder, loc, i));
+  }
+}
+
+Value sparse_tensor::reshapeValuesToLevels(
+    OpBuilder &builder, Location loc, SparseTensorEncodingAttr enc,
+    const SmallVectorImpl<Value> &dimSizes, Value valuesBuffer,
+    Value idxBuffer) {
+  // Use the dstIdx to store the level sizes.
+  unsigned rank = enc.getDimLevelType().size();
+  SmallVector<Value> lvlSizes;
+  for (unsigned i = 0; i < dimSizes.size(); i++)
+    lvlSizes.push_back(dimSizes[toOrigDim(enc, i)]);
+  storeIndices(builder, loc, rank, idxBuffer, lvlSizes);
+  // The memref ReshapeOp requires the sizes buffer to have a static
+  // shape.
+  idxBuffer = builder.create<memref::CastOp>(
+      loc, MemRefType::get({rank}, builder.getIndexType()), idxBuffer);
+  SmallVector<int64_t> shape(rank, ShapedType::kDynamic);
+  Type elemTp = valuesBuffer.getType().cast<MemRefType>().getElementType();
+  return builder.create<memref::ReshapeOp>(loc, MemRefType::get(shape, elemTp),
+                                           valuesBuffer, idxBuffer);
+}
+
 Value sparse_tensor::genToPointers(OpBuilder &builder, Location loc,
                                    Value tensor, uint64_t d) {
   RankedTensorType srcTp = tensor.getType().cast<RankedTensorType>();
   SparseTensorEncodingAttr encSrc = getSparseTensorEncoding(srcTp);
   Type ptrTp = get1DMemRefType(getPointerOverheadType(builder, encSrc),
                                /*withLayout=*/false);
   return builder.create<ToPointersOp>(loc, ptrTp, tensor,
                                       builder.getIndexAttr(d));

 }
 
 Value sparse_tensor::genValMemSize(OpBuilder &builder, Location loc,
                                    Value tensor) {
   SmallVector<Value> fields;
   auto desc = getMutDescriptorFromTensorTuple(tensor, fields);
   return desc.getValMemSize(builder, loc);
 }
 No newline at end of file

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

for (unsigned i = 0; i < rank; i++) {
     idx = builder.create<memref::LoadOp>(loc, ind, idx);
     if (offsetDim == i && offset)
       idx = builder.create<arith::AddIOp>(loc, idx, offset);
     ivs.push_back(idx);
   }
   return ivs;
 }
 
-/// Converts the vector indices and store it into the memory pointed by
-/// `ind`, apply (optional) `offset` on `offsetDim`.
-static void storeIndices(OpBuilder &builder, Location loc, unsigned rank,
-                         Value ind, ValueRange ivs, unsigned offsetDim = 0,
-                         Value offset = Value()) {
-  for (unsigned i = 0; i < rank; i++) {
-    Value idx = ivs[i];
-    if (offsetDim == i && offset)
-      idx = builder.create<arith::AddIOp>(loc, idx, offset);
-    builder.create<memref::StoreOp>(loc, idx, ind,
-                                    constantIndex(builder, loc, i));
-  }
-}
-
 /// Inserts a value stored in `elemPtr` into a dense tensor created by
 /// allocDenseTensor().
 static void insertScalarIntoDenseTensor(OpBuilder &builder, Location loc,
                                         Value elemPtr, Value tensor,
                                         ValueRange ivs) {
   Value elemV = builder.create<memref::LoadOp>(loc, elemPtr);
   builder.create<memref::StoreOp>(loc, elemV, tensor, ivs);
 }
if (encDst) {
       dstIdx = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       if (allDense) {
         dstTensor = dst;
         // Get the values buffer for the sparse tensor and reshape it to the
         // corresponding dense tensor shape.
         dst = genValuesCall(rewriter, loc,
                             MemRefType::get({ShapedType::kDynamic}, elemTp),
                             {dst});
-
         // Use the dstIdx to store the level sizes.
-        SmallVector<Value> lvlSizes;
-        for (unsigned i = 0; i < sizes.size(); i++)
-          lvlSizes.push_back(sizes[toOrigDim(encDst, i)]);
-        storeIndices(rewriter, loc, rank, dstIdx, lvlSizes);
-        // The memref ReshapeOp requires the sizes buffer to have a static
-        // shape.
-        Value typedBuffer = rewriter.create<memref::CastOp>(
-            loc, MemRefType::get({rank}, rewriter.getIndexType()), dstIdx);
-        SmallVector<int64_t> shape(rank, ShapedType::kDynamic);
-        dst = rewriter.create<memref::ReshapeOp>(
-            loc, MemRefType::get(shape, elemTp), dst, typedBuffer);
+        dst = reshapeValuesToLevels(rewriter, loc, encDst, sizes, dst, dstIdx);
       } else {
         dstPerm = params.getDim2LvlMap();
         elemPtr = genAllocaScalar(rewriter, loc, elemTp);
       }
     } else {
       // TODO: Dense buffers should be allocated/deallocated via the callback
       // in BufferizationOptions.
       dst = allocDenseTensor(rewriter, loc, dstTp, sizes);