Thanks for adding this comment!

This function returns three tensors. How do these relate to the inputs and to the result of the projection?

IIUC, for

lhs, rhs, lhs_hash_opt = sparse_mask_like_prepare_sparse_inputs(t, mask)

we have the following invariants

lhs.indices() == t.indices()[lhs_hash]
lhs.values() == t.values()[lhs_hash]
rhs == mask

where lhs_hash = lhs_hash_opt or slice(None), lhs may be a copy or a view of t, rhs is uncoaleced view of mask.

See the comment in #94991 (comment).

What is the advantage of returning mask as uncoalesced?

If mask is uncoalesced, but the other argument is not, the COO intersection kernel will return a tensor with the same indices as mask but will do binary search of mask hashes into the hashes of the other argument's indices all without calls to sort and coalesce. The COO intersection kernel is heavily optimized to take advantage of is_coalesced of either argument, see https://github.com/pytorch/pytorch/pull/92976/files. As such, calling an intersection kernel with arguments (a, b) might produce results with a.indices() or b.indices() depending on which is more performant (i.e. does not sync) based on whether a.is_coalesced or b.is_coalesced. In order to force this kernel to do what we want for sparse_mask, we mark certain arguments as "coalesced" (if need to after sort) and mask as uncoalesced to make sure the result has indices mask.indices().

This (forcing a coalesced input to be uncoalesced) sounds like a convoluted way to control COO intersection kernel functionality (to ensure that the result indices are the input indices).

It still better than writing things from scratch, imho. The COO intersection kernel is not a public function with fixed semantics, but we want it to be fast and without any syncs if possible, so that, say, mul(a, b) and mul(b, a) produce the very same tensor and without any syncs if either a or b is coalesced. In my opinion we can sacrifice some clarity for performance here given that the array of ops that we can implement with this kernel is quite substantial and is performance critical (do you remember @amjames was working on removing calls to coalesce?) Fast sparse on CUDA is what sells it in the first place, and there is still some room for improvement in the intersection kernel as I see it... But I hear you, and will probably modify interface a bit now that I know the use cases better. Current design assumed just mul, I did not know about sparse_mask and its importance back then (and other use cases, like non-symmetric context). That could be a nice follow-up once performance is here...

Ok, so essentially we're starting to write manual fusions. If we had torch.compile support we could presumably generate the code for this pattern. Since this is a hot path this is ok, but it's going to start to become difficult to test correctness.

Since _sparse_mask_projection is a native function you could even write a test that compares it to the pattern that you're fusing.

Is this a common general pattern? Maybe we can apply it in more locations. If we generalize the fused composite native function that you're adding here a bit, maybe it applies to more locations? It can then also subsume all the logic of sparse_mask_like_grad and we can explicitly say "This function represents the fusion of this pattern". That'll be easier to understand to future maintainers.

@cpuhrsch , alternatively, we can make this function composite implicit now that sparse_mask does support backward with COO inputs. This way we can remove the code for backward altogether. The result: much less code, still sync-less backward albeit potentially slower compared to this impl (there is only one way to call backward, there is no back-and-force between grad and input, because of the fixed semantics of sparse_mask). But, yes, it is a generic pattern forced by some "sparse-semantics" functions which might become obsolete with differentiable sparse_mask or something similar but more flexible in enforcing projection direction (aka a public interface for our COO intersection primitive kernel).
Another possibility: we can create a public function which is doing either _sparse_mask_projection and/or sparse_mask (whichever is faster) which is explicitly differentiable and could be used to enforce sparse semantics all while having flexible backward. Then any sparse semantics function will be a composition of this method and the underlying logic it implements.

I think we should go with your first approach for now and keep the amount of code low and then work on torch.compile support to provide these fusions. We can keep what you have for now as a reference and goal for that integration. Otherwise we'll have to undo these manual fusions again once we have torch.compile support.