[API Proposal]: Add TensorPool<T>

[tjwald]

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

[dotnet-policy-service]

Tagging subscribers to this area: @dotnet/area-system-numerics-tensors
See info in area-owners.md if you want to be subscribed.