Reduce time complexity of node replacment in PyTorch frontend

[marty1885]

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).

[marty1885]

@nvukobratTT Poke on this. I've applied the requested changes.

[nvukobratTT]

@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.

[marty1885]

@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.

[nvukobratTT]

Hey @marty1885 I hope you don't mind for late dellays. We're doing our best to merge these changes as soon as possible :)) However, a few personnel dedicated to this merge were on vacation, and currently, there are holidays upcoming. Because of if, the final merge is still pending and I hope it'll get resolved soon.

Once again I hope you don't mind the wait, and we'll push for this change to be merged as soon as possible, as it's beneficial for us and other contributors as well :))) Thanks once again for your and team's effort!

[marty1885]

@nvukobratTT Sure but more then a month to merge a single patch seems excessive. And will become a barrier for future contributors as it might be easier for them to just fork and not bother with upstreaming.

Do you have an estimate on when this could be merged? And any chance the merging process be improved?

[nvukobratTT]

@marty1885 I agree with you! One of the purposes of these bounty hunts was to see where there are loopholes in our current process. We're working on solving them before we release the next set of bounty programs.

The wait time isn't related only to this change. It's small indeed, however, we want to cover other changes as well and push them soon after we solve current process limitations. There are a few more details we need to discuss and cover, but I hope we'll have some feedback for you and the rest of the contributors soon :))

Once again, thanks for the effort!