Repository: pytorch
Issue: Refactor/Cleanup lazy IR codegen
Labels: triaged,module: lazy

This is a place to collect goals/motivations for a refactoring pass of the lazy tensor IR codegen infra.

1) ~~Ok, so the LazyIrSchema object knows "what are the C++ types and names of all arguments in my schema", and now it also knows "which of my arguments were originally scalars before processing", using this separate list.~~

~~I think the other "data-model-y" option you could take would be to represent that info per-argument using a new type (e.g. a LTCArgument that knows both its NamedCType, and whether it was a scalar argument). The advantage there is that your processing functions (node_ctor_arg_rvalue_string) can run on each argument in isolation, instead of having to take in the whole schema. But if this list is used in a really minimal/targeted way then maybe that's overkill.~~
(1) addressed by https://github.com/pytorch/pytorch/pull/73939

2) RE aten interned strings: now that this is all in core, should we just add these symbols in core instead of having to handle them specially in codegen?

3) any reason why we need to use the metadata.kernel to choose how to name the shape functions? For example, we could do something like f.func.name.unambiguous_name() instead.  
Note: 
here is one case where they differ.

ipdb> p metadata.kernel
'arange_out'
ipdb> p f.func.name.unambiguous_name()
'arange_start_out'

what we probably want to have happen (based on what some of the other codegen in core does), is that the shape signatures are part of some API (the "lazy API", or "lazy shape API") that knows how to generate its name. That way at both the place where you define those functions in the codegen, and at the place where you call those functions, you can use the same api call to generate the name (and you don't have to remember whether or not it was actually using metadata.kernel or f.unambiguous_name()

4) another suggestion from @bdhirsh
In [tools/codegen/dest/lazy_ir.py](https://github.com/pytorch/pytorch/pull/73793#discussion_r831239842):

> @@ -180,17 +182,21 @@ def lazy_tensor_decls(value_types: List[NamedCType], tensor_class: str, schema:
         if t.name in schema.wrapped_scalar_names:
             # no lazy tensor wrapper for scalars that are promoted to IR values
             continue
-        if isinstance(t.type, BaseCType):
-            lazy_tensor_decls.append(
-                f"{tensor_class}Ptr lazy_{t.name} = "
-                f"torch::lazy::GetLtcTensorOrCreateForWrappedNumber({t.name}, *common_device);")
+        elif isinstance(t.type, BaseCType):
+            if t.type.type is tensorListValueT:
Going back to Nick's question [above](https://github.com/pytorch/pytorch/pull/73793/files#r822284536), I think Will's comment [here](https://github.com/pytorch/pytorch/pull/73793/files#diff-97a6e7faf841231c6814dbe6eebbb10217c557cf1437ffa241dca5d6cb10cfeeR11) makes sense to me. It looks like you're using the NamedCType here for two purposes:

(1) getting the C++ string to emit (torch::lazy::Value)
(2) pattern matching on the type to figure out extra stuff to emit (it's a TensorList type, so emit extra stuff)

But NamedCType isn't great for pattern matching, and tensorListValueT looks like a good example of why - when we see Tensor and Tensor[] arguments in the schema, we want both to map to a torch::lazy::Value (so they're C++ type is the same), but we still need to handle those cases differently (which you get around by creating a separate CType that you can pattern match on for Tensor lists, even though the C++ type is the same).

This totally doesn't require an immediate change (I saw the PR already landed, just figured I'd add a quick comment since I was tagged :) ). One.. slightly less hacky / easy way to get rid of the tensorListValueT hack would be to make LazyIrSchema hold a Tuple[Binding] instead of a Tuple[NamedCType]. That's still a bit weird since that's not the main purpose of Binding, but if you look at where it's [defined](https://github.com/pytorch/pytorch/blob/aff74cd654571ae49c650d124861821b90660251/tools/codegen/api/types.py#L305), it holds both a NamedCType (where you can store valueT) and an Argument (that can hold the Tensor[] information, which you can use for pattern matching).

The change for LazyIrSchema's constructor would be pretty simple - something like:

...
positional_arg_types.append(Binding(
            nctype=NamedCType(arg.name, process_ir_type(arg.type)),
            name=arg.name,
            default=None,
            argument=arg,
))

cc @bdhirsh  - please feel free to list other stuff here