[mlir][sparse] Factoring out NewCallParams class in SparseTensorConversion.cpp

[wrengr]

Summary

The new class helps encapsulate the arguments to _mlir_ciface_newSparseTensor so that client code doesn't depend on the details of the API. (This makes way for the next differential which significantly alters the API.)

[aartbik]

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

250

I would be okay keeping this method inline as well. I can see why you moved it out, so that the class is perhaps a bit more readable, but having it all in one place makes sense to me too. But your call.

[wrengr]

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

250

I mainly left the definition outside the class in order to minimize the diff. Though since it is a rather substantial definition, I'm not sure moving it into the class would really help legibility (especially since all the other methods are extremely short). I think I'll leave it for now, and if anyone feels strongly they can always move it into the class later :)

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

static Value genDimSizeCall(OpBuilder &builder, Location loc,
   StringRef name = "sparseDimSize";
   SmallVector<Value, 2> params{
       src, constantIndex(builder, loc, toStoredDim(enc, idx))};
   Type iTp = builder.getIndexType();
   return createFuncCall(builder, loc, name, iTp, params, EmitCInterface::Off)
       .getResult(0);
 }
 
-/// Generates a call into the "swiss army knife" method of the sparse runtime
-/// support library for materializing sparse tensors into the computation.
-static Value genNewCall(OpBuilder &builder, Location loc,
-                        ArrayRef<Value> params) {
-  StringRef name = "newSparseTensor";
-  Type pTp = getOpaquePointerType(builder);
-  return createFuncCall(builder, loc, name, pTp, params, EmitCInterface::On)
-      .getResult(0);
-}
-
 /// Compute the size from type (for static sizes) or from an already-converted
 /// opaque pointer source (for dynamic sizes) at the given dimension.
 static Value sizeFromPtrAtDim(OpBuilder &builder, Location loc,
                               SparseTensorEncodingAttr &enc, ShapedType stp,
                               Value src, unsigned dim) {
   auto shape = stp.getShape();
   if (shape[dim] == ShapedType::kDynamicSize)
     return genDimSizeCall(builder, loc, enc, src, dim);
static Value genBuffer(OpBuilder &builder, Location loc, ValueRange values) {
   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;
 }
 
-/// Populates parameters required to call the "swiss army knife" method of the
-/// sparse runtime support library for materializing sparse tensors into the
-/// computation.
-static void newParams(OpBuilder &builder, SmallVector<Value, 8> &params,
-                      Location loc, ShapedType stp,
-                      SparseTensorEncodingAttr &enc, Action action,
-                      ValueRange szs, Value ptr = Value()) {
-  ArrayRef<DimLevelType> dlt = enc.getDimLevelType();
-  unsigned sz = dlt.size();
+/// 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 {
+public:
+  /// Allocates the `ValueRange` for the `func::CallOp` parameters,
+  /// but does not initialize them.
+  NewCallParams(OpBuilder &builder, Location loc)
+      : builder(builder), loc(loc), pTp(getOpaquePointerType(builder)) {}
+
+  /// Initializes all static parameters (i.e., those which indicate
+  /// type-level information such as the encoding and sizes), generating
+  /// MLIR buffers as needed, and returning `this` for method chaining.
+  /// This method does not set the action and pointer arguments, since
+  /// those are handled by `genNewCall` instead.
+  NewCallParams &genBuffers(SparseTensorEncodingAttr enc, ValueRange sizes,
+                            ShapedType stp);
+
+  /// (Re)sets the C++ template type parameters, and returns `this`
+  /// for method chaining.  This is already done as part of `genBuffers`,
+  /// but is factored out so that it can also be called independently
+  /// whenever subsequent `genNewCall` calls want to reuse the same
+  /// buffers but different type parameters.
+  //
+  // TODO: This is only ever used by sparse2sparse-viaCOO `ConvertOp`;
+  // is there a better way to handle that than this one-off setter method?
+  NewCallParams &setTemplateTypes(SparseTensorEncodingAttr enc,
+                                  ShapedType stp) {
+    params[kParamPtrTp] = constantPointerTypeEncoding(builder, loc, enc);
+    params[kParamIndTp] = constantIndexTypeEncoding(builder, loc, enc);
+    params[kParamValTp] =
+        constantPrimaryTypeEncoding(builder, loc, stp.getElementType());
+    return *this;
+  }
+
+  /// Checks whether all the static parameters have been initialized.
+  bool isInitialized() const {
+    for (unsigned i = 0; i < kNumStaticParams; ++i)
+      if (!params[i])
+        return false;
+    return true;
+  }
+
+  /// Gets the dimension-to-level mapping.
+  //
+  // TODO: This is only ever used for passing into `genAddEltCall`;
+  // is there a better way to encapsulate that pattern (both to avoid
+  // this one-off getter, and to avoid potential mixups)?
+  Value getDim2LvlMap() const {
+    assert(isInitialized() && "Must initialize before getDim2LvlMap");
+    return params[kParamDim2Lvl];
+  }
+
+  /// Generates a function call, with the current static parameters
+  /// and the given dynamic arguments.
+  Value genNewCall(Action action, Value ptr = Value()) {
+    assert(isInitialized() && "Must initialize before genNewCall");
+    StringRef name = "newSparseTensor";
+    params[kParamAction] = constantAction(builder, loc, action);
+    params[kParamPtr] = ptr ? ptr : builder.create<LLVM::NullOp>(loc, pTp);
+    return createFuncCall(builder, loc, name, pTp, params, EmitCInterface::On)
+        .getResult(0);
+  }
+
+private:
+  static constexpr unsigned kNumStaticParams = 6;
+  static constexpr unsigned kNumDynamicParams = 2;
+  static constexpr unsigned kNumParams = kNumStaticParams + kNumDynamicParams;
+  static constexpr unsigned kParamLvlTypes = 0;
+  static constexpr unsigned kParamDimSizes = 1;
+  static constexpr unsigned kParamDim2Lvl = 2;
+  static constexpr unsigned kParamPtrTp = 3;
+  static constexpr unsigned kParamIndTp = 4;
+  static constexpr unsigned kParamValTp = 5;
+  static constexpr unsigned kParamAction = 6;
+  static constexpr unsigned kParamPtr = 7;
+
+  OpBuilder &builder;
+  Location loc;
+  Type pTp;
+  Value params[kNumParams];
+};
+
+// TODO: see the note at `_mlir_ciface_newSparseTensor` about how
+// the meaning of the various arguments (e.g., "sizes" vs "shapes")
+// is inconsistent between the different actions.
+NewCallParams &NewCallParams::genBuffers(SparseTensorEncodingAttr enc,
+                                         ValueRange dimSizes, ShapedType stp) {
+  const unsigned lvlRank = enc.getDimLevelType().size();
+  const unsigned dimRank = stp.getRank();
   // Sparsity annotations.
-  SmallVector<Value, 4> attrs;
-  for (unsigned i = 0; i < sz; i++)
-    attrs.push_back(constantDimLevelTypeEncoding(builder, loc, dlt[i]));
-  params.push_back(genBuffer(builder, loc, attrs));
-  // Dimension sizes array of the enveloping tensor. Useful for either
+  SmallVector<Value, 4> lvlTypes;
+  for (auto dlt : enc.getDimLevelType())
+    lvlTypes.push_back(constantDimLevelTypeEncoding(builder, loc, dlt));
+  assert(lvlTypes.size() == lvlRank && "Level-rank mismatch");
+  params[kParamLvlTypes] = genBuffer(builder, loc, lvlTypes);
+  // Dimension-sizes array of the enveloping tensor.  Useful for either
   // verification of external data, or for construction of internal data.
-  params.push_back(genBuffer(builder, loc, szs));
-  // Dimension order permutation array. This is the "identity" permutation by
-  // default, or otherwise the "reverse" permutation of a given ordering, so
-  // that indices can be mapped quickly to the right position.
-  SmallVector<Value, 4> rev(sz);
-  for (unsigned i = 0; i < sz; i++)
-    rev[toOrigDim(enc, i)] = constantIndex(builder, loc, i);
-  params.push_back(genBuffer(builder, loc, rev));
+  assert(dimSizes.size() == dimRank && "Dimension-rank mismatch");
+  params[kParamDimSizes] = genBuffer(builder, loc, dimSizes);
+  // The dimension-to-level mapping.  We must preinitialize `dim2lvl`
+  // so that the true branch below can perform random-access `operator[]`
+  // assignment.
+  SmallVector<Value, 4> dim2lvl(dimRank);
+  auto dimOrder = enc.getDimOrdering();
+  if (dimOrder) {
+    assert(dimOrder.isPermutation());
+    for (unsigned l = 0; l < lvlRank; l++) {
+      // The `d`th source variable occurs in the `l`th result position.
+      uint64_t d = dimOrder.getDimPosition(l);
+      dim2lvl[d] = constantIndex(builder, loc, l);
+    }
+  } else {
+    assert(dimRank == lvlRank && "Rank mismatch");
+    for (unsigned i = 0; i < lvlRank; i++)
+      dim2lvl[i] = constantIndex(builder, loc, i);
+  }
+  params[kParamDim2Lvl] = genBuffer(builder, loc, dim2lvl);
   // Secondary and primary types encoding.
-  Type elemTp = stp.getElementType();
-  params.push_back(constantPointerTypeEncoding(builder, loc, enc));
-  params.push_back(constantIndexTypeEncoding(builder, loc, enc));
-  params.push_back(constantPrimaryTypeEncoding(builder, loc, elemTp));
-  // User action.
-  params.push_back(constantAction(builder, loc, action));
-  // Payload pointer.
-  if (!ptr)
-    ptr = builder.create<LLVM::NullOp>(loc, getOpaquePointerType(builder));
-  params.push_back(ptr);
+  setTemplateTypes(enc, stp);
+  // Finally, make note that initialization is complete.
+  assert(isInitialized() && "Initialization failed");
+  // And return `this` for method chaining.
+  return *this;
 }
 
 /// Generates a call to obtain the values array.
 static Value genValuesCall(OpBuilder &builder, Location loc, ShapedType tp,
                            ValueRange ptr) {
   SmallString<15> name{"sparseValues",
                        primaryTypeFunctionSuffix(tp.getElementType())};
   return createFuncCall(builder, loc, name, tp, ptr, EmitCInterface::On)
genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
   Type elemTp = srcTp.getElementType();
   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(), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
   SmallVector<Value, 4> srcSizes;
-  SmallVector<Value, 8> params;
   sizesFromPtr(rewriter, srcSizes, loc, encSrc, srcTp, adaptor.getSrc());
-  newParams(rewriter, params, loc, srcTp, noPerm, Action::kToIterator, srcSizes,
-            adaptor.getSrc());
-  Value iter = genNewCall(rewriter, loc, params);
+  NewCallParams params(rewriter, loc);
+  Value iter = params.genBuffers(noPerm, srcSizes, srcTp)
+                   .genNewCall(Action::kToIterator, adaptor.getSrc());
   // Start a new COO for the destination tensor.
   SmallVector<Value, 4> dstSizes;
-  params.clear();
   if (dstTp.hasStaticShape()) {
     sizesFromType(rewriter, dstSizes, loc, dstTp);
   } else {
     ArrayRef<int64_t> dstShape = dstTp.getShape();
     genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
                        op.getReassociationIndices());
   }
-  newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, dstSizes);
-  Value coo = genNewCall(rewriter, loc, params);
-  Value dstPerm = params[2];
+  Value coo =
+      params.genBuffers(encDst, dstSizes, dstTp).genNewCall(Action::kEmptyCOO);
+  Value dstPerm = params.getDim2LvlMap();
   // Construct a while loop over the iterator.
   Value srcIdx = genAlloca(rewriter, loc, srcRank, rewriter.getIndexType());
   Value dstIdx = genAlloca(rewriter, loc, dstRank, rewriter.getIndexType());
   Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
   SmallVector<Value> noArgs;
   SmallVector<Type> noTypes;
   auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
   Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
   rewriter.setInsertionPointToEnd(before);
   Value cond = genGetNextCall(rewriter, loc, iter, srcIdx, elemPtr);
   rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
   // Translate indices from source to target and insert. Note that we do
   // not need to store the value in elemPtr, as the value is still there.
   Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
   rewriter.setInsertionPointToStart(after);
   translateIndices(loc, rewriter, op.getReassociationIndices(), dstTp, srcTp,
                    dstIdx, srcIdx, dstSizes, srcSizes);
   genAddEltCall(rewriter, loc, elemTp, coo, elemPtr, dstIdx, dstPerm);
   rewriter.create<scf::YieldOp>(loc);
   // Final call to construct sparse tensor storage and free temporary resources.
   rewriter.setInsertionPointAfter(whileOp);
