Torch compile with custom backend data from host to device for each epoch

Tiru_B: <p>Hi All</p>
<p>I have integrated our custom accelerator to pytorch 2.0. Our integration approach is similar to torch  XLA, which will scale well for Inference. But For Training i have few doubts. Though my  integration approach works but its not scalable.</p>
<p>Here is the issue.</p>
<p>Torch Dynamo give forward graph with output and its activation outputs, which will be used in bacward process.</p>
<p>so  My graph will  be like below  ```<br>
class GraphModule(torch.nn.Module):<br>
def forward(self, primals_1: “f32[64, 100]”, primals_2: “f32[64]”, primals_3: “f32[10, 64]”, primals_4: “f32[10]”, primals_5: “Sym(s0)”, primals_6: “f32[s0, 100]”):<br>
# File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:19 in forward, code: x = self.fc1(x)<br>
permute: “f32[100, 64]” = torch.ops.aten.permute.default(primals_1, [1, 0]);  primals_1 = None<br>
mm: “f32[s0, 64]” = torch.ops.aten.mm.default(primals_6, permute);  permute = None<br>
add: “f32[s0, 64]” = torch.ops.aten.add.Tensor(mm, primals_2);  mm = primals_2 = None</p>
<pre><code>    # File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:20 in forward, code: x = self.relu(x)
    relu: "f32[s0, 64]" = torch.ops.aten.relu.default(add);  add = None
    
    # File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:21 in forward, code: x = self.fc2(x)
    permute_1: "f32[64, 10]" = torch.ops.aten.permute.default(primals_3, [1, 0]);  primals_3 = None
    mm_1: "f32[s0, 10]" = torch.ops.aten.mm.default(relu, permute_1)
    add_1: "f32[s0, 10]" = torch.ops.aten.add.Tensor(mm_1, primals_4);  mm_1 = primals_4 = None
    return [add_1, primals_6, relu, permute_1, primals_5]
</code></pre>
<pre><code class="lang-auto">which has 5 outputs so my hardware also returns 5 outputs, which will be consumed by backward graph.

</code></pre>
<p>class GraphModule(torch.nn.Module):<br>
def forward(self, primals_5: “f32[32, 100]”, relu: “f32[32, 64]”, permute_1: “f32[64, 10]”, tangents_1: “f32[32, 10]”):<br>
# File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:20 in forward, code: x = self.relu(x)<br>
le: “b8[32, 64]” = torch.ops.aten.le.Scalar(relu, 0)<br>
full_default: “f32<span class="chcklst-box fa fa-square-o fa-fw"></span>” = torch.ops.aten.full.default(<span class="chcklst-box fa fa-square-o fa-fw"></span>, 0.0, dtype = torch.float32, layout = torch.strided, device = device(type=‘cpu’), pin_memory = False)</p>
<pre><code>    # File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:21 in forward, code: x = self.fc2(x)
    permute_2: "f32[10, 64]" = torch.ops.aten.permute.default(permute_1, [1, 0]);  permute_1 = None
    mm_2: "f32[32, 64]" = torch.ops.aten.mm.default(tangents_1, permute_2);  permute_2 = None
    
    # File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:20 in forward, code: x = self.relu(x)
    where: "f32[32, 64]" = torch.ops.aten.where.self(le, full_default, mm_2);  le = full_default = mm_2 = None
    
    # File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:19 in forward, code: x = self.fc1(x)
    permute_6: "f32[64, 32]" = torch.ops.aten.permute.default(where, [1, 0])
    mm_4: "f32[64, 100]" = torch.ops.aten.mm.default(permute_6, primals_5);  permute_6 = primals_5 = None
    permute_7: "f32[100, 64]" = torch.ops.aten.permute.default(mm_4, [1, 0]);  mm_4 = None
    permute_8: "f32[64, 100]" = torch.ops.aten.permute.default(permute_7, [1, 0]);  permute_7 = None
    sum_2: "f32[1, 64]" = torch.ops.aten.sum.dim_IntList(where, [0], True);  where = None
    view_1: "f32[64]" = torch.ops.aten.reshape.default(sum_2, [64]);  sum_2 = None
    
    # File: /home/tiru/work/tvm/python/torch_pac/utests/test_MLP_train.py:21 in forward, code: x = self.fc2(x)
    permute_3: "f32[10, 32]" = torch.ops.aten.permute.default(tangents_1, [1, 0])
    mm_3: "f32[10, 64]" = torch.ops.aten.mm.default(permute_3, relu);  permute_3 = relu = None
    permute_4: "f32[64, 10]" = torch.ops.aten.permute.default(mm_3, [1, 0]);  mm_3 = None
    permute_5: "f32[10, 64]" = torch.ops.aten.permute.default(permute_4, [1, 0]);  permute_4 = None
    sum_1: "f32[1, 10]" = torch.ops.aten.sum.dim_IntList(tangents_1, [0], True);  tangents_1 = None
    view: "f32[10]" = torch.ops.aten.reshape.default(sum_1, [10]);  sum_1 = None
    return [permute_8, view_1, permute_5, view, None]
</code></pre>
<pre><code class="lang-auto">
Since the symbols are same between the graph, i can make my hardware to run without copy to host and host to device betwee the graph for each run.

But, runtime wrapper on top  aot_dispatch_base will have check how may inputs are returning from the graph, becuase of this check i have return  those tensor to this run time wrapper. which i needs copy from device. I have to copy back device for backwar.

So Is there a way from the application it self to tell don't copy data between the graph runs, When ever user request to read data from device. then only reads from device.

kind of egar execution. I can do this thing with our own APIS, not depended on Torch flow. But i want similar thing from torch.</code></pre>