Add support for torch.compile in RGCNConv

[akihironitta]

Fixes #8467.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (b3f0eb3) 89.25% compared to head (691b270) 89.26%.

Additional details and impacted files

@@           Coverage Diff           @@
##           master    #8783   +/-   ##
=======================================
  Coverage   89.25%   89.26%           
=======================================
  Files         467      467           
  Lines       29919    29928    +9     
=======================================
+ Hits        26705    26714    +9     
  Misses       3214     3214           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[rusty1s]

Wouldn't we need to reset it somewhere?

[rusty1s]

Sorry for asking so many questions. What was the reason for the graph break in the first place? Computing the heuristic? Can't we just do

if is_compiling():
    use_segment_matmul = False
else:
    use_segment_matmul = torch_geometric.backend.use_segment_matmul
    if use_segment_matmul is None:
        pass

etc

[akihironitta]

I'll keep this PR on hold, as it makes more sense to first update the segment matmul heuristic (what @puririshi98 did in #8615). I see two graph breaks in the test case due to:

Break Reasons:
  Break Reason 1:
    Reason: call_function BuiltinVariable(int) [TensorVariable()] {}
    User Stack:
      <FrameSummary file /home/akihiro/work/github.com/pyg-team/pytorch_geometric/torch_geometric/nn/conv/rgcn_conv.py, line 235 in forward>
  Break Reason 2:
    Reason: call_function BuiltinVariable(bool) [TensorVariable()] {}
    User Stack:
      <FrameSummary file /home/akihiro/work/github.com/pyg-team/pytorch_geometric/torch_geometric/backend.py, line 55 in use_segment_matmul_heuristic>