BUG: free-threaded dynamic loop addition and ideally loop replacements in ufuncs

[seberg]

I didn't think this would get back so quickly, but we should allow dynamic addition of loops in ufuncs.
Previously, this was OK at least when done during promotion (which is to me the interesting part -- other times could add their own lock, but this path cannot).

Now, we could just disable the duplicate addition check and things will probably work, but not sure what is best.
I am not sure there is much of a point to not just lock everything when we do promotion (promotion should be very rare!).

I also would like to allow loop replacement eventually, and this has come around a bit quicker now as well.
Now, that requires clearing the cache, which is already similar to what is happening when growing, but the trivial way to add it requires leaking the cache buckets (which growing also currently does?!).

We can just move back the locking one level up, but I do wonder a bit if all of this is actually worth it compared to an RWlock?

Right now, I added dynamic loops for imag/real, but it is internal, so there is the "ignore that it is already added" flag public already.

CC @kumaraditya303 and @ngoldbaum.

[ngoldbaum]

Can you elaborate with a concrete example of what you want to do that you can't do now? I'm a little confused by the issue description.

[seberg]

There are two things in #30984 I am adding new loops within the promoter function. That is a loop for new DTypes is added during promotion.
While that was a bit broken until that PR, that is a feature that would be nice to have properly, because it is something numba does (compile new loops on the fly depending on the DTypes coming in).

From a cache perspective that is actually OK, as long as we ignore the second registration (which we do). But for adding it to the loops list is currently an error. It could just be ignored there, but I am not sure if it doesn't make sense to just lock things a bit more broadly (i.e. always lock when promoting).

The second thing is related to #31068 we have PyUFunc_ReplaceLoopBySignature. That is mostly unused API (probably?) except for that pnumpy project that Matti was once involved with.
For the sake of making that go away/breaking it, it would be nice to be able to point to an alternative. But that does require invalidating the cache.
(I.e. if one replaces implementations or even adds specialized versions of existing things that include casts, one needs to invalidate the cache in principle. This can make a lot of sense in a JIT context.)

The first is easy, since you could just ignore the double registration, that feels a bit awkward generally, but maybe that is fine. But it would also work to lock everything when promotion happens (so that the loop is only added once if it is dynamically added at promotion time).

The second, I am not sure how easy it is. It could be a really cool thing to try and speed up NumPy with JIT'ed loops, but PyUFunc_ReplaceLoopBySignature is the explicit need, and I admit it is a bit dubious in whether it is used.

[ngoldbaum]

I'm still really confused. You're talking about PyArrayIdentityHash, right? It would really help me if you could point to concrete code diffs where current limitations of PyArrayIdentityHash are causing problems. If you're talking about replacing loops in the dispatch cache, what about Kumar's comment here? #31068 (comment). E.g., using copy-on-write.

We can just move back the locking one level up, but I do wonder a bit if all of this is actually worth it compared to an RWlock?

We moved away from the RWlock because there would still be a mutex that needs to be held for a short amount of time in the fast path. For this sort of read-only access pattern where we are protecting against rare writes, it ends up being a poor trade-off: many simultaneous reads are guaranteed to cause contention, even after the cache is filled.

I don't really have a strong opinion about any of this, as long as we can figure out a way to keep the lock-free fast path.

[seberg]

For the first part the cache is clearly fine since there is no replacement. But, the lock used for the cache could be locked more widely during promotion to avoid races while adding loops. I.e. #30984 adds a loop inside the promoter which is cool for JIT'ing etc.
And that currently only works for private API because it has the flag to ignore that we are trying to add the same thing twice.

Yes, Kumar's comment is part of the second issue, but I wanted a clearer issue on it.

We moved away from the RWlock because there would still be a mutex that needs to be held for a short amount of time in the fast path. [...] as long as we can figure out a way to keep the lock-free fast path.

Sure, I don't care too much if we leak loops when replacing them, e.g. leaking the whole cache seems a bit funny but maybe fine. For the first part, I am asking wondering about a lock for the promotion part, which is outside the fast-path.

For the second, Kumar's comment may be the right thing. I dunno, you don't expect a lot of replacement, etc. so maybe effectively leaking stuff is fine.
(I honestly don't have a clear feeling for the trade-off of removing congestion in some cases vs. not cleaning up here but it is little memory and bounded, even if the speed difference may also be low.)

[ngoldbaum]

And that currently only works for private API because it has the flag to ignore that

Ah, I missed 38ced07 in the discussion for #30984. I'd appreciate a direct ping when you feel the need to add hacks to work around issues you encounter on the free-threaded build.

@kumaraditya303 can I get you to weigh in here? Is there a way to make PyArrayIdentityHash thread-safe if we allow loop replacement?

For the leaks, I'm not entirely sure why you think cleanup of ufuncs means we have to leak things. Are you talking about how we made the array method objects immortal? We should probably just back that out and try to find a different way to reduce reference count contention in casts if they objects aren't actually immortal. I don't think we should intentionally leak anything.

[seberg]

For the leaks, I'm not entirely sure why you think cleanup of ufuncs means we have to leak things

Ah, I was just meaning the way that the cache currently keeps the old buckets around on resize in free-threading mode (IIUC), because it needs a delayed de-allocate otherwise. And I suppose I guessed that clearing the cache would do the same.
But the comment in the other thread already mentions a possible future deferred dealloc too, so...

[kumaraditya303]

can I get you to weigh in here? Is there a way to make PyArrayIdentityHash thread-safe if we allow loop replacement?

I can think of two ways to keep the hashtable lock-free while allowing for replacement:

• The first way is to replace the entire backing buckets on every replacement and keeping the old buckets in the linked-list of old buckets. This would consume a lot of memory if a lot of replacements are performed but if it is rare then it will be simpler to implement. If CPython exposes the PyMem_FreeDelayed as an API then the memory could be freed on a ceval quiescent state. This is same as what I suggested in ENH: Improve ufunc performance by caching the loop function #31068 (comment)

• The second way is to add another level of indirection to the hashtable so that individual entries can be replaced without having to replace the entire buckets. This would consume much less memory as only the replaced entry would be kept in memory for lock-free concurrent readers. Here's a rough outline of the new cache layout:

  struct bucket_entry {
      PyObject *value;
      PyObject *keys[];     /* flexible array, length = key_len */
  };
  
  struct buckets {
      struct buckets *prev;
      npy_intp size;
      npy_intp nelem;
      struct bucket_entry *entries[];  /* array of pointers, NULL = empty slot */
  };

  With this structure you can allocate a new entry for a key and replace it via an atomic store on the entry itself and replace it. You would have keep the old entries alive for concurrent readers but this layout would consume much lesser memory. The downside is that it would cache lookup a bit slower because of one extra data-dependent memory load.