-  params[6] = constantAction(rewriter, loc, Action::kFromCOO);
-  params[7] = coo;
-  Value dst = genNewCall(rewriter, loc, params);
+  Value dst = params.genNewCall(Action::kFromCOO, coo);
   genDelCOOCall(rewriter, loc, elemTp, coo);
   genDelIteratorCall(rewriter, loc, elemTp, iter);
   rewriter.replaceOp(op, dst);
   return success();
 }
 
 // Generates a while loop that iterates over the COO list extracted
 // from `t`, using `bodyBuilder` to build the loop body.
static void genSparseCOOIterationLoop(
   unsigned rank = tensorTp.getRank();
   Type elemTp = tensorTp.getElementType();
 
   // Start an iterator over the tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
       rewriter.getContext(), enc.getDimLevelType(), AffineMap(), AffineMap(),
       enc.getPointerBitWidth(), enc.getIndexBitWidth());
   SmallVector<Value, 4> sizes;
-  SmallVector<Value, 8> params;
   sizesFromPtr(rewriter, sizes, loc, noPerm, tensorTp, t);
-  newParams(rewriter, params, loc, tensorTp, noPerm, Action::kToIterator, sizes,
-            t);
-  Value iter = genNewCall(rewriter, loc, params);
+  Value iter = NewCallParams(rewriter, loc)
+                   .genBuffers(noPerm, sizes, tensorTp)
+                   .genNewCall(Action::kToIterator, t);
 
   // Construct a while loop over the iterator.
   Value srcIdx = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
   Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
   SmallVector<Value> noArgs;
   SmallVector<Type> noTypes;
   auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
   Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
