Performance / optimization with identical array arguments

[JeffGreen]

Hello! Judging from the stablehlo / make_japxr() output of :

some_func.lower(my_arr, my_arr).as_text()

it seems that JAX doesn't recognize that both array arguments are identical. I'm guessing this means that things like optimization of loads, cache performance, etc. is sub-optimal. Is there a way to indicate to JAX that arguments for a particular invocation of a given signature are identical?

[mattjj]

Thanks for the question!

If you want a single argument, you should write jit(lambda x: some_func(x, x)).lower(my_arr, my_arr).as_text().

When you write f.lower(*args), it's really a convenience layer over specifying a type at which you want to trace f. We don't look at the values of args, or their object ids; we effectively replace them with something like args = [jax.ShapeDtypeStruct(x.shape, x.dtype) for x in args].

You could imagine an example-arguments-to-type policy which automatically de-duplicated, by value or object id. But that would have its own surprises: what would the calling convention be? To pass a single argument (even though we called lower with two arguments)? Or to pass two and then check that they're the same object (or the same value)?

When things are ambiguous, we tend to prefer the simplest policies, requiring the user to be explicit. That is, in this case we require you de-duplicate the inputs yourself. Luckily you can always write a wrapper which performs that logic automatically, if you prefer it!

What do you think?

[JeffGreen]

Thanks for the response!

I'm aware of manual de-duplication as in your lambda example, but it the case of many more arguments, some of which may be sets of duplicates, it becomes more complex. As a potential feature, it's kind of the opposite of c's restrict so I was inclined to think / hope that XLA somehow handled it on the compiler / framework level as opposed to pushing it back up to the user.

With respect to your comment about implementation of automatic deduplication (and without knowing the nuances of stablehlo!), I could imagine JAX / XLA annotating each argument with an incrementing integer in the case of a never-seen argument (by identity) or a previously-used integer if it's a duplicate. As this would be transparent to the user, you'd still need to call with two arguments.

With all that being said, it seems like it's possible to do on my end, hopefully even in the more complex case that this trivial example came from.

Thanks again.