Should we use namespace, opname = torch._library.utils.parse_namespace(qualname) ?

but parse_namespace only parses - we'd still need to write the op lookup. we can do something like this

try:
    op_overload = torch._library.utils.lookup_op(qualname)
except (ValueError, AttributeError):
    pass

lookup_op uses parse_namespace internally and handles the full lookup

Does this handle inplace operation correctly? Should we return a mutated argument (e.g. Tensor(a!) out), but not create a new tensor?

_all_reduce returns the input directly

valid point. FallbackKernel tracks mutations via handle_aliasing_and_mutation, but mutates_and_returns_first_arg only works for aten ops. For inplace symm_mem ops, we should either verify the current behavior is correct or handle them like _all_reduce_ does. let me investigate it

I just checked mutable ops returning Tensor(a!) will fail - can_auto_functionalize returns False for non-aten ops with aliased returns. We can use _CollectiveKernel.create_inplace instead. I will test it

Should we also reuse the name parsing logic in torch._library.utils here?

aah, correct. We should

cc @zou3519 mind looking at the library changes ?

I'm not expecting a CollectiveKernel create_inplace refactoring in this pr. Can we separate that to a separate pr ?

agreed, I'll remove the _CollectiveKernel.create_inplace usage from this pr entirely. The symm_mem buffer realization now will happen generically inside ir.FallbackKernel.create via _maybe_realize_symm_mem_args, which will check the registry and realizes the registered args as comm buffers. The per-op custom lowerings and the _create_symm_mem_lowering / _get_mutated_return_arg helpers have been removed. I'll handle the CollectiveKernel.create_inplace refactoring in a follow-up PR.

The maybe_realize_symm_memory handling i'm expecting to just go in ir.FallbackKernel. Can you put it there ?

nit: you can use a tree_map util to recursively do this instead of manual inspect.

Can you please add documentation to these fields?

sorry, why do we need this?

without the bypass, mutable symm_mem ops (e.g., multimem_all_reduce_, two_shot_all_reduce_) hit the NotImplementedError. they have schema.is_mutable = True but can_auto_functionalize returns False for them, and they don't have a Functionalize dispatch key implementation.

could we fix this before land. cc @zou3519 to comment. i dont know why can_auto_functionalize matters in inductor lowering.

@eellison the mutable symm_mem ops (e.g., multimem_all_reduce_, two_shot_all_reduce_) return aliased tensors (Tensor(a!) -> Tensor(a!)), so they fail both mutates_and_returns_first_arg (aten-only gate) and can_auto_functionalize (no aliased returns), hitting the NotImplementedError... the proper fix would be routing these through _CollectiveKernel.create_inplace, but you mentioned earlier that refactoring should be a separate PR. Should I go ahead and do it here instead? @zou3519 any thoughts on whether the can_auto_functionalize guard should be relaxed for this case?

@RohitRathore1 can you change the mutable symm ops to not return the aliased tensor? Instead you can write any code that calls them to just use the input tensor.

Otherwise, we can special case these but we are going to need to work a bit to special case them

@zou3519 because no callers rely on the return values of the mutable symm_mem ops, so changing them to return () is safe. However, that's a schema change to the ops themselves which feels like it should be a separate PR. would it be ok to keep the bypass for now and do the schema change as a follow-up? or would you prefer it in this PR?

We should figure it out now, we risk silent incorrectness otherwise. Either we fix the schema, or we are going to need to special case these operators in functionalization. If we special case the operators then we need to make specific changes to special case the operators.

@kwen2501 would you be okay if we changed the schema for these operators? The TL;DR is that custom operators are not allowed to do something like the following:

def op_(x):
    mutate(x)
    return x

Instead they should be written like the following, to avoid returning the input directly as the output:

def op_(x):
    mutate(x)

If there are BC concerns here, then we can try to special case them in compile.

Thanks @zou3519 . I am fine with breaking the API signatures now.
@RohitRathore1 can you change the APIs for those couple ops? Thanks!

@RohitRathore1 You can also pull a separate PR and base this PR on top. Either way works.

@kwen2501 created #179144 with the schema changes... will rebase #173513 on top of it once this lands.