[mlir][sparse] Moving/renaming genBuffer to allocaBuffer

[wrengr]

Summary

This allows allocaBuffer to be used outside of SparseTensorConversion.cpp, which will be helpful for a some future commits.

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

 /// type, but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`).
 Value genAlloca(OpBuilder &builder, Location loc, unsigned sz, Type tp);
 
 /// Generates an uninitialized temporary buffer with room for one value
 /// of the given type, and returns the `memref<$tp>`.
 Value genAllocaScalar(OpBuilder &builder, Location loc, Type tp);
 
+/// Generates a temporary buffer, initializes it with the given contents,
+/// and returns it as type `memref<? x $tp>` (rather than specifying the
+/// size of the buffer).
+Value allocaBuffer(OpBuilder &builder, Location loc, ValueRange values);
+
 /// Generates code to allocate a buffer of the given type, and zero
 /// initialize it.  If the buffer type has any dynamic sizes, then the
 /// `sizes` parameter should be as filled by sizesFromPtr(); that way
 /// we can reuse the genDimSizeCall() results generated by sizesFromPtr().
 Value allocDenseTensor(OpBuilder &builder, Location loc,
                        RankedTensorType tensorTp, ValueRange sizes);
 
 /// Generates code to deallocate a dense buffer.

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

Value mlir::sparse_tensor::genAlloca(OpBuilder &builder, Location loc, Value sz,
   return builder.create<memref::AllocaOp>(loc, memTp, ValueRange{sz});
 }
 
 Value mlir::sparse_tensor::genAllocaScalar(OpBuilder &builder, Location loc,
                                            Type tp) {
   return builder.create<memref::AllocaOp>(loc, MemRefType::get({}, tp));
 }
 
