[API Proposal]: Add TensorPool<T>
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
- Domain
- machine-learning
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
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
More from dotnet/runtime
-
agentic-workflows untriaged
Difficulty 2/5 1-3 hours Newbie friendliness 76/100
-
area-System.Reflection blocking-clean-ci-optional Known Build Error os-mac-os-x untriaged
Difficulty 2/5 1-3 hours Newbie friendliness 68/100
-
area-CodeGen-coreclr untriaged
Difficulty 1/5 Under an hour Newbie friendliness 92/100
-
agentic-workflows untriaged
Difficulty 1/5 Under an hour Newbie friendliness 78/100
-
area-VM-meta-mono untriaged
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
Similar issues
-
Difficulty 2/5 1-3 hours Newbie friendliness 86/100
-
:watch: Not Triaged 11.0 fundamentals/subsvc
Difficulty 2/5 1-3 hours Newbie friendliness 92/100
dotnet/AspNetCore.Docs#37699 ·
-
Difficulty 2/5 1-3 hours Newbie friendliness 72/100
SubtitleEdit/subtitleedit#15108 · 1 comment ·
-
area/docs-content Bug pulumi/docs
Difficulty 1/5 1-3 hours Newbie friendliness 94/100
-
Create parent directories only after the containment check in InstallHelper.TryExtractToDirectory Open
Difficulty 2/5 1-3 hours Newbie friendliness 78/100
PowerShell/PSResourceGet#2056 ·