Reduce time complexity of node replacment in PyTorch frontend #1
Conversation
marty1885
changed the title
python/tvm/relay/frontend/pytorch.py
Outdated
| @@ -5871,7 +5871,19 @@ def export_c_graph(location, graph): | |||
| fname = os.path.join(location, f"tvm_exported_c_graph_{time_stamp}.txt") | |||
| with open(f"{fname}", "w") as f: | |||
| f.write(str(graph)) | |||
|
|
|||
|
|
|||
| def _binray_search(lst, func): | |||
There was a problem hiding this comment.
Correct spelling to _binary_search(lst, func):
python/tvm/relay/frontend/pytorch.py
Outdated
|
|
||
| relevant_ops = node_inputs_map[orig_node] | ||
| begin_idx = _binary_search(relevant_ops, lambda x: x[0] > node_idx) | ||
| for idx, node in relevant_ops[begin_idx:]: |
There was a problem hiding this comment.
No need for idx, can you replace it with _?
python/tvm/relay/frontend/pytorch.py
Outdated
| @@ -5871,7 +5871,19 @@ def export_c_graph(location, graph): | |||
| fname = os.path.join(location, f"tvm_exported_c_graph_{time_stamp}.txt") | |||
| with open(f"{fname}", "w") as f: | |||
| f.write(str(graph)) | |||
|
|
|||
|
|
|||
| def _binary_search(lst, func): | |||
There was a problem hiding this comment.
Instead of calling arg name func, let's use condition for easier understanding
|
@nvukobratTT Poke on this. I've applied the requested changes. |
Hey @marty1885, sorry for the delayed response. We have to do a few internal checks before final approval and merge. As soon as I get those I'll approve this MR :) Until then, freely continue with model bringup with these changes cherry-picked. I hope this isn't causing too many issues on your end. |
|
@nvukobratTT Poke again. Is there anything I can do to push forward the PR? I feel this patch would be helpful to every user of BUDA and wants it to be included in future releases. I understand your schedule is busy. More then a month to merge a widely beneficial patch seems a bit long. |
Hi,
This is a part of the joined work with @JushBJJ and @JonathanALevine to get LLMs running in BUDA (and the bounties). This PR reduces the time complexity of the node replacement step. This makes time to compile large models tolerable. And pretty much required when compiling really large models like RWKVv4 which has 600K nodes.
Ref. tenstorrent/tt-buda-demos#22
The old code loops through all nodes for each node it wants to replace. Which makes it an O(N^2) algorithm. Now it uses a hash map and reduces the complexity down to O(N).