+Value mlir::sparse_tensor::allocaBuffer(OpBuilder &builder, Location loc,
+                                        ValueRange values) {
+  const unsigned sz = values.size();
+  assert(sz >= 1);
+  Value buffer = genAlloca(builder, loc, sz, values[0].getType());
+  for (unsigned i = 0; i < sz; i++) {
+    Value idx = constantIndex(builder, loc, i);
+    builder.create<memref::StoreOp>(loc, values[i], buffer, idx);
+  }
+  return buffer;
+}
+
 Value mlir::sparse_tensor::allocDenseTensor(OpBuilder &builder, Location loc,
                                             RankedTensorType tensorTp,
                                             ValueRange sizes) {
   Type elemTp = tensorTp.getElementType();
   auto shape = tensorTp.getShape();
   auto memTp = MemRefType::get(shape, elemTp);
   SmallVector<Value> dynamicSizes;
   for (unsigned i = 0, rank = tensorTp.getRank(); i < rank; i++) {

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

 /// but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`). Unlike temporary buffers on the stack,
 /// this buffer must be explicitly deallocated by client.
 static Value genAlloc(RewriterBase &rewriter, Location loc, Value sz, Type tp) {
   auto memTp = MemRefType::get({ShapedType::kDynamic}, tp);
   return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{sz});
 }
 
-/// Generates a temporary buffer of the given type and given contents.
-static Value genBuffer(OpBuilder &builder, Location loc, ValueRange values) {
-  unsigned sz = values.size();
-  assert(sz >= 1);
-  Value buffer = genAlloca(builder, loc, sz, values[0].getType());
-  for (unsigned i = 0; i < sz; i++) {
-    Value idx = constantIndex(builder, loc, i);
-    builder.create<memref::StoreOp>(loc, values[i], buffer, idx);
-  }
-  return buffer;
-}
-
 /// Generates a temporary buffer for the level-types of the given encoding.
 static Value genLvlTypesBuffer(OpBuilder &builder, Location loc,
                                SparseTensorEncodingAttr enc) {
   SmallVector<Value> lvlTypes;
   auto dlts = enc.getDimLevelType();
   lvlTypes.reserve(dlts.size());
   for (auto dlt : dlts)
     lvlTypes.push_back(constantDimLevelTypeEncoding(builder, loc, dlt));
-  return genBuffer(builder, loc, lvlTypes);
+  return allocaBuffer(builder, loc, lvlTypes);
 }
 
 /// This class abstracts over the API of `_mlir_ciface_newSparseTensor`:
 /// the "swiss army knife" method of the sparse runtime support library
 /// for materializing sparse tensors into the computation.  This abstraction
 /// reduces the need to make modifications to client code whenever that
 /// API changes.
 class NewCallParams final {
NewCallParams &NewCallParams::genBuffers(SparseTensorEncodingAttr enc,
                                          ValueRange dimSizes, ShapedType stp) {
   const unsigned lvlRank = enc.getDimLevelType().size();
   const unsigned dimRank = stp.getRank();
   // Sparsity annotations.
   params[kParamLvlTypes] = genLvlTypesBuffer(builder, loc, enc);
   // Dimension-sizes array of the enveloping tensor.  Useful for either
   // verification of external data, or for construction of internal data.
   assert(dimSizes.size() == dimRank && "Dimension-rank mismatch");
-  params[kParamDimSizes] = genBuffer(builder, loc, dimSizes);
+  params[kParamDimSizes] = allocaBuffer(builder, loc, dimSizes);
   // The level-sizes array must be passed as well, since for arbitrary
   // dim2lvl mappings it cannot be trivially reconstructed at runtime.
   // For now however, since we're still assuming permutations, we will
   // initialize this parameter alongside the `dim2lvl` and `lvl2dim`
   // parameters below.  We preinitialize `lvlSizes` for code symmetry.
   SmallVector<Value> lvlSizes(lvlRank);
   // The dimension-to-level mapping and its inverse.  We must preinitialize
   // `dim2lvl` so that the true branch below can perform random-access
if (dimOrder) {
     }
   } else {
     assert(dimRank == lvlRank && "Rank mismatch");
     for (unsigned i = 0; i < lvlRank; i++) {
       dim2lvl[i] = lvl2dim[i] = constantIndex(builder, loc, i);
       lvlSizes[i] = dimSizes[i];
     }
   }
-  params[kParamLvlSizes] = genBuffer(builder, loc, lvlSizes);
-  params[kParamLvl2Dim] = genBuffer(builder, loc, lvl2dim);
+  params[kParamLvlSizes] = allocaBuffer(builder, loc, lvlSizes);
+  params[kParamLvl2Dim] = allocaBuffer(builder, loc, lvl2dim);
   params[kParamDim2Lvl] =
-      dimOrder ? genBuffer(builder, loc, dim2lvl) : params[kParamLvl2Dim];
+      dimOrder ? allocaBuffer(builder, loc, dim2lvl) : params[kParamLvl2Dim];
   // Secondary and primary types encoding.
   setTemplateTypes(enc, stp);
   // Finally, make note that initialization is complete.
   assert(isInitialized() && "Initialization failed");
   // And return `this` for method chaining.
   return *this;
 }
 
matchAndRewrite(NewOp op, OpAdaptor adaptor,
     auto enc = getSparseTensorEncoding(stp);
     if (!enc)
       return failure();
     const unsigned dimRank = stp.getRank();
     const unsigned lvlRank = enc.getDimLevelType().size();
     // Construct the dimShape.
     const auto dimShape = stp.getShape();
     SmallVector<Value> dimShapeValues = getDimShape(rewriter, loc, stp);
-    Value dimShapeBuffer = genBuffer(rewriter, loc, dimShapeValues);
+    Value dimShapeBuffer = allocaBuffer(rewriter, loc, dimShapeValues);
     // Allocate `SparseTensorReader` and perform all initial setup that
     // does not depend on lvlSizes (nor dim2lvl, lvl2dim, etc).
     Type opaqueTp = getOpaquePointerType(rewriter);
     Value valTp =
         constantPrimaryTypeEncoding(rewriter, loc, stp.getElementType());
     Value reader =
         createFuncCall(rewriter, loc, "createCheckedSparseTensorReader",
                        opaqueTp,
if (auto dimOrder = enc.getDimOrdering()) {
         Value dim = constantIndex(rewriter, loc, d);
         dim2lvlValues[d] = lvl;
         lvl2dimValues[l] = dim;
         lvlSizeValues[l] =
             (dimShape[d] == ShapedType::kDynamic)
                 ? rewriter.create<memref::LoadOp>(loc, dimSizesBuffer, dim)
                 : dimShapeValues[d];
       }
-      lvlSizesBuffer = genBuffer(rewriter, loc, lvlSizeValues);
-      lvl2dimBuffer = genBuffer(rewriter, loc, lvl2dimValues);
-      dim2lvlBuffer = genBuffer(rewriter, loc, dim2lvlValues);
+      lvlSizesBuffer = allocaBuffer(rewriter, loc, lvlSizeValues);
+      lvl2dimBuffer = allocaBuffer(rewriter, loc, lvl2dimValues);
+      dim2lvlBuffer = allocaBuffer(rewriter, loc, dim2lvlValues);
     } else {
       assert(dimRank == lvlRank && "Rank mismatch");
       SmallVector<Value> iotaValues;
       iotaValues.reserve(lvlRank);
       for (unsigned i = 0; i < lvlRank; i++)
         iotaValues.push_back(constantIndex(rewriter, loc, i));
       lvlSizesBuffer = dimSizesBuffer ? dimSizesBuffer : dimShapeBuffer;
-      dim2lvlBuffer = lvl2dimBuffer = genBuffer(rewriter, loc, iotaValues);
+      dim2lvlBuffer = lvl2dimBuffer = allocaBuffer(rewriter, loc, iotaValues);
     }
     // Use the `reader` to parse the file.
     SmallVector<Value, 8> params{
         reader,
         lvlSizesBuffer,
         genLvlTypesBuffer(rewriter, loc, enc),
         lvl2dimBuffer,
         dim2lvlBuffer,