[API Proposal]: Add TensorPool<T>

Open
#115,201 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Assessment

Difficulty
5/5
Estimated time
Over a week
Newbie friendliness
30/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Stale
Tech stack
csharp

Research direction

No runtime files or tests are named in the issue; start with the proposed TensorPool<TTensor, T> API and the ArrayPool comparison. Review the linked Example/SentimentInference/Example.SentimentInference.Console usage to understand the pooling scenario, then resolve the open API choices, including clearing tensors on return; done requires an agreed design and implementation scope.

Written by the indexing model from the issue text.

Description

api-suggestion area-System.Numerics.Tensors
Background and motivation

Due to the nature of the usage of tensors they tend to be really big.
In addition, due to usage patterns creating copies with different operations applied to them, many of which can't be in place, applications using Tensors create many temporaries which are large.

In my ML application, there are fixed number of threads running operations on batches of sentences, so there are a constant number of tensors active at each moment.
I am spending ~1% of my time in the GC, mostly on the large 500KB Tensors I am tokenizing and processing.
I have not measured the overhead of the allocations and GC pauses in my server example, but I assume it will be larger since I do not have all of the requests ahead of time and so have to create more tensors.

Bucketing the tensors by size and pooling instances will reduce the allocations done by the server, reduce the GC time, and increase throughput.

This is almost exactly like ArrayPool<T>, and should probably be a wrapper of it.

This is something I discussed with @tannergooding but I didn't see any API proposals.
I will state, that like ArrayPool<T>, this pool might in some situations be overkill and cause more harm than benefit, as the GC is doing an amazing job.

API Proposal
namespace System.Numeric.Tensors.Buffers;

public class TensorPool<TTensor, T> where TTensor: ITensor<TTensor, T>
{
    public TensorPool();
    public TensorPool(int maxFlattenedLength, int maxTensorsPerBucket);  // Configuration Params.

    public TTensor Rent(scoped ReadOnlySpan<nint> lengths);
    public TTensor Rent(scoped ReadOnlySpan<nint> lengths, scoped ReadOnlySpan<nint> strides);

    public void Return(TTensor tensor, bool clearTensor = true);  // should it default to false? 
}
API Usage
int maxFlattenedLength = 1000;
int maxTensorsPerBucket = 10;
int exampleCount = 10000;
var tensorPool = new TensorPool<Tensor<float>, float>(maxFlattenedLength, maxTensorsPerBucket);

var lengths = new nint[] { 100, 100 };
float[] averages = new float[10000];
Parallel.For(0, 10000, i =>
{
    Tensor<float> tensor1 = tensorPool.Rent(lengths);
    Tensor<float> tensor2 = tensorPool.Rent(lengths);
    Tensor<float> tensor3 = tensorPool.Rent(lengths);

    Tensor.FillGaussianNormalDistribution(tensor1.AsTensorSpan());
    tensor2.Fill(i * i);
    Tensor.Add(tensor1.AsReadOnlyTensorSpan(), tensor2, tensor3);
    Tensor.Multiply(tensor1.AsReadOnlyTensorSpan(), tensor3, tensor2);

    averages[i] = Tensor.Average<float>(tensor2.AsTensorSpan());

    tensorPool.Return(tensor1, clearTensor: true);
    tensorPool.Return(tensor2, clearTensor: true);
    tensorPool.Return(tensor3, clearTensor: true);
});

foreach (float average in averages) Console.WriteLine(average);
Alternative Designs
  • Each consumer wanting to pool Tensors will have to write their own wrapper around ArrayPool and deal with converting the lengths and strides to the correct size of backing array.
  • Not pooling Tensors - relying on the GC to do a good job on big chunks of memory.
  • Should the default for TensorPool be to clear the Tensor on return? Usually they are used for calculations, and so will most likely not want leftovers from the previous calculation.
Risks

No response

Dominant language
C#
Stars
18.3k
Forks
5.6k
PR merge metrics
PR metrics pending

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from dotnet/runtime

All issues in dotnet/runtime

Similar issues

More C# issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.