matchAndRewrite(NewOp op, OpAdaptor adaptor,
     Location loc = op.getLoc();
     Type resType = op.getType();
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
     // Generate the call to construct tensor from ptr. The sizes are
     // inferred from the result type of the new operator.
     SmallVector<Value, 4> sizes;
-    SmallVector<Value, 8> params;
     ShapedType stp = resType.cast<ShapedType>();
     sizesFromType(rewriter, sizes, loc, stp);
     Value ptr = adaptor.getOperands()[0];
-    newParams(rewriter, params, loc, stp, enc, Action::kFromFile, sizes, ptr);
-    rewriter.replaceOp(op, genNewCall(rewriter, loc, params));
+    rewriter.replaceOp(op, NewCallParams(rewriter, loc)
+                               .genBuffers(enc, sizes, stp)
+                               .genNewCall(Action::kFromFile, ptr));
     return success();
   }
 };
 
 /// Sparse conversion rule for the alloc operator.
 class SparseTensorAllocConverter
     : public OpConversionPattern<bufferization::AllocTensorOp> {
 public:
for (int64_t i = 0; i < resType.getRank(); ++i) {
         sizes.push_back(adaptor.getOperands()[operandCtr++]);
       } else {
         sizes.push_back(
             rewriter.create<arith::ConstantIndexOp>(loc, op.getStaticSize(i)));
       }
     }
     // Generate the call to construct empty tensor. The sizes are
     // explicitly defined by the arguments to the alloc operator.
