Use weakref in storing tensors as keys (follow-up to #111470)

[pearu]

This PR addresses the discussion items in #111470 (comment), that is,

• use weakref when storing tensors as keys,
• add storage_offset to the key data,
• and revise the description of the TensorAsKey utility.

Stack from ghstack (oldest at bottom):

• Fix scatter_mm kernel failure on non-contiguous tensor arguments #112337
• -> Use weakref in storing tensors as keys (follow-up to #111470) #112076
• Eliminate try-catch block around triton::_triton_bsr_dense_mm_out call. #112154

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/112076

• 📄 Preview Python docs built from this PR
• 📄 Preview C++ docs built from this PR
• ❓ Need help or want to give feedback on the CI? Visit the bot commands wiki or our office hours

Note: Links to docs will display an error until the docs builds have been completed.

✅ No Failures

As of commit 4902f95 with merge base f5088d2 ():
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[albanD]

I'm confused about what this means tbh.

[pearu]

The definition of _bsr_scatter_mm_indices_data function is:

@lru_cache(maxsize=128)
def _bsr_scatter_mm_indices_data(indices_format, M, K, N, Ms, Ks, nbatches, SPLIT_N, crow_indices_as_key, col_indices_as_key):
    ...

which means that a key to LRU cache constitutes of all arguments to the function:

lru_cache_key = (indices_format, M, K, N, Ms, Ks, nbatches, SPLIT_N, crow_indices_as_key, col_indices_as_key)

while the value of the LRU cache is the result of the function. Notice that the argument crow_indices_as_key is TensorAsKey(crow_indices).
On the other hand, the result of the function is a tuple (indices_format, c_indices, r_offsets, q_offsets) where originally c_indices was crow_indices.

So, this means that the tensor crow_indices is stored both in the key as well as in the value of LRU cache. Now, when the crow_indices variable goes out of scope, the corresponding tensor ought to be garbage collected (assuming that crow_indices appears as weakref instance in the key). However, the tensor cannot be garbage collected because it is referenced by the value of the LRU cache. Therefore, the crow_indices tensor must be cloned so that garbage collection would be effective when the crow_indices variable goes out of the scope, and later, when the LRU cache will overfill, the cloned tensor will be cleaned up as well.

[pearu]

Hmm, after thinking about it, I think the crow_indices tensor in the value part of the LRU cache ought to be weakref as well. Then we don't have to wait for an overfill for the tensor to be garbage collected. I'll fix it unless I missed something right now..

[pearu]

This comment is now removed. It turns out that it is not sufficient to use crow_indices and col_indices as keys because these have shorter lifetime than the original sparse compressed tensor. Now the original tensor is wrapped with TensorAsKey to make sure that the cached item exists in the life time of the original tensor.