Avoid memory allocations and deallocations when creating NVTETensor

[ptrendx]

Description

Changed the transformer_engine::Tensor allocation to use static pool. NVTETensor is no longer a pointer to Tensor, but rather an index into the pool.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Introduce TensorAllocator with a static pool capable of storing 20MB of Tensors
• Changed the reinterpret_cast<Tensor*> logic for converting NVTETensors to transformer_engine::Tensor to call a dedicated conversion function instead.

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[ptrendx]

/te-ci

[timmoon10]

The fixed 20 MiB allocation size is a little hacky. How about storing the Tensors in unique_ptrs?

Suggested change

	std::vector<Tensor>memory;
	std::vector<std::unique_ptr<Tensor>> memory;

This way the vector can freely reallocate and move unique_ptrs without disturbing the Tensors' addresses. We could also switch to a more appropriate data structure like unordered_set.

[ptrendx]

I was thinking quite a lot about this actually before settling on that design. There are a few reasons why I chose it:

• with NVTETensor as a pointer (so that you can use the unordered_set) you don't have the ability to check if the parameter given to the function is actually a valid Tensor. Bugs in user code could then lead to crashes/security vulneabilities.
• in the regular usage the NVTETensor values used in a function would be very close together, helping with the memory access pattern (that would be best if we got rid of the std::vector for the shape in Tensor and instead used NVTEShape, making Tensor fully POD. Keep in mind we are fighting for microseconds here.
• I would argue that the 20MB limit is actually way too high. It enables creating over 83k Tensors. There is no real way of getting to that number without memory leak.

[timmoon10]

• Making NVTETensor an index for security reasons seems sensible. I would suggest keeping this PR's index logic, but storing the tensors themselves in unique_ptrs.
• In most cases, unique_ptr should compile down to the same instructions as raw pointers. The only performance impact should be with caching.
• If cache misses become performance critical, we could allocate Tensors in blocks of a few cache lines. This would result in identical performance as a single contiguous block, but with an adaptive buffer size.
• I don't think we should strive to make Tensor a full POD. It would benefit from a redesign for polymorphism (see the visitor pattern prototype in this branch).

[ptrendx]

/te-ci

[ptrendx]

/te-ci

[timmoon10]

Overall LGTM. I still don't like that we allocate Tensors out of a fixed 20 MiB buffer, but it's an implementation detail that's not exposed externally.