-    SmallVector<Value, 8> params;
-    ShapedType stp = resType.cast<ShapedType>();
-    newParams(rewriter, params, loc, stp, enc, Action::kEmpty, sizes);
-    rewriter.replaceOp(op, genNewCall(rewriter, loc, params));
+    rewriter.replaceOp(op,
+                       NewCallParams(rewriter, loc)
+                           .genBuffers(enc, sizes, resType.cast<ShapedType>())
+                           .genNewCall(Action::kEmpty));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
if (encDst && encSrc) {
       // Using the coordinate scheme as an intermediate does not always
       // yield the fastest conversion but avoids the need for a full
       // O(N^2) conversion matrix.
       if (encDst == encSrc) {
         rewriter.replaceOp(op, adaptor.getOperands()); // hidden nop cast
         return success();
       }
       SmallVector<Value, 4> sizes;
-      SmallVector<Value, 8> params;
+      NewCallParams params(rewriter, loc);
       ShapedType stp = srcType.cast<ShapedType>();
       sizesFromPtr(rewriter, sizes, loc, encSrc, stp, src);
       bool useDirectConversion;
       switch (options.sparseToSparseStrategy) {
       case SparseToSparseConversionStrategy::kViaCOO:
         useDirectConversion = false;
         break;
       case SparseToSparseConversionStrategy::kDirect:
         useDirectConversion = true;
         assert(canUseDirectConversion(encDst.getDimLevelType()) &&
                "Unsupported target for direct sparse-to-sparse conversion");
         break;
       case SparseToSparseConversionStrategy::kAuto:
         useDirectConversion = canUseDirectConversion(encDst.getDimLevelType());
         break;
       }
       if (useDirectConversion) {
-        newParams(rewriter, params, loc, stp, encDst, Action::kSparseToSparse,
-                  sizes, src);
-        rewriter.replaceOp(op, genNewCall(rewriter, loc, params));
+        rewriter.replaceOp(op, params.genBuffers(encDst, sizes, stp)
+                                   .genNewCall(Action::kSparseToSparse, src));
       } else { // use via-COO conversion.
         // Set up encoding with right mix of src and dst so that the two
         // method calls can share most parameters, while still providing
         // the correct sparsity information to either of them.
         auto enc = SparseTensorEncodingAttr::get(
             op->getContext(), encDst.getDimLevelType(), encDst.getDimOrdering(),
             encDst.getHigherOrdering(), encSrc.getPointerBitWidth(),
             encSrc.getIndexBitWidth());
-        newParams(rewriter, params, loc, stp, enc, Action::kToCOO, sizes, src);
-        Value coo = genNewCall(rewriter, loc, params);
-        params[3] = constantPointerTypeEncoding(rewriter, loc, encDst);
-        params[4] = constantIndexTypeEncoding(rewriter, loc, encDst);
-        params[6] = constantAction(rewriter, loc, Action::kFromCOO);
-        params[7] = coo;
-        Value dst = genNewCall(rewriter, loc, params);
+        // TODO: This is the only place where `kToCOO` (or `kToIterator`)
+        // is called with a non-identity permutation.  Is there any clean
+        // way to push the permutation over to the `kFromCOO` side instead?
+        Value coo =
+            params.genBuffers(enc, sizes, stp).genNewCall(Action::kToCOO, src);
+        Value dst = params.setTemplateTypes(encDst, stp)
+                        .genNewCall(Action::kFromCOO, coo);
         genDelCOOCall(rewriter, loc, stp.getElementType(), coo);
         rewriter.replaceOp(op, dst);
       }
       return success();
     }
     if (!encDst && encSrc) {
       // This is sparse => dense conversion, which is handled as follows:
       //   dst = new Tensor(0);
       //   iter = new SparseTensorIterator(src);
       //   while (elem = iter->getNext()) {
       //     dst[elem.indices] = elem.value;
       //   }
       //   delete iter;
       RankedTensorType dstTensorTp = resType.cast<RankedTensorType>();
       RankedTensorType srcTensorTp = srcType.cast<RankedTensorType>();
       unsigned rank = dstTensorTp.getRank();
       Type elemTp = dstTensorTp.getElementType();
-      // Fabricate a no-permutation encoding for newParams().
+      // Fabricate a no-permutation encoding for NewCallParams
       // The pointer/index types must be those of `src`.
       // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
           op->getContext(),
           SmallVector<DimLevelType>(rank, DimLevelType::Dense), AffineMap(),
           AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
       SmallVector<Value, 4> sizes;
-      SmallVector<Value, 8> params;
       sizesFromPtr(rewriter, sizes, loc, encSrc, srcTensorTp, src);
-      newParams(rewriter, params, loc, dstTensorTp, encDst, Action::kToIterator,
-                sizes, src);
-      Value iter = genNewCall(rewriter, loc, params);
+      Value iter = NewCallParams(rewriter, loc)
+                       .genBuffers(encDst, sizes, dstTensorTp)
+                       .genNewCall(Action::kToIterator, src);
       Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
       Block *insertionBlock = rewriter.getInsertionBlock();
       // TODO: Dense buffers should be allocated/deallocated via the callback
       // in BufferizationOptions.
       Value dst = allocDenseTensor(rewriter, loc, dstTensorTp, sizes);
       SmallVector<Value> noArgs;
       SmallVector<Type> noTypes;
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     // 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().
     ShapedType stp = resType.cast<ShapedType>();
     unsigned rank = stp.getRank();
     SmallVector<Value, 4> sizes;
-    SmallVector<Value, 8> params;
     sizesFromSrc(rewriter, sizes, loc, src);
-    newParams(rewriter, params, loc, stp, encDst, Action::kEmptyCOO, sizes);
-    Value coo = genNewCall(rewriter, loc, params);
+    NewCallParams params(rewriter, loc);
+    Value coo =
+        params.genBuffers(encDst, sizes, stp).genNewCall(Action::kEmptyCOO);
     Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
-    Value perm = params[2];
+    Value perm = params.getDim2LvlMap();
     Type eltType = stp.getElementType();
     Value elemPtr = genAllocaScalar(rewriter, loc, eltType);
     genDenseTensorOrSparseConstantIterLoop(
         rewriter, loc, src, rank,
         [&](OpBuilder &builder, Location loc, Value val, ValueRange indices) {
           for (unsigned i = 0; i < rank; i++) {
             Value idx = constantIndex(builder, loc, i);
             builder.create<memref::StoreOp>(loc, indices[i], ind, idx);
           }
           builder.create<memref::StoreOp>(loc, val, elemPtr);
           genAddEltCall(builder, loc, eltType, coo, elemPtr, ind, perm);
         });
     // Final call to construct sparse tensor storage.
-    params[6] = constantAction(rewriter, loc, Action::kFromCOO);
-    params[7] = coo;
-    Value dst = genNewCall(rewriter, loc, params);
+    Value dst = params.genNewCall(Action::kFromCOO, coo);
     genDelCOOCall(rewriter, loc, eltType, coo);
     rewriter.replaceOp(op, dst);
     return success();
   }
 
 private:
   /// Options to control sparse code generation.
   SparseTensorConversionOptions options;
