Propagate parameter names #1420
Conversation
There was a problem hiding this comment.
This changes the semantics of the external calling convention of the MLIR module, which is one of the most complicated parts of the compiler (or any compiler). We currently provide guarantees about positional arguments. Is there no way for you to pass this information directly to the lower-level compiler you are calling into, since Torch-MLIR doesn't use this information, and we guarantee the order of positional arguments?
|
Also, conceptually, I don't understand why an LTC backend would need this information. Can you explain what your backend does with it? |
|
@silvasean This will not affect the MLIR graph generated by Torch-MLIR LTC unless the parameter names are explicitly set If not set it will return a null attribute same as before
This is something specifically required for the Cerebras LTC backend. We require the parameter names be present in the MLIR graph in order to be able to generate our internal tensor mappings. |
|
Can you handle this outside Torch-MLIR? E.g. after lowering the Torch-MLIR graph you will add your Cerebras-specific attributes afterward. |
I don't think so. Just from the fact that we need this change to the Torch dialect to be able to lowering it from TS Jit to MLIR Is it such a great concern to add this in Torch-MLIR, seeing as the parameter's name needs to explicitly be set in order for it to show up in the MLIR graph to begin with. This is something we're only doing in the Cerebras backend |
|
Yes, this is a huge concern since it changes the external calling convention of the module. Please find a way to do it outside Conceptually this seems like it should be not very difficult -- it should be possible to build a map associating the parameter names to function argument numbers, and then after you lower the graph with Torch-MLIR, you apply those attributes on the parameters. |
|
@silvasean I was able to attach the parameter name attribute after lowering into MLIR, so I've cleaned up and removed that change. Please review again to see if the external calling conventions are now preserved |
- Purpose of this is to allow modification of func_op by subclass backend
antoniojkim
force-pushed
the
antoniojkim/parameter_names
branch
from
cb84d49 to
3e4ec62
Compare
* Propagate parameter name to MLIR * Add TorchMlirNode Constructor Hook * Make func_op mutable - Purpose of this is to allow modification of func_op by subclass backend * Clean up unnecessary changes * Remove unnecessary attribute case * Address PR comments
We need to propagate parameter names into the MLIR graph in order to create proper tensor mappings for the Cerebras backend.
CC: @ke1337 @glebk-cerebras @vaibhavc-cerebras @prateekgu-cerebras