Repository: pytorch
Issue: [compile] Tracker for `torchrec_dlrm` issues
Labels: triaged,module: custom-operators,module: functionalization,oncall: pt2,module: pt2-dispatcher

First issue
~~~
import torch
from dataclasses import dataclass, field
from typing import Any

@dataclass
class DClass():
   sharding_contexts: Any = field(default_factory=list)
   a: int = 1

def fn(x):
    return DClass().a * x

opt_fn = torch.compile(fn)
x = torch.randn(4)
opt_fn(x)

~~~

-----------------


Second issue

~~~
from math import inf
import torch
from torch import tensor, device
import torch.fx as fx
import torch._dynamo
from torch._dynamo.testing import rand_strided
from torch._dynamo.debug_utils import run_fwd_maybe_bwd

import torch._dynamo.config
import torch._inductor.config
import torch._functorch.config
import fbgemm_gpu


torch._inductor.config.fallback_random = True






from torch.nn import *
class Repro(torch.nn.Module):
    def __init__(self):
        super().__init__()



    def forward(self, L_L_indices_ : torch.Tensor, L_L_offsets_ : torch.Tensor, L_L_self_rows_per_table_ : torch.Tensor, L_L_self_bounds_check_warning_ : torch.Tensor):
        l_l_indices_ = L_L_indices_
        l_l_offsets_ = L_L_offsets_
        l_l_self_rows_per_table_ = L_L_self_rows_per_table_
        l_l_self_bounds_check_warning_ = L_L_self_bounds_check_warning_
        long = l_l_indices_.long();  l_l_indices_ = None
        long_1 = l_l_offsets_.long();  l_l_offsets_ = None
        bounds_check_indices = torch.ops.fbgemm.bounds_check_indices(l_l_self_rows_per_table_, long, long_1, 1, l_l_self_bounds_check_warning_, None);  l_l_self_rows_per_table_ = long = l_l_self_bounds_check_warning_ = None
        return [long_1]


mod = Repro()

def load_args(reader):
    buf0 = reader.storage('b34ae4d796654d6804a16e92852d0395c29d3ea7', 832, device=device(type='cuda', index=0), dtype_hint=torch.int64)
    reader.tensor(buf0, (104,), dtype=torch.int64)  # L_L_indices_
    buf1 = reader.storage('0aca75493e9fed822d70e8524e19324441bfaa4d', 840, device=device(type='cuda', index=0), dtype_hint=torch.int64)
    reader.tensor(buf1, (105,), dtype=torch.int64)  # L_L_offsets_
    buf2 = reader.storage('d0cca63d71b4e5a07f9cb3f2fbd5e3aa1c5e4c35', 208, device=device(type='cuda', index=0), dtype_hint=torch.int64)
    reader.tensor(buf2, (26,), dtype=torch.int64)  # L_L_self_rows_per_table_
    buf3 = reader.storage('05fe405753166f125559e7c9ac558654f107c7e9', 8, device=device(type='cuda', index=0), dtype_hint=torch.int64)
    reader.tensor(buf3, (1,), dtype=torch.int64)  # L_L_self_bounds_check_warning_
load_args._version = 0

if __name__ == '__main__':
    from torch._dynamo.repro.after_dynamo import run_repro
    run_repro(mod, load_args, accuracy=False, command='run',
        save_dir='/scratch/anijain/work/gcc-10/pytorch/checkpoints', autocast=False, backend='aot_eager')

~~~

~~~
RuntimeError: Found a custom (non-ATen) operator that either mutates or its inputs: fbgemm::bounds_check_indices.. Getting these operators to work with functionalization requires some extra work. For mutable ops you need to register a corresponding out-of-place variant of the op, and you also need to register a Functionalization kernel that performs some boilerplate, telling functionalization to map from the mutable op to the out-of-place op. See a more complete example of how to do this at https://gist.github.com/bdhirsh/7dadbf6296f8f7d1abcf4c482f438aaa. Please file a GitHub issue if you run into any problems.
~~~
-----------

Third-issue - Bad Fx Graph
Repro - `TORCHDYNAMO_DEBUG_FUNCTION="invoke" TORCH_LOGS="+dynamo" python benchmarks/dynamo/torchbench.py  --backend=aot_eager --float32 --inference --device cuda  --accuracy --only torchrec_dlrm`

The `TORCHDYNAMO_DEBUG_FUNCTION` flag calls Dynamo only on the relevant function.
 
---------



cc @bdhirsh @ezyang @msaroufim @wconstab @zou3519 @ngimel