Ticket

#1041, #856

Problem description

• We cannot populate the argument types for torch models.

What's changed

• Added custom torch operator torch.ops.tt.mark_argument_attributes which applies a stablehlo.custom_call to a tensor :
  • This operator returns the input tensor.
  • This operator has attributes "argument_type" and "name". (Name is optional)
  • This operation can ONLY be run on an xla device, PyTorch will raise an error if the user attempts to execute this operation on a non-xla device.
  • This op can be captured by the dynamo and is a valid inside a GraphModule
• Added a pass to populate the argument types of the SHLO module using this custom call
  • This pass does something special, If at least one argument is populated as an "input", then all un-populated arguments will be automatically defined as "constant".
    • This works under the assumption that every user input is annotated. Users will not be expected to populate the argument types themselves. It will either be done by our torch.compile backend or by the utility function we provide: mark_model_user_inputs.
• Edited propagateRoleAttribute to simply populate the ttcore.argument_type arg attribute in the function rather than tt.input_role. This way we can bypass PopulateArgumentTypes in tt-mlir entirely as the graph is already populated.
• Deleted the environment variable override for populating argument types as it should no longer be needed.

Checklist

• New/Existing tests provide coverage for changes

I didn't add a test because we will not really be able to take advantage of this until the tt_torch backend is merged: #954, I will add a test after both of these are merged.

I did add an example though, using the "openxla" backend.