matchAndRewrite(ConcatenateOp op, OpAdaptor adaptor,
     // A pointer to the value being inserted (if dense => sparse)
     Value elemPtr;
     // Memory that holds the COO for destination tensor (if sparse out)
     Value dstIdx;
     // The offset applied to the dimenstion to be concated (starting from 0)
     Value offset = constantIndex(rewriter, loc, 0);
 
     SmallVector<Value, 4> sizes;
-    SmallVector<Value, 8> params;
+    NewCallParams params(rewriter, loc);
     concatSizesFromInputs(rewriter, sizes, loc, dstTp, op.getInputs(),
                           concatDim);
 
     if (encDst) {
       // Start a new COO for the destination tensor.
-      newParams(rewriter, params, loc, dstTp, encDst, Action::kEmptyCOO, sizes);
-      dst = genNewCall(rewriter, loc, params);
-      dstPerm = params[2];
+      dst =
+          params.genBuffers(encDst, sizes, dstTp).genNewCall(Action::kEmptyCOO);
+      dstPerm = params.getDim2LvlMap();
       elemPtr = genAllocaScalar(rewriter, loc, elemTp);
       dstIdx = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     } else {
       // TODO: Dense buffers should be allocated/deallocated via the callback
       // in BufferizationOptions.
       dst = allocDenseTensor(rewriter, loc, dstTp, sizes);
     }
     for (auto it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
for (auto it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
       Value curDim = encSrc ? sizeFromPtrAtDim(rewriter, loc, encSrc, srcTp,
                                                adaptedOp, concatDim)
                             : linalg::createOrFoldDimOp(rewriter, loc,
                                                         adaptedOp, concatDim);
 
       offset = rewriter.create<arith::AddIOp>(loc, offset, curDim);
     }
     if (encDst) {
-      params[6] = constantAction(rewriter, loc, Action::kFromCOO);
       // In sparse output case, the destination holds the COO.
       Value coo = dst;
-      params[7] = coo;
-      dst = genNewCall(rewriter, loc, params);
+      dst = params.genNewCall(Action::kFromCOO, coo);
       // Release resources.
       genDelCOOCall(rewriter, loc, elemTp, coo);
       rewriter.replaceOp(op, dst);
     } else {
       rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
     }
     return success();
   }
 };
 
 /// Sparse conversion rule for the output operator.
 class SparseTensorOutConverter : public OpConversionPattern<OutOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(OutOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
     // Convert to default permuted COO.
     Value src = adaptor.getOperands()[0];
     auto encSrc = getSparseTensorEncoding(srcType);
     SmallVector<Value, 4> sizes;
-    SmallVector<Value, 8> params;
     sizesFromPtr(rewriter, sizes, loc, encSrc, srcType, src);
     auto enc = SparseTensorEncodingAttr::get(
         op->getContext(), encSrc.getDimLevelType(), AffineMap(), AffineMap(),
         encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
-    newParams(rewriter, params, loc, srcType, enc, Action::kToCOO, sizes, src);
-    Value coo = genNewCall(rewriter, loc, params);
+    Value coo = NewCallParams(rewriter, loc)
+                    .genBuffers(enc, sizes, srcType)
+                    .genNewCall(Action::kToCOO, src);
     // Then output the tensor to external file with indices in the externally
     // visible lexicographic index order. A sort is required if the source was
     // not in that order yet (note that the sort can be dropped altogether if
     // external format does not care about the order at all, but here we assume
     // it does).
-    bool sort =
-        encSrc.getDimOrdering() && !encSrc.getDimOrdering().isIdentity();
-    params.clear();
-    params.push_back(coo);
-    params.push_back(adaptor.getOperands()[1]);
-    params.push_back(constantI1(rewriter, loc, sort));
+    Value sort = constantI1(rewriter, loc,
+                            encSrc.getDimOrdering() &&
+                                !encSrc.getDimOrdering().isIdentity());
+    SmallVector<Value, 3> outParams{coo, adaptor.getOperands()[1], sort};
     Type eltType = srcType.getElementType();
     SmallString<18> name{"outSparseTensor", primaryTypeFunctionSuffix(eltType)};
-    createFuncCall(rewriter, loc, name, {}, params, EmitCInterface::Off);
+    createFuncCall(rewriter, loc, name, {}, outParams, EmitCInterface::Off);
     genDelCOOCall(rewriter, loc, eltType, coo);
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 } // namespace