Modify insert-gpu-allocs pass to also consider xegpu.store/load

[dchigarev]

Hi @Menooker, these changes in IMEX are required for linalg->xegpu->gpu exe pipeline in graph compiler.

The PR modifies insert-gpu-allocs that copies data between CPU and GPU so that it would also consider xegpu.load/store operations as a load/store operations and trigger the copy between devices.

I've tested the changes locally on the GPU pipeline (via xegpu).

[Menooker]

The patch LGTM. I am happy with merging this to our development branch. However, this looks like a bug fix which is worth upstreaming. I would suggest to submit it to the upstream IMEX repo. @kurapov-peter What do you think?

[dchigarev]

The patch LGTM. I am happy with merging this to our development branch. However, this looks like a bug fix which is worth upstreaming. I would suggest to submit it to the upstream IMEX repo. @kurapov-peter What do you think?

Yes, I also wanted to push this fix to IMEX, but in my understanding the time from PR submission until it's merged would be longer in case of the main IMEX branch (comparing with our dev branch). So I though I'll push this fix to our dev first and then upstream to IMEX in order to save time. Do you think I should do this the other way around (first to main IMEX and the port it to our dev branch)?

[Menooker]

The patch LGTM. I am happy with merging this to our development branch. However, this looks like a bug fix which is worth upstreaming. I would suggest to submit it to the upstream IMEX repo. @kurapov-peter What do you think?

Yes, I also wanted to push this fix to IMEX, but in my understanding the time from PR submission until it's merged would be longer in case of the main IMEX branch (comparing with our dev branch). So I though I'll push this fix to our dev first and then upstream to IMEX in order to save time. Do you think I should do this the other way around (first to main IMEX and the port it to our dev branch)?

Sure we can merge it to the development branch. I would suggest to apply the PR to a new branch "dev", which I have created in this repo just now. I have submitted this branch NewLower to IMEX upstream and there is a PR in that repo to track NewLower branch. We'd better not touch this branch and move future development on the new "dev" branch.